From 9370c82dbd1774941f9d8a05c9eafdac1ecb2e2c Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 28 Aug 2026 01:37:05 +0200 Subject: [PATCH 01/34] Rebase GLM-Next support onto master, and migrate to llama-memory-hybrid-idx --- conversion/__init__.py | 2 + conversion/base.py | 16 +- conversion/glm.py | 212 ++++++++++ conversion/qwen3vl.py | 30 +- gguf-py/gguf/constants.py | 74 ++++ gguf-py/gguf/gguf_writer.py | 9 + gguf-py/gguf/tensor_mapping.py | 32 ++ src/llama-arch.cpp | 8 + src/llama-arch.h | 5 + src/llama-context.cpp | 2 +- src/llama-graph.cpp | 179 ++++++++- src/llama-graph.h | 30 ++ src/llama-hparams.h | 2 + src/llama-memory-hybrid-idx.cpp | 198 ++++++++- src/llama-memory-hybrid-idx.h | 5 + src/llama-model.cpp | 41 ++ src/llama-model.h | 5 + src/llama-quant.cpp | 18 + src/models/deepseek4.cpp | 159 -------- src/models/glm5-next.cpp | 690 ++++++++++++++++++++++++++++++++ src/models/models.h | 55 +-- tools/mtmd/clip-impl.h | 3 + tools/mtmd/clip-model.h | 3 + tools/mtmd/clip.cpp | 25 ++ tools/mtmd/models/glm4v.cpp | 6 +- tools/mtmd/mtmd-image.cpp | 69 ++++ tools/mtmd/mtmd-image.h | 6 + tools/mtmd/mtmd.cpp | 7 + 28 files changed, 1694 insertions(+), 197 deletions(-) create mode 100644 src/models/glm5-next.cpp 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/base.py b/conversion/base.py index daae28e92adc..e9915f5c3f0e 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1416,12 +1416,13 @@ def does_token_look_special(self, token: str | bytes) -> bool: return seems_special # used for GPT-2 BPE and WordPiece vocabs - def get_vocab_base(self) -> tuple[list[str], list[int], str]: + def get_vocab_base(self, tokenizer=None) -> tuple[list[str], list[int], str]: tokens: list[str] = [] toktypes: list[int] = [] - from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(self.dir_model) + if tokenizer is None: + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model) vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) # ty: ignore[unresolved-attribute] assert max(tokenizer.vocab.values()) < vocab_size # ty: ignore[unresolved-attribute] @@ -2167,11 +2168,12 @@ def _set_vocab_glmedge(self): special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|endoftext|>"]) # ty: ignore[unresolved-attribute] special_vocab.add_to_gguf(self.gguf_writer) - def _set_vocab_glm(self): - from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(self.dir_model) + def _set_vocab_glm(self, tokenizer=None): + if tokenizer is None: + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model) special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) - tokens, toktypes, tokpre = self.get_vocab_base() + tokens, toktypes, tokpre = self.get_vocab_base(tokenizer) self.gguf_writer.add_tokenizer_model("gpt2") self.gguf_writer.add_tokenizer_pre(tokpre) self.gguf_writer.add_token_list(tokens) diff --git a/conversion/glm.py b/conversion/glm.py index 7544f850cb22..7402a1a11de6 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -402,3 +402,215 @@ 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_main_layers = self.hparams["num_hidden_layers"] + self.n_nextn_layers = self.hparams.get("num_nextn_predict_layers", 0) + self.skip_mtp = self.no_mtp or self.n_nextn_layers == 0 + + self.block_count = self.n_main_layers + 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): + from transformers import AutoTokenizer + try: + tokenizer = AutoTokenizer.from_pretrained(self.dir_model) + except ValueError: + # the repo ships only a transformers v5 style tokenizer.json, load it directly + from transformers import PreTrainedTokenizerFast + tokenizer = PreTrainedTokenizerFast(tokenizer_file=str(self.dir_model / "tokenizer.json")) + return self._set_vocab_glm(tokenizer) + + 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 + + if name.startswith(("model.visual.", "visual.")): + return None + + assert cls._n_main_layers is not None + m = re.match(r"model\.(?:language_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", + "model.language_model.embed_tokens.weight", "model.language_model.norm.weight", + ): + return None + + return name, gen + + def is_kda_layer(self, il: int) -> bool: + if il >= self.n_main_layers: + return False + return self.hparams["layer_types"][il] == "linear_attention" + + def set_gguf_parameters(self): + hp = self.hparams + n_layer = self.n_main_layers + + # the loader reads this array before it knows about NextN, so cover all + hp["num_key_value_heads"] = [0 if self.is_kda_layer(il) else 1 for il in range(self.block_count)] + + super().set_gguf_parameters() + 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)) + if (indexer_types := hp.get("indexer_types")) is not None: + self.gguf_writer.add_indexer_types([t == "full" for t in indexer_types[:n_layer]]) + + # 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.startswith("model.language_model."): + name = "model." + name[len("model.language_model."):] + + 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 + if (new_name.startswith(("blk.", "output_hc")) and any(k in new_name for k in + ("hc_attn_", "hc_ffn_", "indexer.kpool", "ssm_a", "ssm_dt", "exp_probs_b"))): + 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..c0885f438d0c 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_limit(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 fffbd6745397..703bfe8f3009 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -224,6 +224,8 @@ 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 + KPOOL_SELECT_TAIL = "{arch}.attention.indexer.kpool_select_tail" # GLM5-Next class HyperConnection: COUNT = "{arch}.hyper_connection.count" @@ -381,6 +383,7 @@ class ClipVision: IMAGE_MEAN = "clip.vision.image_mean" IMAGE_STD = "clip.vision.image_std" SPATIAL_MERGE_SIZE = "clip.vision.spatial_merge_size" + SWIGLU_LIMIT = "clip.vision.swiglu_limit" 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" @@ -558,6 +561,7 @@ class MODEL_ARCH(IntEnum): GLM4 = auto() GLM4_MOE = auto() GLM_DSA = auto() + GLM5_NEXT = auto() BITNET = auto() T5 = auto() T5ENCODER = auto() @@ -887,6 +891,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() @@ -1306,6 +1312,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", @@ -1634,6 +1641,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.kpool_gate", + MODEL_TENSOR.INDEXER_KPOOL_APE: "blk.{bid}.indexer.kpool_ape", # vision MODEL_TENSOR.V_MMPROJ: "mm.{bid}", MODEL_TENSOR.V_MMPROJ_FC: "mm.model.fc", @@ -3990,6 +3999,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, @@ -5647,6 +5720,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 b1d161bcb37e..41d4b5c64d49 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -816,6 +816,12 @@ 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_max_alibi_bias(self, bias: float) -> None: self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias) @@ -1370,6 +1376,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_limit(self, value: float) -> None: + self.add_float32(Keys.ClipVision.SWIGLU_LIMIT, 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 5e61f61f7f0d..f26384040112 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,8 @@ 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_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 +681,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.kpool_gate" }, + { LLM_TENSOR_INDEXER_KPOOL_APE, "blk.%d.indexer.kpool_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 +955,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 +1080,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: diff --git a/src/llama-arch.h b/src/llama-arch.h index ca7d55a5fd78..09029c11f941 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,8 @@ 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_OUTPUT_GROUP_COUNT, LLM_KV_ATTENTION_OUTPUT_LORA_RANK, LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, @@ -677,6 +680,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 fb88919f9d67..1968c203a0e2 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2293,7 +2293,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 || diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0ef..1cfb08580ac4 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1779,7 +1779,7 @@ ggml_tensor * llm_graph_context::build_ffn( tmp = ggml_clamp(ctx0, tmp, -limit, limit); cb(tmp, "ffn_up_clamped", il); - 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_clamp(ctx0, cur, -INFINITY, limit); cb(cur, "ffn_gate_clamped", il); cur = ggml_swiglu_split(ctx0, cur, tmp); @@ -2176,7 +2176,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( up = ggml_clamp(ctx0, up, -limit, limit); cb(up, "ffn_moe_up_clamped", il); - 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_clamp(ctx0, cur, -INFINITY, limit); cb(cur, "ffn_moe_gate_clamped", il); cur = ggml_swiglu_split(ctx0, cur, up); @@ -3410,6 +3410,181 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const { return (llm_graph_input_dsv4 *) res->add_input(std::move(inp)); } +// manifold-constrained hyper-connections (mHC), deepseek4 and glm5-next + +static ggml_tensor * hc_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) { + return ggml_view_1d(ctx, t, ne0, ggml_row_size(t->type, i0)); +} + +static ggml_tensor * hc_view_2d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t i0) { + return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], ggml_row_size(t->type, i0)); +} + +static ggml_tensor * hc_affine(ggml_context * ctx, ggml_tensor * x, ggml_tensor * scale, ggml_tensor * base) { + x = ggml_mul(ctx, x, scale); + x = ggml_add(ctx, x, base); + return x; +} + +ggml_tensor * llm_graph_context::build_hc_pre( + ggml_tensor * x, + ggml_tensor * weights, + int il) const { + GGML_ASSERT(x->ne[0] == n_embd); + GGML_ASSERT(x->ne[1] == hparams.dsv4_hc_mult); + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = x->ne[2]; + + if (cparams.fused_dsv4_hc_pre && il >= 0) { + ggml_tensor * result = ggml_dsv4_hc_pre(ctx0, x, weights); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, result, il}); + return result; + } + + ggml_tensor * result = nullptr; + for (int64_t ih = 0; ih < hc; ++ih) { + ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]); + ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]); + ggml_tensor * cur = ggml_mul(ctx0, xh, wh); + result = result ? ggml_add(ctx0, result, cur) : cur; + } + + return result; +} + +ggml_tensor * llm_graph_context::build_hc_sinkhorn( + ggml_tensor * comb, + int il) const { + GGML_UNUSED(il); + + // comb is [dst_hc, src_hc, n_tokens]. Sinkhorn follows the reference: + // row softmax over dst, one column normalization, then repeated row/column normalization. + comb = ggml_soft_max(ctx0, comb); + + ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1); + eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps); + + comb = ggml_add(ctx0, comb, eps); + + auto norm_cols = [&]() { + ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3)); + ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst); + col_sum = ggml_add(ctx0, col_sum, eps); + col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3); + comb = ggml_div(ctx0, comb, col_sum); + }; + + auto norm_rows = [&]() { + ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb); + row_sum = ggml_add(ctx0, row_sum, eps); + comb = ggml_div(ctx0, comb, row_sum); + }; + + norm_cols(); + for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) { + norm_rows(); + norm_cols(); + } + + return comb; +} + +ggml_tensor * llm_graph_context::build_hc_pre( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il) const { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc*n_embd; + const int64_t hc_mix_dim = (2 + hc)*hc; + const int64_t nt = x->ne[2]; + + GGML_ASSERT(hc == 4); + GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim); + + ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt); + ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps); + ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); + cb(mixes, "hc_mixes", il); + + ggml_tensor * scale_pre = hc_view_1d(ctx0, hc_scale, 1, 0); + ggml_tensor * scale_post = hc_view_1d(ctx0, hc_scale, 1, 1); + + ggml_tensor * base_pre = hc_view_1d(ctx0, hc_base, hc, 0); + ggml_tensor * base_post = hc_view_1d(ctx0, hc_base, hc, hc); + + ggml_tensor * pre = hc_view_2d(ctx0, mixes, hc, nt, 0); + pre = hc_affine(ctx0, pre, scale_pre, base_pre); + pre = ggml_sigmoid(ctx0, pre); + pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps); + cb(pre, "hc_pre", il); + + *post = hc_view_2d(ctx0, mixes, hc, nt, hc); + *post = hc_affine(ctx0, *post, scale_post, base_post); + *post = ggml_sigmoid(ctx0, *post); + *post = ggml_scale(ctx0, *post, 2.0f); + cb(*post, "hc_post", il); + + if (cparams.fused_dsv4_hc_comb) { + *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps, + (int32_t) hparams.dsv4_hc_sinkhorn_iters); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il}); + } else { + ggml_tensor * scale_comb = hc_view_1d(ctx0, hc_scale, 1, 2); + ggml_tensor * base_comb = hc_view_1d(ctx0, hc_base, hc*hc, 2*hc); + + *comb = hc_view_2d(ctx0, mixes, hc*hc, nt, 2*hc); + *comb = hc_affine(ctx0, *comb, scale_comb, base_comb); + *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt); + *comb = build_hc_sinkhorn(*comb, il); + } + cb(*comb, "hc_comb", il); + + ggml_tensor * result = build_hc_pre(x, pre, il); + return result; +} + +ggml_tensor * llm_graph_context::build_hc_post( + ggml_tensor * x, + ggml_tensor * residual, + ggml_tensor * post, + ggml_tensor * comb, + int il) const { + GGML_ASSERT(x->ne[0] == n_embd); + GGML_ASSERT(residual->ne[1] == hparams.dsv4_hc_mult); + + if (cparams.fused_dsv4_hc_post) { + ggml_tensor * result = ggml_dsv4_hc_post(ctx0, x, residual, post, comb); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, result, il}); + return result; + } + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = x->ne[1]; + + ggml_tensor * out = nullptr; + for (int64_t dst = 0; dst < hc; ++dst) { + ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst*post->nb[0]); + ggml_tensor * cur = ggml_mul(ctx0, x, post_dst); + + for (int64_t src = 0; src < hc; ++src) { + ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]); + ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2], + dst*comb->nb[0] + src*comb->nb[1]); + cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst)); + } + + cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt); + out = out ? ggml_concat(ctx0, out, cur, 1) : cur; + } + + return out; +} + ggml_tensor * llm_graph_context::build_rs( ggml_tensor * s, ggml_tensor * state_copy_main, diff --git a/src/llama-graph.h b/src/llama-graph.h index b388e028cb53..e30f915197c2 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1339,6 +1339,36 @@ struct llm_graph_context { // hybrid // + // hyper-connections (mHC) + + + // collapse the hc streams with per-stream weights + ggml_tensor * build_hc_pre( + ggml_tensor * x, + ggml_tensor * weights, + int il) const; + + // returns the collapsed input and fills the post / comb weights + ggml_tensor * build_hc_pre( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il) const; + + ggml_tensor * build_hc_post( + ggml_tensor * x, + ggml_tensor * residual, + ggml_tensor * post, + ggml_tensor * comb, + int il) const; + + ggml_tensor * build_hc_sinkhorn( + ggml_tensor * comb, + int il) const; + llm_graph_input_mem_hybrid * build_inp_mem_hybrid() const; llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k() const; diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 1411692a8909..39db7ba37304 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -261,6 +261,8 @@ 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; // 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 d4e59d77e570..0475c149f72d 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" @@ -48,7 +52,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 k-pool indexer of glm5-next caches key | gate per token + hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size * (model.hparams.indexer_kpool > 0 ? 2 : 1); LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); @@ -463,3 +468,194 @@ void llama_memory_hybrid_idx_context::set_input_qsa( } } } + +// k-pool DSA indexer (glm5-next) + +namespace { + +struct kpool_seq { + llama_pos pos_min = 0; + std::vector> cells; // (pos, cell) + std::vector pools; +}; + +// per-sequence sorted (pos, cell) lists and the complete pools among the first n_kv cells +static std::vector kpool_collect(const llama_kv_cells & cells, uint32_t kpool, uint32_t n_kv) { + std::vector res(LLAMA_MAX_SEQ); + + const uint32_t n = std::min(n_kv, cells.size()); + for (uint32_t i = 0; i < n; ++i) { + if (cells.is_empty(i)) { + continue; + } + const llama_pos p = cells.pos_get(i); + for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + if (cells.seq_has(i, s)) { + res[s].cells.emplace_back(p, i); + } + } + } + + for (auto & sq : res) { + 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; + + // a pool is complete when kpool consecutive positions, from pos_min all have a cell + for (size_t j = 0; j + kpool <= sq.cells.size(); ) { + const llama_pos p0 = sq.cells[j].first; + if ((p0 - sq.pos_min) % 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; + } + } + } + + return res; +} + +// the last padded pool is always unused +static uint32_t kpool_pad(uint32_t n_pool) { + return std::max(64u, GGML_PAD(n_pool + 1, 64u)); +} + +} + +uint32_t llama_memory_hybrid_idx_context::get_n_kpool(uint32_t kpool) const { + GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); + GGML_ASSERT(get_n_stream() == 1 && "TODO: k-pool indexer with multiple streams"); + + const auto seqs = kpool_collect(mem->get_mem_idx()->get_cells(0), kpool, get_idx()->get_n_kv()); + + uint32_t n_pool = 0; + for (const auto & sq : seqs) { + n_pool += (uint32_t) sq.pools.size(); + } + + return kpool_pad(n_pool); +} + +void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs, ggml_tensor * cell_pool, + const llama_ubatch * ubatch, uint32_t kpool) 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_idxs->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(pool_mask->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(tail_idxs->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(cell_pool->buffer)); + + const uint32_t n_kv = get_idx()->get_n_kv(); + const auto seqs = kpool_collect(mem->get_mem_idx()->get_cells(0), kpool, n_kv); + + const uint32_t n_tokens = ubatch->n_tokens; + const uint32_t n_pool = pool_idxs->ne[1]; + + GGML_ASSERT(pool_idxs->ne[0] == (int64_t) kpool); + 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(cell_pool->ne[0] == (int64_t) n_kv); + + // the cell of the first ubatch token + uint32_t dummy_cell = 0; + { + const llama_seq_id s = ubatch->seq_id[0][0]; + const auto & sq = 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; + } + + // pools are laid out per sequence + std::vector seq_pool_start(LLAMA_MAX_SEQ, 0); + std::vector pool_end; + pool_end.reserve(n_pool); + + // cells outside any complete pool map to the last (always unused) pool + int32_t * cpool = (int32_t *) cell_pool->data; + std::fill(cpool, cpool + n_kv, (int32_t) n_pool - 1); + + int32_t * pidx = (int32_t *) pool_idxs->data; + for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + const auto & sq = seqs[s]; + seq_pool_start[s] = (uint32_t) pool_end.size(); + for (uint32_t j : sq.pools) { + const uint32_t ip = (uint32_t) pool_end.size(); + GGML_ASSERT(ip + 1 < n_pool); + for (uint32_t k = 0; k < kpool; ++k) { + const uint32_t cell = sq.cells[j + k].second; + pidx[ip*kpool + k] = (int32_t) cell; + cpool[cell] = (int32_t) ip; + } + pool_end.push_back(sq.cells[j + kpool - 1].first); + } + } + const uint32_t n_pool_real = (uint32_t) pool_end.size(); + for (uint32_t ip = n_pool_real; ip < n_pool; ++ip) { + for (uint32_t k = 0; k < kpool; ++k) { + pidx[ip*kpool + k] = (int32_t) dummy_cell; + } + } + + // 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) 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); + } + }; + 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 = seqs[s]; + + const uint32_t n_tail = (uint32_t) ((p - sq.pos_min + 1) % kpool); + + for (uint32_t k = 0; k < kpool - 1; ++k) { + int32_t cell = (int32_t) n_kv; + 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; + } + } + tidx[(size_t) i*(kpool - 1) + k] = cell; + } + } +} diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index e3472646d0f6..35dd5ced9781 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -137,6 +137,11 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // bias F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns] // the caller then adds the attention mask, the only part of the bias that varies within a block + // glm5-next, complete pools of kpool consecutive positions per sequence, scored as whole pools + uint32_t get_n_kpool(uint32_t kpool) const; // padded pool count, the last pool is always unused + void set_input_kpool(ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs, ggml_tensor * cell_pool, + const llama_ubatch * ubatch, uint32_t kpool) 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; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index fc83658dd7ff..26dc9008d5fc 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"; @@ -2301,6 +2304,43 @@ 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); + }; + + 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); @@ -2813,6 +2853,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 0d7352ac500b..748f4a626971 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 c414caa173fa..fb3a0e4b318a 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -328,6 +328,24 @@ 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: the k-pool position bias is added elementwise; the indexer, mHC mixers, + // KDA gates and MLA low-rank paths are small and precision-sensitive (~1 GB total) + quantize &= name.find("indexer.kpool_ape.weight") == std::string::npos; + if (arch == LLM_ARCH_GLM5_NEXT) { + quantize &= name.find("indexer.") == std::string::npos; + quantize &= name.find("hc_attn_fn.weight") == std::string::npos; + quantize &= name.find("hc_ffn_fn.weight") == 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_q_a.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; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index fc816e2aeb43..3f7c62e53ac9 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -284,165 +284,6 @@ static ggml_tensor * dsv4_hc_affine( return x; } -ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( - ggml_tensor * x, - ggml_tensor * weights, - int il) const { - GGML_ASSERT(x->ne[0] == n_embd); - GGML_ASSERT(x->ne[1] == hparams.dsv4_hc_mult); - - const int64_t hc = hparams.dsv4_hc_mult; - const int64_t nt = x->ne[2]; - - if (cparams.fused_dsv4_hc_pre && il >= 0) { - ggml_tensor * result = ggml_dsv4_hc_pre(ctx0, x, weights); - res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, result, il}); - return result; - } - - ggml_tensor * result = nullptr; - for (int64_t ih = 0; ih < hc; ++ih) { - ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]); - ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]); - ggml_tensor * cur = ggml_mul(ctx0, xh, wh); - result = result ? ggml_add(ctx0, result, cur) : cur; - } - - return result; -} - -ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn( - ggml_tensor * comb, - int il) const { - GGML_UNUSED(il); - - // comb is [dst_hc, src_hc, n_tokens]. Sinkhorn follows the reference: - // row softmax over dst, one column normalization, then repeated row/column normalization. - comb = ggml_soft_max(ctx0, comb); - - ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1); - eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps); - - comb = ggml_add(ctx0, comb, eps); - - auto norm_cols = [&]() { - ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3)); - ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst); - col_sum = ggml_add(ctx0, col_sum, eps); - col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3); - comb = ggml_div(ctx0, comb, col_sum); - }; - - auto norm_rows = [&]() { - ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb); - row_sum = ggml_add(ctx0, row_sum, eps); - comb = ggml_div(ctx0, comb, row_sum); - }; - - norm_cols(); - for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) { - norm_rows(); - norm_cols(); - } - - return comb; -} - -ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( - ggml_tensor * x, - ggml_tensor * hc_fn, - ggml_tensor * hc_scale, - ggml_tensor * hc_base, - ggml_tensor ** post, - ggml_tensor ** comb, - int il) const { - const int64_t hc = hparams.dsv4_hc_mult; - const int64_t hc_dim = hc*n_embd; - const int64_t hc_mix_dim = (2 + hc)*hc; - const int64_t nt = x->ne[2]; - - GGML_ASSERT(hc == 4); - GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim); - - ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt); - ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps); - ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); - cb(mixes, "hc_mixes", il); - - ggml_tensor * scale_pre = dsv4_view_1d(ctx0, hc_scale, 1, 0); - ggml_tensor * scale_post = dsv4_view_1d(ctx0, hc_scale, 1, 1); - - ggml_tensor * base_pre = dsv4_view_1d(ctx0, hc_base, hc, 0); - ggml_tensor * base_post = dsv4_view_1d(ctx0, hc_base, hc, hc); - - ggml_tensor * pre = dsv4_view_2d(ctx0, mixes, hc, nt, 0); - pre = dsv4_hc_affine(ctx0, pre, scale_pre, base_pre); - pre = ggml_sigmoid(ctx0, pre); - pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps); - cb(pre, "hc_pre", il); - - *post = dsv4_view_2d(ctx0, mixes, hc, nt, hc); - *post = dsv4_hc_affine(ctx0, *post, scale_post, base_post); - *post = ggml_sigmoid(ctx0, *post); - *post = ggml_scale(ctx0, *post, 2.0f); - cb(*post, "hc_post", il); - - if (cparams.fused_dsv4_hc_comb) { - *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps, - (int32_t) hparams.dsv4_hc_sinkhorn_iters); - res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il}); - } else { - ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2); - ggml_tensor * base_comb = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc); - - *comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc); - *comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb); - *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt); - *comb = build_hc_sinkhorn(*comb, il); - } - cb(*comb, "hc_comb", il); - - ggml_tensor * result = build_hc_pre(x, pre, il); - return result; -} - -ggml_tensor * llama_model_deepseek4::graph::build_hc_post( - ggml_tensor * x, - ggml_tensor * residual, - ggml_tensor * post, - ggml_tensor * comb, - int il) const { - GGML_ASSERT(x->ne[0] == n_embd); - GGML_ASSERT(residual->ne[1] == hparams.dsv4_hc_mult); - - if (cparams.fused_dsv4_hc_post) { - ggml_tensor * result = ggml_dsv4_hc_post(ctx0, x, residual, post, comb); - res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, result, il}); - return result; - } - - const int64_t hc = hparams.dsv4_hc_mult; - const int64_t nt = x->ne[1]; - - ggml_tensor * out = nullptr; - for (int64_t dst = 0; dst < hc; ++dst) { - ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst*post->nb[0]); - ggml_tensor * cur = ggml_mul(ctx0, x, post_dst); - - for (int64_t src = 0; src < hc; ++src) { - ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]); - ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2], - dst*comb->nb[0] + src*comb->nb[1]); - cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst)); - } - - cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt); - out = out ? ggml_concat(ctx0, out, cur, 1) : cur; - } - - return out; -} - ggml_tensor * llama_model_deepseek4::graph::build_hc_head( ggml_tensor * x, ggml_tensor * hc_fn, diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp new file mode 100644 index 000000000000..f9dc0e3ea38a --- /dev/null +++ b/src/models/glm5-next.cpp @@ -0,0 +1,690 @@ +#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; + + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all); + + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + if (i >= hparams.n_layer()) { + hparams.n_head_kv_arr[i] = 1; + } + hparams.is_recr_impl[i] = i < hparams.n_layer() && 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); + 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 (TODO) + 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; + 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, 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 { + return std::make_unique(*this, params); +} + +// Mean over the hyper-connection streams +static ggml_tensor * glm5_hc_mean(ggml_context * ctx, ggml_tensor * x) { + 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); + 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])); + } + return ggml_scale(ctx, acc, 1.0f/hc); +} + +// 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) { + 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); + + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, last_conv_x, + 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), + (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + + 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_idxs, pool_mask, tail_idxs, cell_pool, ubatch, kpool); + } + + 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_idxs->ne[1] == mctx->get_n_kpool(kpool); + res &= pool_mask->ne[1] == params.ubatch.n_tokens; + res &= tail_idxs->ne[1] == params.ubatch.n_tokens; + res &= cell_pool->ne[0] == idx->get_n_kv(); + + return res; + } + + ggml_tensor * k_idxs = nullptr; // I64 [n_tokens] + ggml_tensor * pool_idxs = nullptr; // I32 [kpool, n_pool] + ggml_tensor * pool_mask = nullptr; // F32/F16 [n_pool, n_tokens] + ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens] + ggml_tensor * cell_pool = nullptr; // I32 [n_kv] + + const llama_memory_hybrid_idx_context * mctx; + const uint32_t kpool; +}; + +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(kpool); + const uint32_t n_kv = mctx_idx->get_n_kv(); + + // 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_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); + inp->cell_pool = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_kv); + ggml_set_input(inp->pool_idxs); + ggml_set_input(inp->pool_mask); + ggml_set_input(inp->tail_idxs); + ggml_set_input(inp->cell_pool); + + 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) : + llm_build_delta_net_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 + if (inp_out_ids) { + 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 = glm5_hc_mean(ctx0, inpL); + cb(cur, "hc_head", -1); + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + 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(); + + 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); + 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); + 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); + 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); + + const float eps = hparams.f_norm_rms_eps; + Qcur = ggml_l2_norm(ctx0, Qcur, eps); + Kcur = ggml_l2_norm(ctx0, Kcur, eps); + + auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); + + ggml_tensor * output = ggml_cont(ctx0, attn_out.first); + ggml_tensor * new_state = attn_out.second; + cb(output, "kda_scan_out", il); + + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, new_state, + ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); + + // 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, 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_idxs->ne[1]; + + // queries + ggml_tensor * 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); + + ggml_tensor * packed = ggml_concat(ctx0, ik, ig, 0); + packed = ggml_reshape_3d(ctx0, packed, 2*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]; + k_all = ggml_view_2d(ctx0, k_all, 2*n_embd_indexer, n_kv, k_all->nb[2], 0); + + // Gather the member of every pool + ggml_tensor * rows = ggml_get_rows(ctx0, k_all, ggml_reshape_1d(ctx0, inp_kpool->pool_idxs, kpool*n_pool)); + rows = ggml_reshape_3d(ctx0, rows, 2*n_embd_indexer, kpool, n_pool); + + ggml_tensor * pk = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_pool, rows->nb[1], rows->nb[2], 0); + ggml_tensor * pg = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_pool, 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_pool] + ggml_tensor * probs = ggml_soft_max(ctx0, logits); + + pk = ggml_cont(ctx0, ggml_permute(ctx0, pk, 1, 0, 2, 3)); + ggml_tensor * pooled = ggml_sum_rows(ctx0, ggml_mul(ctx0, probs, pk)); // [1, head_dim, n_pool] + pooled = ggml_reshape_3d(ctx0, pooled, n_embd_indexer, 1, n_pool); + cb(pooled, "indexer_pool_k", il); + + 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_cont(ctx0, ggml_top_k(ctx0, score, n_top_pool)); // [n_top_pool, n_tokens] + cb(top_k, "indexer_top_k", il); + + // Pool-level selection mask, -inf everywhere except the selected pools + ggml_tensor * pool_sel = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_pool, n_tokens); + pool_sel = ggml_fill(ctx0, pool_sel, -INFINITY); + ggml_tensor * zeros = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_top_pool, n_tokens), 0.0f); + pool_sel = ggml_set_rows(ctx0, pool_sel, zeros, ggml_reshape_3d(ctx0, top_k, n_top_pool, n_tokens, 1)); + pool_sel = ggml_reshape_2d(ctx0, pool_sel, n_pool, n_tokens); + + // Expand to the cells + ggml_tensor * pool_sel_t = ggml_cont(ctx0, ggml_transpose(ctx0, pool_sel)); // [n_tokens, n_pool] + ggml_tensor * sel_t = ggml_get_rows(ctx0, pool_sel_t, inp_kpool->cell_pool); // [n_tokens, n_kv] + ggml_tensor * sel = ggml_cont(ctx0, ggml_transpose(ctx0, sel_t)); // [n_kv, n_tokens] + + if (hparams.indexer_kpool_select_tail) { + ggml_tensor * pad = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_tokens), -INFINITY); + sel = ggml_concat(ctx0, sel, pad, 0); // [n_kv + 1, n_tokens] + sel = ggml_reshape_3d(ctx0, sel, 1, n_kv + 1, n_tokens); + ggml_tensor * tzeros = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, kpool - 1, n_tokens), 0.0f); + sel = ggml_set_rows(ctx0, sel, tzeros, ggml_reshape_3d(ctx0, inp_kpool->tail_idxs, kpool - 1, n_tokens, 1)); + sel = ggml_view_2d(ctx0, sel, n_kv, n_tokens, sel->nb[2], 0); + } + 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 * sel = nullptr; + if (hparams.is_indexer_full(il)) { + sel = build_kpool_select(cur, qr, 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)); + + // Combine the causal mask with the indexer selection + ggml_tensor * kq_mask = inp_attn->get_kq_mask(); + ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32); + mask = ggml_add(ctx0, ggml_reshape_2d(ctx0, mask, mask->ne[0], mask->ne[1]), sel); + if (kq_mask->type != GGML_TYPE_F32) { + mask = ggml_cast(ctx0, mask, kq_mask->type); + } + mask = ggml_reshape_4d(ctx0, mask, 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); + + ggml_tensor * 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; +} diff --git a/src/models/models.h b/src/models/models.h index af60764c2f7f..4b58baadf4d0 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1180,22 +1180,6 @@ struct llama_model_deepseek4 : public llama_model_base { graph(const llm_graph_params & params) : llm_graph_context(params) {} graph(const llama_model & model, const llm_graph_params & params); - ggml_tensor * build_hc_pre( - ggml_tensor * x, - ggml_tensor * hc_fn, - ggml_tensor * hc_scale, - ggml_tensor * hc_base, - ggml_tensor ** post, - ggml_tensor ** comb, - int il) const; - - ggml_tensor * build_hc_post( - ggml_tensor * x, - ggml_tensor * residual, - ggml_tensor * post, - ggml_tensor * comb, - int il) const; - ggml_tensor * build_hc_head( ggml_tensor * x, ggml_tensor * hc_fn, @@ -1289,14 +1273,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 { @@ -2489,6 +2465,37 @@ 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) {} + 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; + + struct graph : public llm_build_delta_net_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, 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); + }; + + 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/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index f6045093c637..75a0fad692f9 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_LIMIT "clip.vision.swiglu_limit" #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..79183b179adb 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -93,6 +93,9 @@ struct clip_hparams { float eps = 1e-6; float rope_theta = 0.0; + + // clamp the SwiGLU gate to (-inf, limit] and up to [-limit, limit] when > 0 (glm5-next) + float swiglu_limit = 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..6f5ad0bac395 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.swiglu_limit > 0.0f) { + cur = ggml_clamp(ctx0, cur, -INFINITY, hparams.swiglu_limit); + tmp = ggml_clamp(ctx0, tmp, -hparams.swiglu_limit, hparams.swiglu_limit); + 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,19 @@ 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); + get_f32(KEY_SWIGLU_LIMIT, hparams.swiglu_limit, false); + get_u32(KEY_IMAGE_MIN_PIXELS, hparams.image_min_pixels); + get_u32(KEY_IMAGE_MAX_PIXELS, hparams.image_max_pixels); + hparams.warmup_image_size = static_cast(std::sqrt(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 +2562,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 +4021,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 +4048,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 +4130,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 +4756,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 +5905,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|> From d4ce82595d1c242820c01cada2cb87d9604f0c21 Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 28 Aug 2026 01:51:57 +0200 Subject: [PATCH 02/34] Add initial MTP support --- src/llama-model.cpp | 10 ++++ src/models/glm5-next.cpp | 124 ++++++++++++++++++++++++++++++++++++++- src/models/models.h | 11 ++++ 3 files changed, 143 insertions(+), 2 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 26dc9008d5fc..154038672afb 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2321,6 +2321,16 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, 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, diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index f9dc0e3ea38a..0ab2141ae0a9 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -193,6 +193,9 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) { } std::unique_ptr llama_model_glm5_next::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -415,7 +418,9 @@ llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_p } // narrow to the output tokens, then collapse the streams - if (inp_out_ids) { + // 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); @@ -425,6 +430,14 @@ llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_p 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; @@ -655,7 +668,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( cb(q_absorbed, "q_absorbed", il); ggml_tensor * sel = nullptr; - if (hparams.is_indexer_full(il)) { + 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, layer, mctx_hyb, inp_kpool, il); *prev_sel = sel; } else { @@ -688,3 +701,110 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( 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 4b58baadf4d0..65bef6bbe191 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2491,6 +2491,17 @@ struct llama_model_glm5_next : public llama_model_base { 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) : + llm_build_delta_net_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; From d2fc716c8dbe4ef4898a6b77a227d66137ded1c9 Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 28 Aug 2026 20:02:48 +0200 Subject: [PATCH 03/34] Merge branch optimizations. Reduce allocated compute buffer size, speed up long context decode, fla, and slight MTP improvements. --- common/speculative.cpp | 89 +++++++ src/llama-context.cpp | 126 ++++++++++ src/llama-context.h | 9 + src/llama-ext.h | 4 + src/llama-graph.cpp | 8 + src/llama-graph.h | 4 + src/llama-memory-hybrid-idx.cpp | 402 +++++++++++++++++++++++++------- src/llama-memory-hybrid-idx.h | 36 ++- src/models/glm5-next.cpp | 346 +++++++++++++++++++-------- src/models/models.h | 2 +- 10 files changed, 849 insertions(+), 177 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index b5348ab6f3d5..b464cab8ae1a 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -15,6 +15,7 @@ #include #include #include +#include #include #include #include @@ -1381,6 +1382,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) @@ -1394,6 +1400,12 @@ 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))); + const char * env_share = getenv("LLAMA_GLM5_MTP_INDEX_SHARE"); + dsa_index_share = (env_share == nullptr || atoi(env_share) != 0) && 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", @@ -1459,6 +1471,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 { @@ -1474,6 +1488,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; @@ -1481,6 +1499,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) { @@ -1531,6 +1606,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 @@ -1620,6 +1697,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 @@ -1652,6 +1731,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 @@ -1674,6 +1757,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. @@ -1752,6 +1839,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/src/llama-context.cpp b/src/llama-context.cpp index 1968c203a0e2..77956b921058 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" @@ -971,6 +972,73 @@ 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; + } + + auto * mem = static_cast(memory.get()); + if (mem->get_mtp_dsa_index_share() == enabled) { + return true; + } + + mem->set_mtp_dsa_index_share(enabled); + sched_need_reserve = true; + return true; +} + +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(); @@ -1726,6 +1794,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; @@ -1955,6 +2032,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; @@ -3807,6 +3914,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 1cfb08580ac4..a066a8872b2a 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) { diff --git a/src/llama-graph.h b/src/llama-graph.h index e30f915197c2..30e87f5fc0e4 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-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 0475c149f72d..6db5c3e227ca 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -52,8 +52,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); - // the k-pool indexer of glm5-next caches key | gate per token - hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size * (model.hparams.indexer_kpool > 0 ? 2 : 1); + // 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); LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); @@ -136,9 +136,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 +170,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 +182,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 +192,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 +202,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 +212,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(); } } @@ -312,9 +340,19 @@ 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)) { + // 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(); } @@ -471,145 +509,293 @@ void llama_memory_hybrid_idx_context::set_input_qsa( // k-pool DSA indexer (glm5-next) -namespace { +// Cache the k-pool layout for this 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; -struct kpool_seq { - llama_pos pos_min = 0; - std::vector> cells; // (pos, cell) - std::vector pools; + uint32_t n_pool_real = 0; + uint32_t n_new = 0; + bool cache_safe = true; + + bool have_new = false; // Whether the ubatch-dependent part is filled. + size_t i_ubatch = SIZE_MAX; // The ubatch this state was computed for. }; -// per-sequence sorted (pos, cell) lists and the complete pools among the first n_kv cells -static std::vector kpool_collect(const llama_kv_cells & cells, uint32_t kpool, uint32_t n_kv) { - std::vector res(LLAMA_MAX_SEQ); +namespace { - const uint32_t n = std::min(n_kv, cells.size()); - for (uint32_t i = 0; i < n; ++i) { - if (cells.is_empty(i)) { - continue; - } - const llama_pos p = cells.pos_get(i); +// 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::kpool_state & llama_memory_hybrid_idx_context::kpool_get_state( + uint32_t kpool, 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"); + + if (!kpool_st || kpool_st->i_ubatch != i_cur) { + kpool_st = std::make_unique(); + + auto & st = *kpool_st; + st.i_ubatch = i_cur; + st.seqs.resize(LLAMA_MAX_SEQ); + + const auto & cells = mem->get_mem_idx()->get_cells(0); + + const uint32_t n_kv = get_idx()->get_n_kv(); + const uint32_t n = std::min(n_kv, cells.size()); + + // Scan only active sequences + std::vector active; for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { - if (cells.seq_has(i, s)) { - res[s].cells.emplace_back(p, i); + if (cells.seq_pos_min(s) >= 0) { + active.push_back(s); } } - } - for (auto & sq : res) { - if (sq.cells.empty()) { - continue; - } - if (!std::is_sorted(sq.cells.begin(), sq.cells.end())) { - std::sort(sq.cells.begin(), sq.cells.end()); + for (uint32_t i = 0; i < n; ++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; + } } - sq.pos_min = sq.cells.front().first; - // a pool is complete when kpool consecutive positions, from pos_min all have a cell - for (size_t j = 0; j + kpool <= sq.cells.size(); ) { - const llama_pos p0 = sq.cells[j].first; - if ((p0 - sq.pos_min) % kpool != 0) { - ++j; + for (auto & sq : st.seqs) { + if (sq.cells.empty()) { 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 (!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; } } - if (ok) { - sq.pools.push_back((uint32_t) j); - j += kpool; - } else { - ++j; + + st.n_pool_real += (uint32_t) sq.pools.size(); + } + } + + auto & st = *kpool_st; + + if (ubatch != nullptr && !st.have_new) { + // 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++; + } } } + + st.have_new = true; } - return res; + return st; } -// the last padded pool is always unused -static uint32_t kpool_pad(uint32_t n_pool) { - return std::max(64u, GGML_PAD(n_pool + 1, 64u)); +uint32_t llama_memory_hybrid_idx_context::get_n_kpool(uint32_t kpool) const { + return kpool_pad(kpool_get_state(kpool, nullptr).n_pool_real); } +uint32_t llama_memory_hybrid_idx_context::get_n_kpool_new(uint32_t kpool, const llama_ubatch * ubatch) const { + return kpool_get_state(kpool, ubatch).n_new; } -uint32_t llama_memory_hybrid_idx_context::get_n_kpool(uint32_t kpool) const { - GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); - GGML_ASSERT(get_n_stream() == 1 && "TODO: k-pool indexer with multiple streams"); - - const auto seqs = kpool_collect(mem->get_mem_idx()->get_cells(0), kpool, get_idx()->get_n_kv()); +bool llama_memory_hybrid_idx_context::get_kpool_cache_safe(uint32_t kpool) const { + return kpool_get_state(kpool, nullptr).cache_safe; +} - uint32_t n_pool = 0; - for (const auto & sq : seqs) { - n_pool += (uint32_t) sq.pools.size(); - } +bool llama_memory_hybrid_idx_context::get_mtp_dsa_index_share() const { + return mem != nullptr && mem->get_mtp_dsa_index_share(); +} - return kpool_pad(n_pool); +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_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs, ggml_tensor * cell_pool, +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, uint32_t kpool) 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)); - GGML_ASSERT(ggml_backend_buffer_is_host(cell_pool->buffer)); const uint32_t n_kv = get_idx()->get_n_kv(); - const auto seqs = kpool_collect(mem->get_mem_idx()->get_cells(0), kpool, n_kv); + + const auto & st = kpool_get_state(kpool, ubatch); const uint32_t n_tokens = ubatch->n_tokens; - const uint32_t n_pool = pool_idxs->ne[1]; + const uint32_t n_pool = (uint32_t) pool_cells->ne[0]; + const uint32_t n_new = st.n_new; - GGML_ASSERT(pool_idxs->ne[0] == (int64_t) kpool); + 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(cell_pool->ne[0] == (int64_t) n_kv); + 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); + } + } - // the cell of the first ubatch token + // 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 = seqs[s]; + 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); - // cells outside any complete pool map to the last (always unused) pool - int32_t * cpool = (int32_t *) cell_pool->data; - std::fill(cpool, cpool + n_kv, (int32_t) n_pool - 1); + 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; - int32_t * pidx = (int32_t *) pool_idxs->data; + uint32_t i_new = 0; for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { - const auto & sq = seqs[s]; + const auto & sq = st.seqs[s]; seq_pool_start[s] = (uint32_t) pool_end.size(); - for (uint32_t j : sq.pools) { + 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) { - const uint32_t cell = sq.cells[j + k].second; - pidx[ip*kpool + k] = (int32_t) cell; - cpool[cell] = (int32_t) ip; + 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[ip*kpool + k] = (int32_t) dummy_cell; + pidx[(size_t) ip*kpool + k] = sentinel; } } @@ -627,9 +813,17 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_idxs, g std::fill(row, row + n_pool, drop); const uint32_t p0 = seq_pool_start[s]; - const uint32_t p1 = p0 + (uint32_t) seqs[s].pools.size(); + 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) { @@ -642,20 +836,72 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_idxs, g 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 = seqs[s]; + const auto & sq = st.seqs[s]; - const uint32_t n_tail = (uint32_t) ((p - sq.pos_min + 1) % kpool); + 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 = (int32_t) n_kv; + 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, uint32_t kpool) 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_get_state(kpool, ubatch); + 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 35dd5ced9781..5b979900afdb 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -75,6 +75,15 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_kv_cache * get_mem_idx() const; // nullptr when the model carries no indexer + // Sequence edits invalidate cached relative pools. + bool kpool_is_dirty () const { return kpool_dirty; } + void kpool_clear_dirty() const { 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 @@ -85,6 +94,11 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_hparams hparams_idx; const std::unique_ptr mem_idx; + + // Mutable because it is captured and cleared through const contexts. + mutable 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 { @@ -110,7 +124,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 @@ -138,9 +152,17 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns] // the caller then adds the attention mask, the only part of the bias that varies within a block // glm5-next, complete pools of kpool consecutive positions per sequence, scored as whole pools - uint32_t get_n_kpool(uint32_t kpool) const; // padded pool count, the last pool is always unused - void set_input_kpool(ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs, ggml_tensor * cell_pool, + // Cache sequence-private complete pools and re-pool only changed pools. + uint32_t get_n_kpool (uint32_t kpool) const; // Padded pool count, where the last pool is always unused. + uint32_t get_n_kpool_new(uint32_t kpool, const llama_ubatch * ubatch) const; // Exact count of new pools. + bool get_kpool_cache_safe(uint32_t kpool) 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, uint32_t kpool) const; + void set_input_mtp_dsa_selection(ggml_tensor * sel, ggml_tensor * mask, bool gather, + const llama_ubatch * ubatch, uint32_t kpool) 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, @@ -158,4 +180,12 @@ 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; + + // Cached k-pool layout. + struct kpool_state; + kpool_state & kpool_get_state(uint32_t kpool, const llama_ubatch * ubatch) const; + mutable std::unique_ptr kpool_st; + + // Clear a pending full re-pool only after the first ubatch succeeds + bool kpool_dirty_batch = false; }; diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 0ab2141ae0a9..5b3234d62878 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -255,7 +255,10 @@ class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { void set_input(const llama_ubatch * ubatch) override { mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); - mctx->set_input_kpool(pool_idxs, pool_mask, tail_idxs, cell_pool, ubatch, kpool); + mctx->set_input_kpool(pool_cells, pool_idxs, pool_mask, tail_idxs, gather_mask, gather, new_pool_idxs, new_pool_rep, ubatch, kpool); + if (reuse_sel != nullptr) { + mctx->set_input_mtp_dsa_selection(reuse_sel, gather_mask, gather, ubatch, kpool); + } } bool can_reuse(const llm_graph_params & params) override { @@ -268,23 +271,42 @@ class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { bool res = true; - res &= k_idxs->ne[0] == params.ubatch.n_tokens; - res &= pool_idxs->ne[1] == mctx->get_n_kpool(kpool); - res &= pool_mask->ne[1] == params.ubatch.n_tokens; - res &= tail_idxs->ne[1] == params.ubatch.n_tokens; - res &= cell_pool->ne[0] == idx->get_n_kv(); + res &= k_idxs->ne[0] == params.ubatch.n_tokens; + res &= pool_cells->ne[0] == mctx->get_n_kpool(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(kpool, ¶ms.ubatch); + res &= cache_safe == mctx->get_kpool_cache_safe(kpool); + 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_idxs = nullptr; // I32 [kpool, n_pool] - ggml_tensor * pool_mask = nullptr; // F32/F16 [n_pool, n_tokens] - ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens] - ggml_tensor * cell_pool = nullptr; // I32 [n_kv] + 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) { @@ -294,21 +316,70 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui const uint32_t kpool = hparams.indexer_kpool; const uint32_t n_pool = mctx_hyb->get_n_kpool(kpool); const uint32_t n_kv = mctx_idx->get_n_kv(); + const uint32_t n_new = mctx_hyb->get_n_kpool_new(kpool, &ubatch); + const bool cache_safe = mctx_hyb->get_kpool_cache_safe(kpool); // 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_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); - inp->cell_pool = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_kv); + 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_set_input(inp->cell_pool); + + 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. + { + static const int64_t max_ub = [] { + const char * s = getenv("LLAMA_GLM5_GATHER_UBATCH"); + return s != nullptr ? (int64_t) atoll(s) : (int64_t) 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)); } @@ -498,9 +569,11 @@ ggml_tensor * llama_model_glm5_next::graph::build_kda_layer( 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); - const float eps = hparams.f_norm_rms_eps; - Qcur = ggml_l2_norm(ctx0, Qcur, eps); - Kcur = ggml_l2_norm(ctx0, Kcur, eps); + // 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); auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); @@ -534,7 +607,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kda_layer( // 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, const llama_layer & layer, + 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(); @@ -542,12 +615,15 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( 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_idxs->ne[1]; - - // queries - ggml_tensor * 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); + 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); @@ -557,77 +633,127 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( ggml_tensor * ig = ggml_mul_mat(ctx0, layer.indexer_kpool_gate, cur); cb(ig, "indexer_gate", il); - ggml_tensor * packed = ggml_concat(ctx0, ik, ig, 0); - packed = ggml_reshape_3d(ctx0, packed, 2*n_embd_indexer, 1, n_tokens); + // 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]; - k_all = ggml_view_2d(ctx0, k_all, 2*n_embd_indexer, n_kv, k_all->nb[2], 0); - // Gather the member of every pool - ggml_tensor * rows = ggml_get_rows(ctx0, k_all, ggml_reshape_1d(ctx0, inp_kpool->pool_idxs, kpool*n_pool)); - rows = ggml_reshape_3d(ctx0, rows, 2*n_embd_indexer, kpool, n_pool); + 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_pool, rows->nb[1], rows->nb[2], 0); - ggml_tensor * pg = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_pool, rows->nb[1], rows->nb[2], ggml_row_size(rows->type, n_embd_indexer)); + 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_pool] - ggml_tensor * probs = ggml_soft_max(ctx0, logits); + 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)); - ggml_tensor * pooled = ggml_sum_rows(ctx0, ggml_mul(ctx0, probs, pk)); // [1, head_dim, n_pool] + 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 * 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 * 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_cont(ctx0, ggml_top_k(ctx0, score, n_top_pool)); // [n_top_pool, n_tokens] + 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); - 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}); + 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 { - 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(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; } - 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_cont(ctx0, ggml_top_k(ctx0, score, n_top_pool)); // [n_top_pool, n_tokens] - cb(top_k, "indexer_top_k", il); - - // Pool-level selection mask, -inf everywhere except the selected pools - ggml_tensor * pool_sel = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_pool, n_tokens); - pool_sel = ggml_fill(ctx0, pool_sel, -INFINITY); - ggml_tensor * zeros = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_top_pool, n_tokens), 0.0f); - pool_sel = ggml_set_rows(ctx0, pool_sel, zeros, ggml_reshape_3d(ctx0, top_k, n_top_pool, n_tokens, 1)); - pool_sel = ggml_reshape_2d(ctx0, pool_sel, n_pool, n_tokens); - - // Expand to the cells - ggml_tensor * pool_sel_t = ggml_cont(ctx0, ggml_transpose(ctx0, pool_sel)); // [n_tokens, n_pool] - ggml_tensor * sel_t = ggml_get_rows(ctx0, pool_sel_t, inp_kpool->cell_pool); // [n_tokens, n_kv] - ggml_tensor * sel = ggml_cont(ctx0, ggml_transpose(ctx0, sel_t)); // [n_kv, n_tokens] - - if (hparams.indexer_kpool_select_tail) { - ggml_tensor * pad = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_tokens), -INFINITY); - sel = ggml_concat(ctx0, sel, pad, 0); // [n_kv + 1, n_tokens] - sel = ggml_reshape_3d(ctx0, sel, 1, n_kv + 1, n_tokens); - ggml_tensor * tzeros = ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, kpool - 1, n_tokens), 0.0f); - sel = ggml_set_rows(ctx0, sel, tzeros, ggml_reshape_3d(ctx0, inp_kpool->tail_idxs, kpool - 1, n_tokens, 1)); - sel = ggml_view_2d(ctx0, sel, n_kv, n_tokens, sel->nb[2], 0); + + // 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; @@ -667,9 +793,11 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( 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, layer, mctx_hyb, inp_kpool, il); + 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"); @@ -680,20 +808,48 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( 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)); - // Combine the causal mask with the indexer selection - ggml_tensor * kq_mask = inp_attn->get_kq_mask(); - ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32); - mask = ggml_add(ctx0, ggml_reshape_2d(ctx0, mask, mask->ne[0], mask->ne[1]), sel); - if (kq_mask->type != GGML_TYPE_F32) { - mask = ggml_cast(ctx0, mask, kq_mask->type); - } - mask = ggml_reshape_4d(ctx0, mask, kq_mask->ne[0], kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]); - cb(mask, "kq_mask_dsa", 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 * 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); + 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 * out = build_attn_mha(q_absorbed, k, v, nullptr, mask, nullptr, layer.wv_b, kq_scale, 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); diff --git a/src/models/models.h b/src/models/models.h index 65bef6bbe191..cad251afd212 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2485,7 +2485,7 @@ struct llama_model_glm5_next : public llama_model_base { 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, const llama_layer & layer, + 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, From 98557115ef196b608b7090771309aa9ef628bd1c Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 28 Aug 2026 23:31:01 +0200 Subject: [PATCH 04/34] Review driven changes, remove env vars, protect tensors --- common/speculative.cpp | 3 +-- conversion/glm.py | 3 ++- conversion/qwen3vl.py | 2 +- gguf-py/gguf/constants.py | 7 ++++--- gguf-py/gguf/gguf_writer.py | 7 +++++-- src/llama-arch.cpp | 5 +++-- src/llama-arch.h | 1 + src/llama-context.cpp | 11 ++++++----- src/llama-hparams.h | 1 + src/llama-quant.cpp | 12 +++++++++++- src/models/glm5-next.cpp | 6 ++---- tools/mtmd/clip-impl.h | 2 +- tools/mtmd/clip-model.h | 8 ++++++-- tools/mtmd/clip.cpp | 14 +++++++++----- 14 files changed, 53 insertions(+), 29 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index b464cab8ae1a..6aa03a16bac7 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1400,8 +1400,7 @@ 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))); - const char * env_share = getenv("LLAMA_GLM5_MTP_INDEX_SHARE"); - dsa_index_share = (env_share == nullptr || atoi(env_share) != 0) && llama_set_mtp_dsa_index_share(ctx_dft, true); + dsa_index_share = llama_set_mtp_dsa_index_share(ctx_dft, true); if (dsa_index_share) { dsa_sel.resize(n_seq); } diff --git a/conversion/glm.py b/conversion/glm.py index 7402a1a11de6..ffee9f0a0475 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -511,6 +511,7 @@ def set_gguf_parameters(self): 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[:n_layer]]) @@ -593,7 +594,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter 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 if (new_name.startswith(("blk.", "output_hc")) and any(k in new_name for k in - ("hc_attn_", "hc_ffn_", "indexer.kpool", "ssm_a", "ssm_dt", "exp_probs_b"))): + ("hc_attn_", "hc_ffn_", "indexer_compressor_", "ssm_a", "ssm_dt", "exp_probs_b"))): return gguf.GGMLQuantizationType.F32 return super().tensor_force_quant(name, new_name, bid, n_dims) diff --git a/conversion/qwen3vl.py b/conversion/qwen3vl.py index c0885f438d0c..11ce68515b2c 100644 --- a/conversion/qwen3vl.py +++ b/conversion/qwen3vl.py @@ -264,7 +264,7 @@ def set_gguf_parameters(self): 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_limit(float(limit)) + 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 diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 703bfe8f3009..d225763ea6f4 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -225,6 +225,7 @@ class Indexer: 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: @@ -383,7 +384,7 @@ class ClipVision: IMAGE_MEAN = "clip.vision.image_mean" IMAGE_STD = "clip.vision.image_std" SPATIAL_MERGE_SIZE = "clip.vision.spatial_merge_size" - SWIGLU_LIMIT = "clip.vision.swiglu_limit" + 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" @@ -1641,8 +1642,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.kpool_gate", - MODEL_TENSOR.INDEXER_KPOOL_APE: "blk.{bid}.indexer.kpool_ape", + 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", diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 41d4b5c64d49..505afdcac06c 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -822,6 +822,9 @@ def add_indexer_kpool(self, value: int) -> None: 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) @@ -1376,8 +1379,8 @@ 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_limit(self, value: float) -> None: - self.add_float32(Keys.ClipVision.SWIGLU_LIMIT, value) + 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/src/llama-arch.cpp b/src/llama-arch.cpp index f26384040112..4840bcf6a879 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -287,6 +287,7 @@ static const std::map LLM_KV_NAMES = { { 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" }, @@ -681,8 +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.kpool_gate" }, - { LLM_TENSOR_INDEXER_KPOOL_APE, "blk.%d.indexer.kpool_ape" }, + { 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" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 09029c11f941..0213a632e0de 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -292,6 +292,7 @@ enum llm_kv { 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, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 77956b921058..fb00b7c4dd8e 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -977,14 +977,15 @@ bool llama_context::set_mtp_dsa_index_share(bool enabled) { return false; } + enabled = enabled && model.hparams.indexer_index_share_mtp; + auto * mem = static_cast(memory.get()); - if (mem->get_mtp_dsa_index_share() == enabled) { - return true; + if (mem->get_mtp_dsa_index_share() != enabled) { + mem->set_mtp_dsa_index_share(enabled); + sched_need_reserve = true; } - mem->set_mtp_dsa_index_share(enabled); - sched_need_reserve = true; - return true; + return enabled; } bool llama_context::set_mtp_dsa_selection(const int32_t * data, size_t size) { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 39db7ba37304..c36646f5b0c8 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -263,6 +263,7 @@ struct llama_hparams { 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-quant.cpp b/src/llama-quant.cpp index fb3a0e4b318a..0a342beb3d0a 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -330,9 +330,10 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param // GLM5-Next: the k-pool position bias is added elementwise; the indexer, mHC mixers, // KDA gates and MLA low-rank paths are small and precision-sensitive (~1 GB total) - quantize &= name.find("indexer.kpool_ape.weight") == std::string::npos; if (arch == LLM_ARCH_GLM5_NEXT) { quantize &= name.find("indexer.") == std::string::npos; + quantize &= name.find("indexer_compressor_gate.weight") == std::string::npos; + quantize &= name.find("indexer_compressor_ape.weight") == std::string::npos; quantize &= name.find("hc_attn_fn.weight") == std::string::npos; quantize &= name.find("hc_ffn_fn.weight") == std::string::npos; quantize &= name.find("ssm_f_a.weight") == std::string::npos; @@ -344,6 +345,15 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param 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; + + // the NextN draft head runs once per speculative token, so its error compounds + quantize &= name.find("nextn.eh_proj.weight") == std::string::npos; + + quantize &= name.find("attn_q.weight") == std::string::npos; //This is too strict, relax these later + quantize &= name.find("attn_k.weight") == std::string::npos; + quantize &= name.find("attn_v.weight") == std::string::npos; + quantize &= name.find("attn_q_b.weight") == std::string::npos; + quantize &= name.find("attn_output.weight") == std::string::npos; } // do not quantize RWKV's small yet 2D weights diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 5b3234d62878..f7dd7f00f102 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -48,6 +48,7 @@ void llama_model_glm5_next::load_arch_hparams(llama_model_loader & ml) { 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); @@ -343,10 +344,7 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui // Gather selected latents for small decode batches when n_kv exceeds n_sel. { - static const int64_t max_ub = [] { - const char * s = getenv("LLAMA_GLM5_GATHER_UBATCH"); - return s != nullptr ? (int64_t) atoll(s) : (int64_t) 16; - }(); + 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); diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 75a0fad692f9..28406d8125a6 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -63,7 +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_LIMIT "clip.vision.swiglu_limit" +#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" diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 79183b179adb..7539c82b08e4 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -94,8 +94,12 @@ struct clip_hparams { float eps = 1e-6; float rope_theta = 0.0; - // clamp the SwiGLU gate to (-inf, limit] and up to [-limit, limit] when > 0 (glm5-next) - float swiglu_limit = 0.0f; + 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 6f5ad0bac395..2406b00a64e9 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -648,9 +648,9 @@ ggml_tensor * clip_graph::build_ffn( switch (type_op) { case FFN_SILU: if (gate) { - if (hparams.swiglu_limit > 0.0f) { - cur = ggml_clamp(ctx0, cur, -INFINITY, hparams.swiglu_limit); - tmp = ggml_clamp(ctx0, tmp, -hparams.swiglu_limit, hparams.swiglu_limit); + 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); @@ -1739,10 +1739,14 @@ struct clip_model_loader { hparams.n_merge = 2; hparams.image_resize_algo = RESIZE_ALGO_BICUBIC; get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); - get_f32(KEY_SWIGLU_LIMIT, hparams.swiglu_limit, 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.warmup_image_size = static_cast(std::sqrt(hparams.image_max_pixels)); hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup } break; case PROJECTOR_TYPE_LLAMA4: From b93ed51f4ff7c52a3e58c0a1063056bc478222b3 Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 28 Aug 2026 23:58:43 +0200 Subject: [PATCH 05/34] Strip MTP for initial PR --- common/speculative.cpp | 80 ----------------- src/llama-context.cpp | 126 --------------------------- src/llama-context.h | 8 -- src/llama-ext.h | 3 - src/llama-graph.cpp | 8 -- src/llama-graph.h | 4 - src/llama-memory-hybrid-idx.cpp | 77 ---------------- src/llama-memory-hybrid-idx.h | 11 --- src/models/glm5-next.cpp | 150 ++------------------------------ src/models/models.h | 10 --- 10 files changed, 7 insertions(+), 470 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 6aa03a16bac7..e00eea511d03 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1382,10 +1382,6 @@ 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) @@ -1400,11 +1396,6 @@ 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", @@ -1471,7 +1462,6 @@ 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 { @@ -1487,10 +1477,6 @@ 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; @@ -1498,63 +1484,6 @@ 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) { @@ -1605,7 +1534,6 @@ 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); @@ -1696,7 +1624,6 @@ 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); @@ -1730,9 +1657,6 @@ 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 @@ -1756,9 +1680,6 @@ 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 @@ -1838,7 +1759,6 @@ 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/src/llama-context.cpp b/src/llama-context.cpp index fb00b7c4dd8e..44b5285a4d46 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -972,74 +972,6 @@ 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(); @@ -1795,15 +1727,6 @@ 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; @@ -2033,36 +1956,6 @@ 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; @@ -3915,25 +3808,6 @@ 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 8073d96fb2fa..abca6468deb2 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -100,9 +100,6 @@ 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( @@ -303,11 +300,6 @@ 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 diff --git a/src/llama-ext.h b/src/llama-ext.h index 71f7b72d4ed8..1f284885b8eb 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -107,9 +107,6 @@ 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 a066a8872b2a..1cfb08580ac4 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1323,8 +1323,6 @@ 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); @@ -1371,12 +1369,6 @@ 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) { diff --git a/src/llama-graph.h b/src/llama-graph.h index 30e87f5fc0e4..e30f915197c2 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -900,8 +900,6 @@ 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]; } @@ -937,8 +935,6 @@ 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-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 6db5c3e227ca..6c381df6c4dd 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -136,27 +136,9 @@ 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); } @@ -171,7 +153,6 @@ 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); @@ -183,7 +164,6 @@ 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(); } } @@ -193,7 +173,6 @@ 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(); } } @@ -203,7 +182,6 @@ 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(); } } @@ -213,7 +191,6 @@ 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(); } } @@ -678,14 +655,6 @@ bool llama_memory_hybrid_idx_context::get_kpool_cache_safe(uint32_t kpool) const return kpool_get_state(kpool, nullptr).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, uint32_t kpool) const { @@ -859,49 +828,3 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, } } } -void llama_memory_hybrid_idx_context::set_input_mtp_dsa_selection( - ggml_tensor * sel, ggml_tensor * mask, bool gather, - const llama_ubatch * ubatch, uint32_t kpool) 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_get_state(kpool, ubatch); - 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 5b979900afdb..3602a0567fb4 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -79,10 +79,6 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { bool kpool_is_dirty () const { return kpool_dirty; } void kpool_clear_dirty() const { 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 @@ -97,8 +93,6 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { // Mutable because it is captured and cleared through const contexts. mutable 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 { @@ -156,14 +150,9 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { uint32_t get_n_kpool (uint32_t kpool) const; // Padded pool count, where the last pool is always unused. uint32_t get_n_kpool_new(uint32_t kpool, const llama_ubatch * ubatch) const; // Exact count of new pools. bool get_kpool_cache_safe(uint32_t kpool) 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, uint32_t kpool) const; - void set_input_mtp_dsa_selection(ggml_tensor * sel, ggml_tensor * mask, bool gather, - const llama_ubatch * ubatch, uint32_t kpool) 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; diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index f7dd7f00f102..03e067418f77 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -195,7 +195,7 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) { std::unique_ptr llama_model_glm5_next::build_arch_graph(const llm_graph_params & params) const { if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { - return std::make_unique(*this, params); + throw std::runtime_error("GLM5-Next NextN graph not implemented yet"); } return std::make_unique(*this, params); } @@ -257,9 +257,6 @@ class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { 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, kpool); - if (reuse_sel != nullptr) { - mctx->set_input_mtp_dsa_selection(reuse_sel, gather_mask, gather, ubatch, kpool); - } } bool can_reuse(const llm_graph_params & params) override { @@ -281,11 +278,6 @@ class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { // The new pool path is sized exactly res &= n_new == mctx->get_n_kpool_new(kpool, ¶ms.ubatch); res &= cache_safe == mctx->get_kpool_cache_safe(kpool); - 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; } @@ -296,7 +288,6 @@ class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { 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 @@ -306,7 +297,6 @@ class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { uint32_t n_sel = 0; bool cache_safe = true; bool gather = false; - bool mtp_share = false; uint32_t n_kv = 0; }; @@ -348,24 +338,15 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui 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) { + if (inp->gather) { 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; @@ -616,12 +597,9 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( 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); - } + ggml_tensor * 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); @@ -681,8 +659,8 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( 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 * 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); @@ -717,16 +695,9 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( // 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); @@ -855,110 +826,3 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( 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 cad251afd212..9393304af98a 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2492,16 +2492,6 @@ struct llama_model_glm5_next : public llama_model_base { 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) : - llm_build_delta_net_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; From ba436fd3ced2d6c7d16f94678c29b8d0880dc3e9 Mon Sep 17 00:00:00 2001 From: timkhronos Date: Sat, 29 Aug 2026 00:02:00 +0200 Subject: [PATCH 06/34] Clean up after mtp strip --- common/speculative.cpp | 8 -------- 1 file changed, 8 deletions(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index e00eea511d03..69c9ddf07f6b 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -15,7 +15,6 @@ #include #include #include -#include #include #include #include @@ -1382,7 +1381,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::vector i_last; std::vector> chain_h; - 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) @@ -1463,7 +1461,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { verify_h_rows.assign(n_seq, 0); } - ~common_speculative_impl_draft_mtp() override { auto * ctx_dft = this->params.ctx_dft; for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) { @@ -1534,7 +1531,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { auto * ctx_tgt = this->params.ctx_tgt; auto * ctx_dft = this->params.ctx_dft; - 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 @@ -1624,7 +1620,6 @@ 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; - common_batch_clear(batch); // keep track of which sequences are still drafting @@ -1657,7 +1652,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { int i = 0; while (n_drafting > 0) { - // 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 @@ -1680,7 +1674,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { break; } - // 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. @@ -1759,7 +1752,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { ++i; } - if (chain_heads) { llama_set_nextn_layer_offset(ctx_dft, 0); // restore default for non-draft decodes } From 3ff8d92c7a582592abb9404647ce96e97372d644 Mon Sep 17 00:00:00 2001 From: timkhronos Date: Sat, 29 Aug 2026 00:03:19 +0200 Subject: [PATCH 07/34] Clean up after mtp strip --- common/speculative.cpp | 1 - 1 file changed, 1 deletion(-) diff --git a/common/speculative.cpp b/common/speculative.cpp index 69c9ddf07f6b..a9501a9d1ead 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1459,7 +1459,6 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { verify_h.assign(n_seq, {}); verify_h_rows.assign(n_seq, 0); - } ~common_speculative_impl_draft_mtp() override { auto * ctx_dft = this->params.ctx_dft; From f176c8a3fb448a9a009410c555b09b689388d760 Mon Sep 17 00:00:00 2001 From: timkhronos Date: Sat, 29 Aug 2026 00:03:55 +0200 Subject: [PATCH 08/34] Update speculative.cpp --- common/speculative.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/common/speculative.cpp b/common/speculative.cpp index a9501a9d1ead..b5348ab6f3d5 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -1460,6 +1460,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { verify_h.assign(n_seq, {}); verify_h_rows.assign(n_seq, 0); } + ~common_speculative_impl_draft_mtp() override { auto * ctx_dft = this->params.ctx_dft; for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) backend_chains.size(); ++seq_id) { From c34db5814e26ae32ac3c2428092c15e0824d8f04 Mon Sep 17 00:00:00 2001 From: timkhronos Date: Sat, 29 Aug 2026 00:05:28 +0200 Subject: [PATCH 09/34] Update llama-context.h --- src/llama-context.h | 1 - 1 file changed, 1 deletion(-) diff --git a/src/llama-context.h b/src/llama-context.h index abca6468deb2..bf91daa8b562 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -300,7 +300,6 @@ struct llama_context { // sets llm_graph_result::t_h_nextn buffer_view embd_nextn = {nullptr, 0}; - // host buffers for output layer input embeddings, per layer // populated when cparams.output_layer_inp[il] is true std::vector> embd_layer_inp; From 0aa327b3f7e3183f2dd68d40882aa95b4a97bdbe Mon Sep 17 00:00:00 2001 From: timkhronos Date: Sat, 29 Aug 2026 00:06:09 +0200 Subject: [PATCH 10/34] Clean up after mtp strip --- src/llama-ext.h | 1 - 1 file changed, 1 deletion(-) diff --git a/src/llama-ext.h b/src/llama-ext.h index 1f284885b8eb..92a759b7a0ae 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -107,7 +107,6 @@ 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); - // 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); From ec1cdbf016516c6d0b37c2911f5444c47ee68ee4 Mon Sep 17 00:00:00 2001 From: Chrono Date: Sat, 29 Aug 2026 14:10:17 +0200 Subject: [PATCH 11/34] Fix tokenizer ignore merges --- conversion/base.py | 8 ++++---- conversion/glm.py | 2 +- src/llama-vocab.cpp | 5 +++++ 3 files changed, 10 insertions(+), 5 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index e9915f5c3f0e..1b8db40207dc 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1416,7 +1416,7 @@ def does_token_look_special(self, token: str | bytes) -> bool: return seems_special # used for GPT-2 BPE and WordPiece vocabs - def get_vocab_base(self, tokenizer=None) -> tuple[list[str], list[int], str]: + def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[list[str], list[int], str]: tokens: list[str] = [] toktypes: list[int] = [] @@ -1426,7 +1426,7 @@ def get_vocab_base(self, tokenizer=None) -> tuple[list[str], list[int], str]: vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) # ty: ignore[unresolved-attribute] assert max(tokenizer.vocab.values()) < vocab_size # ty: ignore[unresolved-attribute] - tokpre = self.get_vocab_base_pre(tokenizer) + tokpre = tokpre or self.get_vocab_base_pre(tokenizer) reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} # ty: ignore[unresolved-attribute] added_vocab = tokenizer.get_added_vocab() # ty: ignore[unresolved-attribute] @@ -2168,12 +2168,12 @@ def _set_vocab_glmedge(self): special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|endoftext|>"]) # ty: ignore[unresolved-attribute] special_vocab.add_to_gguf(self.gguf_writer) - def _set_vocab_glm(self, tokenizer=None): + def _set_vocab_glm(self, tokenizer=None, tokpre: str | None = None): if tokenizer is None: from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(self.dir_model) special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) - tokens, toktypes, tokpre = self.get_vocab_base(tokenizer) + tokens, toktypes, tokpre = self.get_vocab_base(tokenizer, tokpre=tokpre) self.gguf_writer.add_tokenizer_model("gpt2") self.gguf_writer.add_tokenizer_pre(tokpre) self.gguf_writer.add_token_list(tokens) diff --git a/conversion/glm.py b/conversion/glm.py index ffee9f0a0475..7e807780bdad 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -436,7 +436,7 @@ def set_vocab(self): # the repo ships only a transformers v5 style tokenizer.json, load it directly from transformers import PreTrainedTokenizerFast tokenizer = PreTrainedTokenizerFast(tokenizer_file=str(self.dir_model / "tokenizer.json")) - return self._set_vocab_glm(tokenizer) + return self._set_vocab_glm(tokenizer, tokpre="glm5") def index_tensors(self, remote_hf_model_id: str | None = None): hp = self.hparams.get("text_config", self.hparams) diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index ff926ceecd17..68633d6b33a5 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2257,6 +2257,11 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "chatglm-bpe") { pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4; special_bos_id = LLAMA_TOKEN_NULL; + } else if ( + 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; From 7152e9bf218ff92baabac61778e99f5a6fe6c966 Mon Sep 17 00:00:00 2001 From: Chrono Date: Sat, 29 Aug 2026 14:17:14 +0200 Subject: [PATCH 12/34] Improve quantization protection selection --- src/llama-quant.cpp | 57 +++++++++++++++++++++++++-------------------- 1 file changed, 32 insertions(+), 25 deletions(-) diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 0a342beb3d0a..ee1217f68b3e 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -328,32 +328,23 @@ 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: the k-pool position bias is added elementwise; the indexer, mHC mixers, - // KDA gates and MLA low-rank paths are small and precision-sensitive (~1 GB total) + // glm5-next if (arch == LLM_ARCH_GLM5_NEXT) { - quantize &= name.find("indexer.") == std::string::npos; - quantize &= name.find("indexer_compressor_gate.weight") == std::string::npos; - quantize &= name.find("indexer_compressor_ape.weight") == std::string::npos; - quantize &= name.find("hc_attn_fn.weight") == std::string::npos; - quantize &= name.find("hc_ffn_fn.weight") == 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_q_a.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; - - // the NextN draft head runs once per speculative token, so its error compounds - quantize &= name.find("nextn.eh_proj.weight") == std::string::npos; - - quantize &= name.find("attn_q.weight") == std::string::npos; //This is too strict, relax these later - quantize &= name.find("attn_k.weight") == std::string::npos; - quantize &= name.find("attn_v.weight") == std::string::npos; - quantize &= name.find("attn_q_b.weight") == std::string::npos; - quantize &= name.find("attn_output.weight") == std::string::npos; + 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 @@ -479,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)) { From fdc54bedfb074fb6d4491beea68a64040b02431d Mon Sep 17 00:00:00 2001 From: Chrono Date: Mon, 31 Aug 2026 13:20:42 +0200 Subject: [PATCH 13/34] Refactor mhc helpers, graph base --- src/llama-graph.cpp | 175 ------------------------------------- src/llama-graph.h | 30 ------- src/models/deepseek4.cpp | 183 +++++++++++++++++++++++++++++++++++++-- src/models/glm5-next.cpp | 15 +--- src/models/models.h | 52 ++++++++++- 5 files changed, 226 insertions(+), 229 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 1cfb08580ac4..40b59fc3644e 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3410,181 +3410,6 @@ llm_graph_input_dsv4 * llm_graph_context::build_inp_dsv4() const { return (llm_graph_input_dsv4 *) res->add_input(std::move(inp)); } -// manifold-constrained hyper-connections (mHC), deepseek4 and glm5-next - -static ggml_tensor * hc_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) { - return ggml_view_1d(ctx, t, ne0, ggml_row_size(t->type, i0)); -} - -static ggml_tensor * hc_view_2d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t i0) { - return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], ggml_row_size(t->type, i0)); -} - -static ggml_tensor * hc_affine(ggml_context * ctx, ggml_tensor * x, ggml_tensor * scale, ggml_tensor * base) { - x = ggml_mul(ctx, x, scale); - x = ggml_add(ctx, x, base); - return x; -} - -ggml_tensor * llm_graph_context::build_hc_pre( - ggml_tensor * x, - ggml_tensor * weights, - int il) const { - GGML_ASSERT(x->ne[0] == n_embd); - GGML_ASSERT(x->ne[1] == hparams.dsv4_hc_mult); - - const int64_t hc = hparams.dsv4_hc_mult; - const int64_t nt = x->ne[2]; - - if (cparams.fused_dsv4_hc_pre && il >= 0) { - ggml_tensor * result = ggml_dsv4_hc_pre(ctx0, x, weights); - res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, result, il}); - return result; - } - - ggml_tensor * result = nullptr; - for (int64_t ih = 0; ih < hc; ++ih) { - ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]); - ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]); - ggml_tensor * cur = ggml_mul(ctx0, xh, wh); - result = result ? ggml_add(ctx0, result, cur) : cur; - } - - return result; -} - -ggml_tensor * llm_graph_context::build_hc_sinkhorn( - ggml_tensor * comb, - int il) const { - GGML_UNUSED(il); - - // comb is [dst_hc, src_hc, n_tokens]. Sinkhorn follows the reference: - // row softmax over dst, one column normalization, then repeated row/column normalization. - comb = ggml_soft_max(ctx0, comb); - - ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1); - eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps); - - comb = ggml_add(ctx0, comb, eps); - - auto norm_cols = [&]() { - ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3)); - ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst); - col_sum = ggml_add(ctx0, col_sum, eps); - col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3); - comb = ggml_div(ctx0, comb, col_sum); - }; - - auto norm_rows = [&]() { - ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb); - row_sum = ggml_add(ctx0, row_sum, eps); - comb = ggml_div(ctx0, comb, row_sum); - }; - - norm_cols(); - for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) { - norm_rows(); - norm_cols(); - } - - return comb; -} - -ggml_tensor * llm_graph_context::build_hc_pre( - ggml_tensor * x, - ggml_tensor * hc_fn, - ggml_tensor * hc_scale, - ggml_tensor * hc_base, - ggml_tensor ** post, - ggml_tensor ** comb, - int il) const { - const int64_t hc = hparams.dsv4_hc_mult; - const int64_t hc_dim = hc*n_embd; - const int64_t hc_mix_dim = (2 + hc)*hc; - const int64_t nt = x->ne[2]; - - GGML_ASSERT(hc == 4); - GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim); - - ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt); - ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps); - ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); - cb(mixes, "hc_mixes", il); - - ggml_tensor * scale_pre = hc_view_1d(ctx0, hc_scale, 1, 0); - ggml_tensor * scale_post = hc_view_1d(ctx0, hc_scale, 1, 1); - - ggml_tensor * base_pre = hc_view_1d(ctx0, hc_base, hc, 0); - ggml_tensor * base_post = hc_view_1d(ctx0, hc_base, hc, hc); - - ggml_tensor * pre = hc_view_2d(ctx0, mixes, hc, nt, 0); - pre = hc_affine(ctx0, pre, scale_pre, base_pre); - pre = ggml_sigmoid(ctx0, pre); - pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps); - cb(pre, "hc_pre", il); - - *post = hc_view_2d(ctx0, mixes, hc, nt, hc); - *post = hc_affine(ctx0, *post, scale_post, base_post); - *post = ggml_sigmoid(ctx0, *post); - *post = ggml_scale(ctx0, *post, 2.0f); - cb(*post, "hc_post", il); - - if (cparams.fused_dsv4_hc_comb) { - *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps, - (int32_t) hparams.dsv4_hc_sinkhorn_iters); - res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il}); - } else { - ggml_tensor * scale_comb = hc_view_1d(ctx0, hc_scale, 1, 2); - ggml_tensor * base_comb = hc_view_1d(ctx0, hc_base, hc*hc, 2*hc); - - *comb = hc_view_2d(ctx0, mixes, hc*hc, nt, 2*hc); - *comb = hc_affine(ctx0, *comb, scale_comb, base_comb); - *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt); - *comb = build_hc_sinkhorn(*comb, il); - } - cb(*comb, "hc_comb", il); - - ggml_tensor * result = build_hc_pre(x, pre, il); - return result; -} - -ggml_tensor * llm_graph_context::build_hc_post( - ggml_tensor * x, - ggml_tensor * residual, - ggml_tensor * post, - ggml_tensor * comb, - int il) const { - GGML_ASSERT(x->ne[0] == n_embd); - GGML_ASSERT(residual->ne[1] == hparams.dsv4_hc_mult); - - if (cparams.fused_dsv4_hc_post) { - ggml_tensor * result = ggml_dsv4_hc_post(ctx0, x, residual, post, comb); - res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, result, il}); - return result; - } - - const int64_t hc = hparams.dsv4_hc_mult; - const int64_t nt = x->ne[1]; - - ggml_tensor * out = nullptr; - for (int64_t dst = 0; dst < hc; ++dst) { - ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst*post->nb[0]); - ggml_tensor * cur = ggml_mul(ctx0, x, post_dst); - - for (int64_t src = 0; src < hc; ++src) { - ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]); - ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2], - dst*comb->nb[0] + src*comb->nb[1]); - cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst)); - } - - cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt); - out = out ? ggml_concat(ctx0, out, cur, 1) : cur; - } - - return out; -} - ggml_tensor * llm_graph_context::build_rs( ggml_tensor * s, ggml_tensor * state_copy_main, diff --git a/src/llama-graph.h b/src/llama-graph.h index e30f915197c2..b388e028cb53 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1339,36 +1339,6 @@ struct llm_graph_context { // hybrid // - // hyper-connections (mHC) - - - // collapse the hc streams with per-stream weights - ggml_tensor * build_hc_pre( - ggml_tensor * x, - ggml_tensor * weights, - int il) const; - - // returns the collapsed input and fills the post / comb weights - ggml_tensor * build_hc_pre( - ggml_tensor * x, - ggml_tensor * hc_fn, - ggml_tensor * hc_scale, - ggml_tensor * hc_base, - ggml_tensor ** post, - ggml_tensor ** comb, - int il) const; - - ggml_tensor * build_hc_post( - ggml_tensor * x, - ggml_tensor * residual, - ggml_tensor * post, - ggml_tensor * comb, - int il) const; - - ggml_tensor * build_hc_sinkhorn( - ggml_tensor * comb, - int il) const; - llm_graph_input_mem_hybrid * build_inp_mem_hybrid() const; llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k() const; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 3f7c62e53ac9..69dc4ee2a6f8 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -263,15 +263,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( @@ -284,6 +284,173 @@ static ggml_tensor * dsv4_hc_affine( return x; } +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_pre( + ggml_tensor * x, + ggml_tensor * weights, + int il) const { + GGML_ASSERT(x->ne[0] == n_embd); + GGML_ASSERT(x->ne[1] == hparams.dsv4_hc_mult); + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = x->ne[2]; + + if (cparams.fused_dsv4_hc_pre && il >= 0) { + ggml_tensor * result = ggml_dsv4_hc_pre(ctx0, x, weights); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, result, il}); + return result; + } + + ggml_tensor * result = nullptr; + for (int64_t ih = 0; ih < hc; ++ih) { + ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]); + ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]); + ggml_tensor * cur = ggml_mul(ctx0, xh, wh); + result = result ? ggml_add(ctx0, result, cur) : cur; + } + + return result; +} + +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_sinkhorn( + ggml_tensor * comb, + int il) const { + GGML_UNUSED(il); + + // comb is [dst_hc, src_hc, n_tokens]. Sinkhorn follows the reference: + // row softmax over dst, one column normalization, then repeated row/column normalization. + comb = ggml_soft_max(ctx0, comb); + + ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1); + eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps); + + comb = ggml_add(ctx0, comb, eps); + + auto norm_cols = [&]() { + ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3)); + ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst); + col_sum = ggml_add(ctx0, col_sum, eps); + col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3); + comb = ggml_div(ctx0, comb, col_sum); + }; + + auto norm_rows = [&]() { + ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb); + row_sum = ggml_add(ctx0, row_sum, eps); + comb = ggml_div(ctx0, comb, row_sum); + }; + + norm_cols(); + for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) { + norm_rows(); + norm_cols(); + } + + return comb; +} + +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_pre( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il) const { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc*n_embd; + const int64_t hc_mix_dim = (2 + hc)*hc; + const int64_t nt = x->ne[2]; + + GGML_ASSERT(hc == 4); + GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim); + + ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt); + ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps); + ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm); + cb(mixes, "hc_mixes", il); + + ggml_tensor * scale_pre = dsv4_view_1d(ctx0, hc_scale, 1, 0); + ggml_tensor * scale_post = dsv4_view_1d(ctx0, hc_scale, 1, 1); + + ggml_tensor * base_pre = dsv4_view_1d(ctx0, hc_base, hc, 0); + ggml_tensor * base_post = dsv4_view_1d(ctx0, hc_base, hc, hc); + + ggml_tensor * pre = dsv4_view_2d(ctx0, mixes, hc, nt, 0); + pre = dsv4_hc_affine(ctx0, pre, scale_pre, base_pre); + pre = ggml_sigmoid(ctx0, pre); + pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps); + cb(pre, "hc_pre", il); + + *post = dsv4_view_2d(ctx0, mixes, hc, nt, hc); + *post = dsv4_hc_affine(ctx0, *post, scale_post, base_post); + *post = ggml_sigmoid(ctx0, *post); + *post = ggml_scale(ctx0, *post, 2.0f); + cb(*post, "hc_post", il); + + if (cparams.fused_dsv4_hc_comb) { + *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps, + (int32_t) hparams.dsv4_hc_sinkhorn_iters); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il}); + } else { + ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2); + ggml_tensor * base_comb = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc); + + *comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc); + *comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb); + *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt); + *comb = build_hc_sinkhorn(*comb, il); + } + cb(*comb, "hc_comb", il); + + ggml_tensor * result = build_hc_pre(x, pre, il); + return result; +} + +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_post( + ggml_tensor * x, + ggml_tensor * residual, + ggml_tensor * post, + ggml_tensor * comb, + int il) const { + GGML_ASSERT(x->ne[0] == n_embd); + GGML_ASSERT(residual->ne[1] == hparams.dsv4_hc_mult); + + if (cparams.fused_dsv4_hc_post) { + ggml_tensor * result = ggml_dsv4_hc_post(ctx0, x, residual, post, comb); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, result, il}); + return result; + } + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = x->ne[1]; + + ggml_tensor * out = nullptr; + for (int64_t dst = 0; dst < hc; ++dst) { + ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst*post->nb[0]); + ggml_tensor * cur = ggml_mul(ctx0, x, post_dst); + + for (int64_t src = 0; src < hc; ++src) { + ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]); + ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2], + dst*comb->nb[0] + src*comb->nb[1]); + cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst)); + } + + cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt); + out = out ? ggml_concat(ctx0, out, cur, 1) : cur; + } + + 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, @@ -1058,7 +1225,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); @@ -1075,7 +1242,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]); } @@ -1157,7 +1324,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 index 03e067418f77..5707ac028add 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -200,17 +200,6 @@ std::unique_ptr llama_model_glm5_next::build_arch_graph(const return std::make_unique(*this, params); } -// Mean over the hyper-connection streams -static ggml_tensor * glm5_hc_mean(ggml_context * ctx, ggml_tensor * x) { - 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); - 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])); - } - return ggml_scale(ctx, acc, 1.0f/hc); -} - // 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, @@ -364,7 +353,7 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui } llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params) : - llm_build_delta_net_base(params), model(model) { + llama_model_deepseek4::graph_base(params), model(model) { ggml_tensor * cur; @@ -476,7 +465,7 @@ llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_p inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs); } - cur = glm5_hc_mean(ctx0, inpL); + cur = build_hc_mean(inpL); cb(cur, "hc_head", -1); cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); diff --git a/src/models/models.h b/src/models/models.h index 9393304af98a..15961170d90d 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1176,8 +1176,53 @@ 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) {} + // 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, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il) const; + + ggml_tensor * build_hc_post( + ggml_tensor * x, + ggml_tensor * residual, + ggml_tensor * post, + 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( @@ -2473,7 +2518,8 @@ struct llama_model_glm5_next : public llama_model_base { // k-pool indexer inputs on top of the generic hybrid input class llm_graph_input_kpool; - struct graph : public llm_build_delta_net_base { + // 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; From 74bb0e39153b2a1569b918a1fd95f7610e645fc9 Mon Sep 17 00:00:00 2001 From: Chrono Date: Mon, 31 Aug 2026 14:56:06 +0200 Subject: [PATCH 14/34] Lint Fixes --- conversion/base.py | 12 ++++++------ conversion/dream.py | 13 +++++++------ conversion/glm.py | 4 ++-- conversion/laguna.py | 4 ++-- conversion/llada.py | 13 +++++++------ 5 files changed, 24 insertions(+), 22 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index 1b8db40207dc..3f96b03d5c86 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1423,15 +1423,15 @@ def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[lis if tokenizer is None: from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(self.dir_model) - vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) # ty: ignore[unresolved-attribute] - assert max(tokenizer.vocab.values()) < vocab_size # ty: ignore[unresolved-attribute] + vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) + assert max(tokenizer.vocab.values()) < vocab_size tokpre = tokpre or self.get_vocab_base_pre(tokenizer) - reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} # ty: ignore[unresolved-attribute] - added_vocab = tokenizer.get_added_vocab() # ty: ignore[unresolved-attribute] + reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} + added_vocab = tokenizer.get_added_vocab() - added_tokens_decoder = tokenizer.added_tokens_decoder # ty: ignore[unresolved-attribute] + added_tokens_decoder = tokenizer.added_tokens_decoder for i in range(vocab_size): if i not in reverse_vocab: @@ -1444,7 +1444,7 @@ def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[lis # To avoid unexpected issues - we make sure to normalize non-normalized tokens if not added_tokens_decoder[i].normalized: previous_token = token - token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) # ty: ignore[unresolved-attribute, invalid-assignment] + token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) if previous_token != token: logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer") diff --git a/conversion/dream.py b/conversion/dream.py index 14f25404d67f..dcd1bb9fd993 100644 --- a/conversion/dream.py +++ b/conversion/dream.py @@ -13,21 +13,22 @@ class DreamModel(TextModel): model_arch = gguf.MODEL_ARCH.DREAM - def get_vocab_base(self) -> tuple[list[str], list[int], str]: + def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[list[str], list[int], str]: tokens: list[str] = [] toktypes: list[int] = [] - from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) + if tokenizer is None: + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) - vocab_dict = tokenizer.get_vocab() # ty: ignore[unresolved-attribute] + vocab_dict = tokenizer.get_vocab() vocab_size = self.hparams.get("vocab_size", len(vocab_dict)) assert max(vocab_dict.values()) < vocab_size - tokpre = self.get_vocab_base_pre(tokenizer) + tokpre = tokpre or self.get_vocab_base_pre(tokenizer) reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in vocab_dict.items()} - added_vocab = tokenizer.get_added_vocab() # ty: ignore[unresolved-attribute] + added_vocab = tokenizer.get_added_vocab() for i in range(vocab_size): if i not in reverse_vocab: diff --git a/conversion/glm.py b/conversion/glm.py index 7e807780bdad..360355c680d6 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -593,8 +593,8 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter 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 - if (new_name.startswith(("blk.", "output_hc")) and any(k in new_name for k in - ("hc_attn_", "hc_ffn_", "indexer_compressor_", "ssm_a", "ssm_dt", "exp_probs_b"))): + 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) diff --git a/conversion/laguna.py b/conversion/laguna.py index 29e0b3d6b318..0757eca12c58 100644 --- a/conversion/laguna.py +++ b/conversion/laguna.py @@ -45,14 +45,14 @@ def set_vocab(self) -> None: self.gguf_writer.add_eot_token_id(extra[0]) logger.info(f"gguf: registered eot_token_id={extra[0]} from eos list {eos_ids}") - def get_vocab_base(self) -> tuple[list[str], list[int], str]: + def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[list[str], list[int], str]: # is the assistant turn-end (registered as eot below). The # HF tokenizer flags it special=false, so the base classifies it as # USER_DEFINED and llama.cpp renders its text into generated content, # leaking "" and breaking response parsing. It is a control # marker, so promote it to CONTROL: llama.cpp then treats it as # end-of-generation and suppresses its text. - tokens, toktypes, tokpre = super().get_vocab_base() + tokens, toktypes, tokpre = super().get_vocab_base(tokenizer, tokpre=tokpre) for i, tok in enumerate(tokens): if tok == "": toktypes[i] = gguf.TokenType.CONTROL diff --git a/conversion/llada.py b/conversion/llada.py index c03607191a76..03ea421cf0a5 100644 --- a/conversion/llada.py +++ b/conversion/llada.py @@ -16,21 +16,22 @@ class LLaDAModel(TextModel): model_arch = gguf.MODEL_ARCH.LLADA undo_permute = True - def get_vocab_base(self) -> tuple[list[str], list[int], str]: + def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[list[str], list[int], str]: tokens: list[str] = [] toktypes: list[int] = [] - from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) + if tokenizer is None: + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) - vocab_dict = tokenizer.get_vocab() # ty: ignore[unresolved-attribute] + vocab_dict = tokenizer.get_vocab() vocab_size = self.hparams.get("vocab_size", len(vocab_dict)) assert max(vocab_dict.values()) < vocab_size - tokpre = self.get_vocab_base_pre(tokenizer) + tokpre = tokpre or self.get_vocab_base_pre(tokenizer) reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in vocab_dict.items()} - added_vocab = tokenizer.get_added_vocab() # ty: ignore[unresolved-attribute] + added_vocab = tokenizer.get_added_vocab() for i in range(vocab_size): if i not in reverse_vocab: From a771613af20f3dc60247e4b6a3d11516f0664673 Mon Sep 17 00:00:00 2001 From: timkhronos Date: Mon, 31 Aug 2026 16:33:33 +0200 Subject: [PATCH 15/34] Apply suggestions from code review MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Sigbjørn Skjæret --- conversion/glm.py | 8 ++------ src/models/glm5-next.cpp | 2 +- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/conversion/glm.py b/conversion/glm.py index 360355c680d6..6336c59303d2 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -453,14 +453,13 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return None assert cls._n_main_layers is not None - m = re.match(r"model\.(?:language_model\.)?layers\.(\d+)\.", name) + 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", - "model.language_model.embed_tokens.weight", "model.language_model.norm.weight", ): return None @@ -532,10 +531,7 @@ def set_gguf_parameters(self): 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.startswith("model.language_model."): - name = "model." + name[len("model.language_model."):] - - if name == "lm_head.weight" and self.hparams.get("tie_word_embeddings", False): + if name == "lm_head.weight" and self.hparams.get("tie_word_embeddings", False): return # routed experts diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 5707ac028add..5462d111a30c 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -120,7 +120,7 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) { 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, i), {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); From 5728a4bfc886ceb0a1b4e2c59995c9e34c6de046 Mon Sep 17 00:00:00 2001 From: Chrono Date: Mon, 31 Aug 2026 22:27:41 +0200 Subject: [PATCH 16/34] Skip glm5-next in model saver, fix CRLF --- conversion/glm.py | 4 ++-- src/llama-model-saver.cpp | 1 + src/models/glm5-next.cpp | 2 +- 3 files changed, 4 insertions(+), 3 deletions(-) diff --git a/conversion/glm.py b/conversion/glm.py index 6336c59303d2..c341a440f14c 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -453,7 +453,7 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return None assert cls._n_main_layers is not None - m = re.match(r"model\.layers\.(\d+)\.", name) + 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: @@ -531,7 +531,7 @@ def set_gguf_parameters(self): 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): + if name == "lm_head.weight" and self.hparams.get("tie_word_embeddings", False): return # routed experts 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/models/glm5-next.cpp b/src/models/glm5-next.cpp index 5462d111a30c..49b6cafc3a86 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -120,7 +120,7 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) { 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_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); From c35bddd016c622e0539f80b7cb6e9405e1c51a34 Mon Sep 17 00:00:00 2001 From: Chrono Date: Mon, 31 Aug 2026 22:54:06 +0200 Subject: [PATCH 17/34] Skip glm5-next in sweep --- tests/test-llama-archs.cpp | 3 +++ 1 file changed, 3 insertions(+) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 35a3286e4a1b..e1554638ec1a 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -501,6 +501,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. } From 3bdb2d84109dc460bab27a3b31842de6ecee19c5 Mon Sep 17 00:00:00 2001 From: Chrono Date: Tue, 1 Sep 2026 01:10:11 +0200 Subject: [PATCH 18/34] Remove T4 fallback --- conversion/base.py | 30 ++++++++++++++---------------- conversion/dream.py | 13 ++++++------- conversion/glm.py | 10 ++-------- conversion/laguna.py | 4 ++-- conversion/llada.py | 13 ++++++------- src/llama-vocab.cpp | 2 +- 6 files changed, 31 insertions(+), 41 deletions(-) diff --git a/conversion/base.py b/conversion/base.py index 3f96b03d5c86..daae28e92adc 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1416,22 +1416,21 @@ def does_token_look_special(self, token: str | bytes) -> bool: return seems_special # used for GPT-2 BPE and WordPiece vocabs - def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[list[str], list[int], str]: + def get_vocab_base(self) -> tuple[list[str], list[int], str]: tokens: list[str] = [] toktypes: list[int] = [] - if tokenizer is None: - from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(self.dir_model) - vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) - assert max(tokenizer.vocab.values()) < vocab_size + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model) + vocab_size = self.hparams.get("vocab_size", len(tokenizer.vocab)) # ty: ignore[unresolved-attribute] + assert max(tokenizer.vocab.values()) < vocab_size # ty: ignore[unresolved-attribute] - tokpre = tokpre or self.get_vocab_base_pre(tokenizer) + tokpre = self.get_vocab_base_pre(tokenizer) - reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} - added_vocab = tokenizer.get_added_vocab() + reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in tokenizer.vocab.items()} # ty: ignore[unresolved-attribute] + added_vocab = tokenizer.get_added_vocab() # ty: ignore[unresolved-attribute] - added_tokens_decoder = tokenizer.added_tokens_decoder + added_tokens_decoder = tokenizer.added_tokens_decoder # ty: ignore[unresolved-attribute] for i in range(vocab_size): if i not in reverse_vocab: @@ -1444,7 +1443,7 @@ def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[lis # To avoid unexpected issues - we make sure to normalize non-normalized tokens if not added_tokens_decoder[i].normalized: previous_token = token - token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) + token = tokenizer.decode(tokenizer.encode(token, add_special_tokens=False)) # ty: ignore[unresolved-attribute, invalid-assignment] if previous_token != token: logger.info(f"{repr(previous_token)} is encoded and decoded back to {repr(token)} using AutoTokenizer") @@ -2168,12 +2167,11 @@ def _set_vocab_glmedge(self): special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|endoftext|>"]) # ty: ignore[unresolved-attribute] special_vocab.add_to_gguf(self.gguf_writer) - def _set_vocab_glm(self, tokenizer=None, tokpre: str | None = None): - if tokenizer is None: - from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(self.dir_model) + def _set_vocab_glm(self): + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model) special_vocab = gguf.SpecialVocab(self.dir_model, load_merges=True) - tokens, toktypes, tokpre = self.get_vocab_base(tokenizer, tokpre=tokpre) + tokens, toktypes, tokpre = self.get_vocab_base() self.gguf_writer.add_tokenizer_model("gpt2") self.gguf_writer.add_tokenizer_pre(tokpre) self.gguf_writer.add_token_list(tokens) diff --git a/conversion/dream.py b/conversion/dream.py index dcd1bb9fd993..14f25404d67f 100644 --- a/conversion/dream.py +++ b/conversion/dream.py @@ -13,22 +13,21 @@ class DreamModel(TextModel): model_arch = gguf.MODEL_ARCH.DREAM - def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[list[str], list[int], str]: + def get_vocab_base(self) -> tuple[list[str], list[int], str]: tokens: list[str] = [] toktypes: list[int] = [] - if tokenizer is None: - from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) - vocab_dict = tokenizer.get_vocab() + vocab_dict = tokenizer.get_vocab() # ty: ignore[unresolved-attribute] vocab_size = self.hparams.get("vocab_size", len(vocab_dict)) assert max(vocab_dict.values()) < vocab_size - tokpre = tokpre or self.get_vocab_base_pre(tokenizer) + tokpre = self.get_vocab_base_pre(tokenizer) reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in vocab_dict.items()} - added_vocab = tokenizer.get_added_vocab() + added_vocab = tokenizer.get_added_vocab() # ty: ignore[unresolved-attribute] for i in range(vocab_size): if i not in reverse_vocab: diff --git a/conversion/glm.py b/conversion/glm.py index c341a440f14c..5bc15ab9ac7c 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -429,14 +429,8 @@ def __init__(self, *args, **kwargs): self.hparams.pop("head_dim", None) def set_vocab(self): - from transformers import AutoTokenizer - try: - tokenizer = AutoTokenizer.from_pretrained(self.dir_model) - except ValueError: - # the repo ships only a transformers v5 style tokenizer.json, load it directly - from transformers import PreTrainedTokenizerFast - tokenizer = PreTrainedTokenizerFast(tokenizer_file=str(self.dir_model / "tokenizer.json")) - return self._set_vocab_glm(tokenizer, tokpre="glm5") + # 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) diff --git a/conversion/laguna.py b/conversion/laguna.py index 0757eca12c58..29e0b3d6b318 100644 --- a/conversion/laguna.py +++ b/conversion/laguna.py @@ -45,14 +45,14 @@ def set_vocab(self) -> None: self.gguf_writer.add_eot_token_id(extra[0]) logger.info(f"gguf: registered eot_token_id={extra[0]} from eos list {eos_ids}") - def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[list[str], list[int], str]: + def get_vocab_base(self) -> tuple[list[str], list[int], str]: # is the assistant turn-end (registered as eot below). The # HF tokenizer flags it special=false, so the base classifies it as # USER_DEFINED and llama.cpp renders its text into generated content, # leaking "" and breaking response parsing. It is a control # marker, so promote it to CONTROL: llama.cpp then treats it as # end-of-generation and suppresses its text. - tokens, toktypes, tokpre = super().get_vocab_base(tokenizer, tokpre=tokpre) + tokens, toktypes, tokpre = super().get_vocab_base() for i, tok in enumerate(tokens): if tok == "": toktypes[i] = gguf.TokenType.CONTROL diff --git a/conversion/llada.py b/conversion/llada.py index 03ea421cf0a5..c03607191a76 100644 --- a/conversion/llada.py +++ b/conversion/llada.py @@ -16,22 +16,21 @@ class LLaDAModel(TextModel): model_arch = gguf.MODEL_ARCH.LLADA undo_permute = True - def get_vocab_base(self, tokenizer=None, tokpre: str | None = None) -> tuple[list[str], list[int], str]: + def get_vocab_base(self) -> tuple[list[str], list[int], str]: tokens: list[str] = [] toktypes: list[int] = [] - if tokenizer is None: - from transformers import AutoTokenizer - tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) + from transformers import AutoTokenizer + tokenizer = AutoTokenizer.from_pretrained(self.dir_model, trust_remote_code=True) - vocab_dict = tokenizer.get_vocab() + vocab_dict = tokenizer.get_vocab() # ty: ignore[unresolved-attribute] vocab_size = self.hparams.get("vocab_size", len(vocab_dict)) assert max(vocab_dict.values()) < vocab_size - tokpre = tokpre or self.get_vocab_base_pre(tokenizer) + tokpre = self.get_vocab_base_pre(tokenizer) reverse_vocab = {id_: encoded_tok for encoded_tok, id_ in vocab_dict.items()} - added_vocab = tokenizer.get_added_vocab() + added_vocab = tokenizer.get_added_vocab() # ty: ignore[unresolved-attribute] for i in range(vocab_size): if i not in reverse_vocab: diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 68633d6b33a5..ebd46aa7082e 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2253,11 +2253,11 @@ 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; From 81f95af49427695153b85086b80237455cf3088e Mon Sep 17 00:00:00 2001 From: Chrono Date: Tue, 1 Sep 2026 12:01:31 +0200 Subject: [PATCH 19/34] Review cleanup --- conversion/glm.py | 20 ++++++++------------ 1 file changed, 8 insertions(+), 12 deletions(-) diff --git a/conversion/glm.py b/conversion/glm.py index 5bc15ab9ac7c..29d0881db06c 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -417,14 +417,12 @@ class Glm5NextModel(TextModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - self.n_main_layers = self.hparams["num_hidden_layers"] self.n_nextn_layers = self.hparams.get("num_nextn_predict_layers", 0) self.skip_mtp = self.no_mtp or self.n_nextn_layers == 0 - self.block_count = self.n_main_layers 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.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) self.hparams.pop("head_dim", None) @@ -459,17 +457,15 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return name, gen - def is_kda_layer(self, il: int) -> bool: - if il >= self.n_main_layers: - return False - return self.hparams["layer_types"][il] == "linear_attention" - def set_gguf_parameters(self): hp = self.hparams - n_layer = self.n_main_layers - # the loader reads this array before it knows about NextN, so cover all - hp["num_key_value_heads"] = [0 if self.is_kda_layer(il) else 1 for il in range(self.block_count)] + 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, since the generic loader validates this array's length against the full count before NextN is known + # the NextN entry's value itself is never read + hp["num_key_value_heads"] = n_kv_heads + [1] * (self.block_count - len(n_kv_heads)) super().set_gguf_parameters() self.gguf_writer.add_vocab_size(hp["vocab_size"]) @@ -506,7 +502,7 @@ def set_gguf_parameters(self): 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[:n_layer]]) + self.gguf_writer.add_indexer_types([t == "full" for t in indexer_types]) # mHC assert hp.get("mhc", True) From 1eca274ed67274f9ca52c846438dc07863988172 Mon Sep 17 00:00:00 2001 From: Chrono Date: Tue, 1 Sep 2026 14:05:19 +0200 Subject: [PATCH 20/34] Review suggestions --- conversion/glm.py | 7 ++----- src/models/glm5-next.cpp | 3 --- 2 files changed, 2 insertions(+), 8 deletions(-) diff --git a/conversion/glm.py b/conversion/glm.py index 29d0881db06c..63712e8c7bdc 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -458,16 +458,13 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca 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, since the generic loader validates this array's length against the full count before NextN is known - # the NextN entry's value itself is never read - hp["num_key_value_heads"] = n_kv_heads + [1] * (self.block_count - len(n_kv_heads)) - - super().set_gguf_parameters() + 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) diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 49b6cafc3a86..6fdc76926b5c 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -18,9 +18,6 @@ void llama_model_glm5_next::load_arch_hparams(llama_model_loader & ml) { // the MLA cache holds the compressed latent hparams.n_embd_head_v_full = hparams.n_lora_kv; - ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); - GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all); - for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { if (i >= hparams.n_layer()) { hparams.n_head_kv_arr[i] = 1; From 7de5a8e39be48fab0240df5cc2fa22214507288f Mon Sep 17 00:00:00 2001 From: Chrono Date: Tue, 1 Sep 2026 15:21:18 +0200 Subject: [PATCH 21/34] Defer separate MTP gguf handling to MTP PR, drop filter --- conversion/glm.py | 3 --- src/models/glm5-next.cpp | 7 +++---- 2 files changed, 3 insertions(+), 7 deletions(-) diff --git a/conversion/glm.py b/conversion/glm.py index 63712e8c7bdc..e0901d0c340a 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -441,9 +441,6 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return None name, gen = titem - if name.startswith(("model.visual.", "visual.")): - return None - 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 diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 6fdc76926b5c..f9445f3892ae 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -68,10 +68,9 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) { 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 (TODO) - 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; + // the NextN block is loaded but only used by the MTP graph. + // Separated trunk_only/mtp_only handling TODO with DECODER_MTP graph in the MTP follow up + int mtp_flags = 0; if (!ml.load_mtp) { mtp_flags |= TENSOR_SKIP; } From 9fe9fd71f874f56c2307ead6f63fa50580df7d2f Mon Sep 17 00:00:00 2001 From: Chrono Date: Tue, 1 Sep 2026 19:12:36 +0200 Subject: [PATCH 22/34] Repad n_head_kv --- conversion/glm.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/conversion/glm.py b/conversion/glm.py index e0901d0c340a..5b68f9ac6a88 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -461,6 +461,8 @@ def set_gguf_parameters(self): 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) From a386cd7610f66089c885e2bd2c7ad7eabefc83cc Mon Sep 17 00:00:00 2001 From: Chrono Date: Tue, 1 Sep 2026 19:22:21 +0200 Subject: [PATCH 23/34] kpool init apply --- src/llama-memory-hybrid-idx.cpp | 276 +++++++++++++++++--------------- src/llama-memory-hybrid-idx.h | 33 ++-- src/models/glm5-next.cpp | 14 +- 3 files changed, 172 insertions(+), 151 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index a57e3420eb63..8d75312246b8 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -66,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 @@ -610,6 +611,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) {} @@ -621,7 +647,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, @@ -644,6 +674,7 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( 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)) { + 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(); } @@ -672,6 +703,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; } @@ -700,178 +739,152 @@ void llama_memory_hybrid_idx_context::set_input_qsa( // k-pool DSA indexer (glm5-next) -// Cache the k-pool layout for this 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; - - bool have_new = false; // Whether the ubatch-dependent part is filled. - size_t i_ubatch = SIZE_MAX; // The ubatch this state was computed for. -}; - -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::kpool_state & llama_memory_hybrid_idx_context::kpool_get_state( - uint32_t kpool, const llama_ubatch * ubatch) const { +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"); - if (!kpool_st || kpool_st->i_ubatch != i_cur) { - kpool_st = std::make_unique(); + const uint32_t kpool = mem->get_kpool(); - auto & st = *kpool_st; - st.i_ubatch = i_cur; - st.seqs.resize(LLAMA_MAX_SEQ); + kpool_state st; + st.seqs.resize(LLAMA_MAX_SEQ); - const auto & cells = mem->get_mem_idx()->get_cells(0); + const auto & cells = mem->get_mem_idx()->get_cells(0); - const uint32_t n_kv = get_idx()->get_n_kv(); - const uint32_t n = std::min(n_kv, cells.size()); + // 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); + } + } - // 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()); + } - for (uint32_t i = 0; i < n; ++i) { - if (cells.is_empty(i)) { + 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; } - 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; + 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 (n_seq_cell > 1) { - st.cache_safe = false; + if (ok) { + sq.pools.push_back((uint32_t) j); + j += kpool; + } else { + ++j; } } - 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(); + } - 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); } - } - auto & st = *kpool_st; + return st; + } - if (ubatch != nullptr && !st.have_new) { - // Shared cells cannot cache sequence relative pools. - const bool all_new = !st.cache_safe || (kpool_dirty_batch && i_cur == 0); + // 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]); - } + 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 (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]; + 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; - } + 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; - } + 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; + 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++; - } + 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++; } } - - st.have_new = true; } return st; } -uint32_t llama_memory_hybrid_idx_context::get_n_kpool(uint32_t kpool) const { - return kpool_pad(kpool_get_state(kpool, nullptr).n_pool_real); +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(uint32_t kpool, const llama_ubatch * ubatch) const { - return kpool_get_state(kpool, ubatch).n_new; +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(uint32_t kpool) const { - return kpool_get_state(kpool, nullptr).cache_safe; +bool llama_memory_hybrid_idx_context::get_kpool_cache_safe() const { + return kpool_cur().cache_safe; } 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, uint32_t kpool) const { + 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)); @@ -879,9 +892,10 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, GGML_ASSERT(ggml_backend_buffer_is_host(pool_mask->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(tail_idxs->buffer)); - const uint32_t n_kv = get_idx()->get_n_kv(); + const uint32_t kpool = mem->get_kpool(); + const uint32_t n_kv = get_idx()->get_n_kv(); - const auto & st = kpool_get_state(kpool, ubatch); + const auto & st = kpool_cur(); const uint32_t n_tokens = ubatch->n_tokens; const uint32_t n_pool = (uint32_t) pool_cells->ne[0]; diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index dbeaa45c7713..5b6f8804d11a 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -87,9 +87,12 @@ 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() const { kpool_dirty = false; } + void kpool_clear_dirty() { kpool_dirty = false; } 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 @@ -102,8 +105,9 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { const std::unique_ptr mem_idx; - // Mutable because it is captured and cleared through const contexts. - mutable bool kpool_dirty = false; + const uint32_t n_kpool; + + bool kpool_dirty = false; }; class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { @@ -148,20 +152,19 @@ 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 - // Cache sequence-private complete pools and re-pool only changed pools. - uint32_t get_n_kpool (uint32_t kpool) const; // Padded pool count, where the last pool is always unused. - uint32_t get_n_kpool_new(uint32_t kpool, const llama_ubatch * ubatch) const; // Exact count of new pools. - bool get_kpool_cache_safe(uint32_t kpool) 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; 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, uint32_t kpool) const; + 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 @@ -173,10 +176,14 @@ 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; - // Cached k-pool layout. + // K-pool layouts struct kpool_state; - kpool_state & kpool_get_state(uint32_t kpool, const llama_ubatch * ubatch) const; - mutable std::unique_ptr kpool_st; + 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/models/glm5-next.cpp b/src/models/glm5-next.cpp index f9445f3892ae..65ecee297cc1 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -241,7 +241,7 @@ class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { 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, kpool); + mctx->set_input_kpool(pool_cells, pool_idxs, pool_mask, tail_idxs, gather_mask, gather, new_pool_idxs, new_pool_rep, ubatch); } bool can_reuse(const llm_graph_params & params) override { @@ -255,14 +255,14 @@ class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { bool res = true; res &= k_idxs->ne[0] == params.ubatch.n_tokens; - res &= pool_cells->ne[0] == mctx->get_n_kpool(kpool); + 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(kpool, ¶ms.ubatch); - res &= cache_safe == mctx->get_kpool_cache_safe(kpool); + res &= n_new == mctx->get_n_kpool_new(); + res &= cache_safe == mctx->get_kpool_cache_safe(); return res; } @@ -290,10 +290,10 @@ llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::bui GGML_ASSERT(mctx_idx != nullptr); const uint32_t kpool = hparams.indexer_kpool; - const uint32_t n_pool = mctx_hyb->get_n_kpool(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(kpool, &ubatch); - const bool cache_safe = mctx_hyb->get_kpool_cache_safe(kpool); + 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; From 120eb9e4c02081c45829b09105208a7df8165e8f Mon Sep 17 00:00:00 2001 From: Chrono Date: Wed, 2 Sep 2026 00:03:59 +0200 Subject: [PATCH 24/34] Order by descending score --- src/models/glm5-next.cpp | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 65ecee297cc1..562c45666b80 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -669,7 +669,16 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( 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_cont(ctx0, ggml_top_k(ctx0, score, n_top_pool)); // [n_top_pool, n_tokens] + 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, From d0c5589128f9eebd2951ebc1724b081455d3462a Mon Sep 17 00:00:00 2001 From: Chrono Date: Wed, 2 Sep 2026 11:30:18 +0200 Subject: [PATCH 25/34] Drop guard --- src/models/glm5-next.cpp | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 562c45666b80..fecb7ae07a0c 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -19,10 +19,7 @@ void llama_model_glm5_next::load_arch_hparams(llama_model_loader & ml) { hparams.n_embd_head_v_full = hparams.n_lora_kv; for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { - if (i >= hparams.n_layer()) { - hparams.n_head_kv_arr[i] = 1; - } - hparams.is_recr_impl[i] = i < hparams.n_layer() && hparams.n_head_kv(i) == 0; + hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; } ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); From 8c2893991524fd3c91b02dee0d1005aceefcf0b8 Mon Sep 17 00:00:00 2001 From: Chrono Date: Wed, 2 Sep 2026 14:01:01 +0200 Subject: [PATCH 26/34] read kpool from hparams, clarify kpool cache flags, remove kpool_build_state(nullptr) --- src/llama-memory-hybrid-idx.cpp | 53 ++++++++++++++++++--------------- src/llama-memory-hybrid-idx.h | 18 +++++------ 2 files changed, 38 insertions(+), 33 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 8d75312246b8..0347aaf28431 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -66,8 +66,7 @@ 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 @@ -158,7 +157,7 @@ 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; + kpool_cache_stale = true; } return get_mem_attn()->seq_rm(seq_id, p0, p1); @@ -169,7 +168,7 @@ 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; + kpool_cache_stale = true; } } @@ -178,7 +177,7 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { if (mem_idx) { mem_idx->seq_keep(seq_id); - kpool_dirty = true; + kpool_cache_stale = true; } } @@ -187,7 +186,7 @@ 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; + kpool_cache_stale = true; } } @@ -196,7 +195,7 @@ 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; + kpool_cache_stale = true; } } @@ -649,7 +648,7 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hy ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : 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)); + kpool_states.push_back(kpool_build_layout()); } } @@ -676,15 +675,15 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( 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(); + kpool_stale_batch = mem->kpool_cache_is_stale(); } 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 (i_cur == 0 && kpool_stale_batch && mem != nullptr) { + mem->kpool_cache_clear(); } if (ctx_idx) { @@ -708,7 +707,7 @@ bool llama_memory_hybrid_idx_context::apply() { GGML_ASSERT(i_cur <= kpool_states.size()); kpool_states.resize(i_cur); - kpool_states.push_back(kpool_build_state(&get_ubatch())); + kpool_states.push_back(kpool_build_state(get_ubatch())); } return res; @@ -739,8 +738,8 @@ void llama_memory_hybrid_idx_context::set_input_qsa( // 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 { +// Layout only, used by the full cache context so get_n_kpool() works during graph reserve. +llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kpool_build_layout() const { GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); GGML_ASSERT(get_n_stream() == 1 && "TODO: k-pool indexer with multiple streams"); @@ -811,22 +810,28 @@ llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kp 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; + for (auto & sq : st.seqs) { + sq.is_new.assign(sq.pools.size(), 0); } + return st; +} + +// Layout of the cells as of this ubatch plus which pools it completes or rewrites. +llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kpool_build_state( + const llama_ubatch & ubatch) const { + kpool_state st = kpool_build_layout(); + + const uint32_t kpool = mem->get_kpool(); + // 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); + const bool all_new = !st.cache_safe || (kpool_stale_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 (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) { diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 5b6f8804d11a..09c43d82e97f 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -88,11 +88,12 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { bool blk_bias) const; // The model's indexer pool size. - uint32_t get_kpool() const { return n_kpool; } + uint32_t get_kpool() const { return hparams_idx.indexer_kpool; } - // Sequence edits invalidate cached relative pools. - bool kpool_is_dirty () const { return kpool_dirty; } - void kpool_clear_dirty() { kpool_dirty = false; } + // The pooled keys persist in the idx cache across batches. + // Sequence edits shift the pool grid and stale the cached values. + bool kpool_cache_is_stale() const { return kpool_cache_stale; } + void kpool_cache_clear () { kpool_cache_stale = false; } 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 @@ -105,9 +106,7 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { const std::unique_ptr mem_idx; - const uint32_t n_kpool; - - bool kpool_dirty = false; + bool kpool_cache_stale = false; }; class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { @@ -178,7 +177,8 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // K-pool layouts struct kpool_state; - kpool_state kpool_build_state(const llama_ubatch * ubatch) const; + kpool_state kpool_build_layout() const; + kpool_state kpool_build_state(const llama_ubatch & ubatch) const; const kpool_state & kpool_cur() const; std::vector kpool_states; @@ -186,5 +186,5 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { bool kpool_track = false; // Clear a pending full re-pool only after the first ubatch succeeds - bool kpool_dirty_batch = false; + bool kpool_stale_batch = false; }; From a2f1d20c0da8167d4ea6271b0d139180f9d5a605 Mon Sep 17 00:00:00 2001 From: Chrono Date: Wed, 2 Sep 2026 14:16:01 +0200 Subject: [PATCH 27/34] Add glm5-next support to model saver and add arch test fixture --- src/llama-arch.cpp | 1 + src/llama-model-saver.cpp | 4 +++- tests/test-llama-archs.cpp | 26 ++++++++++++++++++-------- 3 files changed, 22 insertions(+), 9 deletions(-) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 9f91cef60897..cc1256629029 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1151,6 +1151,7 @@ bool llm_arch_supports_sm_tensor(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_QWEN3TTS: case LLM_ARCH_QWEN4EXP: // TODO: fix test-llama-archs return false; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 4fb9d556b151..9eace6805601 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -32,7 +32,6 @@ 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; @@ -298,6 +297,9 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size); + add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, hparams.indexer_kpool); + add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, hparams.indexer_kpool_select_tail); + add_kv(LLM_KV_ATTENTION_INDEXER_INDEX_SHARE_MTP, hparams.indexer_index_share_mtp); add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks); add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true); add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 21f05dbe1acb..2304242b8f59 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -118,6 +118,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 + || arch == LLM_ARCH_GLM5_NEXT || arch == LLM_ARCH_MISTRAL4) { n_embd = 128; n_head = 1; @@ -165,14 +166,19 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE || arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR || - arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3) { + arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_GLM5_NEXT) { GGML_ASSERT(n_layer >= 2); std::vector n_head_per_layer; n_head_per_layer.reserve(n_layer); for (uint32_t il = 0; il < n_layer; il++) { n_head_per_layer.push_back(il == 1 ? 0 : n_head); } - ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head_per_layer); + // GLM5 next KDA heads come from the uniform head count, only head_count_kv is per layer. + if (arch == LLM_ARCH_GLM5_NEXT) { + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head); + } else { + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head_per_layer); + } ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer); } else { ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head); @@ -191,10 +197,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 + || arch == LLM_ARCH_GLM5_NEXT || arch == LLM_ARCH_MISTRAL4) { - ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(576)); + // GLM5 next MLA is nope only, the cache row is the compressed latent alone. + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, arch == LLM_ARCH_GLM5_NEXT ? uint32_t(512) : uint32_t(576)); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512)); - ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); + ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, arch == LLM_ARCH_GLM5_NEXT ? uint32_t(0) : uint32_t(64)); ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192)); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128)); if (arch == LLM_ARCH_DOTS3NOTE) { @@ -249,8 +257,10 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the // indexer head count is independent of the main attention head count. - if (arch == LLM_ARCH_QWEN4EXP) { + if (arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_GLM5_NEXT) { ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8)); // without this the QSA layers fall back to dense and go uncovered ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector(n_layer, 4)); @@ -288,6 +298,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, uint32_t(4)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, true); ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4})); @@ -481,6 +493,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_MIMO2: case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_KIMI_K3: + case LLM_ARCH_GLM5_NEXT: case LLM_ARCH_STEP35: case LLM_ARCH_MISTRAL4: case LLM_ARCH_MELLUM: @@ -525,9 +538,6 @@ 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. } From fe3187de7a7d132742282f2602cf1f8fc3db1d5d Mon Sep 17 00:00:00 2001 From: Chrono Date: Wed, 2 Sep 2026 15:53:02 +0200 Subject: [PATCH 28/34] Review cleanup --- src/llama-context.cpp | 1 - src/llama-memory-hybrid-idx.cpp | 23 +++++++++++++---------- src/llama-memory-hybrid-idx.h | 9 +++++++-- 3 files changed, 20 insertions(+), 13 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index c4273e1e112d..0240c54aa470 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -7,7 +7,6 @@ #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" diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 0347aaf28431..fc8327ddc26f 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -647,8 +647,9 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hy 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())) { - if (mem->get_mem_idx() != nullptr && mem->get_kpool() > 0) { - kpool_states.push_back(kpool_build_layout()); + if (kpool_track()) { + kpool_st = std::make_unique(kpool_build_layout()); + i_kpool = 0; } } @@ -673,7 +674,6 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( 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)) { - kpool_track = mem->get_mem_idx() != nullptr && mem->get_kpool() > 0; // Sequence edits require a full re-pool. kpool_stale_batch = mem->kpool_cache_is_stale(); } @@ -703,16 +703,19 @@ bool llama_memory_hybrid_idx_context::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())); + if (res && kpool_track()) { + kpool_st = std::make_unique(kpool_build_state(get_ubatch())); + i_kpool = i_cur; } return res; } +bool llama_memory_hybrid_idx_context::kpool_track() const { + // Derived from mem instead of being cached. + return mem != nullptr && mem->get_mem_idx() != nullptr && mem->get_kpool() > 0 && !ns_ubatch.empty(); +} + const llama_kv_cache_context * llama_memory_hybrid_idx_context::get_idx() const { return static_cast(ctx_idx.get()); } @@ -870,9 +873,9 @@ llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kp } 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()"); + GGML_ASSERT(kpool_st != nullptr && i_kpool == i_cur && "k-pool state read before apply()"); - return kpool_states[i_cur]; + return *kpool_st; } uint32_t llama_memory_hybrid_idx_context::get_n_kpool() const { diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 09c43d82e97f..c27cadb635ae 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -180,10 +180,15 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { kpool_state kpool_build_layout() const; kpool_state kpool_build_state(const llama_ubatch & ubatch) const; const kpool_state & kpool_cur() const; - std::vector kpool_states; + + // unique_ptr because kpool_state is incomplete here. + std::unique_ptr kpool_st; + + // The ubatch kpool_st was built for, guards against reads before apply. + size_t i_kpool = SIZE_MAX; // Whether this context tracks k-pool states. - bool kpool_track = false; + bool kpool_track() const; // Clear a pending full re-pool only after the first ubatch succeeds bool kpool_stale_batch = false; From 2b533e0950b6b57decf0da01a79521db3612c22d Mon Sep 17 00:00:00 2001 From: Chrono Date: Wed, 2 Sep 2026 19:53:38 +0200 Subject: [PATCH 29/34] Kpool pooled caching clarify --- src/llama-memory-hybrid-idx.cpp | 25 ++++++++++++++++--------- src/llama-memory-hybrid-idx.h | 8 ++++---- src/models/glm5-next.cpp | 1 + 3 files changed, 21 insertions(+), 13 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index fc8327ddc26f..69082cb99aa1 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -157,7 +157,7 @@ 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_cache_stale = true; + mem_idx_stale = true; } return get_mem_attn()->seq_rm(seq_id, p0, p1); @@ -168,7 +168,7 @@ 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_cache_stale = true; + mem_idx_stale = true; } } @@ -177,7 +177,7 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { if (mem_idx) { mem_idx->seq_keep(seq_id); - kpool_cache_stale = true; + mem_idx_stale = true; } } @@ -186,7 +186,7 @@ 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_cache_stale = true; + mem_idx_stale = true; } } @@ -195,7 +195,7 @@ 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_cache_stale = true; + mem_idx_stale = true; } } @@ -675,15 +675,15 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) { // Sequence edits require a full re-pool. - kpool_stale_batch = mem->kpool_cache_is_stale(); + mem_idx_stale_batch = mem->mem_idx_is_stale(); } 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_stale_batch && mem != nullptr) { - mem->kpool_cache_clear(); + if (i_cur == 0 && mem_idx_stale_batch && mem != nullptr) { + mem->mem_idx_stale_clear(); } if (ctx_idx) { @@ -821,6 +821,13 @@ llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kp } // Layout of the cells as of this ubatch plus which pools it completes or rewrites. +// Pool cache lifecycle: +// 1. cpy_k writes each token's key | gate into its idx cache row, pooled slot are zeroed. +// 2. This scan derives the current grouping from mem_idx and marks the pools the ubatch touches or completes as is_new, during decode that's one pool every kpool tokens, zero elsewise. +// 3. The graph pools only the is_new pools and set_rows each result into the pooled slot of the pool's last member row. +// 4. All pools are gathered in one get_rows via pool_cells, fresh ones just written, older ones from whatever batch last wrote them. +// Any seq_* edit regroups the pools, so it sets mem_idx_stale and the first ubatch of the next batch re-pools everything from the still-valid key | gate rows, rewriting the (possibly different) rep rows. +// Orphaned pooled slots are never cleared, a slot is only ever read through pool_cells, which is derived from the current grouping every ubatch. llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kpool_build_state( const llama_ubatch & ubatch) const { kpool_state st = kpool_build_layout(); @@ -828,7 +835,7 @@ llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kp const uint32_t kpool = mem->get_kpool(); // Pools touched by this ubatch are re-pooled, shared cells cannot cache sequence relative pools. - const bool all_new = !st.cache_safe || (kpool_stale_batch && i_cur == 0); + const bool all_new = !st.cache_safe || (mem_idx_stale_batch && i_cur == 0); std::vector> upos(LLAMA_MAX_SEQ); if (!all_new) { diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index c27cadb635ae..27069d8bd46e 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -92,8 +92,8 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { // The pooled keys persist in the idx cache across batches. // Sequence edits shift the pool grid and stale the cached values. - bool kpool_cache_is_stale() const { return kpool_cache_stale; } - void kpool_cache_clear () { kpool_cache_stale = false; } + bool mem_idx_is_stale() const { return mem_idx_stale; } + void mem_idx_stale_clear () { mem_idx_stale = false; } 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 @@ -106,7 +106,7 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { const std::unique_ptr mem_idx; - bool kpool_cache_stale = false; + bool mem_idx_stale = false; }; class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { @@ -191,5 +191,5 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { bool kpool_track() const; // Clear a pending full re-pool only after the first ubatch succeeds - bool kpool_stale_batch = false; + bool mem_idx_stale_batch = false; }; diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index fecb7ae07a0c..4806918859ed 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -602,6 +602,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( 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); + // View into the persistent pooled slots of the idx cache. Guarded by mem_idx_stale. 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)); From ff6be954fff86243c6413bcb79dfcf1f4c3e5d69 Mon Sep 17 00:00:00 2001 From: Chrono Date: Thu, 3 Sep 2026 12:22:03 +0200 Subject: [PATCH 30/34] Add multi stream support --- src/llama-kv-cache.cpp | 10 ++++ src/llama-kv-cache.h | 6 ++ src/llama-memory-hybrid-idx.cpp | 98 +++++++++++++++++++++++---------- src/llama-model.cpp | 3 - src/models/glm5-next.cpp | 22 ++++---- 5 files changed, 95 insertions(+), 44 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index fd7ce0bb6e4e..e45b28142490 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1248,6 +1248,12 @@ const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const { return v_cells[seq_to_stream[seq_id]]; } +uint32_t llama_kv_cache::get_stream(llama_seq_id seq_id) const { + GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size()); + + return seq_to_stream[seq_id]; +} + uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const { uint32_t result = 0; @@ -2795,6 +2801,10 @@ ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il) cons return kv->get_k(ctx, il, n_kv, sinfos[i_cur]); } +ggml_tensor * llama_kv_cache_context::get_k_storage(int32_t il) const { + return kv->get_k_storage(il); +} + ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) const { return kv->get_v(ctx, il, n_kv, sinfos[i_cur]); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index c4d8699def12..b6eb47e3a0f4 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -168,6 +168,9 @@ class llama_kv_cache : public llama_memory_i { const llama_kv_cells & get_cells(llama_seq_id seq_id) const; + // The stream holding seq_id's cells. + uint32_t get_stream(llama_seq_id seq_id) const; + // state_read, plus the cells the restored tokens were placed in // a cache that mirrors another one (the qwen4exp indexer) must not search for its own cells: two searches agree only by luck // sinfos_out: if set, filled with the layout used; a stream with no cells leaves an empty entry @@ -398,6 +401,9 @@ class llama_kv_cache_context : public llama_memory_context_i { ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il) const; + // The full K storage tensor of the layer, spanning all streams. + ggml_tensor * get_k_storage(int32_t il) const; + // store k_cur and v_cur in the cache based on the provided head location // note: the heads in k_cur and v_cur should be laid out contiguously in memory // - k_cur [n_embd_head_k, n_head_k, n_tokens] diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 69082cb99aa1..8c9bbf1815eb 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -614,7 +614,8 @@ static std::vector llama_memory_hybrid_idx_ns(const llama_kv_cache::sl 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 + uint32_t strm = 0; // Stream holding this sequence's cells + std::vector> cells; // Position and stream local cell pairs, sorted by position. std::vector pools; std::vector is_new; }; @@ -744,37 +745,56 @@ void llama_memory_hybrid_idx_context::set_input_qsa( // Layout only, used by the full cache context so get_n_kpool() works during graph reserve. llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kpool_build_layout() 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); + const auto * kv = mem->get_mem_idx(); + const uint32_t n_stream_kv = kv->get_n_stream(); - // 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); - } - } + if (n_stream_kv == 1) { + const auto & cells = kv->get_cells(0); - for (uint32_t i = 0; i < cells.size(); ++i) { - if (cells.is_empty(i)) { - continue; + // 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); + } } - 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; + + 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; } } - if (n_seq_cell > 1) { - st.cache_safe = false; + } else { + // When kv is non unified, one stream per sequence, so streams never share cells. Cell indices stay stream-local. + for (llama_seq_id s = 0; s < (llama_seq_id) n_stream_kv; ++s) { + const auto & cells = kv->get_cells(s); + if (cells.seq_pos_min(s) < 0) { + continue; + } + auto & sq = st.seqs[s]; + sq.strm = kv->get_stream(s); + for (uint32_t i = 0; i < cells.size(); ++i) { + if (!cells.is_empty(i) && cells.seq_has(i, s)) { + sq.cells.emplace_back(cells.pos_get(i), i); + } + } } } @@ -901,7 +921,6 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, 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)); @@ -933,14 +952,29 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, } } + const uint32_t kv_size = mem->get_mem_idx()->get_size(); + const uint32_t n_stream_kv = mem->get_mem_idx()->get_n_stream(); + + auto gcell = [&](const kpool_state::seq & sq, uint32_t cell) { + return (int64_t) sq.strm*kv_size + cell; + }; + + // Sequences present in this ubatch, pools of absent sequences must fall on the scatter sentinel row. + std::vector seq_in_ub(LLAMA_MAX_SEQ, 0); + for (uint32_t i = 0; i < n_tokens; ++i) { + for (int32_t k = 0; k < ubatch->n_seq_id[i]; ++k) { + seq_in_ub[ubatch->seq_id[i][k]] = 1; + } + } + // Use the first ubatch cell for padded gathers. - uint32_t dummy_cell = 0; + int64_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; + dummy_cell = gcell(sq, it->second); } // Gather maps padding to a real cell and masks it separately. @@ -974,6 +1008,9 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, 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(); + + const bool inert = !gather && n_stream_kv > 1 && !seq_in_ub[s]; + 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(); @@ -981,19 +1018,20 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, // 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; + pcell[ip] = (int32_t) gcell(sq, rep); for (uint32_t k = 0; k < kpool; ++k) { - pidx[(size_t) ip*kpool + k] = (int32_t) sq.cells[j + k].second; + pidx[(size_t) ip*kpool + k] = inert ? sentinel : + (int32_t) (gather ? gcell(sq, sq.cells[j + k].second) : (int64_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; + nidx[(size_t) i_new*kpool + k] = (int32_t) gcell(sq, sq.cells[j + k].second); } if (nrep != nullptr) { - nrep[i_new] = (int64_t) rep; + nrep[i_new] = gcell(sq, rep); } ++i_new; } @@ -1005,7 +1043,7 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, 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; + pcell[ip] = (int32_t) dummy_cell; // pool_cells always addresses the K storage for (uint32_t k = 0; k < kpool; ++k) { pidx[(size_t) ip*kpool + k] = sentinel; } @@ -1059,7 +1097,7 @@ void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, 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; + cell = (int32_t) (gather ? gcell(sq, it->second) : (int64_t) it->second); real = true; } } diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 1166ee3e2187..6e194ec99512 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2313,9 +2313,6 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } 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) { diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 4806918859ed..74abad4c5527 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -597,14 +597,15 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( 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 * k_store = mctx_lid->get_k_storage(il); // [3*n_embd_indexer, kv_size, n_stream] + GGML_ASSERT(k_store->ne[0] == 3*n_embd_indexer); + const int64_t n_cells = k_store->ne[1]*k_store->ne[2]; + const int64_t n_kv = mctx_lid->get_n_kv(); - ggml_tensor * kg_all = ggml_view_2d(ctx0, k_all, 2*n_embd_indexer, n_kv, k_all->nb[2], 0); + ggml_tensor * kg_all = ggml_view_2d(ctx0, k_store, 2*n_embd_indexer, n_cells, k_store->nb[1], 0); // View into the persistent pooled slots of the idx cache. Guarded by mem_idx_stale. - 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_all = ggml_view_2d(ctx0, k_store, n_embd_indexer, n_cells, k_store->nb[1], + ggml_row_size(k_store->type, 2*n_embd_indexer)); ggml_tensor * pooled_new = nullptr; // Pool only entries completed by this ubatch. @@ -713,7 +714,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( 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); + GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[1]*kq_mask->ne[2]*kq_mask->ne[3] == n_tokens); sel = ggml_add(ctx0, sel, ggml_reshape_2d(ctx0, kq_mask, n_kv, n_tokens)); cb(sel, "indexer_sel", il); @@ -778,11 +779,10 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( 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 * k = mctx_mla->get_k_storage(il); // [kv_lora_rank, kv_size, n_stream] + GGML_ASSERT(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 * rows = ggml_view_2d(ctx0, k, k->ne[0], k->ne[1]*k->ne[2], k->nb[1], 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); From 1b564d2dd826cfc9c5b9f84e025f22c154680818 Mon Sep 17 00:00:00 2001 From: Chrono Date: Thu, 3 Sep 2026 13:23:01 +0200 Subject: [PATCH 31/34] Finish Rebase --- src/models/glm5-next.cpp | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 74abad4c5527..1e4407b29787 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -22,7 +22,7 @@ void llama_model_glm5_next::load_arch_hparams(llama_model_loader & ml) { 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_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); 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); @@ -160,7 +160,7 @@ void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) { 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_ff_exp = hparams.n_ff_exp(i); 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); @@ -809,7 +809,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( 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); + out = build_attn_mha(q_absorbed, k, v, nullptr, mask, nullptr, layer.wv_b, 0, kq_scale, il); } cb(out, "kqv_out", il); From 99fdaab443ac9311430e625fb99d66f935a2d842 Mon Sep 17 00:00:00 2001 From: Chrono Date: Thu, 3 Sep 2026 14:15:38 +0200 Subject: [PATCH 32/34] Sparse FA fir DSA prefill --- src/models/glm5-next.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 1e4407b29787..e3780032860d 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -809,7 +809,7 @@ ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( 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, 0, kq_scale, il); + out = build_attn_mha(q_absorbed, k, v, nullptr, mask, nullptr, layer.wv_b, inp_kpool->n_sel, kq_scale, il); } cb(out, "kqv_out", il); From 5c4bd50b1fb1daabe2644f9cc05e8da0fdb61884 Mon Sep 17 00:00:00 2001 From: Chrono Date: Fri, 4 Sep 2026 01:37:16 +0200 Subject: [PATCH 33/34] Const --- tools/mtmd/mtmd-image.cpp | 2 +- tools/mtmd/mtmd-image.h | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index d8f2d887c33f..a11215509323 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -792,7 +792,7 @@ mtmd_image_preproc_out mtmd_image_preprocessor_dyn_size::preprocess(const clip_i // 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) { +mtmd_image_preproc_out mtmd_image_preprocessor_glm5v::preprocess(const clip_image_u8 & img) const { GGML_ASSERT(hparams.image_min_pixels > 0 && hparams.image_max_pixels > 0); const int64_t factor = hparams.patch_size * hparams.n_merge; diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index a074e4eed331..4fa6207d03be 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -126,7 +126,7 @@ struct mtmd_image_preprocessor_dyn_size : mtmd_image_preprocessor { // 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; + mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const 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 From 8134115f88ed8018474e7db69afcfe97fb097fc4 Mon Sep 17 00:00:00 2001 From: timkhronos Date: Fri, 4 Sep 2026 23:25:35 +0200 Subject: [PATCH 34/34] Update llama-model.cpp to fix rebase error --- src/llama-model.cpp | 1 + 1 file changed, 1 insertion(+) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 7942458e4bac..fa5cb0b02a80 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2371,6 +2371,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* filter_attn */ std::move(filter_attn), /* filter_recr */ std::move(filter_recr), /* filter_idx */ std::move(filter_idx)); + } break; case LLM_ARCH_HY_V4: { if (hparams.indexer_top_k == 0) {