diff --git a/common/speculative.cpp b/common/speculative.cpp index 851a47b9a584..92f9cafac64c 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1357,6 +1357,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::vector i_last; std::vector> chain_h; + bool dsa_index_share = false; + size_t dsa_sel_width = 0; + std::vector> dsa_sel; + std::vector 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,6 +1462,10 @@ 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; @@ -1457,6 +1473,63 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { 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 } diff --git a/conversion/__init__.py b/conversion/__init__.py index a5632fcc4bb9..eefb619cd327 100644 --- a/conversion/__init__.py +++ b/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", diff --git a/conversion/glm.py b/conversion/glm.py index 7544f850cb22..5b68f9ac6a88 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -402,3 +402,198 @@ def set_vocab(self): special_vocab._set_special_token("unk", tokenizer.get_added_vocab()[""]) # 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}") diff --git a/conversion/qwen3vl.py b/conversion/qwen3vl.py index 4fec708c9ff3..11ce68515b2c 100644 --- a/conversion/qwen3vl.py +++ b/conversion/qwen3vl.py @@ -228,10 +228,12 @@ class Qwen3ASRMmprojModel(Qwen3OmniMmprojModel): @ModelBase.register("Glm4vForConditionalGeneration", "Glm4vMoeForConditionalGeneration", "GlmOcrForConditionalGeneration") @ModelBase.example("zai-org/GLM-4.1V-9B-Thinking", "zai-org/GLM-4.5V") class Glm4VVisionModel(Qwen3VLVisionModel): + projector_type = gguf.VisionProjectorType.GLM4V + def set_gguf_parameters(self): MmprojModel.set_gguf_parameters(self) # skip Qwen3VLVisionModel parameters assert self.hparams_vision is not None - self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.GLM4V) + self.gguf_writer.add_clip_projector_type(self.projector_type) hidden_act = str(self.hparams_vision.get("hidden_act", "")).lower() if hidden_act == "gelu": @@ -249,6 +251,32 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(data_torch, name, bid) +@ModelBase.register("Glm5NextForConditionalGeneration") +@ModelBase.example("zai-org/GLM-5.3-Flash") +class Glm5NextVisionModel(Glm4VVisionModel): + # GLM-5.3-Flash vision tower. glm4v layout with per-head qk-norm, no post-conv norm and no learned position embeddings. + # Images are placed on a ceil aligned canvas with padding. + + projector_type = gguf.VisionProjectorType.GLM5V + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_vision is not None + self.gguf_writer.add_vision_spatial_merge_size(int(self.hparams_vision.get("spatial_merge_size", 2))) + if (limit := self.hparams_vision.get("swiglu_limit")) is not None: + self.gguf_writer.add_vision_swiglu_clamp(float(limit)) + + # image token budget from the processor, stored as single-frame pixel counts + pc = self.preprocessor_config + patch = int(pc.get("patch_size", 14)) + merge = int(pc.get("merge_size", 2)) + pixels_per_token = (patch * merge) ** 2 + if (min_tok := pc.get("min_image_tokens")) is not None: + self.gguf_writer.add_vision_min_pixels(int(min_tok) * pixels_per_token) + if (max_tok := pc.get("max_image_tokens")) is not None: + self.gguf_writer.add_vision_max_pixels(int(max_tok) * pixels_per_token) + + @ModelBase.register("Qwen3VLForConditionalGeneration") @ModelBase.example("Qwen/Qwen3-VL-4B-Instruct") class Qwen3VLTextModel(Qwen3Model): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index c99feb3c795c..5e12bc71f9e6 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -225,6 +225,9 @@ class Indexer: BLOCK_SIZE = "{arch}.attention.indexer.block_size" # MSA LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" # MSA TYPES = "{arch}.attention.indexer.types" + KPOOL = "{arch}.attention.indexer.kpool" # GLM5-Next + INDEX_SHARE_MTP = "{arch}.attention.indexer.index_share_mtp" # GLM5-Next + KPOOL_SELECT_TAIL = "{arch}.attention.indexer.kpool_select_tail" # GLM5-Next class HyperConnection: COUNT = "{arch}.hyper_connection.count" @@ -382,6 +385,7 @@ class ClipVision: IMAGE_MEAN = "clip.vision.image_mean" IMAGE_STD = "clip.vision.image_std" SPATIAL_MERGE_SIZE = "clip.vision.spatial_merge_size" + SWIGLU_CLAMP = "clip.vision.swiglu_clamp" EXPERT_COUNT_PER_LAYER = "clip.vision.expert_count_per_layer" # dots3note pyramid MoE, 0 = dense layer EXPERT_USED_COUNT = "clip.vision.expert_used_count" USE_GELU = "clip.use_gelu" @@ -559,6 +563,7 @@ class MODEL_ARCH(IntEnum): GLM4 = auto() GLM4_MOE = auto() GLM_DSA = auto() + GLM5_NEXT = auto() BITNET = auto() T5 = auto() T5ENCODER = auto() @@ -888,6 +893,8 @@ class MODEL_TENSOR(IntEnum): INDEXER_COMPRESSOR_WGATE = auto() INDEXER_COMPRESSOR_APE = auto() INDEXER_COMPRESSOR_NORM = auto() + INDEXER_KPOOL_GATE = auto() + INDEXER_KPOOL_APE = auto() # vision V_MMPROJ = auto() V_MMPROJ_FC = auto() @@ -1307,6 +1314,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GLM4: "glm4", MODEL_ARCH.GLM4_MOE: "glm4moe", MODEL_ARCH.GLM_DSA: "glm-dsa", + MODEL_ARCH.GLM5_NEXT: "glm5-next", MODEL_ARCH.BITNET: "bitnet", MODEL_ARCH.T5: "t5", MODEL_ARCH.T5ENCODER: "t5encoder", @@ -1635,6 +1643,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: "blk.{bid}.indexer_compressor_gate", MODEL_TENSOR.INDEXER_COMPRESSOR_APE: "blk.{bid}.indexer_compressor_ape", MODEL_TENSOR.INDEXER_COMPRESSOR_NORM: "blk.{bid}.indexer_compressor_norm", + MODEL_TENSOR.INDEXER_KPOOL_GATE: "blk.{bid}.indexer_compressor_gate", + MODEL_TENSOR.INDEXER_KPOOL_APE: "blk.{bid}.indexer_compressor_ape", # vision MODEL_TENSOR.V_MMPROJ: "mm.{bid}", MODEL_TENSOR.V_MMPROJ_FC: "mm.model.fc", @@ -3991,6 +4001,70 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.GLM5_NEXT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + # mHC + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_SCALE, + # KDA (linear attention) layers + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_F_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_G_A, + MODEL_TENSOR.SSM_G_B, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + # MLA (nope) + DSA layers + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.INDEXER_KPOOL_GATE, + MODEL_TENSOR.INDEXER_KPOOL_APE, + # FFN + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], MODEL_ARCH.BITNET: [ MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, @@ -5648,6 +5722,7 @@ class VisionProjectorType: LFM2A = "lfm2a" # audio MUSIC_FLAMINGO = "musicflamingo" # audio GLM4V = "glm4v" + GLM5V = "glm5v" YOUTUVL = "youtuvl" NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index d95fe9b1ac3c..049a79bae9a0 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -821,6 +821,15 @@ def add_indexer_types(self, value: Sequence[bool]) -> None: key = Keys.Attention.Indexer.TYPES.format(arch=self.arch) self.add_array(key, value) + def add_indexer_kpool(self, value: int) -> None: + self.add_uint32(Keys.Attention.Indexer.KPOOL.format(arch=self.arch), value) + + def add_indexer_kpool_select_tail(self, value: bool) -> None: + self.add_bool(Keys.Attention.Indexer.KPOOL_SELECT_TAIL.format(arch=self.arch), value) + + def add_indexer_index_share_mtp(self, value: bool) -> None: + self.add_bool(Keys.Attention.Indexer.INDEX_SHARE_MTP.format(arch=self.arch), value) + def add_max_alibi_bias(self, bias: float) -> None: self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias) @@ -1378,6 +1387,9 @@ def add_vision_image_mean(self, values: Sequence[float]) -> None: def add_vision_image_std(self, values: Sequence[float]) -> None: self.add_array(Keys.ClipVision.IMAGE_STD, values) + def add_vision_swiglu_clamp(self, value: float) -> None: + self.add_float32(Keys.ClipVision.SWIGLU_CLAMP, value) + def add_vision_spatial_merge_size(self, value: int) -> None: self.add_uint32(Keys.ClipVision.SPATIAL_MERGE_SIZE, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 861acfe181fe..b89d93d2b451 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1317,6 +1317,38 @@ class TensorNameMap: "model.layers.{bid}.self_attn.indexer.wq_b", # DSA ), + MODEL_TENSOR.INDEXER_KPOOL_GATE: ( + "model.layers.{bid}.self_attn.indexer.index_kpool_compress_gate", # glm5-next + ), + + MODEL_TENSOR.INDEXER_KPOOL_APE: ( + "model.layers.{bid}.self_attn.indexer.index_kpool_compress_ape", # glm5-next + ), + + MODEL_TENSOR.HC_ATTN_FN: ( + "model.layers.{bid}.hc_attn_fn", # glm5-next + ), + + MODEL_TENSOR.HC_ATTN_BASE: ( + "model.layers.{bid}.hc_attn_base", # glm5-next + ), + + MODEL_TENSOR.HC_ATTN_SCALE: ( + "model.layers.{bid}.hc_attn_scale", # glm5-next + ), + + MODEL_TENSOR.HC_FFN_FN: ( + "model.layers.{bid}.hc_ffn_fn", # glm5-next + ), + + MODEL_TENSOR.HC_FFN_BASE: ( + "model.layers.{bid}.hc_ffn_base", # glm5-next + ), + + MODEL_TENSOR.HC_FFN_SCALE: ( + "model.layers.{bid}.hc_ffn_scale", # glm5-next + ), + MODEL_TENSOR.INDEXER_Q_PROJ: ( "model.layers.{bid}.self_attn.index_q_proj", # MSA ), diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 7adb87411abd..df94431ae663 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -149,6 +149,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_MAINCODER, "maincoder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, { LLM_ARCH_KIMI_K3, "kimi-k3" }, + { LLM_ARCH_GLM5_NEXT, "glm5-next" }, { LLM_ARCH_TALKIE, "talkie" }, { LLM_ARCH_MELLUM, "mellum" }, { LLM_ARCH_NANBEIGE, "nanbeige" }, @@ -284,6 +285,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" }, { LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" }, { LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" }, + { LLM_KV_ATTENTION_INDEXER_KPOOL, "%s.attention.indexer.kpool" }, + { LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, "%s.attention.indexer.kpool_select_tail" }, + { LLM_KV_ATTENTION_INDEXER_INDEX_SHARE_MTP, "%s.attention.indexer.index_share_mtp" }, { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" }, { LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" }, { LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" }, @@ -678,6 +682,8 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" }, { LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" }, { LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "blk.%d.indexer_compressor_norm" }, + { LLM_TENSOR_INDEXER_KPOOL_GATE, "blk.%d.indexer_compressor_gate" }, + { LLM_TENSOR_INDEXER_KPOOL_APE, "blk.%d.indexer_compressor_ape" }, { LLM_TENSOR_FFN_GATE_TID2EID, "blk.%d.ffn_gate_tid2eid" }, { LLM_TENSOR_MASKED_EMBD_CENTROIDS, "masked_embd_centroids" }, { LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" }, @@ -950,6 +956,8 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, {LLM_TENSOR_INDEXER_COMPRESSOR_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_INDEXER_KPOOL_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_INDEXER_KPOOL_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, {LLM_TENSOR_FFN_GATE_TID2EID, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, {LLM_TENSOR_NEXTN_PROJ_PRE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_NEXTN_PROJ_POST, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, @@ -1073,6 +1081,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_KIMI_K3: + case LLM_ARCH_GLM5_NEXT: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_QWEN4EXP: @@ -1107,6 +1116,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { case LLM_ARCH_LFM2: case LLM_ARCH_LFM2MOE: case LLM_ARCH_BAILINGMOE3: + case LLM_ARCH_GLM5_NEXT: return true; default: return false; diff --git a/src/llama-arch.h b/src/llama-arch.h index ca7d55a5fd78..0213a632e0de 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -150,6 +150,7 @@ enum llm_arch { LLM_ARCH_MAINCODER, LLM_ARCH_KIMI_LINEAR, LLM_ARCH_KIMI_K3, + LLM_ARCH_GLM5_NEXT, LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, @@ -289,6 +290,9 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, LLM_KV_ATTENTION_INDEXER_TYPES, + LLM_KV_ATTENTION_INDEXER_KPOOL, + LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, + LLM_KV_ATTENTION_INDEXER_INDEX_SHARE_MTP, LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, LLM_KV_ATTENTION_OUTPUT_LORA_RANK, LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, @@ -677,6 +681,8 @@ enum llm_tensor { LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, LLM_TENSOR_INDEXER_COMPRESSOR_APE, LLM_TENSOR_INDEXER_COMPRESSOR_NORM, + LLM_TENSOR_INDEXER_KPOOL_GATE, // glm5-next: k-pool gate scores + LLM_TENSOR_INDEXER_KPOOL_APE, // glm5-next: k-pool position bias LLM_TENSOR_FFN_GATE_TID2EID, LLM_TENSOR_NEXTN_PROJ_PRE, LLM_TENSOR_NEXTN_PROJ_POST, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index f286bd1da2a5..1a44bba19e2e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -7,6 +7,7 @@ #include "llama-batch.h" #include "llama-io.h" #include "llama-memory.h" +#include "llama-memory-hybrid-idx.h" #include "llama-mmap.h" #include "llama-model.h" #include "llama-ext.h" @@ -979,6 +980,74 @@ float * llama_context::get_embeddings_nextn_ith(int32_t i) { } } +bool llama_context::set_mtp_dsa_index_share(bool enabled) { + if (model.arch != LLM_ARCH_GLM5_NEXT || cparams.ctx_type != LLAMA_CONTEXT_TYPE_MTP || memory == nullptr) { + return false; + } + + enabled = enabled && model.hparams.indexer_index_share_mtp; + + auto * mem = static_cast(memory.get()); + if (mem->get_mtp_dsa_index_share() != enabled) { + mem->set_mtp_dsa_index_share(enabled); + sched_need_reserve = true; + } + + return enabled; +} + +bool llama_context::set_mtp_dsa_selection(const int32_t * data, size_t size) { + if (model.arch != LLM_ARCH_GLM5_NEXT || cparams.ctx_type != LLAMA_CONTEXT_TYPE_MTP || memory == nullptr) { + return false; + } + + auto * mem = static_cast(memory.get()); + if (!mem->get_mtp_dsa_index_share()) { + return false; + } + + mem->set_mtp_dsa_selection(data, size); + return true; +} + +const int32_t * llama_context::get_mtp_dsa_selection(size_t * size) { + if (size != nullptr) { + *size = 0; + } + if (mtp_dsa_sel_raw.empty() || mtp_dsa_sel_raw.size() != mtp_dsa_sel_mask.size() || mtp_dsa_sel_width == 0) { + return nullptr; + } + GGML_ASSERT(memory != nullptr && model.arch == LLM_ARCH_GLM5_NEXT); + GGML_ASSERT(mtp_dsa_sel_raw.size() == mtp_dsa_sel_width*mtp_dsa_sel_seq.size()); + + auto * mem = static_cast(memory.get()); + auto * mem_idx = mem->get_mem_idx(); + GGML_ASSERT(mem_idx != nullptr); + + mtp_dsa_sel.resize(mtp_dsa_sel_raw.size()); + for (size_t row = 0; row < mtp_dsa_sel_seq.size(); ++row) { + const llama_seq_id seq_id = mtp_dsa_sel_seq[row]; + GGML_ASSERT(seq_id >= 0); + const auto & cells = mem_idx->get_cells(seq_id); + + for (size_t j = 0; j < mtp_dsa_sel_width; ++j) { + const size_t i = row*mtp_dsa_sel_width + j; + const int32_t cell = mtp_dsa_sel_raw[i]; + if (mtp_dsa_sel_mask[i] != 0.0f) { + mtp_dsa_sel[i] = -1; + continue; + } + GGML_ASSERT(cell >= 0 && (uint32_t) cell < cells.size()); + GGML_ASSERT(!cells.is_empty((uint32_t) cell) && cells.seq_has((uint32_t) cell, seq_id)); + mtp_dsa_sel[i] = cells.pos_get((uint32_t) cell); + } + } + if (size != nullptr) { + *size = mtp_dsa_sel.size(); + } + return mtp_dsa_sel.data(); +} + float * llama_context::get_embeddings_layer_inp(uint32_t lid) { output_reorder(); @@ -1736,6 +1805,15 @@ int llama_context::decode(const llama_batch & batch_inp) { output_swaps.clear(); + if (!mtp_dsa_sel_raw.empty()) { + synchronize(); + } + mtp_dsa_sel_raw.clear(); + mtp_dsa_sel_mask.clear(); + mtp_dsa_sel_seq.clear(); + mtp_dsa_sel.clear(); + mtp_dsa_sel_width = 0; + sched_reserve(); bool did_optimize = false; @@ -1965,6 +2043,36 @@ int llama_context::decode(const llama_batch & batch_inp) { } } + auto * t_mtp_sel = res->get_mtp_dsa_sel(); + auto * t_mtp_mask = res->get_mtp_dsa_mask(); + if (t_mtp_sel != nullptr || t_mtp_mask != nullptr) { + GGML_ASSERT(t_mtp_sel != nullptr && t_mtp_mask != nullptr); + GGML_ASSERT(t_mtp_sel->type == GGML_TYPE_I32 && t_mtp_mask->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(t_mtp_sel) && ggml_is_contiguous(t_mtp_mask)); + GGML_ASSERT(t_mtp_sel->ne[0] == t_mtp_mask->ne[0] && t_mtp_sel->ne[1] == (int64_t) ubatch.n_tokens); + GGML_ASSERT(ggml_nelements(t_mtp_sel) == ggml_nelements(t_mtp_mask)); + + const size_t width = (size_t) t_mtp_sel->ne[0]; + if (mtp_dsa_sel_width == 0) { + mtp_dsa_sel_width = width; + mtp_dsa_sel_raw.resize(width*n_tokens_all); + mtp_dsa_sel_mask.resize(width*n_tokens_all); + mtp_dsa_sel_seq.resize(n_tokens_all, -1); + } + GGML_ASSERT(width == mtp_dsa_sel_width); + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + GGML_ASSERT(ubatch.n_seq_id[i] == 1); + mtp_dsa_sel_seq[(size_t) n_tokens_prev + i] = ubatch.seq_id[i][0]; + } + + const size_t offset = width*(size_t) n_tokens_prev; + ggml_backend_t backend_sel = ggml_backend_sched_get_tensor_backend(sched.get(), t_mtp_sel); + ggml_backend_t backend_mask = ggml_backend_sched_get_tensor_backend(sched.get(), t_mtp_mask); + GGML_ASSERT(backend_sel != nullptr && backend_mask != nullptr); + ggml_backend_tensor_get_async(backend_sel, t_mtp_sel, mtp_dsa_sel_raw.data() + offset, 0, ggml_nbytes(t_mtp_sel)); + ggml_backend_tensor_get_async(backend_mask, t_mtp_mask, mtp_dsa_sel_mask.data() + offset, 0, ggml_nbytes(t_mtp_mask)); + } + if (has_samplers) { const auto stride = n_vocab; @@ -2303,7 +2411,7 @@ void llama_context::output_reorder() { uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { uint32_t res; - if (model.arch == LLM_ARCH_KIMI_K3) { + if (model.arch == LLM_ARCH_KIMI_K3 || model.arch == LLM_ARCH_GLM5_NEXT) { // the n_tokens*40 budget below is exhausted at ubatch 3840 res = std::max(n_tokens * 160, 64u * model.n_tensors()); } else if (model.arch == LLM_ARCH_QWEN3NEXT || @@ -3906,6 +4014,25 @@ float * llama_get_embeddings_nextn_ith(llama_context * ctx, int32_t i) { return ctx->get_embeddings_nextn_ith(i); } +bool llama_set_mtp_dsa_index_share(llama_context * ctx, bool enabled) { + return ctx != nullptr && ctx->set_mtp_dsa_index_share(enabled); +} + +bool llama_set_mtp_dsa_selection(llama_context * ctx, const int32_t * data, size_t size) { + return ctx != nullptr && ctx->set_mtp_dsa_selection(data, size); +} + +const int32_t * llama_get_mtp_dsa_selection(llama_context * ctx, size_t * size) { + if (ctx == nullptr) { + if (size != nullptr) { + *size = 0; + } + return nullptr; + } + ctx->synchronize(); + return ctx->get_mtp_dsa_selection(size); +} + float * llama_get_embeddings_layer_inp(llama_context * ctx, uint32_t lid) { ctx->synchronize(); diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b562..8073d96fb2fa 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -100,6 +100,9 @@ struct llama_context { size_t get_sampled_probs_count(int32_t idx); const llama_token * get_sampled_candidates_ith(int32_t idx); + bool set_mtp_dsa_index_share(bool enabled); + bool set_mtp_dsa_selection(const int32_t * data, size_t size); + const int32_t * get_mtp_dsa_selection(size_t * size); size_t get_sampled_candidates_count(int32_t idx); void attach_threadpool( @@ -300,6 +303,12 @@ struct llama_context { // sets llm_graph_result::t_h_nextn buffer_view embd_nextn = {nullptr, 0}; + std::vector mtp_dsa_sel_raw; + std::vector mtp_dsa_sel_mask; + std::vector mtp_dsa_sel_seq; + std::vector mtp_dsa_sel; + size_t mtp_dsa_sel_width = 0; + // host buffers for output layer input embeddings, per layer // populated when cparams.output_layer_inp[il] is true std::vector> embd_layer_inp; diff --git a/src/llama-ext.h b/src/llama-ext.h index 92a759b7a0ae..71f7b72d4ed8 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -107,6 +107,10 @@ LLAMA_API float * llama_get_embeddings_nextn(struct llama_context * ctx); // LLAMA_API float * llama_get_embeddings_ith(struct llama_context * ctx, int32_t i); LLAMA_API float * llama_get_embeddings_nextn_ith(struct llama_context * ctx, int32_t i); +LLAMA_API bool llama_set_mtp_dsa_index_share(struct llama_context * ctx, bool enabled); +LLAMA_API bool llama_set_mtp_dsa_selection(struct llama_context * ctx, const int32_t * data, size_t size); +LLAMA_API const int32_t * llama_get_mtp_dsa_selection(struct llama_context * ctx, size_t * size); + // Set whether the context outputs the input embeddings of a specific layer LLAMA_API void llama_set_embeddings_layer_inp(struct llama_context * ctx, uint32_t lid, bool value); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 72db486cacea..5369521c26cf 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1323,6 +1323,8 @@ void llm_graph_result::reset() { t_embd = nullptr; t_embd_pooled = nullptr; t_h_nextn = nullptr; + t_mtp_dsa_sel = nullptr; + t_mtp_dsa_mask = nullptr; t_layer_inp.resize(LLAMA_MAX_LAYERS + 1); std::fill(t_layer_inp.begin(), t_layer_inp.end(), nullptr); @@ -1369,6 +1371,12 @@ void llm_graph_result::set_outputs(const llm_graph_params & params) { if (t_h_nextn != nullptr) { ggml_set_output(t_h_nextn); } + if (t_mtp_dsa_sel != nullptr) { + ggml_set_output(t_mtp_dsa_sel); + } + if (t_mtp_dsa_mask != nullptr) { + ggml_set_output(t_mtp_dsa_mask); + } { const auto & embeddings_layer_inp = params.cparams.embeddings_layer_inp; for (size_t il = 0; il < embeddings_layer_inp.size(); ++il) { @@ -1776,7 +1784,7 @@ ggml_tensor * llm_graph_context::build_ffn( const float limit = hparams.swiglu_clamp_shexp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_swiglu_clamp(ctx0, cur, tmp, limit); } else { tmp = ggml_clamp(ctx0, tmp, -limit, limit); @@ -2170,7 +2178,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( const float limit = hparams.swiglu_clamp_exp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_swiglu_clamp(ctx0, cur, up, limit); } else { up = ggml_clamp(ctx0, up, -limit, limit); diff --git a/src/llama-graph.h b/src/llama-graph.h index b388e028cb53..11cd087aeeff 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -900,6 +900,8 @@ class llm_graph_result { ggml_tensor * get_embd() const { return t_embd; } ggml_tensor * get_embd_pooled() const { return t_embd_pooled; } ggml_tensor * get_h_nextn() const { return t_h_nextn; } + ggml_tensor * get_mtp_dsa_sel() const { return t_mtp_dsa_sel; } + ggml_tensor * get_mtp_dsa_mask() const { return t_mtp_dsa_mask; } ggml_tensor * get_layer_inp(int il) const { return t_layer_inp[il]; } @@ -935,6 +937,8 @@ class llm_graph_result { ggml_tensor * t_embd = nullptr; ggml_tensor * t_embd_pooled = nullptr; ggml_tensor * t_h_nextn = nullptr; // [n_embd, n_outputs] hidden state before final output norm + ggml_tensor * t_mtp_dsa_sel = nullptr; + ggml_tensor * t_mtp_dsa_mask = nullptr; std::vector t_layer_inp; diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 1411692a8909..c36646f5b0c8 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -261,6 +261,9 @@ struct llama_hparams { uint32_t indexer_n_head = 0; uint32_t indexer_head_size = 0; uint32_t indexer_top_k = 0; + uint32_t indexer_kpool = 0; // k-pool size + bool indexer_kpool_select_tail = true; + bool indexer_index_share_mtp = false; // MTP iterations reuse one indexer selection // MSA uint32_t indexer_block_size = 0; uint32_t indexer_local_blocks = 0; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 93b468784a33..1220725bbf79 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -1,5 +1,9 @@ #include "llama-memory-hybrid-idx.h" +#include +#include +#include + #include "llama-impl.h" #include "llama-batch.h" #include "llama-io.h" @@ -49,7 +53,8 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( mem_idx(filter_idx == nullptr ? nullptr : [&] { // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1); - hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; + // The glm5 next indexer caches key, gate and pooled values per token + hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size * (model.hparams.indexer_kpool > 0 ? 3 : 1); // the cached indexer keys are raw, rotation happens after pooling at read time, so a // K-shift must not rotate them while the stream copies in the same update still apply @@ -61,7 +66,8 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( model, hparams_idx, type_k, type_v, v_trans, offload, unified, kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_idx, nullptr, nullptr, "idx_"); - }()) {} + }()), + n_kpool(model.hparams.indexer_kpool) {} llama_memory_context_ptr llama_memory_hybrid_idx::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { // note: repeats llama_memory_hybrid::init_batch, as the indexer needs the attention slot infos that the base context hides @@ -136,9 +142,27 @@ llama_memory_context_ptr llama_memory_hybrid_idx::init_update(llama_context * lc return std::make_unique(this, lctx, optimize); } +void llama_memory_hybrid_idx::set_mtp_dsa_index_share(bool enabled) { + mtp_dsa_index_share = enabled; + if (!enabled) { + mtp_dsa_selection.clear(); + } +} + +void llama_memory_hybrid_idx::set_mtp_dsa_selection(const int32_t * data, size_t size) { + if (data == nullptr) { + GGML_ASSERT(size == 0); + mtp_dsa_selection.clear(); + return; + } + mtp_dsa_selection.assign(data, data + size); +} + void llama_memory_hybrid_idx::clear(bool data) { llama_memory_hybrid::clear(data); + mtp_dsa_selection.clear(); + if (mem_idx) { mem_idx->clear(data); } @@ -152,6 +176,8 @@ bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_po if (mem_idx) { mem_idx->seq_rm(seq_id, p0, p1); + kpool_dirty = true; + mtp_dsa_selection.clear(); } return get_mem_attn()->seq_rm(seq_id, p0, p1); @@ -162,6 +188,8 @@ void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_i if (mem_idx) { mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1); + kpool_dirty = true; + mtp_dsa_selection.clear(); } } @@ -170,6 +198,8 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { if (mem_idx) { mem_idx->seq_keep(seq_id); + kpool_dirty = true; + mtp_dsa_selection.clear(); } } @@ -178,6 +208,8 @@ void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_p if (mem_idx) { mem_idx->seq_add(seq_id, p0, p1, shift); + kpool_dirty = true; + mtp_dsa_selection.clear(); } } @@ -186,6 +218,8 @@ void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_p if (mem_idx) { mem_idx->seq_div(seq_id, p0, p1, d); + kpool_dirty = true; + mtp_dsa_selection.clear(); } } @@ -600,6 +634,31 @@ static std::vector llama_memory_hybrid_idx_ns(const llama_kv_cache::sl return res; } +// The kpool layout of one ubatch. +struct llama_memory_hybrid_idx_context::kpool_state { + struct seq { + llama_pos pos_min = 0; + std::vector> cells; // Position and cell pairs, sorted by position + std::vector pools; + std::vector is_new; + }; + + std::vector seqs; + + uint32_t n_pool_real = 0; + uint32_t n_new = 0; + bool cache_safe = true; +}; + +namespace { + +// The last padded pool is always unused. +uint32_t kpool_pad(uint32_t n_pool) { + return std::max(64u, GGML_PAD(n_pool + 1, 64u)); +} + +} + llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_status status) : llama_memory_hybrid_context(status) {} @@ -611,7 +670,11 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hy ns_ubatch(mem->get_mem_idx() == nullptr ? std::vector() : std::vector{ mem->get_mem_idx()->get_n_stream() }), ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : - new llama_kv_cache_context(mem->get_mem_idx())) {} + new llama_kv_cache_context(mem->get_mem_idx())) { + if (mem->get_mem_idx() != nullptr && mem->get_kpool() > 0) { + kpool_states.push_back(kpool_build_state(nullptr)); + } +} llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( llama_memory_hybrid_idx * mem, @@ -633,9 +696,20 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( mem(mem), ns_ubatch(llama_memory_hybrid_idx_ns(sinfos_idx)), ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : - new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {} + new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) { + kpool_track = mem->get_mem_idx() != nullptr && mem->get_kpool() > 0; + // Sequence edits require a full re-pool. + kpool_dirty_batch = mem->kpool_is_dirty(); +} + +llama_memory_hybrid_idx_context::~llama_memory_hybrid_idx_context() = default; bool llama_memory_hybrid_idx_context::next() { + // Clear only after a successful ubatch. + if (i_cur == 0 && kpool_dirty_batch && mem != nullptr) { + mem->kpool_clear_dirty(); + } + if (ctx_idx) { ctx_idx->next(); } @@ -652,6 +726,14 @@ bool llama_memory_hybrid_idx_context::apply() { res = res & ctx_idx->apply(); } + // Fix the pool layout of this ubatch. + if (res && kpool_track) { + GGML_ASSERT(i_cur <= kpool_states.size()); + kpool_states.resize(i_cur); + + kpool_states.push_back(kpool_build_state(&get_ubatch())); + } + return res; } @@ -677,3 +759,377 @@ void llama_memory_hybrid_idx_context::set_input_qsa( mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); } + +// k-pool DSA indexer (glm5-next) + +llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kpool_build_state( + const llama_ubatch * ubatch) const { + GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); + GGML_ASSERT(get_n_stream() == 1 && "TODO: k-pool indexer with multiple streams"); + + const uint32_t kpool = mem->get_kpool(); + + kpool_state st; + st.seqs.resize(LLAMA_MAX_SEQ); + + const auto & cells = mem->get_mem_idx()->get_cells(0); + + // Scan only active sequences + std::vector active; + for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + if (cells.seq_pos_min(s) >= 0) { + active.push_back(s); + } + } + + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.is_empty(i)) { + continue; + } + const llama_pos p = cells.pos_get(i); + uint32_t n_seq_cell = 0; + for (const llama_seq_id s : active) { + if (cells.seq_has(i, s)) { + st.seqs[s].cells.emplace_back(p, i); + ++n_seq_cell; + } + } + if (n_seq_cell > 1) { + st.cache_safe = false; + } + } + + for (auto & sq : st.seqs) { + if (sq.cells.empty()) { + continue; + } + if (!std::is_sorted(sq.cells.begin(), sq.cells.end())) { + std::sort(sq.cells.begin(), sq.cells.end()); + } + + sq.pos_min = sq.cells.front().first; + + // Pools start at the first valid token + for (size_t j = 0; j + kpool <= sq.cells.size(); ) { + const llama_pos p0 = sq.cells[j].first; + if ((p0 - sq.pos_min) % (llama_pos) kpool != 0) { + ++j; + continue; + } + bool ok = true; + for (uint32_t k = 1; k < kpool; ++k) { + if (sq.cells[j + k].first != p0 + (llama_pos) k) { + ok = false; + break; + } + } + if (ok) { + sq.pools.push_back((uint32_t) j); + j += kpool; + } else { + ++j; + } + } + + st.n_pool_real += (uint32_t) sq.pools.size(); + } + + if (ubatch == nullptr) { + for (auto & sq : st.seqs) { + sq.is_new.assign(sq.pools.size(), 0); + } + + return st; + } + + // Pools touched by this ubatch are re-pooled, shared cells cannot cache sequence relative pools. + const bool all_new = !st.cache_safe || (kpool_dirty_batch && i_cur == 0); + + std::vector> upos(LLAMA_MAX_SEQ); + if (!all_new) { + for (uint32_t i = 0; i < ubatch->n_tokens; ++i) { + for (int32_t k = 0; k < ubatch->n_seq_id[i]; ++k) { + upos[ubatch->seq_id[i][k]].push_back(ubatch->pos[i]); + } + } + for (auto & v : upos) { + if (!std::is_sorted(v.begin(), v.end())) { + std::sort(v.begin(), v.end()); + } + } + } + + for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + auto & sq = st.seqs[s]; + + sq.is_new.assign(sq.pools.size(), all_new ? 1 : 0); + if (all_new) { + st.n_new += (uint32_t) sq.pools.size(); + continue; + } + + const auto & up = upos[s]; + if (up.empty()) { + continue; + } + + for (size_t pi = 0; pi < sq.pools.size(); ++pi) { + const llama_pos p0 = sq.cells[sq.pools[pi]].first; + + auto it = std::lower_bound(up.begin(), up.end(), p0); + if (it != up.end() && *it < p0 + (llama_pos) kpool) { + sq.is_new[pi] = 1; + st.n_new++; + } + } + } + + return st; +} + +const llama_memory_hybrid_idx_context::kpool_state & llama_memory_hybrid_idx_context::kpool_cur() const { + GGML_ASSERT(i_cur < kpool_states.size() && "k-pool state read before apply()"); + + return kpool_states[i_cur]; +} + +uint32_t llama_memory_hybrid_idx_context::get_n_kpool() const { + return kpool_pad(kpool_cur().n_pool_real); +} + +uint32_t llama_memory_hybrid_idx_context::get_n_kpool_new() const { + return kpool_cur().n_new; +} + +bool llama_memory_hybrid_idx_context::get_kpool_cache_safe() const { + return kpool_cur().cache_safe; +} + +bool llama_memory_hybrid_idx_context::get_mtp_dsa_index_share() const { + return mem != nullptr && mem->get_mtp_dsa_index_share(); +} + +size_t llama_memory_hybrid_idx_context::get_mtp_dsa_selection_size() const { + return mem != nullptr ? mem->get_mtp_dsa_selection().size() : 0; +} + +void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs, + ggml_tensor * gather_mask, bool gather, ggml_tensor * new_pool_idxs, ggml_tensor * new_pool_rep, + const llama_ubatch * ubatch) const { + GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); + GGML_ASSERT(get_n_stream() == 1 && "TODO: k-pool indexer with multiple streams"); + GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(pool_idxs->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(pool_mask->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(tail_idxs->buffer)); + + const uint32_t kpool = mem->get_kpool(); + const uint32_t n_kv = get_idx()->get_n_kv(); + + const auto & st = kpool_cur(); + + const uint32_t n_tokens = ubatch->n_tokens; + const uint32_t n_pool = (uint32_t) pool_cells->ne[0]; + const uint32_t n_new = st.n_new; + + GGML_ASSERT(n_pool == kpool_pad(st.n_pool_real)); + GGML_ASSERT(pool_mask->ne[0] == (int64_t) n_pool && pool_mask->ne[1] == (int64_t) n_tokens); + GGML_ASSERT(tail_idxs->ne[0] == (int64_t) kpool - 1 && tail_idxs->ne[1] == (int64_t) n_tokens); + GGML_ASSERT(pool_idxs->ne[0] == (int64_t) kpool && pool_idxs->ne[1] == (int64_t) n_pool); + GGML_ASSERT((n_new == 0) == (new_pool_idxs == nullptr)); + GGML_ASSERT(st.cache_safe || new_pool_rep == nullptr); + GGML_ASSERT(!st.cache_safe || (n_new == 0) == (new_pool_rep == nullptr)); + + if (n_new > 0) { + GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_idxs->buffer)); + GGML_ASSERT(new_pool_idxs->ne[0] == (int64_t) kpool && new_pool_idxs->ne[1] == (int64_t) n_new); + if (new_pool_rep != nullptr) { + GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_rep->buffer)); + GGML_ASSERT(new_pool_rep->ne[0] == (int64_t) n_new); + } + } + + // Use the first ubatch cell for padded gathers. + uint32_t dummy_cell = 0; + { + const llama_seq_id s = ubatch->seq_id[0][0]; + const auto & sq = st.seqs[s]; + auto it = std::lower_bound(sq.cells.begin(), sq.cells.end(), std::make_pair(ubatch->pos[0], 0u)); + GGML_ASSERT(it != sq.cells.end() && it->first == ubatch->pos[0]); + dummy_cell = it->second; + } + + // Gather maps padding to a real cell and masks it separately. + const int32_t sentinel = gather ? (int32_t) dummy_cell : (int32_t) n_kv; + + float * gm = nullptr; + uint32_t n_sel = 0; + uint32_t n_top = 0; // Pools per token in the selection. + if (gather_mask != nullptr) { + GGML_ASSERT(ggml_backend_buffer_is_host(gather_mask->buffer)); + GGML_ASSERT(gather_mask->type == GGML_TYPE_F32); + GGML_ASSERT(gather_mask->ne[3] == (int64_t) n_tokens && gather_mask->ne[1] == 1 && gather_mask->ne[2] == 1); + n_sel = (uint32_t) gather_mask->ne[0]; + // The tail slots, when selected, are the n_sel % kpool != 0 remainder. + n_top = n_sel / kpool; + GGML_ASSERT(n_sel % kpool == 0 || n_sel % kpool == kpool - 1); + gm = (float *) gather_mask->data; + } + + // pools are laid out per sequence + std::vector seq_pool_start(LLAMA_MAX_SEQ, 0); + std::vector pool_end; + pool_end.reserve(n_pool); + + int32_t * pcell = (int32_t *) pool_cells->data; + int32_t * pidx = (int32_t *) pool_idxs->data; + int32_t * nidx = n_new > 0 ? (int32_t *) new_pool_idxs->data : nullptr; + int64_t * nrep = new_pool_rep != nullptr ? (int64_t *) new_pool_rep->data : nullptr; + + uint32_t i_new = 0; + for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + const auto & sq = st.seqs[s]; + seq_pool_start[s] = (uint32_t) pool_end.size(); + for (size_t pi = 0; pi < sq.pools.size(); ++pi) { + const uint32_t j = sq.pools[pi]; + const uint32_t ip = (uint32_t) pool_end.size(); + GGML_ASSERT(ip + 1 < n_pool); + + // The pooled key lives in the last member's row. + const uint32_t rep = sq.cells[j + kpool - 1].second; + pcell[ip] = (int32_t) rep; + + for (uint32_t k = 0; k < kpool; ++k) { + pidx[(size_t) ip*kpool + k] = (int32_t) sq.cells[j + k].second; + } + + if (sq.is_new[pi]) { + GGML_ASSERT(i_new < n_new); + for (uint32_t k = 0; k < kpool; ++k) { + nidx[(size_t) i_new*kpool + k] = (int32_t) sq.cells[j + k].second; + } + if (nrep != nullptr) { + nrep[i_new] = (int64_t) rep; + } + ++i_new; + } + + pool_end.push_back(sq.cells[j + kpool - 1].first); + } + } + GGML_ASSERT(i_new == n_new); + + const uint32_t n_pool_real = (uint32_t) pool_end.size(); + for (uint32_t ip = n_pool_real; ip < n_pool; ++ip) { + pcell[ip] = (int32_t) dummy_cell; + for (uint32_t k = 0; k < kpool; ++k) { + pidx[(size_t) ip*kpool + k] = sentinel; + } + } + + // a pool is visible when it belongs to the token's sequence and ends at or before it + auto fill_mask = [&](auto * data) { + using T = std::remove_pointer_t; + const T keep = llama_cast(0.0f); + const T drop = llama_cast(-INFINITY); + + for (uint32_t i = 0; i < n_tokens; ++i) { + const llama_seq_id s = ubatch->seq_id[i][0]; + const llama_pos p = ubatch->pos[i]; + + T * row = data + (size_t) i*n_pool; + std::fill(row, row + n_pool, drop); + + const uint32_t p0 = seq_pool_start[s]; + const uint32_t p1 = p0 + (uint32_t) st.seqs[s].pools.size(); + const uint32_t nv = (uint32_t) (std::upper_bound(pool_end.begin() + p0, pool_end.begin() + p1, p) - (pool_end.begin() + p0)); + std::fill(row + p0, row + p0 + nv, keep); + + // Finite visible pools occupy the first min(nv, n_top) ranked slots. + if (gm != nullptr) { + const uint32_t nvc = std::min(nv, n_top); + float * grow = gm + (size_t) i*n_sel; + std::fill(grow, grow + (size_t) nvc*kpool, 0.0f); + std::fill(grow + (size_t) nvc*kpool, grow + (size_t) n_top*kpool, -INFINITY); + } + } + }; + if (pool_mask->type == GGML_TYPE_F16) { + fill_mask((ggml_fp16_t *) pool_mask->data); + } else { + fill_mask((float *) pool_mask->data); + } + + int32_t * tidx = (int32_t *) tail_idxs->data; + for (uint32_t i = 0; i < n_tokens; ++i) { + const llama_seq_id s = ubatch->seq_id[i][0]; + const llama_pos p = ubatch->pos[i]; + const auto & sq = st.seqs[s]; + + const uint32_t n_tail = (uint32_t) ((p - sq.pos_min + 1) % (llama_pos) kpool); + + for (uint32_t k = 0; k < kpool - 1; ++k) { + int32_t cell = sentinel; + bool real = false; + if (k < n_tail) { + const llama_pos pt = p - (llama_pos) k; + auto it = std::lower_bound(sq.cells.begin(), sq.cells.end(), std::make_pair(pt, 0u)); + if (it != sq.cells.end() && it->first == pt) { + cell = (int32_t) it->second; + real = true; + } + } + tidx[(size_t) i*(kpool - 1) + k] = cell; + + if (gm != nullptr && n_sel % kpool != 0) { + gm[(size_t) i*n_sel + (size_t) n_top*kpool + k] = real ? 0.0f : -INFINITY; + } + } + } +} +void llama_memory_hybrid_idx_context::set_input_mtp_dsa_selection( + ggml_tensor * sel, ggml_tensor * mask, bool gather, + const llama_ubatch * ubatch) const { + GGML_ASSERT(mem != nullptr && ctx_idx != nullptr); + GGML_ASSERT(sel != nullptr && mask != nullptr && ubatch != nullptr); + GGML_ASSERT(sel->type == GGML_TYPE_I32 && mask->type == GGML_TYPE_F32); + + const auto & saved = mem->get_mtp_dsa_selection(); + const size_t n = (size_t) ggml_nelements(sel); + const size_t width = (size_t) sel->ne[0]; + GGML_ASSERT(saved.size() == n && (size_t) ggml_nelements(mask) == n); + GGML_ASSERT(sel->ne[1] == (int64_t) ubatch->n_tokens); + + const auto & st = kpool_cur(); + const int32_t n_kv = (int32_t) get_idx()->get_n_kv(); + GGML_ASSERT(n_kv > 0); + + std::vector mapped(n); + std::vector valid(n); + for (uint32_t i = 0; i < ubatch->n_tokens; ++i) { + GGML_ASSERT(ubatch->n_seq_id[i] == 1); + const llama_seq_id seq_id = ubatch->seq_id[i][0]; + GGML_ASSERT(seq_id >= 0 && seq_id < LLAMA_MAX_SEQ); + const auto & cells = st.seqs[seq_id].cells; + + for (size_t j = 0; j < width; ++j) { + const size_t k = (size_t) i*width + j; + const llama_pos pos = saved[k]; + auto it = pos >= 0 ? std::lower_bound(cells.begin(), cells.end(), std::make_pair(pos, 0u)) : cells.end(); + const bool found = it != cells.end() && it->first == pos; + + if (found) { + mapped[k] = (int32_t) it->second; + valid[k] = 0.0f; + } else { + mapped[k] = gather ? 0 : n_kv; + valid[k] = -INFINITY; + } + } + } + + ggml_backend_tensor_set(sel, mapped.data(), 0, mapped.size()*sizeof(mapped[0])); + ggml_backend_tensor_set(mask, valid.data(), 0, valid.size()*sizeof(valid[0])); +} + + diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 705189e7eb58..dea2d65e1f57 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -87,6 +87,18 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const; + // The model's indexer pool size. + uint32_t get_kpool() const { return n_kpool; } + + // Sequence edits invalidate cached relative pools. + bool kpool_is_dirty () const { return kpool_dirty; } + void kpool_clear_dirty() { kpool_dirty = false; } + + void set_mtp_dsa_index_share(bool enabled); + bool get_mtp_dsa_index_share() const { return mtp_dsa_index_share; } + void set_mtp_dsa_selection(const int32_t * data, size_t size); + const std::vector & get_mtp_dsa_selection() const { return mtp_dsa_selection; } + private: // forget seq_id (all of it if seq_id < 0) in every cache at once, so a failed restore cannot leave the caches out of step // seq_id < 0 drops the whole context, as the caches themselves do on a failed restore @@ -97,6 +109,13 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_hparams hparams_idx; const std::unique_ptr mem_idx; + + const uint32_t n_kpool; + + bool kpool_dirty = false; + + bool mtp_dsa_index_share = false; + std::vector mtp_dsa_selection; }; class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { @@ -122,7 +141,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { slot_info_vec_t sinfos_idx, std::vector ubatches); - ~llama_memory_hybrid_idx_context() = default; + ~llama_memory_hybrid_idx_context(); // Defined out of line because kpool_state is incomplete here. // // llama_memory_context_i @@ -141,12 +160,24 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // streams in the current slot info, the `ns` of get_k/get_v; 1 if unified uint32_t get_n_stream() const; + // glm5-next, complete pools of kpool consecutive positions per sequence, scored as whole pools. + uint32_t get_n_kpool () const; // Padded pool count, where the last pool is always unused. + uint32_t get_n_kpool_new() const; // Exact count of pools completed by the current ubatch. + bool get_kpool_cache_safe() const; + bool get_mtp_dsa_index_share() const; + size_t get_mtp_dsa_selection_size() const; + void set_input_kpool(ggml_tensor * pool_cells, ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs, + ggml_tensor * gather_mask, bool gather, ggml_tensor * new_pool_idxs, ggml_tensor * new_pool_rep, + const llama_ubatch * ubatch) const; + void set_input_mtp_dsa_selection(ggml_tensor * sel, ggml_tensor * mask, bool gather, + const llama_ubatch * ubatch) const; + void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const; private: - const llama_memory_hybrid_idx * mem = nullptr; + llama_memory_hybrid_idx * mem = nullptr; // streams per ubatch, read from the slot infos before ctx_idx takes them // declared first, so it is initialised while sinfos_idx is still intact @@ -157,4 +188,16 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // mirrors the base class's ubatch cursor, which is private there size_t i_cur = 0; + + // K-pool layouts + struct kpool_state; + kpool_state kpool_build_state(const llama_ubatch * ubatch) const; + const kpool_state & kpool_cur() const; + std::vector kpool_states; + + // Whether this context tracks k-pool states. + bool kpool_track = false; + + // Clear a pending full re-pool only after the first ubatch succeeds + bool kpool_dirty_batch = false; }; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 57919accf095..e5885ae72905 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -80,6 +80,8 @@ llama_memory_recurrent::llama_memory_recurrent( continue; } + has_state = true; + const char * dev_name = "CPU"; ggml_backend_buffer_type_t buft = ggml_backend_cpu_buffer_type(); @@ -190,16 +192,25 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos if (tail_id >= 0) { auto & cell = cells[tail_id]; - // partial rollback via per-token snapshot index (bounded by n_rs_seq) + // Partial removal at the end of the sequence. if (0 < p0 && p0 <= cell.pos && p1 > cell.pos) { + // An empty recurrent cache tracks positions only, so this is trivially valid. + if (!has_state) { + cell.pos = p0 - 1; + return true; + } + + // Partial rollback via the per token snapshot planes, bounded by n_rs_seq. const llama_pos rollback = cell.pos - (p0 - 1); - // pending rollback is single-use + // A pending rollback is single use. const bool pending = rs_idx[seq_id] != 0; if (!pending && rollback >= 1 && rollback <= (llama_pos) n_rs_seq) { set_rs_idx(seq_id, (uint32_t) rollback); cell.pos = p0 - 1; return true; } + LLAMA_LOG_DEBUG("%s: refused partial removal: seq=%d p0=%d tail.pos=%d rollback=%d n_rs_seq=%u pending=%d\n", + __func__, seq_id, p0, cell.pos, rollback, n_rs_seq, (int) pending); return false; } // invalidate tails which will be cleared diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index 4abb3f5cf5c0..4f7b16c8d9a7 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -70,6 +70,8 @@ class llama_memory_recurrent : public llama_memory_i { uint32_t size = 0; // total number of cells, shared across all sequences uint32_t used = 0; // used cells (i.e. at least one seq_id) + bool has_state = false; + // number of recurrent-state snapshots per seq for rollback; tensors are widened to (1 + n_rs_seq) groups uint32_t n_rs_seq = 0; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 8860bd3f4342..4fb9d556b151 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -32,6 +32,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_LAGUNA: case LLM_ARCH_GRANITE_SWA: case LLM_ARCH_DOTS3NOTE: // TODO: need to handle SWA pattern and MLA+SWA config + case LLM_ARCH_GLM5_NEXT: return false; default: return true; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index bfce09de0ca4..1166ee3e2187 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -334,6 +334,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_kimi_linear(params); case LLM_ARCH_KIMI_K3: return new llama_model_kimi_k3(params); + case LLM_ARCH_GLM5_NEXT: + return new llama_model_glm5_next(params); case LLM_ARCH_STEP35: return new llama_model_step35(params); default: @@ -964,6 +966,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_685B_A37B: return "685B.A37B"; case LLM_TYPE_744B_A40B: return "744B.A40B"; case LLM_TYPE_2_8T_A50B: return "2.8T.A50B"; + case LLM_TYPE_320B_A18B: return "320B.A18B"; case LLM_TYPE_E2B: return "E2B"; case LLM_TYPE_E4B: return "E4B"; default: return "?B"; @@ -2308,6 +2311,53 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr); } } break; + case LLM_ARCH_GLM5_NEXT: + { + if (!cparams.kv_unified && cparams.n_seq_max > 1) { + throw std::runtime_error("GLM5-Next requires a unified KV cache for multiple sequences, use --kv-unified"); + } + // KDA layers are recurrent, the DSA layers use a K-only MLA cache plus an indexer cache. + // tThe Nextn block is never attended by the trunk graph + llama_memory_hybrid_idx::layer_filter_cb filter_attn = [&](uint32_t il) { + return il < hparams.n_layer() && !hparams.is_recr(il); + }; + llama_memory_hybrid_idx::layer_filter_cb filter_idx = [&](uint32_t il) { + return il < hparams.n_layer() && !hparams.is_recr(il) && hparams.is_indexer_full(il); + }; + llama_memory_hybrid_idx::layer_filter_cb filter_recr = [&](uint32_t il) { + return il < hparams.n_layer() && hparams.is_recr(il); + }; + + // the draft head is a single DSA layer + if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) { + if (hparams.n_layer_nextn == 0) { + throw std::runtime_error("GLM5-Next MTP requires the NextN block, convert without --no-mtp"); + } + filter_attn = [&](uint32_t il) { return il >= hparams.n_layer(); }; + filter_idx = [&](uint32_t il) { return il >= hparams.n_layer(); }; + filter_recr = [&](uint32_t) { return false; }; + } + + res = new llama_memory_hybrid_idx( + /* model */ *this, + /* attn_type_k */ params.type_k, + /* attn_type_v */ params.type_v, + /* attn_v_trans */ !cparams.flash_attn, + /* attn_kv_size */ cparams.n_ctx_seq, + /* attn_n_pad */ 1, + /* attn_n_swa */ hparams.n_swa, + /* attn_swa_type */ hparams.swa_type, + /* recurrent_type_r */ GGML_TYPE_F32, + /* recurrent_type_s */ GGML_TYPE_F32, + /* recurrent_rs_size */ std::max((uint32_t) 1, cparams.n_seq_max), + /* n_seq_max */ cparams.n_seq_max, + /* n_rs_seq */ cparams.n_rs_seq, + /* offload */ cparams.offload_kqv, + /* unified */ cparams.kv_unified, + /* filter_attn */ std::move(filter_attn), + /* filter_recr */ std::move(filter_recr), + /* filter_idx */ std::move(filter_idx)); + } break; case LLM_ARCH_DOTS3NOTE: { GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); @@ -2820,6 +2870,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_KIMI_K3: + case LLM_ARCH_GLM5_NEXT: return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values diff --git a/src/llama-model.h b/src/llama-model.h index 38066538ed10..ab29b4eee520 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -149,6 +149,7 @@ enum llm_type { LLM_TYPE_685B_A37B, // DeepSeek V3.2 LLM_TYPE_744B_A40B, // GLM-5 LLM_TYPE_2_8T_A50B, // Kimi-K3 + LLM_TYPE_320B_A18B, // GLM-5.3-Flash LLM_TYPE_E2B, LLM_TYPE_E4B, }; @@ -555,6 +556,10 @@ struct llama_layer { struct ggml_tensor * indexer_attn_k = nullptr; struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias + // glm5-next k-pool indexer + struct ggml_tensor * indexer_kpool_gate = nullptr; + struct ggml_tensor * indexer_kpool_ape = nullptr; + // MSA struct ggml_tensor * index_q_proj = nullptr; struct ggml_tensor * index_k_proj = nullptr; diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 34ff25db57e6..aa92c361b051 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -328,6 +328,25 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param quantize &= name.find("indexer.k_proj.weight") == std::string::npos; quantize &= name.find("indexer.q_proj.weight") == std::string::npos; + // glm5-next + if (arch == LLM_ARCH_GLM5_NEXT) { + quantize &= name.find("hc_") == std::string::npos; + quantize &= name.find("indexer.attn_q_b") == std::string::npos; + quantize &= name.find("indexer.attn_k") == std::string::npos; + quantize &= name.find("indexer.proj") == std::string::npos; + quantize &= name.find("indexer_compressor_gate") == std::string::npos; + quantize &= name.find("indexer_compressor_ape") == std::string::npos; + quantize &= name.find("hc_") == std::string::npos; + quantize &= name.find("ssm_f_a.weight") == std::string::npos; + quantize &= name.find("ssm_f_b.weight") == std::string::npos; + quantize &= name.find("ssm_g_a.weight") == std::string::npos; + quantize &= name.find("ssm_g_b.weight") == std::string::npos; + quantize &= name.find("ssm_beta.weight") == std::string::npos; + quantize &= name.find("attn_kv_a_mqa.weight") == std::string::npos; + quantize &= name.find("attn_k_b.weight") == std::string::npos; + quantize &= name.find("attn_v_b.weight") == std::string::npos; + } + // do not quantize RWKV's small yet 2D weights quantize &= name.find("time_mix_first.weight") == std::string::npos; quantize &= name.find("time_mix_w0.weight") == std::string::npos; @@ -451,6 +470,22 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type return std::make_pair(i_layer, n_layer); }; + // by default, for glm5-next, don't let these tensors be quantized below Q8_0 + if (arch == LLM_ARCH_GLM5_NEXT && ( + name.find("attn_q_a") != std::string::npos || + name.find("attn_q_b") != std::string::npos || + name.find("nextn.eh_proj") != std::string::npos)) + { + switch (new_type) { + case GGML_TYPE_F32: + case GGML_TYPE_BF16: + case GGML_TYPE_F16: + break; + default: + return GGML_TYPE_Q8_0; + } + } + // for arches that share the same tensor between the token embeddings and the output, we quantize the token embeddings // with the quantization of the output tensor if (category == tensor_category::OUTPUT || (qs.has_tied_embeddings && category == tensor_category::TOKEN_EMBD)) { diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index ff926ceecd17..ebd46aa7082e 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2253,10 +2253,15 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { pre_type = LLAMA_VOCAB_PRE_TYPE_PORO; clean_spaces = false; } else if ( - tokenizer_pre == "glm4" || tokenizer_pre == "chatglm-bpe") { pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4; special_bos_id = LLAMA_TOKEN_NULL; + } else if ( + tokenizer_pre == "glm4" || + tokenizer_pre == "glm5") { + pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4; + special_bos_id = LLAMA_TOKEN_NULL; + ignore_merges = true; } else if ( tokenizer_pre == "viking") { pre_type = LLAMA_VOCAB_PRE_TYPE_VIKING; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 0157ce7051bd..e66c6bebf4b0 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -261,15 +261,15 @@ static dsv4_state_tensors dsv4_build_state_snapshot( static constexpr int64_t DSV4_CSA_RATIO = 4; static constexpr int64_t DSV4_HCA_RATIO = 128; -// mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens] -static ggml_tensor * dsv4_hc_mean(ggml_context * ctx, ggml_tensor * x) { +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_mean(ggml_tensor * x) const { const int64_t hc = x->ne[1]; - ggml_tensor * acc = ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0); + ggml_tensor * acc = ggml_view_2d(ctx0, x, x->ne[0], x->ne[2], x->nb[2], 0); for (int64_t s = 1; s < hc; ++s) { - acc = ggml_add(ctx, acc, ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1])); + acc = ggml_add(ctx0, acc, ggml_view_2d(ctx0, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1])); } - return ggml_scale(ctx, acc, 1.0f/hc); + return ggml_scale(ctx0, acc, 1.0f/hc); } static ggml_tensor * dsv4_hc_affine( @@ -282,7 +282,8 @@ static ggml_tensor * dsv4_hc_affine( return x; } -ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_pre( ggml_tensor * x, ggml_tensor * weights, int il) const { @@ -309,7 +310,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( return result; } -ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn( +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_sinkhorn( ggml_tensor * comb, int il) const { GGML_UNUSED(il); @@ -346,7 +348,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn( return comb; } -ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_pre( ggml_tensor * x, ggml_tensor * hc_fn, ggml_tensor * hc_scale, @@ -404,7 +407,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( return result; } -ggml_tensor * llama_model_deepseek4::graph::build_hc_post( +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_post( ggml_tensor * x, ggml_tensor * residual, ggml_tensor * post, @@ -441,6 +445,10 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_post( return out; } +// instantiate the mHC helpers for deepseek4 (and dflash) and glm5-next +template struct llama_model_deepseek4::graph_base; +template struct llama_model_deepseek4::graph_base; + ggml_tensor * llama_model_deepseek4::graph::build_hc_head( ggml_tensor * x, ggml_tensor * hc_fn, @@ -1215,7 +1223,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl( } llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_params & params) : - llm_graph_context(params) { + graph_base<>(params) { ggml_tensor * cur; ggml_tensor * inp = build_inp_embd(model.tok_embd); @@ -1232,7 +1240,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p for (int il = 0; il < n_layer; ++il) { if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) { - res->t_layer_inp[il] = dsv4_hc_mean(ctx0, inpL); + res->t_layer_inp[il] = build_hc_mean(inpL); cb(res->t_layer_inp[il], "layer_inp", il); ggml_build_forward_expand(gf, res->t_layer_inp[il]); } @@ -1314,7 +1322,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p } if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) { - res->t_layer_inp[n_layer] = dsv4_hc_mean(ctx0, inpL); + res->t_layer_inp[n_layer] = build_hc_mean(inpL); cb(res->t_layer_inp[n_layer], "layer_inp", n_layer); ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); } diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp new file mode 100644 index 000000000000..7124546591ad --- /dev/null +++ b/src/models/glm5-next.cpp @@ -0,0 +1,967 @@ +#include "models.h" +#include "llama-memory-hybrid-idx.h" + +// GLM5-Next (GLM-5.3-Flash): hybrid KDA (linear) + nope MLA with a k-pool DSA indexer, +// mHC residual streams, DeepSeek-style MoE. + +void llama_model_glm5_next::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps, false); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false); + + // the MLA cache holds the compressed latent + hparams.n_embd_head_v_full = hparams.n_lora_kv; + + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; + } + + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); + if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) { + hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID; + } + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false); + if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false)) { + hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp; + } + + // DSA indexer with k-pool compression + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KPOOL, hparams.indexer_kpool); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, hparams.indexer_kpool_select_tail, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_INDEX_SHARE_MTP, hparams.indexer_index_share_mtp, false); + GGML_ASSERT(hparams.indexer_kpool > 1 && hparams.indexer_top_k % hparams.indexer_kpool == 0); + std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1); + ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false); + + // mHC + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); + ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); + GGML_ASSERT(hparams.dsv4_hc_mult == 4 && "mHC with hc_mult != 4 is not supported"); + + switch (hparams.n_layer()) { + case 45: type = LLM_TYPE_320B_A18B; break; // GLM-5.3-Flash + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_mix_dim = (2 + hc)*hc; + + // The NextN block is loaded but only used by the MTP graph. + const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; + const bool trunk_only = (n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); + int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; + mtp_ready = n_layer_nextn > 0 && !trunk_only && ml.load_mtp; + if (trunk_only) { + LLAMA_LOG_INFO("%s: trunk only GGUF, the MTP draft head is unavailable\n", __func__); + } + if (!ml.load_mtp) { + mtp_flags |= TENSOR_SKIP; + } + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + for (int i = 0; i < n_layer_all; ++i) { + auto & layer = layers[i]; + + const int flags = (i >= n_layer) ? mtp_flags : 0; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); + + if (i < n_layer) { + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, 0); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, 0); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0); + } + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = head_dim * n_head; + + if (hparams.is_recr(i)) { + auto conv = [&](llm_tensor tid) { + ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED); + return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0); + }; + layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q); + layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K); + layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V); + + create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0); + + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0); + layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); + + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, 0); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0); + + layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", i), {n_embd, head_dim}, 0); + layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", i), {head_dim, d_inner}, 0); + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0); + } else { + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim; + + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags); + + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, flags); + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, flags); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, flags); + + const int64_t n_indexer_head = hparams.indexer_n_head; + const int64_t n_embd_indexer = hparams.indexer_head_size; + const int64_t kpool = hparams.indexer_kpool; + + const bool full = i >= n_layer || hparams.is_indexer_full(i); + const int iflags = flags | (full ? 0 : TENSOR_NOT_REQUIRED); + + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_embd_indexer}, iflags); + layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {n_embd_indexer}, iflags); + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, n_indexer_head}, iflags); + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, n_embd_indexer}, iflags); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_indexer_head * n_embd_indexer}, iflags); + layer.indexer_kpool_gate = create_tensor(tn(LLM_TENSOR_INDEXER_KPOOL_GATE, "weight", i), {n_embd, n_embd_indexer}, iflags); + layer.indexer_kpool_ape = create_tensor(tn(LLM_TENSOR_INDEXER_KPOOL_APE, "weight", i), {n_embd_indexer, kpool}, iflags); + } + + if (i < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } else { + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_expert_shared = hparams.n_expert_shared; + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + } + + if (i >= n_layer) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, flags); + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, flags | TENSOR_NOT_REQUIRED); + } + } +} + +std::unique_ptr llama_model_glm5_next::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + if (!mtp_ready) { + throw std::runtime_error("MTP graph requested but the NextN tensors are not loaded"); + } + return std::make_unique(*this, params); + } + return std::make_unique(*this, params); +} + +// Causal conv1d over one of Q/K/V +static ggml_tensor * glm5_conv1d(ggml_cgraph * gf, ggml_context * ctx0, + ggml_tensor * conv_states_all, ggml_tensor * conv_state_all, + int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, + int64_t d_conv, int64_t head_dim, int64_t n_head, + int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head, + int64_t mem_size, int64_t n_rs_seq) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t n_embd_r_total = 3 * conv_state_size; + + ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + n_embd_r_total * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0); + + // Snapshot plane s holds the conv window s tokens back. + const int64_t n_planes = n_rs_seq + 1; + for (int64_t t = 1; t <= n_planes; ++t) { + const int64_t s_idx = std::max(0, n_seq_tokens - n_planes + t); + const int64_t s_slot = n_planes - t; + + ggml_tensor * conv_window = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], s_idx * conv_x->nb[0]); + + ggml_tensor * conv_update = ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + n_embd_r_total * ggml_element_size(conv_states_all), + ((s_slot * mem_size + kv_head) * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all)); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, conv_window, conv_update)); + } + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight); + Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens); + Xcur = ggml_silu(ctx0, Xcur); + + return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs); +} + + +// K-pool indexer inputs +class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { +public: + llm_graph_input_kpool(const llama_memory_hybrid_idx_context * mctx, uint32_t kpool) : mctx(mctx), kpool(kpool) {} + virtual ~llm_graph_input_kpool() = default; + + void set_input(const llama_ubatch * ubatch) override { + mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); + mctx->set_input_kpool(pool_cells, pool_idxs, pool_mask, tail_idxs, gather_mask, gather, new_pool_idxs, new_pool_rep, ubatch); + if (reuse_sel != nullptr) { + mctx->set_input_mtp_dsa_selection(reuse_sel, gather_mask, gather, ubatch); + } + } + + bool can_reuse(const llm_graph_params & params) override { + mctx = static_cast(params.mctx); + + const auto * idx = mctx->get_idx(); + if (idx == nullptr) { + return false; + } + + bool res = true; + + res &= k_idxs->ne[0] == params.ubatch.n_tokens; + res &= pool_cells->ne[0] == mctx->get_n_kpool(); + res &= pool_mask->ne[1] == params.ubatch.n_tokens; + res &= tail_idxs->ne[1] == params.ubatch.n_tokens; + // The scatter mask shape follows n_kv. + res &= n_kv == idx->get_n_kv(); + // The new pool path is sized exactly + res &= n_new == mctx->get_n_kpool_new(); + res &= cache_safe == mctx->get_kpool_cache_safe(); + const bool share = params.cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && mctx->get_mtp_dsa_index_share(); + const size_t saved = mctx->get_mtp_dsa_selection_size(); + const bool reuse = share && saved == (size_t) n_sel*params.ubatch.n_tokens; + res &= mtp_share == share; + res &= (reuse_sel != nullptr) == reuse; + + return res; + } + + ggml_tensor * k_idxs = nullptr; // I64 [n_tokens] + ggml_tensor * pool_cells = nullptr; // I32 [n_pool] cell caching each pool's pooled key + ggml_tensor * pool_idxs = nullptr; // I32 [kpool, n_pool] member cells per pool, n_kv sentinel for the padded pools + ggml_tensor * pool_mask = nullptr; // F32/F16 [n_pool, n_tokens] + ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens] + ggml_tensor * gather_mask = nullptr; // F32 [n_sel, 1, 1, n_tokens] + ggml_tensor * reuse_sel = nullptr; // I32 [n_sel, n_tokens] + ggml_tensor * new_pool_idxs = nullptr; // I32 [kpool, n_new] members of the pools completed this ubatch + ggml_tensor * new_pool_rep = nullptr; // I64 [n_new] cell to write each new pooled key into + + const llama_memory_hybrid_idx_context * mctx; + const uint32_t kpool; + uint32_t n_new = 0; + uint32_t n_sel = 0; + bool cache_safe = true; + bool gather = false; + bool mtp_share = false; + uint32_t n_kv = 0; +}; + +llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::build_inp_kpool(const llama_memory_hybrid_idx_context * mctx_hyb) { + const auto * mctx_idx = mctx_hyb->get_idx(); + GGML_ASSERT(mctx_idx != nullptr); + + const uint32_t kpool = hparams.indexer_kpool; + const uint32_t n_pool = mctx_hyb->get_n_kpool(); + const uint32_t n_kv = mctx_idx->get_n_kv(); + const uint32_t n_new = mctx_hyb->get_n_kpool_new(); + const bool cache_safe = mctx_hyb->get_kpool_cache_safe(); + + // the fused lightning indexer wants an f16 mask + const auto type_mask = cparams.fused_lid ? GGML_TYPE_F16 : GGML_TYPE_F32; + + auto inp = std::make_unique(mctx_hyb, kpool); + + inp->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + inp->pool_cells = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pool); + inp->pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_pool); + inp->pool_mask = ggml_new_tensor_2d(ctx0, type_mask, n_pool, n_tokens); + inp->tail_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool - 1, n_tokens); + ggml_set_input(inp->pool_cells); + ggml_set_input(inp->pool_idxs); + ggml_set_input(inp->pool_mask); + ggml_set_input(inp->tail_idxs); + + ggml_build_forward_expand(gf, inp->pool_cells); + ggml_build_forward_expand(gf, inp->pool_idxs); + ggml_build_forward_expand(gf, inp->pool_mask); + ggml_build_forward_expand(gf, inp->tail_idxs); + + inp->n_kv = n_kv; + + // Gather selected latents for small decode batches when n_kv exceeds n_sel. + { + constexpr int64_t max_ub = 16; + + const int64_t n_top_pool = std::min(n_pool, hparams.indexer_top_k / kpool); + const int64_t n_sel = kpool*n_top_pool + (hparams.indexer_kpool_select_tail ? kpool - 1 : 0); + const bool mtp_share = cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && mctx_hyb->get_mtp_dsa_index_share(); + + inp->n_sel = (uint32_t) n_sel; + inp->gather = (int64_t) n_tokens <= max_ub && (int64_t) n_kv > n_sel; + inp->mtp_share = mtp_share; + + if (inp->gather || mtp_share) { + inp->gather_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens); + ggml_set_input(inp->gather_mask); + // Keep the mask allocated even when no op reads it, because set_input_kpool always fills it. + ggml_build_forward_expand(gf, inp->gather_mask); + } + + const size_t saved = mctx_hyb->get_mtp_dsa_selection_size(); + if (mtp_share && saved == (size_t) n_sel*n_tokens) { + inp->reuse_sel = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_sel, n_tokens); + ggml_set_input(inp->reuse_sel); + } + } + + inp->n_new = n_new; + inp->cache_safe = cache_safe; + if (n_new > 0) { + inp->new_pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_new); + ggml_set_input(inp->new_pool_idxs); + if (cache_safe) { + inp->new_pool_rep = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new); + ggml_set_input(inp->new_pool_rep); + } + } + + return (llm_graph_input_kpool *) res->add_input(std::move(inp)); +} + +llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params) : + llama_model_deepseek4::graph_base(params), model(model) { + + ggml_tensor * cur; + + ggml_tensor * inp = build_inp_embd(model.tok_embd); + cb(inp, "inp_embd", -1); + + // recurrent state + K-only MLA cache through the generic hybrid input, plus the indexer cache + const auto * mctx_hyb = static_cast(mctx); + + auto * inp_hyb = build_inp_mem_hybrid_k(); + auto * inp_rs = inp_hyb->get_recr(); + auto * inp_attn = inp_hyb->get_attn(); + auto * inp_kpool = build_inp_kpool(mctx_hyb); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head_kda = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = n_head_kda * head_dim; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + const int64_t hc = hparams.dsv4_hc_mult; + ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens); + inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1); + cb(inpL, "hc_init", -1); + + ggml_tensor * prev_sel = nullptr; + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * residual = inpL; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + + cur = build_hc_pre(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, il); + cb(cur, "hc_attn_pre", il); + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recr(il)) { + cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda, + d_inner, n_seq_tokens, n_seqs, il); + } else { + cur = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, &prev_sel, il); + } + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "hc_attn_post", il); + + residual = inpL; + cur = build_hc_pre(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, il); + cb(cur, "hc_ffn_pre", il); + + cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = build_ffn(cur, + layer.ffn_up, nullptr, nullptr, + layer.ffn_gate, nullptr, nullptr, + layer.ffn_down, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + } else { + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.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); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * ffn_shexp = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, 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); + } + + inpL = build_hc_post(cur, residual, post, comb, il); + inpL = build_cvec(inpL, il); + cb(inpL, "l_out", il); + } + + // narrow to the output tokens, then collapse the streams + // Unmasked nextn embeddings need all rows. + const bool narrow_early = inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked); + if (narrow_early) { + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + flat = ggml_get_rows(ctx0, flat, inp_out_ids); + inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs); + } + + cur = build_hc_mean(inpL); + cb(cur, "hc_head", -1); + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + + // the post-norm hidden state feeds the draft head + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (inp_out_ids && !narrow_early) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +// KDA layer, g_a/g_b output gate + +ggml_tensor * llama_model_glm5_next::graph::build_kda_layer( + ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs, + int64_t d_conv, int64_t head_dim, int64_t n_head_kda, + int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) { + + const auto * mctx_cur = inp_rs->mctx; + const auto kv_head = mctx_cur->get_head(); + const auto mem_size = mctx_cur->get_size(); + const auto n_rs_seq = (int64_t) cparams.n_rs_seq; + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * Qcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, n_rs_seq); + ggml_tensor * Kcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, n_rs_seq); + ggml_tensor * Vcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, n_rs_seq); + cb(Qcur, "kda_q_conv", il); + cb(Kcur, "kda_k_conv", il); + cb(Vcur, "kda_v_conv", il); + + // Decay gate, ssm_a holds -exp(A_log) + ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a); + g1 = ggml_add(ctx0, g1, layer.ssm_dt_b); + + ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1); + + if (hparams.kda_gate_lower_bound > -INFINITY) { + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); + g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f)); + g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound); + } else { + g1 = ggml_softplus(ctx0, g1); + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); + } + cb(g1, "kda_g1", il); + + g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs); + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "kda_beta", il); + + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs); + + // Match FLA l2 norm + constexpr float l2_eps = 1e-6f; + const float l2_scale = 1.0f / std::sqrt((float) head_dim); + Qcur = ggml_scale(ctx0, ggml_rms_norm(ctx0, Qcur, l2_eps / (float) head_dim), l2_scale); + Kcur = ggml_scale(ctx0, ggml_rms_norm(ctx0, Kcur, l2_eps / (float) head_dim), l2_scale); + + ggml_tensor * output = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g1, beta, state, il); + output = ggml_cont(ctx0, output); + cb(output, "kda_scan_out", il); + + // output gate, then RMSNorm(o) * Sigmoid(g2) + ggml_tensor * g_a = ggml_mul_mat(ctx0, layer.ssm_g_a, cur); + ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g_b, g_a); + g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_tokens); + + ggml_tensor * o = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_tokens); + ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + cb(g2, "kda_g2", il); + cb(normed, "kda_normed", il); + ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2)); + + gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens); + cur = ggml_mul_mat(ctx0, layer.wo, gated); + cb(cur, "kda_out", il); + + return cur; +} + +// Scores pools of kpool consecutive tokens, expands the selected pools and the incomplete tail into an additive mask + +ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( + ggml_tensor * cur, ggml_tensor * qr, ggml_tensor * kq_mask, const llama_layer & layer, + const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il) { + + const auto * mctx_lid = mctx_hyb->get_idx(); + + const int64_t n_indexer_head = hparams.indexer_n_head; + const int64_t n_embd_indexer = hparams.indexer_head_size; + const int64_t kpool = hparams.indexer_kpool; + const int64_t n_pool = inp_kpool->pool_cells->ne[0]; + const int64_t n_new = inp_kpool->n_new; + + ggml_tensor * iq = nullptr; + if (inp_kpool->reuse_sel == nullptr) { + iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr); + iq = ggml_reshape_3d(ctx0, iq, n_embd_indexer, n_indexer_head, n_tokens); + cb(iq, "indexer_q", il); + } + + // Per-token key and pool gate scores, cached together + ggml_tensor * ik = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur); + ik = build_norm(ik, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il); + cb(ik, "indexer_k", il); + + ggml_tensor * ig = ggml_mul_mat(ctx0, layer.indexer_kpool_gate, cur); + cb(ig, "indexer_gate", il); + + // Cache rows store key | gate | pooled + ggml_tensor * pzero = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_indexer, n_tokens), 0.0f); + ggml_tensor * packed = ggml_concat(ctx0, ggml_concat(ctx0, ik, ig, 0), pzero, 0); + packed = ggml_reshape_3d(ctx0, packed, 3*n_embd_indexer, 1, n_tokens); + ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, packed, inp_kpool->k_idxs, il)); + + ggml_tensor * k_all = mctx_lid->get_k(ctx0, il); + GGML_ASSERT(k_all->ne[3] == 1 && "TODO: k-pool indexer with multiple streams"); + const int64_t n_kv = k_all->ne[2]; + + ggml_tensor * kg_all = ggml_view_2d(ctx0, k_all, 2*n_embd_indexer, n_kv, k_all->nb[2], 0); + ggml_tensor * pooled_all = ggml_view_2d(ctx0, k_all, n_embd_indexer, n_kv, k_all->nb[2], + ggml_row_size(k_all->type, 2*n_embd_indexer)); + + ggml_tensor * pooled_new = nullptr; + // Pool only entries completed by this ubatch. + if (n_new > 0) { + ggml_tensor * rows = ggml_get_rows(ctx0, kg_all, ggml_reshape_1d(ctx0, inp_kpool->new_pool_idxs, kpool*n_new)); + rows = ggml_reshape_3d(ctx0, rows, 2*n_embd_indexer, kpool, n_new); + + ggml_tensor * pk = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_new, rows->nb[1], rows->nb[2], 0); + ggml_tensor * pg = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_new, rows->nb[1], rows->nb[2], ggml_row_size(rows->type, n_embd_indexer)); + + ggml_tensor * logits = ggml_add(ctx0, pg, layer.indexer_kpool_ape); + logits = ggml_cont(ctx0, ggml_permute(ctx0, logits, 1, 0, 2, 3)); // [kpool, head_dim, n_new] + ggml_tensor * probs = ggml_soft_max(ctx0, logits); + + pk = ggml_cont(ctx0, ggml_permute(ctx0, pk, 1, 0, 2, 3)); + pooled_new = ggml_sum_rows(ctx0, ggml_mul(ctx0, probs, pk)); // [1, head_dim, n_new] + pooled_new = ggml_reshape_2d(ctx0, pooled_new, n_embd_indexer, n_new); + cb(pooled_new, "indexer_pool_k_new", il); + + if (inp_kpool->cache_safe) { + // Write before the pool gather. + ggml_build_forward_expand(gf, ggml_set_rows(ctx0, pooled_all, pooled_new, inp_kpool->new_pool_rep)); + } + } + + ggml_tensor * pooled = nullptr; + if (inp_kpool->cache_safe) { + pooled = ggml_get_rows(ctx0, pooled_all, inp_kpool->pool_cells); + } else { + GGML_ASSERT(n_new <= n_pool); + ggml_tensor * pad = ggml_fill(ctx0, + ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_indexer, n_pool - n_new), 0.0f); + pooled = n_new > 0 ? ggml_concat(ctx0, pooled_new, pad, 1) : pad; + } + pooled = ggml_reshape_3d(ctx0, pooled, n_embd_indexer, 1, n_pool); + cb(pooled, "indexer_pool_k", il); + + ggml_tensor * sel_idx = inp_kpool->reuse_sel; + if (sel_idx == nullptr) { + ggml_tensor * weights = ggml_mul_mat(ctx0, layer.indexer_proj, cur); + weights = ggml_scale(ctx0, weights, 1.0f / sqrtf(float(n_embd_indexer * n_indexer_head))); + cb(weights, "indexer_weights", il); + + ggml_tensor * score = nullptr; + if (cparams.fused_lid) { + score = ggml_lightning_indexer(ctx0, iq, pooled, weights, inp_kpool->pool_mask); + res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il}); + } else { + ggml_tensor * q_p = ggml_permute(ctx0, iq, 0, 2, 1, 3); // [head_dim, n_tokens, n_head] + ggml_tensor * k_p = ggml_permute(ctx0, pooled, 0, 2, 1, 3); // [head_dim, n_pool, 1] + + ggml_tensor * kq = ggml_mul_mat(ctx0, k_p, q_p); // [n_pool, n_tokens, n_head] + kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); // [n_head, n_tokens, n_pool] + score = ggml_relu(ctx0, kq); + score = ggml_mul(ctx0, score, weights); + score = ggml_sum_rows(ctx0, score); // [1, n_tokens, n_pool] + score = ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3)); // [n_pool, n_tokens, 1] + score = ggml_add(ctx0, score, inp_kpool->pool_mask); + } + cb(score, "indexer_score", il); + + const int64_t n_top_pool = std::min(n_pool, hparams.indexer_top_k / kpool); + ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], UNORDERED + + // The gather mask marks the first min(nv, n_top_pool) slots as the visible pools, so order the set by descending score. + ggml_tensor * sel_score = ggml_get_rows(ctx0, + ggml_reshape_3d(ctx0, score, 1, n_pool, n_tokens), top_k); // [1, n_top_pool, n_tokens] + ggml_tensor * sel_order = ggml_argsort(ctx0, + ggml_reshape_2d(ctx0, sel_score, n_top_pool, n_tokens), GGML_SORT_ORDER_DESC); + top_k = ggml_get_rows(ctx0, + ggml_reshape_3d(ctx0, ggml_cast(ctx0, top_k, GGML_TYPE_F32), 1, n_top_pool, n_tokens), sel_order); + top_k = ggml_cast(ctx0, ggml_cont(ctx0, ggml_reshape_2d(ctx0, top_k, n_top_pool, n_tokens)), GGML_TYPE_I32); + cb(top_k, "indexer_top_k", il); + + sel_idx = ggml_get_rows(ctx0, inp_kpool->pool_idxs, + ggml_reshape_1d(ctx0, top_k, n_top_pool*n_tokens)); // [kpool, n_top_pool*n_tokens] + sel_idx = ggml_reshape_2d(ctx0, sel_idx, kpool*n_top_pool, n_tokens); + + if (hparams.indexer_kpool_select_tail) { + // Append the incomplete tail with n_kv for missing cells. + sel_idx = ggml_concat(ctx0, sel_idx, inp_kpool->tail_idxs, 0); + } + } else { + cb(sel_idx, "indexer_sel_reuse", il); + } + const int64_t n_sel = sel_idx->ne[0]; + + if (inp_kpool->mtp_share && il >= (int) hparams.n_layer() && inp_kpool->reuse_sel == nullptr) { + res->t_mtp_dsa_sel = sel_idx; + res->t_mtp_dsa_mask = inp_kpool->gather_mask; + } + + // Gather returns selected cell indices and masks padding separately. + if (inp_kpool->gather) { + GGML_ASSERT(inp_kpool->gather_mask->ne[0] == n_sel && inp_kpool->gather_mask->ne[3] == n_tokens); + cb(sel_idx, "indexer_sel_idx", il); + return sel_idx; + } + + // Tie scatter storage lifetime to this layer's selected indices. + ggml_tensor * seed = ggml_cast(ctx0, ggml_view_1d(ctx0, sel_idx, 1, 0), GGML_TYPE_F32); + + ggml_tensor * mask_seed = kq_mask->type == GGML_TYPE_F32 ? seed : ggml_cast(ctx0, seed, kq_mask->type); + mask_seed = ggml_fill(ctx0, mask_seed, -INFINITY); + ggml_tensor * mask_all = ggml_repeat_4d(ctx0, mask_seed, 1, n_kv + 1, n_tokens, 1); + mask_all = ggml_reshape_3d(ctx0, mask_all, 1, n_kv + 1, n_tokens); + + ggml_tensor * zero_seed = ggml_fill(ctx0, seed, 0.0f); + ggml_tensor * zeros = ggml_repeat_4d(ctx0, zero_seed, 1, n_sel, n_tokens, 1); + zeros = ggml_reshape_3d(ctx0, zeros, 1, n_sel, n_tokens); + + ggml_tensor * sel = ggml_set_rows(ctx0, mask_all, zeros, ggml_reshape_3d(ctx0, sel_idx, n_sel, n_tokens, 1)); + sel = ggml_view_2d(ctx0, sel, n_kv, n_tokens, sel->nb[2], 0); + + // Fold causal visibility before shared-indexer reuse. + GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tokens && kq_mask->ne[2]*kq_mask->ne[3] == 1); + sel = ggml_add(ctx0, sel, ggml_reshape_2d(ctx0, kq_mask, n_kv, n_tokens)); + cb(sel, "indexer_sel", il); + + return sel; +} + +// Nope MLA layer with sparse attention over the indexer selection + +ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( + ggml_tensor * cur, const llama_layer & layer, + const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_attn_k * inp_attn, + llm_graph_input_kpool * inp_kpool, ggml_tensor ** prev_sel, int il) { + + const auto * mctx_mla = mctx_hyb->get_attn(); + + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - hparams.n_rot(); + const float kq_scale = 1.0f / sqrtf((float) n_embd_head_k_mla); + + GGML_ASSERT(hparams.n_rot() == 0 && "GLM5-Next MLA is nope-only"); + + ggml_tensor * qr = ggml_mul_mat(ctx0, layer.wq_a, cur); + qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(qr, "q_resid", il); + + ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr); + q = ggml_reshape_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens); + + ggml_tensor * kv_cmpr = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + cb(kv_cmpr, "kv_cmpr", il); + + // absorb wk_b so the cache holds only the latent + ggml_tensor * q_absorbed = ggml_permute(ctx0, q, 0, 2, 1, 3); + q_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_absorbed); + q_absorbed = ggml_permute(ctx0, q_absorbed, 0, 2, 1, 3); + cb(q_absorbed, "q_absorbed", il); + + ggml_tensor * kq_mask = inp_attn->get_kq_mask(); + + ggml_tensor * sel = nullptr; + if (il >= (int) hparams.n_layer() || hparams.is_indexer_full(il)) { // the NextN block always has a full indexer + sel = build_kpool_select(cur, qr, kq_mask, layer, mctx_hyb, inp_kpool, il); + *prev_sel = sel; + } else { + GGML_ASSERT(*prev_sel != nullptr && "shared indexer layer must follow a full indexer layer"); + sel = *prev_sel; + } + + ggml_build_forward_expand(gf, q_absorbed); + ggml_build_forward_expand(gf, kv_cmpr); + ggml_build_forward_expand(gf, mctx_mla->cpy_k(ctx0, kv_cmpr, inp_attn->get_k_idxs(), il)); + + ggml_tensor * out = nullptr; + if (inp_kpool->gather) { + // Attend over gathered latents with the token dimension in ne[3]. + + ggml_build_forward_expand(gf, kq_mask); + + ggml_tensor * sel_idx = sel; // I32 [n_sel, n_tokens] + const int64_t n_sel = sel_idx->ne[0]; + + ggml_tensor * k = mctx_mla->get_k(ctx0, il); + GGML_ASSERT(k->ne[3] == 1 && "TODO: gathered DSA with multiple streams"); + GGML_ASSERT(k->ne[1] == 1 && k->ne[0] == kv_lora_rank && "GLM5-Next MLA cache holds a single latent head"); + + ggml_tensor * rows = ggml_view_2d(ctx0, k, k->ne[0], k->ne[2], k->nb[2], 0); + ggml_tensor * k_g = ggml_get_rows(ctx0, rows, ggml_reshape_1d(ctx0, sel_idx, n_sel*n_tokens)); + k_g = ggml_reshape_4d(ctx0, k_g, k->ne[0], n_sel, 1, n_tokens); // F32 [kv_lora_rank, n_sel, 1, n_tokens] + cb(k_g, "kv_gathered", il); + + ggml_tensor * q_g = ggml_permute(ctx0, q_absorbed, 0, 2, 3, 1); // [kv_lora_rank, 1, n_head, n_tokens] + + ggml_tensor * kq = ggml_mul_mat(ctx0, k_g, q_g); // [n_sel, 1, n_head, n_tokens] + ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + kq = ggml_soft_max_ext(ctx0, kq, inp_kpool->gather_mask, kq_scale, 0.0f); + cb(kq, "kq_soft_max_gathered", il); + + ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, k_g)); // [n_sel, kv_lora_rank, 1, n_tokens] + ggml_tensor * kqv = ggml_mul_mat(ctx0, v_t, kq); // [kv_lora_rank, 1, n_head, n_tokens] + kqv = ggml_mul_mat(ctx0, layer.wv_b, kqv); // [n_embd_head_v, 1, n_head, n_tokens] + cb(kqv, "kqv_gathered", il); + + out = ggml_cont(ctx0, ggml_permute(ctx0, kqv, 0, 2, 1, 3)); // [n_embd_head_v, n_head, 1, n_tokens] + out = ggml_reshape_2d(ctx0, out, kqv->ne[0]*n_head, n_tokens); + } else { + // The scatter selection already includes the causal mask. + ggml_tensor * mask = ggml_reshape_4d(ctx0, sel, kq_mask->ne[0], kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]); + cb(mask, "kq_mask_dsa", il); + + ggml_tensor * k = mctx_mla->get_k(ctx0, il); + ggml_tensor * v = ggml_view_4d(ctx0, k, kv_lora_rank, k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); + + out = build_attn_mha(q_absorbed, k, v, nullptr, mask, nullptr, layer.wv_b, kq_scale, il); + } + cb(out, "kqv_out", il); + + out = ggml_mul_mat(ctx0, layer.wo, out); + cb(out, "attn_out", il); + + return out; +} + +// Nextn draft head. The Nextn block is a DSA layer. + +llama_model_glm5_next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) + : graph(model, params, no_trunk_t{}) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM5-Next MTP supports a single NextN block"); + + const int il = hparams.n_layer() + cparams.nextn_layer_offset; + GGML_ASSERT(cparams.nextn_layer_offset == 0); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm && "MTP block tensors missing, convert without --no-mtp"); + + // token and previous hidden state inputs + auto inp = std::make_unique(hparams.n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * tok_embd; + if (ubatch.token) { + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); + } else { + tok_embd = inp->embd; + } + cb(tok_embd, "mtp_tok_embd", il); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + ggml_tensor * h_embd = inp->h; + + res->add_input(std::move(inp)); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // K-only MLA cache plus the indexer cache, no recurrent layers here + const auto * mctx_hyb = static_cast(mctx); + + auto * inp_hyb = build_inp_mem_hybrid_k(); + auto * inp_attn = inp_hyb->get_attn(); + auto * inp_kpool = build_inp_kpool(mctx_hyb); + + ggml_build_forward_expand(gf, inp_hyb->get_recr()->s_copy); + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * cur = ggml_mul_mat(ctx0, layer.nextn.eh_proj, ggml_concat(ctx0, e_norm, h_norm, 0)); + cb(cur, "mtp_eh_proj", il); + + ggml_tensor * inpSA = cur; + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + ggml_tensor * prev_sel = nullptr; + cur = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, &prev_sel, il); + cb(cur, "mtp_attn_out", il); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "mtp_ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_ffn_norm", il); + + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.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); + cb(moe_out, "mtp_ffn_moe_out", il); + + ggml_tensor * ffn_shexp = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "mtp_ffn_shexp", il); + + cur = ggml_add(ctx0, ggml_add(ctx0, moe_out, ffn_shexp), ffn_inp); + cb(cur, "mtp_post_ffn", il); + + // shared_head.norm, then the post-norm hidden state. + ggml_tensor * head_norm_w = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm; + cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1); + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + cur = ggml_mul_mat(ctx0, head_w, cur); + cb(cur, "result_output", -1); + + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index 9b87a40d5af9..dc0a6e4bc8ca 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1176,10 +1176,30 @@ struct llama_model_deepseek4 : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; - struct graph : public llm_graph_context { - graph(const llm_graph_params & params) : llm_graph_context(params) {} - graph(const llama_model & model, const llm_graph_params & params); + // manifold-constrained hyper-connections (mHC), shared by deepseek4 and derived model graphs like glm5-next. + template + struct graph_base : public Base { + graph_base(const llm_graph_params & params) : Base(params) {} + + // members of the dependent base used by the mHC helpers + using Base::ctx0; + using Base::res; + using Base::hparams; + using Base::cparams; + using Base::n_embd; + using Base::norm_rms_eps; + using Base::cb; + + // collapse the hc streams with per-stream weights + ggml_tensor * build_hc_pre( + ggml_tensor * x, + ggml_tensor * weights, + int il) const; + // mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens] + ggml_tensor * build_hc_mean(ggml_tensor * x) const; + + // returns the collapsed input and fills the post / comb weights ggml_tensor * build_hc_pre( ggml_tensor * x, ggml_tensor * hc_fn, @@ -1196,6 +1216,15 @@ struct llama_model_deepseek4 : public llama_model_base { ggml_tensor * comb, int il) const; + ggml_tensor * build_hc_sinkhorn( + ggml_tensor * comb, + int il) const; + }; + + struct graph : public graph_base<> { + graph(const llm_graph_params & params) : graph_base<>(params) {} + graph(const llama_model & model, const llm_graph_params & params); + ggml_tensor * build_hc_head( ggml_tensor * x, ggml_tensor * hc_fn, @@ -1289,14 +1318,6 @@ struct llama_model_deepseek4 : public llama_model_base { float kq_scale, int il) const; - ggml_tensor * build_hc_pre( - ggml_tensor * x, - ggml_tensor * weights, - int il) const; - - ggml_tensor * build_hc_sinkhorn( - ggml_tensor * comb, - int il) const; }; struct graph_mtp : public graph { @@ -2492,6 +2513,52 @@ struct llama_model_kimi_k3 : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +struct llama_model_glm5_next : public llama_model_base { + llama_model_glm5_next(const struct llama_model_params & params) : llama_model_base(params) {} + + // False for trunk only GGUFs and no_mtp loads. + bool mtp_ready = false; + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + // k-pool indexer inputs on top of the generic hybrid input + class llm_graph_input_kpool; + + // mHC helpers from deepseek4, stacked on the delta net helpers + struct graph : public llama_model_deepseek4::graph_base { + graph(const llama_model & model, const llm_graph_params & params); + + const llama_model & model; + + llm_graph_input_kpool * build_inp_kpool(const llama_memory_hybrid_idx_context * mctx_hyb); + + ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_rs * inp_rs, + int64_t d_conv, int64_t head_dim, int64_t n_head_kda, + int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il); + + ggml_tensor * build_kpool_select(ggml_tensor * cur, ggml_tensor * qr, ggml_tensor * kq_mask, const llama_layer & layer, + const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il); + + ggml_tensor * build_dsa_layer(ggml_tensor * cur, const llama_layer & layer, + const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_attn_k * inp_attn, + llm_graph_input_kpool * inp_kpool, ggml_tensor ** prev_sel, int il); + + protected: + // for graph_mtp + struct no_trunk_t {}; + graph(const llama_model & model, const llm_graph_params & params, no_trunk_t) : + llama_model_deepseek4::graph_base(params), model(model) {} + }; + + // Draft head, the Nextn block as a non hyper connected DSA layer + struct graph_mtp : public graph { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_kimi_linear : public llama_model_base { llama_model_kimi_linear(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b2ea245ab846..21f05dbe1acb 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -525,6 +525,9 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_GRANITE_SWITCH) { return false; // FIXME adapter fixture } + if (arch == LLM_ARCH_GLM5_NEXT) { + return false; // FIXME fixture for KDA + k-pool DSA + mhc hparams + } if (arch == LLM_ARCH_LLAMA_EMBED || arch == LLM_ARCH_GEMMA_EMBEDDING || arch == LLM_ARCH_T5ENCODER) { return false; // FIXME Embedding (?) models produce inconsistent results. } diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index f6045093c637..28406d8125a6 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -63,6 +63,7 @@ #define KEY_PROJ_SAMPLE_WINDOW_SIDE "clip.vision.projector.window_side" #define KEY_PROJ_SPATIAL_OFFSETS "clip.vision.projector.spatial_offsets" #define KEY_SPATIAL_MERGE_SIZE "clip.vision.spatial_merge_size" +#define KEY_SWIGLU_CLAMP "clip.vision.swiglu_clamp" #define KEY_MM_PATCH_MERGE_TYPE "clip.vision.mm_patch_merge_type" #define KEY_IMAGE_GRID_PINPOINTS "clip.vision.image_grid_pinpoints" @@ -482,6 +483,7 @@ enum projector_type { PROJECTOR_TYPE_DEEPSEEKOCR2, PROJECTOR_TYPE_LFM2A, PROJECTOR_TYPE_GLM4V, + PROJECTOR_TYPE_GLM5V, PROJECTOR_TYPE_YOUTUVL, PROJECTOR_TYPE_YASA2, PROJECTOR_TYPE_KIMIK25, @@ -546,6 +548,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_DEEPSEEKOCR2, "deepseekocr2"}, { PROJECTOR_TYPE_LFM2A, "lfm2a"}, { PROJECTOR_TYPE_GLM4V, "glm4v"}, + { PROJECTOR_TYPE_GLM5V, "glm5v"}, { PROJECTOR_TYPE_YOUTUVL, "youtuvl"}, { PROJECTOR_TYPE_YASA2, "yasa2"}, { PROJECTOR_TYPE_KIMIK25, "kimik25"}, diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 060938d86e3e..7539c82b08e4 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -93,6 +93,13 @@ struct clip_hparams { float eps = 1e-6; float rope_theta = 0.0; + + std::pair swiglu_clamp_gate = {0.0f, 0.0f}; + std::pair swiglu_clamp_up = {0.0f, 0.0f}; + + bool has_swiglu_clamp() const { + return swiglu_clamp_gate.second > 0.0f || swiglu_clamp_up.second > 0.0f; + } int32_t n_expert_used = 0; std::vector feature_layers; int32_t attn_window_size = 0; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 90de1957586f..2406b00a64e9 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -648,6 +648,11 @@ ggml_tensor * clip_graph::build_ffn( switch (type_op) { case FFN_SILU: if (gate) { + if (hparams.has_swiglu_clamp()) { + cur = ggml_clamp(ctx0, cur, hparams.swiglu_clamp_gate.first, hparams.swiglu_clamp_gate.second); + tmp = ggml_clamp(ctx0, tmp, hparams.swiglu_clamp_up.first, hparams.swiglu_clamp_up.second); + cb(cur, "ffn_gate_clamped", il); + } cur = ggml_swiglu_split(ctx0, cur, tmp); cb(cur, "ffn_swiglu", il); } else { @@ -1078,6 +1083,7 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const builder = std::make_unique(ctx, img); } break; case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: { builder = std::make_unique(ctx, img); } break; @@ -1726,6 +1732,23 @@ struct clip_model_loader { hparams.set_limit_image_tokens(8, 4096); hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup } break; + case PROJECTOR_TYPE_GLM5V: + { + // glm4v tower with clamped SwiGLU, ceil-aligned resize and its own token budget + hparams.rope_theta = 10000.0f; + hparams.n_merge = 2; + hparams.image_resize_algo = RESIZE_ALGO_BICUBIC; + get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); + float swiglu_clamp = 0.0f; + get_f32(KEY_SWIGLU_CLAMP, swiglu_clamp, true); + if (swiglu_clamp > 0.0f) { + hparams.swiglu_clamp_gate = { -INFINITY, swiglu_clamp }; + hparams.swiglu_clamp_up = { -swiglu_clamp, swiglu_clamp }; + } + get_u32(KEY_IMAGE_MIN_PIXELS, hparams.image_min_pixels); + get_u32(KEY_IMAGE_MAX_PIXELS, hparams.image_max_pixels); + hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup + } break; case PROJECTOR_TYPE_LLAMA4: { hparams.rope_theta = 10000.0f; @@ -2543,6 +2566,7 @@ struct clip_model_loader { } } break; case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: { model.mm_fc_w = get_tensor(string_format(TN_MM_PROJECTOR, "weight")); model.mm_ffn_up_w = get_tensor(string_format(TN_MM_UP, "weight")); @@ -4001,6 +4025,7 @@ int clip_n_output_tokens_x(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_EXAONE4_5: case PROJECTOR_TYPE_MIMOVL: case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: case PROJECTOR_TYPE_PADDLEOCR: case PROJECTOR_TYPE_HUNYUANVL: case PROJECTOR_TYPE_YOUTUVL: @@ -4027,6 +4052,7 @@ int clip_n_output_tokens_y(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_EXAONE4_5: case PROJECTOR_TYPE_MIMOVL: case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: case PROJECTOR_TYPE_PADDLEOCR: case PROJECTOR_TYPE_HUNYUANVL: case PROJECTOR_TYPE_YOUTUVL: @@ -4108,6 +4134,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_MIMOVL: case PROJECTOR_TYPE_MINIMAX_M3: case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: case PROJECTOR_TYPE_YOUTUVL: case PROJECTOR_TYPE_MUSE_GLIMMER: { @@ -4733,6 +4760,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { case PROJECTOR_TYPE_QWEN2VL: case PROJECTOR_TYPE_QWEN3VL: case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: { const int merge_ratio = hparams.n_merge; const int pw = image_size_width / patch_size; @@ -5881,6 +5909,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { case PROJECTOR_TYPE_GRANITE4_VISION: return ctx->model.qf_proj_blocks.size() * ctx->model.hparams.projection_dim; case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: return ctx->model.mm_ffn_down_w->ne[1]; case PROJECTOR_TYPE_MIMO_AUDIO: return ctx->model.mm_2_w->ne[1]; diff --git a/tools/mtmd/models/glm4v.cpp b/tools/mtmd/models/glm4v.cpp index 0e1d596b41bb..1cbcea942a39 100644 --- a/tools/mtmd/models/glm4v.cpp +++ b/tools/mtmd/models/glm4v.cpp @@ -41,8 +41,10 @@ ggml_cgraph * clip_graph_glm4v::build() { inp = ggml_add(ctx0, inp, model.patch_bias); cb(inp, "patch_bias", -1); - // pos-conv norm - inp = build_norm(inp, model.norm_embd_w, model.norm_embd_b, norm_t, eps, -1); + // pos-conv norm (absent in GLM-5.3-Flash) + if (model.norm_embd_w) { + inp = build_norm(inp, model.norm_embd_w, model.norm_embd_b, norm_t, eps, -1); + } ggml_tensor * learned_pos_embd = nullptr; // Note: GLM-OCR does not have learned position embeddings diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 0dda8770f292..7d27fb931216 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -786,6 +786,75 @@ mtmd_image_preproc_out mtmd_image_preprocessor_dyn_size::preprocess(const clip_i return output; } +// +// mtmd_image_preprocessor_glm5v +// + +// The canvas is ceil-aligned to patch_size*n_merge and fitted to the token budget. +// Only rescaled to meet the budget and sits top-left, with black padding on the right and bottom +mtmd_image_preproc_out mtmd_image_preprocessor_glm5v::preprocess(const clip_image_u8 & img) { + GGML_ASSERT(hparams.image_min_pixels > 0 && hparams.image_max_pixels > 0); + + const int64_t factor = hparams.patch_size * hparams.n_merge; + const int64_t min_px = hparams.image_min_pixels; // single-frame pixel counts + const int64_t max_px = hparams.image_max_pixels; + const int64_t height = img.get_size().height; + const int64_t width = img.get_size().width; + + auto align = [factor](int64_t v) { return (v + factor - 1) / factor * factor; }; + + // aligned canvas within the budget + int64_t canvas_h = align(height); + int64_t canvas_w = align(width); + + if (canvas_h * canvas_w < min_px) { + const double scale = std::sqrt((double) min_px / (double) (height * width)); + canvas_h = align(std::max(1, (int64_t) std::ceil(height * scale))); + canvas_w = align(std::max(1, (int64_t) std::ceil(width * scale))); + } + + if (canvas_h * canvas_w > max_px) { + // largest content height whose aligned canvas fits the budget + int64_t lo = 1, hi = height; + int64_t best_h = factor, best_w = factor; + while (lo <= hi) { + const int64_t ch = (lo + hi) / 2; + const int64_t cw = std::max(1, width * ch / height); + const int64_t ah = align(ch); + const int64_t aw = align(cw); + if (ah * aw <= max_px) { + best_h = ah; + best_w = aw; + lo = ch + 1; + } else { + hi = ch - 1; + } + } + canvas_h = best_h; + canvas_w = best_w; + } + + // Scaled to fit the canvas, and never upscaled, unless below the min budget + double scale = std::min((double) canvas_h / height, (double) canvas_w / width); + if (height * width >= min_px) { + scale = std::min(1.0, scale); + } + const int content_h = (int) std::max(1, std::min(canvas_h, (int64_t) std::floor(height * scale))); + const int content_w = (int) std::max(1, std::min(canvas_w, (int64_t) std::floor(width * scale))); + + clip_image_u8 content; + img_tool::resize(img, content, clip_image_size{content_w, content_h}, hparams.image_resize_algo, PAD_NONE); + + clip_image_u8 canvas; + canvas.set_size(clip_image_size{(int) canvas_w, (int) canvas_h}, img.is_placeholder()); + img_tool::fill(canvas, {0, 0, 0}); + img_tool::composite(canvas, content, 0, 0); + + mtmd_image_preproc_out output; + output.append(hparams, canvas, true); + return output; +} + // // mtmd_image_preprocessor_longest_edge // diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index 732e27379d27..b4fd3f3285c5 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -123,6 +123,12 @@ struct mtmd_image_preprocessor_dyn_size : mtmd_image_preprocessor { mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; }; +// GLM 5.3 flash, similar to dyn_size, but each edge is aligned up to patch_size*n_merge, and max budget is met by a search over the height with the width scaled proportionally +struct mtmd_image_preprocessor_glm5v : mtmd_image_preprocessor { + mtmd_image_preprocessor_glm5v(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} + mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; +}; + // similar to mtmd_image_preprocessor_dyn_size, but resize the image to have longest edge equal to hparams.image_longest_edge, while preserving aspect ratio struct mtmd_image_preprocessor_longest_edge : mtmd_image_preprocessor { mtmd_image_preprocessor_longest_edge(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 5b306180d62f..2a761031068b 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -860,6 +860,13 @@ struct mtmd_context { img_end = "<|end_of_image|>"; image_preproc = std::make_unique(ctx_v); } break; + case PROJECTOR_TYPE_GLM5V: + { + // <|begin_of_image|> ... (image embeddings) ... <|end_of_image|> + img_beg = "<|begin_of_image|>"; + img_end = "<|end_of_image|>"; + image_preproc = std::make_unique(ctx_v); + } break; case PROJECTOR_TYPE_PADDLEOCR: { // <|IMAGE_START|> ... (image embeddings) ... <|IMAGE_END|>