Add MTP for GLM-5.3-Flash

common/speculative.cpp

@@ -1357,6 +1357,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
     std::vector<int>                i_last;
     std::vector<std::vector<float>> chain_h;
 
+    bool dsa_index_share = false;
+    size_t dsa_sel_width = 0;
+    std::vector<std::vector<int32_t>> dsa_sel;
+    std::vector<int32_t> dsa_sel_batch;
+
     common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq)
         : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq, params.draft.n_max)
         , params(params.draft)
@@ -1370,6 +1375,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                 "MTP input row width must match the target h_nextn width");
         n_mtp_layers = std::max(1, (int) llama_model_n_layer_nextn(llama_get_model(ctx_dft)));
 
+        dsa_index_share = llama_set_mtp_dsa_index_share(ctx_dft, true);
+        if (dsa_index_share) {
+            dsa_sel.resize(n_seq);
+        }
+
         SPC_TRC("%s", "adding speculative implementation 'draft-mtp'\n");
         SPC_TRC("- n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling);
         SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n",
@@ -1435,6 +1445,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
 
         verify_h.assign(n_seq, {});
         verify_h_rows.assign(n_seq, 0);
+
+        SPC_TRC("- dsa_index_share=%d\n", (int) dsa_index_share);
     }
 
     ~common_speculative_impl_draft_mtp() override {
@@ -1450,13 +1462,74 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
         }
         backend_chains.clear();
 
+        if (dsa_index_share && ctx_dft != nullptr) {
+            llama_set_mtp_dsa_index_share(ctx_dft, false);
+        }
+
         if (batch.token != nullptr) {
             free(batch.token);
             batch.token = nullptr;
         }
         llama_batch_free(batch);
     }
 
+    void reset_dsa_index_share() {
+        if (!dsa_index_share) {
+            return;
+        }
+        llama_set_mtp_dsa_selection(params.ctx_dft, nullptr, 0);
+        dsa_sel_width = 0;
+        dsa_sel_batch.clear();
+        for (auto & row : dsa_sel) {
+            row.clear();
+        }
+    }
+
+    bool capture_dsa_index_share(const llama_batch & current) {
+        size_t n = 0;
+        const int32_t * sel = llama_get_mtp_dsa_selection(params.ctx_dft, &n);
+        if (sel == nullptr || current.n_tokens <= 0 || n == 0 || n % (size_t) current.n_tokens != 0) {
+            return false;
+        }
+
+        const size_t width = n / (size_t) current.n_tokens;
+        for (auto & row : dsa_sel) {
+            row.clear();
+        }
+        for (int32_t k = 0; k < current.n_tokens; ++k) {
+            if (current.n_seq_id[k] != 1) {
+                return false;
+            }
+            const llama_seq_id seq_id = current.seq_id[k][0];
+            if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq || !dsa_sel[seq_id].empty()) {
+                return false;
+            }
+            dsa_sel[seq_id].assign(sel + (size_t) k*width, sel + (size_t) (k + 1)*width);
+        }
+        dsa_sel_width = width;
+        return true;
+    }
+
+    bool stage_dsa_index_share(const llama_batch & current) {
+        if (dsa_sel_width == 0 || current.n_tokens <= 0) {
+            return false;
+        }
+
+        dsa_sel_batch.resize(dsa_sel_width*(size_t) current.n_tokens);
+        for (int32_t k = 0; k < current.n_tokens; ++k) {
+            if (current.n_seq_id[k] != 1) {
+                return false;
+            }
+            const llama_seq_id seq_id = current.seq_id[k][0];
+            if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq || dsa_sel[seq_id].size() != dsa_sel_width) {
+                return false;
+            }
+            std::copy(dsa_sel[seq_id].begin(), dsa_sel[seq_id].end(), dsa_sel_batch.begin() + (size_t) k*dsa_sel_width);
+        }
+
+        return llama_set_mtp_dsa_selection(params.ctx_dft, dsa_sel_batch.data(), dsa_sel_batch.size());
+    }
+
     void begin(llama_seq_id seq_id, const llama_tokens & prompt) override {
         const int32_t N = (int32_t) prompt.size();
         if (N <= 0) {
@@ -1507,6 +1580,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
         auto * ctx_tgt = this->params.ctx_tgt;
         auto * ctx_dft = this->params.ctx_dft;
 
+        reset_dsa_index_share();
+
         const size_t row_bytes = (size_t) n_embd * sizeof(float);
 
         // if kv is shared with target (e.g Gemma4), then we can skip this catch-up decode
@@ -1596,6 +1671,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
     void draft(common_speculative_draft_params_vec & dparams) override {
         auto & ctx_dft = params.ctx_dft;
 
+        reset_dsa_index_share();
+
         common_batch_clear(batch);
 
         // keep track of which sequences are still drafting
@@ -1628,6 +1705,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
         int i = 0;
 
         while (n_drafting > 0) {
+            if (dsa_index_share && i > 0 && dsa_sel_width > 0 && !stage_dsa_index_share(batch)) {
+                reset_dsa_index_share();
+            }
+
             // each step decodes under a different head, i.e. a different decoder layer, and
             // KV is per layer. process() filled this layer's KV only for positions < n_past
             // (prompt + accepted prefix) — nothing in the draft region yet. so reset the
@@ -1650,6 +1731,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
                 break;
             }
 
+            if (dsa_index_share && i == 0 && !capture_dsa_index_share(batch)) {
+                reset_dsa_index_share();
+            }
+
             // rebuild the batch for the next step: the growing-KV paths re-add only the
             // new token (the KV already holds the prefix), while chained heads re-add the
             // whole prefix at the next head. dropped sequences are simply not re-added.
@@ -1728,6 +1813,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl {
             ++i;
         }
 
+        reset_dsa_index_share();
+
         if (chain_heads) {
             llama_set_nextn_layer_offset(ctx_dft, 0); // restore default for non-draft decodes
         }

conversion/__init__.py

@@ -104,6 +104,7 @@
     "Glm4MoeLiteForCausalLM": "glm",
     "Glm4vForConditionalGeneration": "glm",
     "Glm4vMoeForConditionalGeneration": "glm",
+    "Glm5NextForConditionalGeneration": "glm",
     "GlmForCausalLM": "chatglm",
     "GlmMoeDsaForCausalLM": "glm",
     "GlmOcrForConditionalGeneration": "glm",
@@ -296,6 +297,7 @@
     "Gemma4UnifiedForConditionalGeneration": "gemma",
     "Glm4vForConditionalGeneration": "qwen3vl",
     "Glm4vMoeForConditionalGeneration": "qwen3vl",
+    "Glm5NextForConditionalGeneration": "qwen3vl",
     "Glm5vForConditionalGeneration": "kimivl",
     "GlmOcrForConditionalGeneration": "qwen3vl",
     "GlmasrModel": "ultravox",

conversion/glm.py

@@ -402,3 +402,198 @@ def set_vocab(self):
         special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<unk>"])  # ty: ignore[unresolved-attribute]
         special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|startoftext|>"])  # ty: ignore[unresolved-attribute]
         special_vocab.add_to_gguf(self.gguf_writer)
+
+
+@ModelBase.register("Glm5NextForConditionalGeneration")
+@ModelBase.example("zai-org/GLM-5.3-Flash")
+class Glm5NextModel(TextModel):
+
+    model_arch = gguf.MODEL_ARCH.GLM5_NEXT
+    supports_mtp_export = True
+
+    _experts: list[dict[str, Tensor]] | None = None
+    _n_main_layers: int | None = None
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        self.n_nextn_layers = self.hparams.get("num_nextn_predict_layers", 0)
+        self.skip_mtp = self.no_mtp or self.n_nextn_layers == 0
+
+        if not self.skip_mtp:
+            self.block_count += self.n_nextn_layers
+            self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+        self.hparams.pop("head_dim", None)
+
+    def set_vocab(self):
+        # requires transformers >= 5, tokpre hash-resolves to glm4
+        return self._set_vocab_glm()
+
+    def index_tensors(self, remote_hf_model_id: str | None = None):
+        hp = self.hparams.get("text_config", self.hparams)
+        type(self)._n_main_layers = hp["num_hidden_layers"]
+        return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
+
+    @classmethod
+    def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+        if (titem := super().filter_tensors(item)) is None:
+            return None
+        name, gen = titem
+
+        assert cls._n_main_layers is not None
+        m = re.match(r"model\.layers\.(\d+)\.", name)
+        is_mtp = m is not None and int(m.group(1)) >= cls._n_main_layers
+
+        if is_mtp and cls.no_mtp:
+            return None
+        if cls.mtp_only and not is_mtp and name not in (
+            "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight",
+        ):
+            return None
+
+        return name, gen
+
+    def set_gguf_parameters(self):
+        super().set_gguf_parameters()
+        hp = self.hparams
+
+        layer_types = hp["layer_types"]
+        n_kv_heads = [0 if t == "linear_attention" else 1 for t in layer_types]
+        assert len(n_kv_heads) == hp["num_hidden_layers"]
+        # Pad to block_count
+        n_kv_heads += [1] * (self.block_count - len(n_kv_heads))
+        self.gguf_writer.add_head_count_kv(n_kv_heads)
+        self.gguf_writer.add_vocab_size(hp["vocab_size"])
+        self.gguf_writer.add_layer_norm_eps(1e-6)
+
+        if not self.skip_mtp:
+            self.gguf_writer.add_nextn_predict_layers(self.n_nextn_layers)
+
+        # KDA
+        lin = hp["linear_attn_config"]
+        assert lin["num_heads"] == hp["num_attention_heads"]
+        self.gguf_writer.add_ssm_conv_kernel(lin["short_conv_kernel_size"])
+        self.gguf_writer.add_kda_head_dim(lin["head_dim"])
+        if (lb := lin.get("gate_lower_bound")) is not None:
+            self.gguf_writer.add_kda_gate_lower_bound(lb)
+
+        # MLA (nope only)
+        assert hp.get("mla_use_nope") and hp["qk_rope_head_dim"] == 0, "expected nope-only MLA"
+        kv_lora_rank = hp["kv_lora_rank"]
+        qk_rope = hp["qk_rope_head_dim"]
+        self.gguf_writer.add_q_lora_rank(hp["q_lora_rank"])
+        self.gguf_writer.add_kv_lora_rank(kv_lora_rank)
+        self.gguf_writer.add_rope_dimension_count(qk_rope)
+        self.gguf_writer.add_key_length(kv_lora_rank + qk_rope)
+        self.gguf_writer.add_value_length(kv_lora_rank)
+        self.gguf_writer.add_key_length_mla(hp["qk_nope_head_dim"] + qk_rope)
+        self.gguf_writer.add_value_length_mla(hp["v_head_dim"])
+
+        # DSA indexer with k-pool compression
+        self.gguf_writer.add_indexer_head_count(hp["index_n_heads"])
+        self.gguf_writer.add_indexer_key_length(hp["index_head_dim"])
+        self.gguf_writer.add_indexer_top_k(hp["index_topk"])
+        self.gguf_writer.add_indexer_kpool(hp["index_kpool"])
+        self.gguf_writer.add_indexer_kpool_select_tail(hp.get("index_kpool_always_select_tail", True))
+        self.gguf_writer.add_indexer_index_share_mtp(hp.get("index_share_for_mtp_iteration", False))
+        if (indexer_types := hp.get("indexer_types")) is not None:
+            self.gguf_writer.add_indexer_types([t == "full" for t in indexer_types])
+
+        # mHC
+        assert hp.get("mhc", True)
+        self.gguf_writer.add_hyper_connection_count(hp["hc_mult"])
+        self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hp["hc_sinkhorn_iters"])
+        self.gguf_writer.add_hyper_connection_epsilon(hp["hc_eps"])
+
+        # MoE
+        self.gguf_writer.add_leading_dense_block_count(hp["first_k_dense_replace"])
+        self.gguf_writer.add_expert_feed_forward_length(hp["moe_intermediate_size"])
+        self.gguf_writer.add_expert_shared_count(hp["n_shared_experts"])
+        self.gguf_writer.add_expert_weights_scale(hp["routed_scaling_factor"])
+        self.gguf_writer.add_expert_weights_norm(hp["norm_topk_prob"])
+        if (limit := hp.get("swiglu_limit")) is not None:
+            self.gguf_writer.add_swiglu_clamp_exp([limit] * self.block_count)
+            self.gguf_writer.add_swiglu_clamp_shexp([limit] * self.block_count)
+
+    def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+        if name == "lm_head.weight" and self.hparams.get("tie_word_embeddings", False):
+            return
+
+        # routed experts
+        if ".mlp.experts." in name:
+            n_experts = self.hparams["n_routed_experts"]
+            assert bid is not None
+            if self._experts is None:
+                self._experts = [{} for _ in range(self.block_count)]
+            self._experts[bid][name] = data_torch
+            if len(self._experts[bid]) < n_experts * 3:
+                return
+            for w_name in ("down_proj", "gate_proj", "up_proj"):
+                datas: list[Tensor] = []
+                for xid in range(n_experts):
+                    ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
+                    datas.append(self._experts[bid].pop(ename))
+                merged = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
+                yield from super().modify_tensors(torch.stack(datas, dim=0), merged, bid)
+            return
+
+        # MLA absorption
+        if name.endswith("kv_b_proj.weight"):
+            n_head = self.hparams["num_attention_heads"]
+            v_head_dim = self.hparams["v_head_dim"]
+            qk_nope_head_dim = self.hparams["qk_nope_head_dim"]
+            assert data_torch.shape[0] == n_head * (v_head_dim + qk_nope_head_dim)
+            kv_b = data_torch.view(n_head, v_head_dim + qk_nope_head_dim, data_torch.shape[-1])
+            k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1)
+            yield from super().modify_tensors(k_b.transpose(1, 2), name.replace("kv_b_proj", "k_b_proj"), bid)
+            yield from super().modify_tensors(v_b, name.replace("kv_b_proj", "v_b_proj"), bid)
+            return
+
+        # KDA conv1d
+        if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")):
+            if data_torch.ndim == 3:
+                d_inner, _, d_conv = data_torch.shape
+            elif data_torch.ndim == 2:
+                d_inner, d_conv = data_torch.shape
+            else:
+                raise ValueError(f"unexpected conv1d rank {data_torch.ndim} for {name}")
+            data_torch = data_torch.reshape(1, d_inner, 1, d_conv)
+
+        if name.endswith(".A_log"):
+            n_head = self.hparams["num_attention_heads"]
+            data_torch = -torch.exp(data_torch.float().flatten()[:n_head])
+
+        if name.endswith(".dt_bias"):
+            name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"
+
+        if re.search(r"\.(hc_(?:attn|ffn)_(?:fn|base|scale)|index_kpool_compress_(?:ape|gate))$", name):
+            yield self.map_tensor_name(name) + ".weight", data_torch
+            return
+
+        yield from super().modify_tensors(data_torch, name, bid)
+
+    def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
+        # keep the small mHC / gating parameters exact
+        exact_keys = ("hc_attn_", "hc_ffn_", "indexer_compressor_", "ssm_a", "ssm_dt", "exp_probs_b")
+        if new_name.startswith(("blk.", "output_hc")) and any(k in new_name for k in exact_keys):
+            return gguf.GGMLQuantizationType.F32
+        return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+    def prepare_metadata(self, vocab_only: bool):
+        from_dir = self.fname_out.is_dir()
+        super().prepare_metadata(vocab_only=vocab_only)
+        if not self.mtp_only or not from_dir:
+            return
+        output_type: str = self.ftype.name.partition("_")[2]
+        fname_default: str = gguf.naming_convention(
+            self.metadata.name, self.metadata.basename, self.metadata.finetune,
+            self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+        self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
+    def prepare_tensors(self):
+        super().prepare_tensors()
+        if self._experts is not None:
+            leftover = [k for d in self._experts for k in d.keys()]
+            if leftover:
+                raise ValueError(f"Unprocessed experts: {leftover}")