Feat/glm4 mtp recursive

[SamuelOliveirads]

Note: Please be aware that this PR shows a high commit count because it was rebased against the latest master after a long period of inactivity.

Recursive MTP Drafting & Graph Optimization for GLM-4.5

This is a follow-up PR to F1LM1/llama.cpp#5, implementing major architectural improvements for Multi-Token Prediction (MTP).

Summary of Changes

1. Multi-Token Support: Enabled recursive drafting loop (Eagle-style) to predict N tokens per step.
2. Graph Optimization: Activated graph reuse and fixed callback handling for MTP operations, significantly reducing CPU/scheduler overhead.
3. Rebase: Updated the branch to the latest master (including server-loop optimizations from server: improve speed of speculative decoding ggml-org/llama.cpp#17808).

---

Detailed Changes

1. Recursive MTP Drafting

Previously, MTP was limited to a single draft token per loop, capping the theoretical maximum performance. This PR implements a recursive generation loop based on the architecture found in sglang#8224.

Architecture:
The MTP draft is not a parallel tensor operation but a recursive loop. The draft inputs require the embedding of the current token and the fixed hidden state from the main model's last interaction.

1. Main Model: Processes Prompt ($t_0...t_n$) $\rightarrow$ Outputs Hidden State ($H_{main}$) and Token ($t_{n+1}$).
2. MTP Loop (Draft 1): Input ($$Emb(t_{n+1})$$) + Context ($H_{main}$) $\rightarrow$ MTP Layer $\rightarrow$ Token ($t_{n+2}$).
3. MTP Loop (Draft 2): Input ($$Emb(t_{n+2})$$) + Context ($H_{main}$) $\rightarrow$ MTP Layer $\rightarrow$ Token ($t_{n+3}$).
4. Verification: Validates the sequence against the main model.

Controls:
I reused common_speculative_params:

• --draft-max N: Sets the maximum number of recursive drafts.
• --draft-p-min P: Confidence threshold. The loop exits early if the MTP model is not confident, saving compute resources.

Recommended Test Command:

-mtp --draft-max 3 --draft-p-min 0.85

(Note: I found better results by limiting draft-max and increasing p-min to avoid low-quality speculation).

2. Rebase & Modernization

I updated the codebase to include the latest commits from master (approx. 5 months of updates). This includes the recent server-loop refactor (ggml-org#17808). While the server refactor itself didn't yield massive gains in my specific MTP tests, the overall codebase improvements combined with my MTP fixes have boosted baseline performance.

3. Graph Reuse Optimization

A major bottleneck was the lack of graph reuse during MTP switching.

• Problem: Switching between Main Model computation (~10ms) and MTP computation (~9.6ms) forced a full graph rebuild and allocation every step.
• Fix: Added proper cb() callbacks to the MTP tensors in build_mtp_tail. This allows the scheduler to correctly offload these tensors and reuse the compute graph structure.
• Result: Subsequent drafts in the loop are significantly faster (e.g., ~2.4ms vs 9.6ms) because the graph overhead is amortized.

---

Performance

Tested on a Windows workstation (Threadripper 5965WX + 2x RTX 3090).
Settings: Small prompt (8 tokens), average of 5 interactions.

Old Branch:

• Baseline (No MTP): ~12.19 t/s
• MTP (1 draft): ~10.52 t/s (Regression)
• MTP (2 drafts): ~10.63 t/s

Updated Branch (This PR):

• Baseline (No MTP): ~12.89 t/s
• MTP (1 draft): ~13.14 t/s (Gain)
• MTP (2 drafts): ~13.54 t/s (Gain)
• MTP (3 drafts): ~12.56 t/s (Diminishing returns due to rejection)

Update (12/13):

After further testing following recent reports, I re-evaluated the performance. I was unable to replicate the regression reported, but my results differ from the initial run.

It is worth noting that this branch (even without MTP enabled) performs similar than master in my tests.

GLM-4.5 Air:

• Master: ~17.60 t/s
• Branch (No MTP): ~18.10 t/s
• MTP (1 Draft): ~15.64 t/s
• MTP (2 Drafts): ~15.63 t/s
• MTP (3 Drafts): ~15.97 t/s

GLM-4.6:

• Master: ~5.18 t/s
• Branch (No MTP): ~5.47 t/s
• MTP (1 Draft): ~5.91 t/s
• MTP (2 Drafts): ~5.78 t/s
• MTP (3 Drafts): ~5.41 t/s

Note: I updated the launch command to explicitly offload the MTP layer (layer 46 for Air, 92 for 4.6) to the GPU, which provided a small performance boost.

Commands Used:

.\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 1 --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|46)\.ffn_.*=CUDA1" ^
    --override-tensor exps=CPU ^
    --no-mmap  ^
    --numa distribute ^
    -mtp --draft-max 3 --draft-p-min 0.85 ^
    --threads 24 --threads-batch 36 ^
    --host 127.0.0.1 --port 8080

.\build\bin\release\llama-server.exe ^
    --model "F:\llm_models\glm-4.6\GLM-4.6-UD-IQ1_S-00001-of-00002.gguf" ^
    --alias GLM-4.6 ^
    --ctx-size 36864 ^
    -ctk q8_0 -ctv q8_0 ^
    -fa 1 --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|19)\.ffn_.*=CUDA0" ^
    -ot "blk\.(20|21|22|23|24|25|26|27|28|29|30|31|32|33|92)\.ffn_.*=CUDA1" ^
    --override-tensor exps=CPU ^
    --no-mmap  ^
    --numa distribute ^
    -mtp --draft-max 3 --draft-p-min 0.85 ^
    --threads 24 --threads-batch 36 ^
    --host 127.0.0.1 --port 8080

Conclusion

This update finally achieves net positive performance over the baseline. I am personally eager to try this on GLM-4.6 with better quantization. I invite others to test this configuration and report findings, as MTP logic is complex and may behave differently across hardware setups.

Also tagging @InfernalDread, as you previously mentioned interest in trying MTP once the performance improved. This update finally yields positive gains over the baseline, so I would love to hear your feedback if you have time to test it.

Transparency Note: I used an LLM to help organize and polish the description of this PR to ensure the architectural changes were explained clearly.

[Stealt91]

Awesome! Thank you for the effort to make this actually worth using! Can't wait for further improvements!

[InfernalDread]

Hello again! Thank you for the mention! I will definitely be trying out these new improvements, can't wait to see what's to come!

[timkhronos]

Hello! I have been meaning to try out these new improvements, however I get a segmentation fault during the MTP pass.

Any idea what could be causing this? I am using a single gpu + cpu linux system, and a q4 quant of 4.6.

I also get the following warnings when building.

[SamuelOliveirads]

Any idea what could be causing this? I am using a single gpu + cpu linux system, and a q4 quant of 4.6.

@timkhronos I was able to replicate the problem. It occurs because you are using GLM-4.6, which wasn't supported in this PR (most of the work was done on an older branch before 4.6 existed).

I have downloaded a GLM-4.6 quant to test and will look deeper to understand and fix the incompatibility. In theory, the architecture should be the same (or very similar) to the 4.5 variant, so I should be able to patch it.

By the way, if anyone wants to test MTP right now, please use the GLM-4.5 variant, at least the Air works okay, probably the full version also.

[SamuelOliveirads]

GLM-4.6 has a small architectural difference compared to the 4.5 variants, specifically in the MTP layer (shared weights vs separate). I just pushed a fix for that and it's working on my end now. Anyone who wants to try MTP with GLM-4.6 should be able to do so without issues

[timkhronos]

First of all, thank you.

After testing, this pr seems to suffer large performance regressions, both in TG and PP, even when not using '-mtp' in the launch arguments. On the latest mainline build I get around 6.2 tok/s TG and 100~120 PP. On this Pr without -mtp, I get around 2.2 tok/s TG and 30 PP. With -mtp I get around 2.7 tok/s TG, if I offload the MTP tensor{92} to the GPU.

[SamuelOliveirads]

After testing, this pr seems to suffer large performance regressions, both in TG and PP, even when not using '-mtp' in the launch arguments. On the latest mainline build I get around 6.2 tok/s TG and 100~120 PP. On this Pr without -mtp, I get around 2.2 tok/s TG and 30 PP. With -mtp I get around 2.7 tok/s TG, if I offload the MTP tensor{92} to the GPU.

That's curious. I ran a couple of tests using main (2fbe3b7) and this branch across different commits to see if I had introduced changes that degrade performance, but I couldn't replicate the problem.

In my tests using GLM-4.5-Air and GLM-4.6, both main and this branch (without MTP enabled) gave statistically identical performance. I have updated the PR description to show my latest results.

I did find some interesting behavior with the Air model: the speed varies quite a bit between runs, even on the same commit. I don't know exactly what causes this—perhaps it's related to how the graph is constructed each time the model loads or even because I'm testing to few examples—but I didn't have this issue with the 4.6 version, which was very consistent.

I also want to thank you for the idea of fully loading the MTP layer on the GPU. It gave me a small but measurable performance boost (jumping from around 14.40 t/s to between 15.6 and 16.0 t/s).