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

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

Adds GLM-5.2 NextN/MTP as a --spec-type draft-mtp target: nextn tensor
loading via the qwen35moe/step35-style presence probe, a graph_mtp
builder (enorm/hnorm/eh_proj + dense MLA + sigmoid-gated MoE with
shared expert + shared head with fallbacks, _s scale tensors passed
for NVFP4), t_h_nextn extraction in the trunk graph, and MTP-context
KV setup: the draft head runs dense MLA, so the MTP context uses a
plain attention KV cache holding only the nextn layer(s) (same
pattern as the hybrid Qwen3.5 MTP context) while the main context
keeps the DSA cache, now filtered to trunk layers only.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

* convert : support --mtp/--no-mtp export for GlmMoeDsaForCausalLM (GLM-5.2)

Opt GLM-5.2 into the supports_mtp_export contract (post-#25641 shape,
mirroring HYV3Model/Step35Model): --no-mtp drops the appended NextN
block (blk.78) and its nextn_predict_layers KV; --mtp keeps only the
NextN block plus shared embeddings/norm/lm_head. Default (bundled)
output is unchanged.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

---------

Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
Satinder Grewal
2026-07-29 14:02:31 +08:00
committed by GitHub
co-authored by Claude Fable 5
parent e9fa0781f1
commit 7be2c65dc9
4 changed files with 366 additions and 15 deletions
+51 -2
View File
@@ -2072,7 +2072,6 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
res = nullptr;
} break;
case LLM_ARCH_DEEPSEEK32:
case LLM_ARCH_GLM_DSA:
{
res = new llama_kv_cache_dsa(
*this,
@@ -2089,6 +2088,56 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
nullptr,
nullptr);
} break;
case LLM_ARCH_GLM_DSA:
{
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) {
// The NextN/MTP draft head runs dense MLA (no DSA indexer), so the
// MTP context uses a plain attention KV cache holding only the
// nextn layer(s) - same pattern as the hybrid Qwen3.5 MTP context.
llama_kv_cache::layer_filter_cb filter =
[&](uint32_t il) { return il >= hparams.n_layer(); };
res = new llama_kv_cache(
*this,
hparams,
params.type_k,
params.type_v,
!cparams.flash_attn,
cparams.offload_kqv,
cparams.kv_unified,
cparams.n_ctx_seq,
cparams.n_seq_max,
1,
hparams.n_swa,
hparams.swa_type,
nullptr,
filter,
nullptr,
nullptr);
} else {
// Main context: DSA cache for the trunk layers only - the nextn
// layer(s) are never attended by the trunk graph.
llama_kv_cache::layer_filter_cb filter = nullptr;
if (hparams.n_layer_nextn > 0) {
filter = [&](uint32_t il) { return il < hparams.n_layer(); };
}
res = new llama_kv_cache_dsa(
*this,
params.type_k,
params.type_v,
!cparams.flash_attn,
cparams.offload_kqv,
cparams.kv_unified,
cparams.n_ctx_seq,
cparams.n_seq_max,
1,
hparams.n_swa,
hparams.swa_type,
filter,
nullptr);
}
} break;
// Models that need standard caching should rely on recurrent/hybrid
// checks
default:
@@ -2194,7 +2243,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
}
if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3) && hparams.n_layer_nextn > 0) {
if ((arch == LLM_ARCH_STEP35 || arch == LLM_ARCH_HY_V3 || arch == LLM_ARCH_GLM_DSA) && hparams.n_layer_nextn > 0) {
if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
filter = [&](uint32_t il) { return il >= hparams.n_layer(); };
} else {