Fix/improve mtp performance

[SamuelOliveirads]

Hey @F1LM1, it's been my turn to be busy lately, so I haven't made much progress. However, this gave me time to rethink the future of this PR, and I decided to clean up the code for a proper review. Here are my thoughts:

• The core idea of the original PR was the MTP implementation. Although performance isn't currently better than having it disabled, it is functional, which I see as a solid contribution to the project.
• This PR has been open for almost four months, and the next set of improvements could take just as long to finish.

I took the latest improvements we made and tested each one with small prompts. This new PR is the result. With my configuration, I'm getting about 84% of the tokens per second compared to running without MTP. I also reviewed the entire project to remove unused code and references, and finally added a command-line argument so everyone can easily try it with or without MTP.

My suggestion is as follows:

1. You review this cleanup, and if it looks good, merge it into the original PR branch.
2. Mark the original PR as ready for review. If necessary, I can tag some of the frequent reviewers to take a look.
3. I will be available to fix any issues found in the original PR and focus on getting it merged. In parallel, I will keep working on the server loop optimization and multi-token MTP in PR Glm4 mtp optimizations #4.

[F1LM1]

Sounds good to me, I'll look over it more formally over the next 2-3 days.

[wishstudio]

I was working on the server loop optimization and just noticed that ngxson has already created a PR lately: ggml-org#17808

It looks very good and hopefully can be merged soon. After that perhaps you can start rebasing this entire repo. There has been a lot of traffic since the creation of this repo so just merging the changes will already need some work.

By the way, could you share your test method (cmdline, gguf, system configuration)? Since I already have a positive gain using this branch it puzzles me that it is still a regression on your system. If you could share the info I can try to reproduce and see if I can find anything suspicious.

[SamuelOliveirads]

(The name of the branch bothered me, but unfortunately the PR closes if I try to rename it... well, it's going to be MT for now)

By the way, could you share your test method (cmdline, gguf, system configuration)? Since I already have a positive gain using this branch it puzzles me that it is still a regression on your system. If you could share the info I can try to reproduce and see if I can find anything suspicious.

Sure. I usually don't test under best-case scenarios. For development, I use Windows, which incurs a slight performance drop. Right now, my specs are a Threadripper 5965WX and two RTX 3090s. I tried months ago with a Ryzen 7 5700X and noticed the same pattern. I was expecting a offload system to no be great, but it still bothers me that performance without MTP is higher. Here are the commands:

cmake -B build -DCMAKE_BUILD_TYPE=Debug -DLLAMA_CURL=OFF -DGGML_CUDA=ON

cmake --build build --config Release -j

.\build\bin\release\llama-server.exe ^
    --model "F:\llm_models\glm-4.5-air_Q4_general\GLM-4.5-Air-IQ4_XS-00001-of-00002.gguf" ^
    --alias GLM-4.5-Air ^
    --ctx-size 36864 ^
    -ctk q8_0 -ctv q8_0 ^
    -fa --verbose ^
   --n-gpu-layers 99 ^
    -b 2048 -ub 1500 ^
   -ot "blk\.(3|4|5|6|7|8|9|10|11|12|13|14|15|16|17|18)\.ffn_.*=CUDA0" ^
    -ot "blk\.(20|21|22|23|24|25|26|27|28|29|30|31|32|33|34)\.ffn_.*=CUDA1" ^
   --override-tensor exps=CPU ^
    -mtp ^
    --threads 24 --threads-batch 36 ^
    --host 127.0.0.1 ^
    --port 8080

I haven't tried using Linux for this branch or with GLM 4.5/4.6 yet, and I also didn't fully optimize for NUMA. That's why I'm eager to see how it performs for others.
llama-bench is not the best tool for testing, at least for me, so I typically try with small prompts and sometimes with code generation or creative writing. On average, small prompts gave 12.19 T/s. Using this branch originally gave me 10.52 T/s, and now it reaches up to 11.05 T/s, which is about 90% of the original performance.

I was working on the server loop optimization and just noticed that ngxson has already created a PR lately: ggml-org#17808

It looks very good and hopefully can be merged soon. After that perhaps you can start rebasing this entire repo. There has been a lot of traffic since the creation of this repo so just merging the changes will already need some work.

That's great to hear! I made negligible improvements to the server loop—nothing that yielded significant gains. I will update the branch, take a look at the mentioned PR, and then create a new one that will allow more than one draft per loop.

[wishstudio]

Thank you for the information! I will take a look. Here are my first intuitions:

In a --cpu-moe setup typically the CPU part is memory bound. IIRC the IQ quants are computational heavier than regular Qx_0 or Qx_K quants. So using that may cause it to shift to computational bound which incur more performance penalties.

To validate this, you can either:

• Try a computational lighter quant, like Q4_K, or better Q8_0.
• Log the time spent in MoE FFNs and check if it is memory bound. You can easily measure by calculating the time spent in the loop for (int cur_a = 0; cur_a < n_as; ++cur_a) in ggml_compute_forward_mul_mat_id function in ggml-cpu.c. You can then compare of the time with the theoretical time of a FFN block of your used quant (calculate the size in bytes of the corresponding weight matrix, then divide it by your actual memory bandwidth, which can be measured by intel MLC tool). You may need to also log the dimensions of ids tensors to do it correctly since there will be some shared experts in consecutive tokens. In my experiments for my gaming rig, the actual time spent in FFN blocks are quite close to theoretical time.

You mentioned NUMA which I think is not the major problem here as although generally speaking llama.cpp is not NUMA aware, the inter CCD bandwidth in a single CPU is still much higher than typical inter socket bandwidth.

The major additional bottleneck with Windows to my knowledge is much higher kernel launch overhead. CUDA graphs can be used to eliminate this but it is currently disabled in split graphs. Anyway it's irrelevant here because you basically have the same overhead with or without MTP.

Another important limitation in Windows is it only allows half of physical RAM to be used as pinned memory. That typically hurt prefill performance in --cpu-moe but not tg. But I guess it's also irrelevant to MTP performance comparison. You can try if adding --no-mmap makes any difference, without the option it will not use pinned memory for model weights.

[F1LM1]

I will need to test more carefully but I'm also using a cpu MoE setup and my tg/s goes from ~10.2 off to ~11.5 on. Haven't yet tried to optimize settings for the on case.

One very loose thought is

ctx-size 36864

This is a weird number; does that mean you're usually absolutely maxing out RAM/VRAM in base configuration? Loading the MTP layer will require a bit of memory so maybe you're filling up VRAM or something?

[SamuelOliveirads]

• Try a computational lighter quant, like Q4_K, or better Q8_0.
• Log the time spent in MoE FFNs and check if it is memory bound. You can easily measure by calculating the time spent in the loop for (int cur_a = 0; cur_a < n_as; ++cur_a) in ggml_compute_forward_mul_mat_id function in ggml-cpu.c. You can then compare of the time with the theoretical time of a FFN block of your used quant (calculate the size in bytes of the corresponding weight matrix, then divide it by your actual memory bandwidth, which can be measured by intel MLC tool). You may need to also log the dimensions of ids tensors to do it correctly since there will be some shared experts in consecutive tokens. In my experiments for my gaming rig, the actual time spent in FFN blocks are quite close to theoretical time.

That's good to know, I will measure it after I finish the next PR. At least right now for me the performance using the Air variant doesn't bother me. My go-to is still the GLM 4.6 in Linux, I just prefer to code using Windows. Still, it's good to see if there is something hurting my performance that doesn't allow MTP to increase higher than without. Also I'm actually curious to try a Q1 just to fit fully in GPU and see how it works.

Another important limitation in Windows is it only allows half of physical RAM to be used as pinned memory. That typically hurt prefill performance in --cpu-moe but not tg. But I guess it's also irrelevant to MTP performance comparison. You can try if adding --no-mmap makes any difference, without the option it will not use pinned memory for model weights.

True, and I saw that when I used GLM 4.5 for the first week; it was around 3 tokens that jumped to 5 by disabling mmap.

This is a weird number; does that mean you're usually absolutely maxing out RAM/VRAM in base configuration? Loading the MTP layer will require a bit of memory so maybe you're filling up VRAM or something?

Yes, when the model was released I used it quite a lot and that context filled the entire VRAM. Today I don't need it anymore, but because I borrowed the old args I didn't change them, and as far as I remember I removed some layers from both GPUs to free memory. But I need to double check if it is offloading to memory again and hurting the performance.

[SamuelOliveirads]

@F1LM1, how is the review going? I hope everything is alright. I’ve submitted #5 and #6 , which should hopefully finalize the MTP implementation—or at least provide a solid foundation for future contributions. If you need any clarification or further changes, please let me know and I will gladly help!

[F1LM1]

@F1LM1, how is the review going? I hope everything is alright. I’ve submitted #5 and #6 , which should hopefully finalize the MTP implementation—or at least provide a solid foundation for future contributions. If you need any clarification or further changes, please let me know and I will gladly help!

Looks really good, super polished compared to what we had before, and I'm also getting reproducible uplift when drafting N = 1 tokens, in the range of 6-10% over baseline (and notably the baseline is better than what it was on the old base). For whatever reason #5 was a bit buggy for me but #6 is not. In most of my test prompts I see the 1st, 2nd, and 3rd speculative token with around 75%, 50%, and 25% acceptance, respectively. I'm not sure if it has to do with overhead in CPU-only or hybrid inference modes but I'm a bit surprised that N > 1 doesn't perform better than it does (on my machine) especially given your overhead optimizations. But maybe it's a sign that things can still improve from here.

Anyway, I have two hopefully very easy suggestions. The first is more important:

• When MTP is off, we still load the MTP layer, which is of course just a waste of memory for anyone who doesn't want to use MTP. And if someone is running their setup "on the edge" in terms of RAM/VRAM this might just tip them over the edge which is a bit annoying; best that we don't accidentally cause any regressions.
• At least on my machine and test prompts, the added overhead of heavy sampling (when drafting) takes this feature from a net positive to a negative again, while in most cases "always greedy sampling" doesn't seem to drop draft acceptance rates by too much. Maybe let's add an override option to allow using the greedy sampler even when reptition/presence/etc penalties are turned on. Honestly maybe even default to using the greedy sampler since in the vast majority of my test cases, either the greedy sampler is better, or (with extreme sampler settings) draft acceptance rates are so bad that there's no point in using MTP to begin with.

[SamuelOliveirads]

Glad to hear that it's working well for your tests!

• When MTP is off, we still load the MTP layer, which is of course just a waste of memory for anyone who doesn't want to use MTP. And if someone is running their setup "on the edge" in terms of RAM/VRAM this might just tip them over the edge which is a bit annoying; best that we don't accidentally cause any regressions.
• At least on my machine and test prompts, the added overhead of heavy sampling (when drafting) takes this feature from a net positive to a negative again, while in most cases "always greedy sampling" doesn't seem to drop draft acceptance rates by too much. Maybe let's add an override option to allow using the greedy sampler even when reptition/presence/etc penalties are turned on. Honestly maybe even default to using the greedy sampler since in the vast majority of my test cases, either the greedy sampler is better, or (with extreme sampler settings) draft acceptance rates are so bad that there's no point in using MTP to begin with.

I always had the feeling that sampling might not give enough benefits for us, but at least in my testing using GLM 4.5 Air and 4.6, it didn't give me detrimental performance. For Air it's inconclusive (basically the same), but using version 4.6 it's an improvement of 18% in draft rate and almost 6% better t/s. Still, because not everyone can get good performance, I preferred to remove it.

To make things easier, I applied both changes in the last PR.

Looks really good, super polished compared to what we had before, and I'm also getting reproducible uplift when drafting N = 1 tokens, in the range of 6-10% over baseline (and notably the baseline is better than what it was on the old base). For whatever reason #5 was a bit buggy for me but #6 is not. In most of my test prompts I see the 1st, 2nd, and 3rd speculative token with around 75%, 50%, and 25% acceptance, respectively. I'm not sure if it has to do with overhead in CPU-only or hybrid inference modes but I'm a bit surprised that N > 1 doesn't perform better than it does (on my machine) especially given your overhead optimizations. But maybe it's a sign that things can still improve from here.

Many things can impact performance. For me, 2 tokens go around a 66% rate, which gives me better performance than 1. Still, MTP is a new architecture in llama.cpp, so it's expected to lack minor but important improvements. Think about the differences between the old and newer baselines; after the GLM 4.5 implementation, many folks made a lot of small improvements that helped us get better performance in #6. This will definitely replicate in MTP when ggml-org#15225 becomes merged.

I finished my to-do list for the MTP implementation and I hope that this will be enough to make the original PR merged.