model: add NextN/MTP speculative decoding support for GLM_DSA (GLM-5.2)

[satindergrewal]

Overview

Adds GLM-5.2's multi-token-prediction (NextN) head as a draft-mtp speculative target for the existing GLM_DSA architecture. This is the fifth in-tree draft-mtp implementation (after qwen35, qwen35moe, step35, cohere2moe). It complements #25407 (GLM 5.2 Indexer); see "Relationship to #24868" below for the GLM-family code-sharing question.

• src/models/glm-dsa.cpp: nextn tensor loading is un-skipped via the qwen35moe/step35-style presence probe (mtp_only / trunk_only detection, so trunk-only GGUFs keep skipping). Adds the full graph_mtp builder: the qwen35moe frame (embd_h input, enorm/hnorm, concat, eh_proj, shared_head fallbacks), deepseek2's MLA attention (absorbed wk_b/wv_b, exact YaRN mscale^2 kq_scale), and deepseek2's sigmoid-gated MoE plus shared-expert FFN. The MTP layer runs dense MLA; the DSA lightning indexer is not used there, matching the GLM_DSA trunk graph. Also fixes the 744B_A40B type string (previously "?B").
• src/models/models.h: graph_mtp declaration.
• src/models/deepseek2.cpp: sets res->t_h_nextn after output_norm in the trunk graph (shared by DEEPSEEK2/MISTRAL4, guarded so MTP-off graphs stay bit-identical to master).
• src/llama-model.cpp: KV filter clause extended to GLM_DSA, so the MTP context allocates KV only for the nextn layer.
• Converter (conversion/glm.py): GlmMoeDsaModel opts into supports_mtp_export (post-convert_hf_to_gguf.py: Refactor how --mtp / --no-mtp is enabled for a model #25641 class-flag shape, mirroring HYV3Model and Step35Model). --no-mtp strips the NextN block and its nextn_predict_layers KV; --mtp emits a draft-only GGUF (NextN block plus shared embeddings/norm/lm_head). Default bundled output is unchanged.
• No new speculative-driver code: the existing --spec-type draft-mtp machinery runs it unchanged.

GLM-5.2's GGUF block_count=79 includes the nextn layer, so n_layer()=78 and the nextn tensors live at blk.78. Existing community GGUFs already carry these tensors and load unmodified.

Validation

Hardware: 2x RTX PRO 6000 Blackwell (96 GB each), CUDA 12.9, SM120a, -sm layer. Model: GLM-5.2 744B-A40B, IQ1_S (the quant that fits fully in VRAM on this rig). Benchmark format: am17an's mtp-bench.py.

Full audit, benchmark scripts, raw outputs, charts, and the determinism control run are published here: https://huggingface.co/datasets/satgeze/glm-dsa-mtp-audit

mtp-bench.py house format, baseline vs MTP (server llama-server -m GLM-5.2-IQ1_S.gguf -ngl 99 -sm layer -c 32768 -np 1 -fa on --jinja [--spec-type draft-mtp]):

Config	Aggregate acceptance	Aggregate speed
baseline (no spec)	n/a	57.7 tok/s
MTP, stochastic sampling (n_max=3, seed 42)	0.5619	up to 77.7 tok/s per task
MTP, greedy ungated (temp 0, top_k 1, p_min 0)	0.6621	57.7 to ~79 tok/s, about 1.37x

Per-position acceptance (greedy, ungated, common_speculative_print_stats): 0.834 / 0.653 / 0.494, mean accepted length 2.98. Acceptance is healthy at all three draft positions, so the framework defaults (n_max=3, p_min=0) are the correct configuration for this model and no per-arch override is proposed. (Contrast hy_v3 #25395, whose single-depth head needed p_min=0.75.)

Real-workload field evidence: a 2-hour agentic coding session (320K ctx configured, q4_0 KV, -np 1) reached 264K context with thousands of draft/verify cycles and zero server-side failures, including clean cancel-on-client-disconnect. Final-task server log: draft acceptance = 0.86207 (150 accepted / 174 generated), mean len = 3.59.

Determinism note. Greedy MTP-on and MTP-off outputs are quality-equivalent but not byte-identical: batched speculative verification evaluates N tokens' logits in one forward pass, whose floating-point reduction order differs from sequential decode and flips argmax at near-tie tokens. This is framework-level behavior, not specific to this port. The same A/B on an already-merged draft-mtp architecture (hy_v3, greedy top_k=1 temp=0 seed=42) diverges identically, with that arch's own MTP genuinely engaged. Raw control logs are in the linked dataset.

Relationship to #24868 (GLM 4.x MTP)

#24868 adds llama_model_deepseek2::graph_mtp and plugs GLM-4.7-Flash (Glm4MoeLite) into draft-mtp. This PR is a different architecture (GLM_DSA: MLA with absorbed wk_b/wv_b plus DSA; the MTP layer runs dense MLA with the exact YaRN mscale^2 kq_scale) with a standalone glm_dsa::graph_mtp. It is not a duplicate, but there is a real code-sharing question: if maintainers prefer consolidating GLM-family MTP builders on deepseek2-reuse (the route #24868 and out-of-tree implementations took), I am glad to restructure onto whichever lands first.

Additional information

Reviewer notes, honestly flagged:

1. Acceptance is healthy at all three draft positions greedy and ungated (0.834 / 0.653 / 0.494, mean length 2.98), so the framework defaults are correct for this model. No per-arch default override is proposed.
2. VRAM: the MTP context roughly doubles per-context compute-buffer overhead. Measured ceiling on the 96 GB x2 rig with the IQ1_S quant is about 128K F16 KV with MTP on (512K q4_0 KV with MTP off). Documented so long-context users are not surprised. (Related: Fix memory allocation for MTP layers for avoid cudaMalloc error #25574 touches MTP-layer allocation.)
3. Acceptance was measured on IQ1_S, the only quant that fits this rig fully in VRAM. Higher quants should only improve acceptance (draft and target share weights). Perplexity versus the HF/vLLM reference was not run (weights are 744B; contributor hardware is workstation-class).
4. CUDA-only testing (2x SM120a). Other backends are unexercised beyond CI.
5. MTP prompt-processing overhead exists as documented by am17an in llama + spec: MTP Support  #22673, acknowledged and not specific to this architecture.
6. Over-context and OOM handling is graceful: an oversized prompt returns HTTP 400 exceed_context_size_error; generation past n_ctx context-shifts and finishes with finish_reason=length; a context too large for VRAM fails the KV allocation and exits cleanly with an error message. The null-context path is guarded by server : properly handle null llama_context #25868, which this branch includes.
7. Quantize: no change needed. The Bad layer 78 abort for GLM-5.2 nextn was already fixed upstream by quant : fix quantizing moe with mtp #24986 (n_layer_all FFN-counter seeding); blk.78 nextn tensors quantize through the standard path. quant: include MTP/NextN block when counting FFN layers (GLM-5.2) #24832 was a duplicate of that one-liner, closed after unrelated changes.

Wiring was verified operation-by-operation against vLLM's GLM MTP reference (which routes glm_moe_dsa through deepseek_mtp.py): concat order [embed, hidden], enorm/hnorm placement, eh_proj, attention dims and ranks and scale, MoE gating, shared-head order, and next-step hidden recycle. 8 of 9 items are exact matches. The one flagged divergence: the MTP layer runs dense MLA, matching this project's own GLM_DSA trunk, where vLLM uses the sparse DSA lightning indexer. This is a pre-existing whole-model design choice in llama.cpp, not an MTP wiring decision. The full audit is in the linked dataset.

Credits: ik_llama.cpp pioneered out-of-tree GLM-DSA MTP (build_deepseek2_mtp); this port validates against its behavior and follows mainline's draft-mtp framework (#18886 / #22673 / #23269 / #23287 / #23643). Landscape also surveyed: xj85770's reference patch and the Mesh-LLM sidecar implementation.

Regression:

• Behavior change for existing bundled GLM-5.2 GGUFs, flagged deliberately: nextn tensors flip from never-loaded to loaded-when-present (same presence-probe pattern as hy_v3 model: add Hy3 (hy_v3) support with MTP speculative decoding #25395, qwen35moe, step35). For a bundled GGUF this costs about 1 MoE layer of weights even when --spec-type draft-mtp is not used. Mitigations shipped here: --no-mtp at conversion strips the block, and --mtp produces a draft-only split. This matches the merged hy_v3 behavior; if maintainers prefer a load-time gate instead, happy to discuss.
• Generation behavior without --spec-type is unchanged (graph identical; nextn tensors are loaded but not executed).
• MTP-off graphs for DEEPSEEK2/MISTRAL4 (which share the trunk builder) stay bit-identical to master via the t_h_nextn guard.
• Other draft-mtp architectures are untouched.
• test-llama-archs: 436 OK / 0 failures, glm_dsa included.

AI usage disclosure

Yes. Full transparency: this implementation was written predominantly by an AI system (Anthropic Claude) working under the direction of @satindergrewal, who commissioned, reviewed, and hardware-tested the resulting builds over multi-hour real workloads, and who takes responsibility for the contribution. The port maps an existing implementation (ik_llama.cpp's glm-dsa MTP) onto mainline's existing draft-mtp framework and verifies wiring against vLLM's reference semantics rather than inventing new design. We understand this project restricts predominantly AI-generated contributions and we defer to the maintainers' judgment: if this cannot be accepted under policy, the PR still serves as a documented, validated reference for whoever implements it by hand, and we are glad to assist them.

[am17an]

C++ part looks okay to me.

PS - @satindergrewal while your work is reasonable your comments/descriptions are not, you should write your messages yourself rather than routing everything through Claude. I'm not sure you understand any of this code but I take responsibility for maintaining it

[satindergrewal]

@am17an You are right.
I won't pretend to understand all of this C++ code, as I come from mostly JavaScript/Golang code experience, but my work is very thorough, always.
I go deep into the subject and try to do my best to gather and do the work as much quality wise as I can. Do test on my hardware, rent cloud servers to do benchmarking and other work to finish off with proper measurements, so my work is not half-cooked presentation. That's my own personal high quality bar I kind of try to maintain, not just for this work, but almost everything I do in general.

Really appreciate you taking responsibility of this code. Much appreciate it. You are free to reject it if it still doesn't meet your standard of quality. However I'll be happy to amend and fix if notified about anything needs change or fixing.

I understand as maintainers it's hard to keep up with a flood of AI code pushes to the projects you maintain, that's fair to expect to have human worded communication.

I go through PR, issues, mainter's and other contributer's forks and their branches and git commits and make sure I am not pushing anything duplicate, and also try to offer help in existing PR and issues with my measurements, without even expecting credit in back.

[am17an]

Yes as I said your work is good, however it would be much better if you wrote your own words as it would deepen your understanding of the codebase. Not to mention the contributing guidelines explicitly forbid comments written by AI.

While I understand your stance is take or leave it, I'm suggesting a middle ground which would help you as well as the project. In the end there is no compulsion for a maintainer to accept a PR even if it is not AI written.

[satindergrewal]

Absolutely, and I have seen maintainers closing other PRs without merge 😅
Experienced that working in open source for many years.

But I has to take help from AI, because to be honest, I am not as competent at this as this work needs to communicate that level of detail. I will catch up, but too slow, and it's much faster if I give control to AI, and make sure it did not mess up. That's why I'm also holding back from pushing more things to project as I suspect folks here don't fully like AI and specially the communication back and forth. So, I'm sort of mixed mind if I push more PRs next time or not here.

[am17an]

If you plan to invest time long-term in this project, read the following - otherwise ignore.

My suggestion would be not to push anything you don't understand, as it doesn't really benefit you or the project longer term. You would be unable to make progress on higher levels of abstractions in the code-base without really understanding how it all works, even if line by line understanding of the code is not important.

and it's much faster if I give control to AI, and make sure it did not mess up

There is no harm in going slower and taking time to understand. In fact, that will allow you to go much faster in the future PRs.

[satindergrewal]

thank you for your kind words @am17an. I'll try my best 🙏

[ChicoPinto70]

Thank you, @satindergrewal !!!!

I ran this PR with SixVolts/GLM-5.2-Q3_K_M in a old Dual Xeon with a lot of tinkering (:-)) and I got a TG speed increase from ~6 T/s to ~7.1 T/s (in code generation).

[fairydreaming]

@satindergrewal Lightning indexer changes were merged recently, it looks like there are some conflicts now with the master to resolve.

Also please add _s tensor scale arguments in build_ffn() and build_moe_ffn() calls like in the current glm-dsa.cpp:

llama.cpp/src/models/glm-dsa.cpp

Lines 482 to 522 in 720d7fa

	if ((uint32_t) il < hparams.n_layer_dense_lead) {
	cur = build_ffn(cur,
	model.layers[il].ffn_up, NULL, model.layers[il].ffn_up_s,
	model.layers[il].ffn_gate, NULL, model.layers[il].ffn_gate_s,
	model.layers[il].ffn_down, NULL, model.layers[il].ffn_down_s,
	NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
	cb(cur, "ffn_out", il);
	} else {
	// MoE branch
	ggml_tensor * moe_out = build_moe_ffn(cur,
	model.layers[il].ffn_gate_inp,
	model.layers[il].ffn_up_exps,
	model.layers[il].ffn_gate_exps,
	model.layers[il].ffn_down_exps,
	model.layers[il].ffn_exp_probs_b,
	n_expert, n_expert_used,
	LLM_FFN_SILU, hparams.expert_weights_norm,
	hparams.expert_weights_scale,
	(llama_expert_gating_func_type) hparams.expert_gating_func,
	il,
	nullptr,
	model.layers[il].ffn_gate_up_exps,
	model.layers[il].ffn_up_exps_s,
	model.layers[il].ffn_gate_exps_s,
	model.layers[il].ffn_down_exps_s);
	cb(moe_out, "ffn_moe_out", il);
	// FFN shared expert
	{
	ggml_tensor * ffn_shexp =
	build_ffn(cur,
	model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s,
	model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s,
	model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s,
	NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
	cb(ffn_shexp, "ffn_shexp", il);
	cur = ggml_add(ctx0, moe_out, ffn_shexp);
	cb(cur, "ffn_out", il);
	}
	}

AFAIK it's needed for NVFP4 quant to work.

I tried this PR and it does seem to speed up code generation a bit (~10 t/s -> ~12 t/s). For English prose only --spec-draft-n-max 1 makes sense, for higher values acceptance rate tanks miserably to 40-50%. But for code even with default 3 drafted tokens it's around 80%.

Edit: In addition to GLM-5.2 I tested this with GLM-5.1 and it also works without issues.

[satindergrewal]

Thanks for testing, and for the GLM-5.1 confirmation @fairydreaming.

Update pushed: rebased past the indexer merge (which restructured GLM_DSA's graph and KV cache, so the MTP wiring needed re-plumbing on the new structure, not just conflict fixes), and added the _s scale args to both FFN calls in the MTP graph for NVFP4.

Re-verified after the changes: greedy acceptance 0.6623 (per-position 0.838/0.650/0.489), same as before the rebase.