From a505a75878db9ec69e5cf48da4002529e911d840 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 26 Aug 2026 06:30:38 +0000 Subject: [PATCH 01/36] llama: load glm5next hparams and tensors Metadata and tensor loading only. The graph entry point throws, as qwen4exp did at the same stage. kda.gate_lower_bound is read as required: kimi-k3 selects the softplus branch when it is absent, which is a different function rather than a missing clamp. The absorbed MLA projections are 3D, so glm5next joins bailingmoe3 in the MXFP4 carve-out that would otherwise quantize them as expert tensors. --- conversion/__init__.py | 3 + conversion/glm5next.py | 258 +++++++++++++++++++++++++++++++++ gguf-py/gguf/constants.py | 67 +++++++++ gguf-py/gguf/gguf_writer.py | 3 + gguf-py/gguf/tensor_mapping.py | 28 ++++ src/llama-arch.cpp | 4 + src/llama-arch.h | 2 + src/llama-hparams.h | 2 + src/llama-model.cpp | 6 +- src/llama-model.h | 1 + src/llama-quant.cpp | 10 +- src/models/glm5next.cpp | 214 +++++++++++++++++++++++++++ src/models/models.h | 8 + 13 files changed, 600 insertions(+), 6 deletions(-) create mode 100644 conversion/glm5next.py create mode 100644 src/models/glm5next.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 8de97e95969..6f23240deb1 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -103,6 +103,8 @@ "Glm4MoeLiteForCausalLM": "glm", "Glm4vForConditionalGeneration": "glm", "Glm4vMoeForConditionalGeneration": "glm", + "Glm5NextForCausalLM": "glm5next", + "Glm5NextForConditionalGeneration": "glm5next", "GlmForCausalLM": "chatglm", "GlmMoeDsaForCausalLM": "glm", "GlmOcrForConditionalGeneration": "glm", @@ -293,6 +295,7 @@ "Gemma4UnifiedForConditionalGeneration": "gemma", "Glm4vForConditionalGeneration": "qwen3vl", "Glm4vMoeForConditionalGeneration": "qwen3vl", + "Glm5NextForConditionalGeneration": "glm5next", "Glm5vForConditionalGeneration": "kimivl", "GlmOcrForConditionalGeneration": "qwen3vl", "GlmasrModel": "ultravox", diff --git a/conversion/glm5next.py b/conversion/glm5next.py new file mode 100644 index 00000000000..24290041e6d --- /dev/null +++ b/conversion/glm5next.py @@ -0,0 +1,258 @@ +from __future__ import annotations + +import re + +from typing import Callable, Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import LazyTorchTensor, ModelBase, TextModel, gguf +from .qwen3vl import Glm4VVisionModel + + +@ModelBase.register("Glm5NextForConditionalGeneration", "Glm5NextForCausalLM") +# [TAG_HF_EXAMPLE_MISSING] +class Glm5NextModel(TextModel): + """GLM-5.3-Flash text tower: hybrid KDA + DSA attention, nope-only MLA, mHC + hyper-connections, and a NextN block with its own DSA attention and indexer. + """ + + model_arch = gguf.MODEL_ARCH.GLM5NEXT + supports_mtp_export = True + + _experts: list[dict[str, Tensor]] | None = None + _main_layers: int | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + nextn_layers = 0 if self.no_mtp else (self.hparams.get("num_nextn_predict_layers", 0) or 0) + self.block_count = self.hparams["num_hidden_layers"] + nextn_layers + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + # two independent spellings of the same partition; disagreement means an + # unexpected config + from_types = {il for il, t in enumerate(self.hparams["layer_types"]) if t == "deepseek_sparse_attention"} + from_list = set(self.hparams["linear_attn_config"]["full_attn_layers"]) + if from_types != from_list: + raise ValueError(f"layer_types picks DSA layers {sorted(from_types)} but full_attn_layers says {sorted(from_list)}") + self._full_attn_layers = from_types + + dense_lead = self.hparams["first_k_dense_replace"] + expected = ["dense"] * dense_lead + ["sparse"] * (self.hparams["num_hidden_layers"] - dense_lead) + if self.hparams["mlp_layer_types"] != expected: + raise ValueError("mlp_layer_types does not match first_k_dense_replace") + + def index_tensors(self, remote_hf_model_id: str | None = None): + # runs before TextModel.__init__ has hoisted text_config to the root + hp = self.hparams.get("text_config", self.hparams) + type(self)._main_layers = hp["num_hidden_layers"] + return super().index_tensors(remote_hf_model_id=remote_hf_model_id) + + def set_vocab(self): + self._set_vocab_glm() + + def is_full_attention(self, bid: int) -> bool: + return bid >= self.hparams["num_hidden_layers"] or bid in self._full_attn_layers + + # -- metadata --------------------------------------------------------- + + def set_gguf_parameters(self): + hp = self.hparams + linear_cfg = hp["linear_attn_config"] + + # checked here, not in the loader: head_count_kv is overwritten below with + # the per-layer 1/0 recurrence marker + if hp["num_attention_heads"] != hp.get("num_key_value_heads"): + raise ValueError("glm5next expects MHA-shaped head counts before MLA absorption") + if hp["qk_rope_head_dim"] != 0 or not hp.get("mla_use_nope"): + raise ValueError("glm5next is nope-only: qk_rope_head_dim must be 0 and mla_use_nope true") + if linear_cfg["num_heads"] != hp["num_attention_heads"]: + raise ValueError("glm5next KDA and full attention are expected to share a head count") + if not hp.get("mhc"): + raise ValueError("glm5next without mHC is not supported") + if hp["index_topk"] % hp["index_kpool"] != 0: + raise ValueError("glm5next index_topk must be a whole number of kpool pools") + + # no GGUF key carries these and the graph cannot express them off, so refuse + # rather than write a silently wrong file + if not hp.get("index_kpool_compress"): + raise ValueError("glm5next without the indexer kpool compressor is not supported") + if not hp.get("index_kpool_always_select_tail"): + raise ValueError("glm5next without always-select-tail kpool is not supported") + if not hp.get("indexer_rope_interleave"): + raise ValueError("glm5next without interleaved indexer rope is not supported") + if set(hp["indexer_types"]) != {"full"}: + raise ValueError("glm5next expects every indexer to be full") + + # drop both: head_dim is 0 in the config, head_count_kv is written as a + # per-layer array below + hp.pop("head_dim", None) + hp.pop("num_key_value_heads", None) + + super().set_gguf_parameters() + + self.gguf_writer.add_vocab_size(hp["vocab_size"]) + + # n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-k3 and bailingmoe3 + self.gguf_writer.add_head_count_kv( + [1 if self.is_full_attention(il) else 0 for il in range(self.block_count)]) + + # --- MLA --- + kv_lora_rank = hp["kv_lora_rank"] + qk_rope_head_dim = 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_head_dim) + self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim) + self.gguf_writer.add_value_length(kv_lora_rank) + self.gguf_writer.add_key_length_mla(hp["qk_nope_head_dim"] + qk_rope_head_dim) + self.gguf_writer.add_value_length_mla(hp["v_head_dim"]) + + # indexer k_norm is a LayerNorm with bias at a fixed 1e-6, not the model's + # RMS eps. glm-dsa omits this key and runs that norm at eps 0 + self.gguf_writer.add_layer_norm_eps(1e-6) + + # --- KDA --- + self.gguf_writer.add_ssm_conv_kernel(linear_cfg["short_conv_kernel_size"]) + self.gguf_writer.add_kda_head_dim(linear_cfg["head_dim"]) + # not a clamp: scales the sigmoid decay gate. required, a missing key + # silently selects the softplus branch instead + self.gguf_writer.add_kda_gate_lower_bound(linear_cfg["gate_lower_bound"]) + + # --- DSA indexer --- + 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"]) + + # --- mHC --- + 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 --- + n_ff_exp = hp["moe_intermediate_size"] + self.gguf_writer.add_expert_feed_forward_length(n_ff_exp) + self.gguf_writer.add_expert_shared_feed_forward_length(n_ff_exp * hp["n_shared_experts"]) + self.gguf_writer.add_expert_shared_count(hp["n_shared_experts"]) + self.gguf_writer.add_leading_dense_block_count(hp["first_k_dense_replace"]) + self.gguf_writer.add_expert_weights_scale(hp["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(hp["norm_topk_prob"]) + + # one limit for the whole model. no dense-FFN clamp key exists, so the + # expert arrays are sized for every layer to cover the leading dense ones + swiglu_limit = float(hp["swiglu_limit"]) + self.gguf_writer.add_swiglu_clamp_exp([swiglu_limit] * self.block_count) + self.gguf_writer.add_swiglu_clamp_shexp([swiglu_limit] * self.block_count) + + if not self.no_mtp and (nextn_layers := hp.get("num_nextn_predict_layers", 0)): + self.gguf_writer.add_nextn_predict_layers(nextn_layers) + + # -- tensors ---------------------------------------------------------- + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + if (titem := super().filter_tensors(item)) is None: + return None + name, gen = titem + + assert cls._main_layers is not None + is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._main_layers + + if is_mtp and cls.no_mtp: + return None + if cls.mtp_only and not is_mtp and name not in ( + "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", + ): + return None + + return name, gen + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # --- KDA conv1d: HF [d_inner, 1, d_conv] -> ggml ne [d_conv, 1, d_inner, 1] --- + if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")): + d_inner = data_torch.shape[0] + d_conv = data_torch.shape[-1] + data_torch = data_torch.reshape(1, d_inner, 1, d_conv) + + # ssm_a holds -exp(A_log), the kimi-k3 convention (bailingmoe3 stores + # +exp(A_log)); the wrong sign turns decay into an unchecked growing state + if name.endswith(".A_log"): + # eager: the sign is the point of the check, and A_log is one per head + decay = LazyTorchTensor.to_eager(torch.exp(data_torch.float())) + if not bool(torch.isfinite(decay).all() and (decay > 0).all()): + raise ValueError(f"{name}: exp(A_log) must be finite and positive") + data_torch = -decay + + # dt_bias -> the name SSM_DT's mapping expects + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + + # bare tensors in the checkpoint, but the GGUF names carry .weight + if re.search(r"\.hc_(attn|ffn)_(fn|base|scale)$", name) or name.endswith( + (".index_kpool_compress_gate", ".index_kpool_compress_ape")): + name += ".weight" + + # --- 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 weight_name in ("down_proj", "gate_proj", "up_proj"): + tensors = [] + for expert_id in range(n_experts): + expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{weight_name}.weight" + tensors.append(self._experts[bid].pop(expert_name)) + merged_name = f"model.layers.{bid}.mlp.experts.{weight_name}.weight" + yield from super().modify_tensors(torch.stack(tensors, dim=0), merged_name, bid) + return + + # --- MLA absorption --- + if name.endswith(".kv_b_proj.weight"): + assert bid is not None + 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), self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_B, bid), bid) + yield from super().modify_tensors(v_b, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V_B, bid), bid) + return + + yield from super().modify_tensors(data_torch, name, bid) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + # learned position table, one row per pooled key; pinned for the same + # reason POS_EMBD is in base.py + if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_APE, bid): + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts is not None: + experts = [name for layer in self._experts for name in layer] + if experts: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("Glm5NextForConditionalGeneration") +# [TAG_HF_EXAMPLE_MISSING] +class Glm5NextVisionModel(Glm4VVisionModel): + """The GLM-4.5V ViT under a `model.visual.` prefix, registered so the mmproj + can be produced. The clip graph is not ported yet: glm4v.cpp normalises the + patch embeddings unconditionally but this tower has no post_conv_layernorm, + and vision_config.swiglu_limit has no key. + """ diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f236a5d2c98..bc38ba87a85 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -220,6 +220,7 @@ 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" # glm5next class HyperConnection: COUNT = "{arch}.hyper_connection.count" @@ -540,6 +541,7 @@ class MODEL_ARCH(IntEnum): GLM4 = auto() GLM4_MOE = auto() GLM_DSA = auto() + GLM5NEXT = auto() BITNET = auto() T5 = auto() T5ENCODER = auto() @@ -1263,6 +1265,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GLM4: "glm4", MODEL_ARCH.GLM4_MOE: "glm4moe", MODEL_ARCH.GLM_DSA: "glm-dsa", + MODEL_ARCH.GLM5NEXT: "glm5next", MODEL_ARCH.BITNET: "bitnet", MODEL_ARCH.T5: "t5", MODEL_ARCH.T5ENCODER: "t5encoder", @@ -3871,6 +3874,70 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.GLM5NEXT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.FFN_NORM, + # mHC, layered on top of the per-layer norms above + 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.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_G_A, + MODEL_TENSOR.SSM_G_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + # DSA (MLA full-attention layers) + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_OUT, + # DSA indexer, with the kpool key compressor + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE, + MODEL_TENSOR.INDEXER_COMPRESSOR_APE, + # FFN: dense on the leading blocks, MoE elsewhere + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + 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, + # NextN/MTP, a full DSA decoder layer with its own indexer + 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, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index d8a96a27bdd..276432623a4 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -812,6 +812,9 @@ def add_indexer_block_size(self, block_size: int) -> None: def add_indexer_local_blocks(self, local_blocks: int) -> None: self.add_uint32(Keys.Attention.Indexer.LOCAL_BLOCKS.format(arch=self.arch), local_blocks) + def add_indexer_kpool(self, kpool: int) -> None: + self.add_uint32(Keys.Attention.Indexer.KPOOL.format(arch=self.arch), kpool) + def add_indexer_types(self, value: Sequence[bool]) -> None: key = Keys.Attention.Indexer.TYPES.format(arch=self.arch) self.add_array(key, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index ef580518e97..050d96391ef 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2680,6 +2680,34 @@ class TensorNameMap: "model.layers.{bid}.post_attention_layernorm", ), }, + MODEL_ARCH.GLM5NEXT: { + # the converter appends the .weight the checkpoint omits, so these + # match through the usual try_suffixes path + MODEL_TENSOR.HC_ATTN_FN: ( + "model.layers.{bid}.hc_attn_fn", + ), + MODEL_TENSOR.HC_ATTN_BASE: ( + "model.layers.{bid}.hc_attn_base", + ), + MODEL_TENSOR.HC_ATTN_SCALE: ( + "model.layers.{bid}.hc_attn_scale", + ), + MODEL_TENSOR.HC_FFN_FN: ( + "model.layers.{bid}.hc_ffn_fn", + ), + MODEL_TENSOR.HC_FFN_BASE: ( + "model.layers.{bid}.hc_ffn_base", + ), + MODEL_TENSOR.HC_FFN_SCALE: ( + "model.layers.{bid}.hc_ffn_scale", + ), + MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: ( + "model.layers.{bid}.self_attn.indexer.index_kpool_compress_gate", + ), + MODEL_TENSOR.INDEXER_COMPRESSOR_APE: ( + "model.layers.{bid}.self_attn.indexer.index_kpool_compress_ape", + ), + }, } mapping: dict[str, tuple[MODEL_TENSOR, str]] diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index eecf444fcf3..0dbb04fe7bd 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -83,6 +83,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GLM4, "glm4" }, { LLM_ARCH_GLM4_MOE, "glm4moe" }, { LLM_ARCH_GLM_DSA, "glm-dsa" }, + { LLM_ARCH_GLM5NEXT, "glm5next" }, { LLM_ARCH_BITNET, "bitnet" }, { LLM_ARCH_T5, "t5" }, { LLM_ARCH_T5ENCODER, "t5encoder" }, @@ -283,6 +284,7 @@ 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_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" }, @@ -1011,6 +1013,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_MINIMAX_01: + case LLM_ARCH_GLM5NEXT: return true; default: return false; @@ -1074,6 +1077,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_GLM5NEXT: case LLM_ARCH_QWEN3TTS: return false; default: diff --git a/src/llama-arch.h b/src/llama-arch.h index 7159e23bf7a..cadbdfff7ef 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -88,6 +88,7 @@ enum llm_arch { LLM_ARCH_GLM4, LLM_ARCH_GLM4_MOE, LLM_ARCH_GLM_DSA, + LLM_ARCH_GLM5NEXT, LLM_ARCH_BITNET, LLM_ARCH_T5, LLM_ARCH_T5ENCODER, @@ -288,6 +289,7 @@ 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_OUTPUT_GROUP_COUNT, LLM_KV_ATTENTION_OUTPUT_LORA_RANK, LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, diff --git a/src/llama-hparams.h b/src/llama-hparams.h index c3c14292c32..bacef90a66e 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -252,6 +252,8 @@ struct llama_hparams { uint32_t indexer_n_head = 0; uint32_t indexer_head_size = 0; uint32_t indexer_top_k = 0; + // glm5next: the indexer scores pools of this many keys instead of single keys + uint32_t indexer_kpool = 0; // MSA uint32_t indexer_block_size = 0; uint32_t indexer_local_blocks = 0; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c34700ff563..462982d1bac 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -201,6 +201,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_deepseek4(params); case LLM_ARCH_GLM_DSA: return new llama_model_glm_dsa(params); + case LLM_ARCH_GLM5NEXT: + return new llama_model_glm5next(params); case LLM_ARCH_MISTRAL4: return new llama_model_mistral4(params); case LLM_ARCH_CHATGLM: @@ -941,6 +943,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_288B_A19B: return "288B.A19B"; case LLM_TYPE_300B_A47B: return "300B.A47B"; case LLM_TYPE_310B_A15B: return "310B.A15B"; + case LLM_TYPE_313B_A17B: return "313B.A17B"; case LLM_TYPE_355B_A32B: return "355B.A32B"; case LLM_TYPE_397B_A17B: return "397B.A17B"; case LLM_TYPE_685B_A37B: return "685B.A37B"; @@ -2441,7 +2444,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_recr = [&](uint32_t il) { return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; - } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01) { + } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01 || arch == LLM_ARCH_GLM5NEXT) { filter_attn = [&](uint32_t il) { return il < hparams.n_layer() && !hparams.is_recr(il); }; @@ -2758,6 +2761,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_GLM5NEXT: 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 44bd9675754..1654b54d33c 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -143,6 +143,7 @@ enum llm_type { LLM_TYPE_288B_A19B, // dots3-note LLM_TYPE_300B_A47B, // Ernie MoE big LLM_TYPE_310B_A15B, // /MiMo-V2-Flash + LLM_TYPE_313B_A17B, // GLM-5.3-Flash LLM_TYPE_355B_A32B, // GLM-4.5 LLM_TYPE_397B_A17B, // Qwen3.5 LLM_TYPE_685B_A37B, // DeepSeek V3.2 diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 20252815d5c..f93f95fddab 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -475,11 +475,11 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type // MoE tensors -> MXFP4 // other tensors -> Q8_0 // MLA projection tensors are also 3D, so match expert tensor roles explicitly. - const bool is_bailingmoe3_expert = arch == LLM_ARCH_BAILINGMOE3 && - (category == tensor_category::FFN_UP || - category == tensor_category::FFN_GATE || - category == tensor_category::FFN_DOWN); - if (tensor->ne[2] > 1 && (arch != LLM_ARCH_BAILINGMOE3 || is_bailingmoe3_expert)) { + const bool has_3d_mla = arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_GLM5NEXT; + const bool is_expert = category == tensor_category::FFN_UP || + category == tensor_category::FFN_GATE || + category == tensor_category::FFN_DOWN; + if (tensor->ne[2] > 1 && (!has_3d_mla || is_expert)) { new_type = GGML_TYPE_MXFP4; } else { new_type = GGML_TYPE_Q8_0; diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp new file mode 100644 index 00000000000..cdb6dde1cb3 --- /dev/null +++ b/src/models/glm5next.cpp @@ -0,0 +1,214 @@ +#include "models.h" + +// +// GLM-5.3-Flash: hybrid KDA (linear) + DSA (nope-only MLA) attention, mHC +// hyper-connections, and a NextN block that is a full DSA decoder layer. +// +// ssm_a holds -exp(A_log), the kimi-k3 convention, so the decay gate reads +// exp(A_log) back as -ssm_a; bailingmoe3 stores +exp(A_log). indistinguishable at +// load time, so conversion/glm5next.py is the only place the sign is checked. +// + +void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + // indexer k_norm is a LayerNorm with bias; without this key it runs at eps 0 + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); + + 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_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); + GGML_ASSERT(hparams.n_lora_q > 0 && "glm5next requires a q LoRA"); + GGML_ASSERT(hparams.n_rot() == 0 && "glm5next MLA is nope-only"); + + // KDA + 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); + // required: absent, kimi-k3 selects the softplus branch, a different + // function, not a missing clamp + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); + GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f); + + // DSA indexer + 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); + GGML_ASSERT(hparams.indexer_kpool > 0); + GGML_ASSERT(hparams.indexer_top_k % hparams.indexer_kpool == 0); + + // 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 > 0); + + // MoE + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); + 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); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false); + + if (hparams.n_ff_shexp == 0) { + hparams.n_ff_shexp = hparams.n_ff_exp * std::max(1u, hparams.n_expert_shared); + } + + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all); + + // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-k3 and bailingmoe3. + // a scalar head_count_kv would make every layer look like DSA, so require both + uint32_t n_recr = 0; + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0; + n_recr += hparams.is_recr_impl[il]; + } + GGML_ASSERT(n_recr > 0 && n_recr < hparams.n_layer() && "glm5next needs a per-layer attention.head_count_kv array"); + + // every glm5next indexer is full; glm-dsa gates its indexer on this + // predicate and the generic loader only zero-fills the array + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + hparams.is_indexer_full_impl[il] = !hparams.is_recr_impl[il]; + } + + switch (hparams.n_layer()) { + case 45: type = hparams.n_embd == 4096 && hparams.n_expert == 288 ? LLM_TYPE_313B_A17B : LLM_TYPE_UNKNOWN; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = head_dim * n_head; + const int64_t d_conv = hparams.ssm_d_conv; + + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + + const int64_t n_embd_indexer = hparams.indexer_head_size; + const int64_t kpool = hparams.indexer_kpool; + + const int64_t hc_dim = (int64_t) hparams.dsv4_hc_mult * n_embd; + const int64_t hc_mix_dim = (2 + (int64_t) hparams.dsv4_hc_mult) * hparams.dsv4_hc_mult; + + // the trunk and the NextN block can be split across two GGUFs in either direction + const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); + 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); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + 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}, TENSOR_NOT_REQUIRED); + if (!output) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int il = 0; il < n_layer_all; ++il) { + auto & layer = layers[il]; + const int flags = il < n_layer ? trunk_flags : mtp_flags; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), {n_embd}, flags); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), {n_embd}, flags); + + // the NextN block keeps the plain residual, so it has no mHC mixer + if (il < n_layer) { + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", il), {hc_dim, hc_mix_dim}, flags); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", il), {hc_mix_dim}, flags); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", il), {3}, flags); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", il), {hc_dim, hc_mix_dim}, flags); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", il), {hc_mix_dim}, flags); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", il), {3}, flags); + } + + if (hparams.is_recr(il)) { + create_tensor_qkv(layer, il, n_embd, d_inner, d_inner, d_inner, flags); + + layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", il), {d_conv, 1, d_inner, 1}, flags); + layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", il), {d_conv, 1, d_inner, 1}, flags); + layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", il), {d_conv, 1, d_inner, 1}, flags); + + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", il), {n_embd, head_dim}, flags); + layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", il), {head_dim, d_inner}, flags); + layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", il), {n_embd, head_dim}, flags); + layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", il), {head_dim, d_inner}, flags); + + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), {n_embd, n_head}, flags); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, il), {n_head}, flags); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), {d_inner}, flags); + + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), {head_dim}, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), {d_inner, n_embd}, flags); + } else { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), {n_embd, q_lora_rank}, flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), {q_lora_rank}, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), {q_lora_rank, n_head * qk_head_dim}, flags); + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), {n_embd, kv_lora_rank}, flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), {kv_lora_rank}, flags); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), {qk_head_dim, kv_lora_rank, n_head}, flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), {kv_lora_rank, v_head_dim, n_head}, flags); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), {n_head * v_head_dim, n_embd}, flags); + + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), {n_embd_indexer}, flags); + layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", il), {n_embd_indexer}, flags); + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", il), {n_embd, hparams.indexer_n_head}, flags); + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", il), {n_embd, n_embd_indexer}, flags); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", il), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, flags); + + // key pooling: DeepSeek-V4 doubles the compressor width, GLM-5.3 does not + layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", il), {n_embd, n_embd_indexer}, flags); + layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", il), {n_embd_indexer, kpool}, flags); + } + + if (il < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd, n_ff}, flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), {n_embd, n_ff}, flags); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), {n_ff, n_embd}, flags); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), {n_embd, n_expert}, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), {n_expert}, flags); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), {n_embd, hparams.n_ff_exp, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), {n_embd, hparams.n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), {hparams.n_ff_exp, n_embd, n_expert}, flags); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), {n_embd, hparams.n_ff_shexp}, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), {n_embd, hparams.n_ff_shexp}, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), {hparams.n_ff_shexp, n_embd}, flags); + } + + if (il >= n_layer) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), {2 * n_embd, n_embd}, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), {n_embd}, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), {n_embd}, flags); + + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), {n_embd}, flags); + // absent in the checkpoint: NextN shares the trunk's embeddings and + // lm_head. only accepted if an export adds them + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); + } + } +} + +std::unique_ptr llama_model_glm5next::build_arch_graph(const llm_graph_params & params) const { + GGML_UNUSED(params); + throw std::runtime_error("glm5next: graph not implemented yet"); +} diff --git a/src/models/models.h b/src/models/models.h index 969429e3b6f..cd223882484 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1331,6 +1331,14 @@ struct llama_model_glm_dsa : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +struct llama_model_glm5next : public llama_model_base { + llama_model_glm5next(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; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_eagle3 : public llama_model_base { llama_model_eagle3(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; From 9e1cc208226cd49699c332ab4a3b1911dc92024d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 26 Aug 2026 07:18:49 +0000 Subject: [PATCH 02/36] llama: add glm5next mHC wide residual glm5next's mHC is DeepSeek-V4's hyper-connection block: same wide residual, same 24-row mixer split, same two activations, same Sinkhorn. Only the final collapse differs, so the graph derives from llama_model_deepseek4::graph and reuses build_hc_pre / build_hc_post / build_hc_sinkhorn rather than restating them, as graph_dsv4 already does in dflash.cpp. dsv4_hc_mean becomes a static member so both archs can reach it; the body and both deepseek4 call sites are otherwise untouched. The generated code for deepseek4 is unchanged apart from the endbr64 landing pad the helper now needs as a global symbol. The four streams start as exact copies of the token embedding and collapse to an unweighted mean after the last layer: this checkpoint has no hc_head. KDA, DSA and the MoE land in later commits, so the two sublayers throw. The mHC wiring around them is final. --- src/llama-context.cpp | 1 + src/models/deepseek4.cpp | 6 +- src/models/glm5next.cpp | 128 ++++++++++++++++++++++++++++++++++++++- src/models/models.h | 22 +++++++ 4 files changed, 152 insertions(+), 5 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0402044da6b..5cbc76afd2a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2302,6 +2302,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_GLM5NEXT || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_01 || diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index fc816e2aeb4..692e25a8db9 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -264,7 +264,7 @@ 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) { +ggml_tensor * llama_model_deepseek4::graph::build_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); @@ -1234,7 +1234,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(ctx0, inpL); cb(res->t_layer_inp[il], "layer_inp", il); ggml_build_forward_expand(gf, res->t_layer_inp[il]); } @@ -1316,7 +1316,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(ctx0, 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/glm5next.cpp b/src/models/glm5next.cpp index cdb6dde1cb3..1e70fcea744 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -43,6 +43,10 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); GGML_ASSERT(hparams.dsv4_hc_mult > 0); + // trunk residual is hc_mult streams wide (deepseek4); lm_head still sees + // n_embd, the streams are averaged first + hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd; + // MoE ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); @@ -208,7 +212,127 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { } } +// sublayers stubbed until the KDA and DSA commits, which must widen these +// signatures and add the position and memory inputs the ctor does not build yet +ggml_tensor * llama_model_glm5next::graph::build_layer_attn( + const llama_model & model, + ggml_tensor * cur, + int il) const { + GGML_UNUSED(model); + GGML_UNUSED(cur); + + throw std::runtime_error(hparams.is_recr(il) + ? "glm5next: KDA attention not implemented yet" + : "glm5next: DSA attention not implemented yet"); +} + +ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( + const llama_model & model, + ggml_tensor * cur, + int il) const { + GGML_UNUSED(model); + GGML_UNUSED(cur); + GGML_UNUSED(il); + + throw std::runtime_error("glm5next: feed-forward not implemented yet"); +} + +llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_params & params) : + llama_model_deepseek4::graph(params) { + ggml_tensor * cur; + + ggml_tensor * inp = build_inp_embd(model.tok_embd); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t hc = hparams.dsv4_hc_mult; + + // hc_mult exact copies of the embedding: no scaling, no one-hot into stream 0 + 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); + + 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] = build_hc_mean(ctx0, inpL); + cb(res->t_layer_inp[il], "layer_inp", il); + ggml_build_forward_expand(gf, res->t_layer_inp[il]); + } + + ggml_tensor * residual = inpL; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + + cur = build_hc_pre(inpL, + model.layers[il].hc_attn_fn, + model.layers[il].hc_attn_scale, + model.layers[il].hc_attn_base, + &post, &comb, il); + cb(cur, "hc_attn_pre", il); + + cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + cur = build_layer_attn(model, cur, il); + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "hc_attn_post", il); + + residual = inpL; + cur = build_hc_pre(inpL, + model.layers[il].hc_ffn_fn, + model.layers[il].hc_ffn_scale, + model.layers[il].hc_ffn_base, + &post, &comb, il); + cb(cur, "hc_ffn_pre", il); + + // expand before the sublayer so op offload does not pull the mHC state + // onto the expert weights' backend, as in deepseek4 + ggml_build_forward_expand(gf, residual); + ggml_build_forward_expand(gf, post); + ggml_build_forward_expand(gf, comb); + + cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_layer_ffn(model, cur, il); + cb(cur, "ffn_out", il); + + inpL = build_hc_post(cur, residual, post, comb, il); + inpL = build_cvec(inpL, il); + cb(inpL, "l_last", il); + } + + if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) { + res->t_layer_inp[n_layer] = build_hc_mean(ctx0, inpL); + cb(res->t_layer_inp[n_layer], "layer_inp", n_layer); + ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); + } + + if (inp_out_ids) { + // flattened: get_rows needs one token's streams to be one contiguous row + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + inpL = ggml_reshape_3d(ctx0, ggml_get_rows(ctx0, flat, inp_out_ids), n_embd, hc, n_outputs); + } + + // no hc_head tensor here: unweighted mean, not DeepSeek-V4's learned gated head + cur = build_hc_mean(ctx0, inpL); + cb(cur, "hc_mean", -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); +} + std::unique_ptr llama_model_glm5next::build_arch_graph(const llm_graph_params & params) const { - GGML_UNUSED(params); - throw std::runtime_error("glm5next: graph not implemented yet"); + // llama_init_from_model accepts an MTP context whenever n_layer_nextn > 0, + // which every glm5next checkpoint has; without this it silently runs the trunk + GGML_ASSERT(params.gtype != LLM_GRAPH_TYPE_DECODER_MTP && "glm5next NextN graph not implemented yet"); + + return std::make_unique(*this, params); } diff --git a/src/models/models.h b/src/models/models.h index cd223882484..2ad98d6fa60 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1294,6 +1294,11 @@ struct llama_model_deepseek4 : public llama_model_base { ggml_tensor * build_hc_sinkhorn( ggml_tensor * comb, int il) const; + + // unweighted mean over the streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens] + static ggml_tensor * build_hc_mean( + ggml_context * ctx, + ggml_tensor * x); }; struct graph_mtp : public graph { @@ -1336,6 +1341,23 @@ struct llama_model_glm5next : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; + // glm5next's mHC is DeepSeek-V4's hyper-connection block (same wide residual, + // 24-row mixer split, activations, Sinkhorn); only the final collapse differs + // (unweighted mean, not a learned gated head), so derive rather than restate + struct graph : public llama_model_deepseek4::graph { + graph(const llama_model & model, const llm_graph_params & params); + + ggml_tensor * build_layer_attn( + const llama_model & model, + ggml_tensor * cur, + int il) const; + + ggml_tensor * build_layer_ffn( + const llama_model & model, + ggml_tensor * cur, + int il) const; + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; From 20a807707eed6d384395411fc134becb88b3e42c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 07:54:41 +0000 Subject: [PATCH 03/36] llama: add glm5next KDA linear attention Copy-adapts kimi-k3's KDA layer rather than kimi-linear's or bailingmoe3's: it already matches on the recurrence ordering, the bounded-sigmoid decay gate and its branch selection, dt_bias added per channel before the reshape, per-head A broadcast, SiLU after the conv, f/g/beta read from the pre-convolution hidden states, and the gated output RMSNorm with a plain weight. Three differences from kimi-k3. The output gate is low rank, g_b(g_a(x)) as in kimi-linear, which is what PR 1's converter emits. The q/k L2 eps is a literal 1e-6, the reference's own constant, not f_norm_rms_eps; ggml_l2_norm implements max(sqrt(sum), eps) rather than sqrt(sum + eps), which at head_dim 128 differs by about eps/(2*sum) and never trips the clamp, so it is close but not bit-exact. And the cross-layer residual, latent MoE, situ activation and MLA output gate have no counterpart here. The conv follows the reference and convolves q|k|v as one depthwise kernel, which keeps the conv state a single contiguous block so build_conv_state can snapshot it. That plus build_recurrent_attn is what makes the layer safe under recurrent-state rollback, so the arch joins llm_arch_supports_rs_rollback; without that entry the guard in llama_context silently clamps n_rs_seq to 0. build_delta_net_autoregressive reshaped a per-channel KDA gate onto ne1, but ne0 is the key axis everywhere else in that function, so it decayed along the value axis. Invisible for GDN, where the gate is scalar and both spellings produce the same [1, 1, H_v, n_seqs], and invisible to the shape checks because S_k == S_v. Fixed rather than asserted around, since glm5next reaches that path on any backend without the fused operator. llama_model_deepseek4::graph now derives from llm_build_delta_net_base so glm5next, which derives from it for the mHC residual, can reach build_delta_net. The base is a method-only mixin over llm_graph_context with no data members and no virtuals beyond the destructor llm_graph_context already has; deepseek4.cpp, dflash.cpp and kimi-k3.cpp compile to byte-identical instructions across the change. graph_max_nodes moves the arch to kimi-k3's tier. Measured on the Tiny fixture with the chunked fallback: 182 nodes plus 15/16 per token for each KDA layer and 46 per layer for the mHC mixers, so the 45-layer model needs 8.3k + 31.9 per token before DSA or the MoE are counted, which overruns the n_tokens*40 budget. test-llama-archs synthesised no MLA, hyper-connection, kpool or expert-weight keys for glm5next, so PR 1's required get_key calls threw out of the sweep and truncated it at 75 of 143 architectures. The fixture is complete now and the row is skipped explicitly while the DSA and feed-forward sublayers still throw. --- src/llama-arch.cpp | 1 + src/llama-context.cpp | 2 +- src/models/deepseek4.cpp | 2 +- src/models/delta-net-base.cpp | 9 ++- src/models/glm5next.cpp | 141 +++++++++++++++++++++++++++++++--- src/models/models.h | 18 ++++- tests/test-llama-archs.cpp | 5 +- 7 files changed, 160 insertions(+), 18 deletions(-) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 0dbb04fe7bd..5a145f9c1ac 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1042,6 +1042,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { case LLM_ARCH_LFM2: case LLM_ARCH_LFM2MOE: case LLM_ARCH_BAILINGMOE3: + case LLM_ARCH_GLM5NEXT: return true; default: return false; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5cbc76afd2a..df10311ba70 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2297,12 +2297,12 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { // 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 || + if (model.arch == LLM_ARCH_GLM5NEXT) { model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_BAILINGMOE3 || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_DEEPSEEK4 || - model.arch == LLM_ARCH_GLM5NEXT || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_01 || diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 692e25a8db9..b1a8e828bbf 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -1217,7 +1217,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) { + llm_build_delta_net_base(params) { ggml_tensor * cur; ggml_tensor * inp = build_inp_embd(model.tok_embd); diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ad661264773..4b6aec8046d 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -330,9 +330,12 @@ std::pair llm_build_delta_net_base::build_delta_ne cb(b, "b_in", il); cb(g, "g_in", il); - // GDA: [1, 1, H_v, n_seqs] - // KDA: [1, S_k, H_v, n_seqs] - g = ggml_reshape_4d(ctx0, g, 1, g->ne[0], H_v, n_seqs); + // the state is indexed [key, value]: ne0 is the key axis, as the k/q reductions + // below rely on, so KDA's per-key-channel decay must broadcast over ne1, not ne0. + // for GDA g->ne[0] is 1 and both spellings agree + // GDA: [1, 1, H_v, n_seqs] + // KDA: [S_k, 1, H_v, n_seqs] + g = ggml_reshape_4d(ctx0, g, g->ne[0], 1, H_v, n_seqs); b = ggml_reshape_4d(ctx0, b, 1, 1, H_v, n_seqs); // [S_v, S_v, H_v, n_seqs] diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 1e70fcea744..8a1594de8c5 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -1,5 +1,7 @@ #include "models.h" +#include "llama-memory-recurrent.h" + // // GLM-5.3-Flash: hybrid KDA (linear) + DSA (nope-only MLA) attention, mHC // hyper-connections, and a NextN block that is a full DSA decoder layer. @@ -21,9 +23,13 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(hparams.n_lora_q > 0 && "glm5next requires a q LoRA"); GGML_ASSERT(hparams.n_rot() == 0 && "glm5next MLA is nope-only"); - // KDA + // KDA. no GGUF key for linear_num_heads, so the KDA head count is + // attention.head_count, which also sizes the recurrent state via n_embd_r/s(). + // conversion/glm5next.py refuses a checkpoint where the two differ 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); + GGML_ASSERT(hparams.ssm_d_conv > 1); + GGML_ASSERT(hparams.n_embd_head_kda > 0); // required: absent, kimi-k3 selects the softplus branch, a different // function, not a missing clamp ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); @@ -212,18 +218,128 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { } } -// sublayers stubbed until the KDA and DSA commits, which must widen these -// signatures and add the position and memory inputs the ctor does not build yet +// +// KDA layer +// +// one depthwise conv over the concatenated q|k|v channels, as in the reference: it +// leaves the conv state one contiguous block, which is what build_conv_state needs to +// snapshot for recurrent-state rollback. three separate convs would be numerically +// identical but would restate the rollback write three times +// +ggml_tensor * llama_model_glm5next::graph::build_kda_layer( + const llama_layer & layer, + llm_graph_input_rs * inp_rs, + ggml_tensor * cur, + int il) { + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = head_dim * n_head; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + const auto * mctx_cur = inp_rs->mctx; + + // f, g and beta read the layer input, NOT the convolved q/k/v + ggml_tensor * inp = cur; + + ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.wq, inp); + ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.wk, inp); + ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.wv, inp); + + ggml_tensor * qkv = ggml_concat(ctx0, ggml_concat(ctx0, Qcur, Kcur, 0), Vcur, 0); + qkv = ggml_reshape_3d(ctx0, qkv, 3*d_inner, n_seq_tokens, n_seqs); + cb(qkv, "kda_qkv", il); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_in = build_conv_state(inp_rs, conv_states_all, qkv, d_conv, 3*d_inner, il); + + // stored separately (kimi-linear, kimi-k3), stacked back into the single kernel + ggml_tensor * conv_w = ggml_concat(ctx0, + ggml_concat(ctx0, + ggml_reshape_2d(ctx0, layer.ssm_q_conv, d_conv, d_inner), + ggml_reshape_2d(ctx0, layer.ssm_k_conv, d_conv, d_inner), 1), + ggml_reshape_2d(ctx0, layer.ssm_v_conv, d_conv, d_inner), 1); + + // SiLU is applied to the conv output, not to the projections + ggml_tensor * conv_out = ggml_silu(ctx0, ggml_ssm_conv(ctx0, conv_in, conv_w)); + cb(conv_out, "kda_conv", il); + + const size_t nb_qkv = ggml_row_size(conv_out->type, 3*d_inner); + const size_t nb_head = ggml_row_size(conv_out->type, head_dim); + + Qcur = ggml_view_4d(ctx0, conv_out, head_dim, n_head, n_seq_tokens, n_seqs, + nb_head, nb_qkv, nb_qkv*n_seq_tokens, 0); + Kcur = ggml_view_4d(ctx0, conv_out, head_dim, n_head, n_seq_tokens, n_seqs, + nb_head, nb_qkv, nb_qkv*n_seq_tokens, ggml_row_size(conv_out->type, d_inner)); + Vcur = ggml_view_4d(ctx0, conv_out, head_dim, n_head, n_seq_tokens, n_seqs, + nb_head, nb_qkv, nb_qkv*n_seq_tokens, ggml_row_size(conv_out->type, 2*d_inner)); + + // 1e-6 is the reference's own constant, not the model's norm eps. ggml_l2_norm + // divides by max(sqrt(sum), eps) where the reference uses sqrt(sum + eps); at + // head_dim 128 the clamp never binds, so close but not bit-exact + Qcur = ggml_l2_norm(ctx0, Qcur, 1e-6f); + Kcur = ggml_l2_norm(ctx0, Kcur, 1e-6f); + cb(Qcur, "kda_q_norm", il); + cb(Kcur, "kda_k_norm", il); + + // the 1/sqrt(head_dim) query scale is applied inside build_delta_net, after this norm + + // forget gate. gate_lower_bound is a multiplicative scale, not a clamp: + // g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias)) + // ssm_a holds -exp(A_log), so exp(A_log) * y == -(y * ssm_a) + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ssm_f_b, ggml_mul_mat(ctx0, layer.ssm_f_a, inp)); + g = ggml_add(ctx0, g, layer.ssm_dt_b); + g = ggml_reshape_3d(ctx0, g, head_dim, n_head, n_tokens); + g = ggml_mul(ctx0, g, ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1)); + g = ggml_sigmoid(ctx0, ggml_scale(ctx0, g, -1.0f)); + g = ggml_scale(ctx0, g, hparams.kda_gate_lower_bound); + g = ggml_reshape_4d(ctx0, g, head_dim, n_head, n_seq_tokens, n_seqs); + cb(g, "kda_gate", il); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, inp); + beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs)); + 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, n_seqs); + + ggml_tensor * out = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g, beta, state, il); + + // the fallbacks return a permuted view, the fused op a contiguous one; cont + // either way rather than depend on which ran + ggml_tensor * o = ggml_cont_3d(ctx0, out, head_dim, n_head, n_tokens); + cb(o, "kda_scan_out", il); + + // low-rank output gate (kimi-k3 has a single full-rank ssm_g instead) + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_g_b, ggml_mul_mat(ctx0, layer.ssm_g_a, inp)); + gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens); + + // RMS over head_dim only, one weight shared by every head, then a plain sigmoid + // gate: not the SiLU that FusedRMSNormGated defaults to + ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, gate)); + cb(gated, "kda_normed", il); + + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, gated, d_inner, n_tokens)); + cb(cur, "kda_out", il); + + return cur; +} + +// DSA is stubbed until its own commit; the mHC wiring around both sublayers is final ggml_tensor * llama_model_glm5next::graph::build_layer_attn( const llama_model & model, + llm_graph_input_rs * inp_rs, ggml_tensor * cur, - int il) const { - GGML_UNUSED(model); + int il) { + if (hparams.is_recr(il)) { + return build_kda_layer(model.layers[il], inp_rs, cur, il); + } + GGML_UNUSED(cur); - throw std::runtime_error(hparams.is_recr(il) - ? "glm5next: KDA attention not implemented yet" - : "glm5next: DSA attention not implemented yet"); + throw std::runtime_error("glm5next: DSA attention not implemented yet"); } ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( @@ -244,6 +360,13 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp = build_inp_embd(model.tok_embd); ggml_tensor * inp_out_ids = build_inp_out_ids(); + // recurrent half of the hybrid memory; the attention half is unused until DSA + llm_graph_input_rs * inp_rs = build_inp_mem_hybrid()->get_recr(); + + GGML_ASSERT(ubatch.n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == ubatch.n_seq_tokens * ubatch.n_seqs); + const int64_t hc = hparams.dsv4_hc_mult; // hc_mult exact copies of the embedding: no scaling, no one-hot into stream 0 @@ -272,7 +395,7 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); - cur = build_layer_attn(model, cur, il); + cur = build_layer_attn(model, inp_rs, cur, il); inpL = build_hc_post(cur, residual, post, comb, il); cb(inpL, "hc_attn_post", il); diff --git a/src/models/models.h b/src/models/models.h index 2ad98d6fa60..8670eda4179 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1173,8 +1173,12 @@ 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) {} + // llm_build_delta_net_base is a method-only mixin over llm_graph_context (no data + // members, no new virtuals), so this graph's layout and behaviour are unchanged. + // deepseek4 has no recurrent layers; it is here so glm5next, which derives from + // this graph for the mHC residual, can reach build_delta_net for its KDA layers + struct graph : public llm_build_delta_net_base { + graph(const llm_graph_params & params) : llm_build_delta_net_base(params) {} graph(const llama_model & model, const llm_graph_params & params); ggml_tensor * build_hc_pre( @@ -1347,10 +1351,18 @@ struct llama_model_glm5next : public llama_model_base { struct graph : public llama_model_deepseek4::graph { graph(const llama_model & model, const llm_graph_params & params); + // not const: the delta-net helpers append to the graph through the base ggml_tensor * build_layer_attn( const llama_model & model, + llm_graph_input_rs * inp_rs, ggml_tensor * cur, - int il) const; + int il); + + ggml_tensor * build_kda_layer( + const llama_layer & layer, + llm_graph_input_rs * inp_rs, + ggml_tensor * cur, + int il); ggml_tensor * build_layer_ffn( const llama_model & model, diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b8fd66ccae5..8b5d4f732c4 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -165,7 +165,7 @@ 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_GLM5NEXT) { GGML_ASSERT(n_layer >= 2); std::vector n_head_per_layer; n_head_per_layer.reserve(n_layer); @@ -213,6 +213,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types); } + } else if (arch == LLM_ARCH_GLM5NEXT) { } else if (arch == LLM_ARCH_MINIMAX_M3) { // partial rotary: n_rot must not exceed the indexer key length (64) ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); @@ -445,6 +446,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_GLM5NEXT: case LLM_ARCH_STEP35: case LLM_ARCH_MISTRAL4: case LLM_ARCH_MELLUM: @@ -505,6 +507,7 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_DEEPSEEK2OCR) { return false; } + if (arch == LLM_ARCH_GLM5NEXT) { // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE) { From 1d99a5ac4e9fda3269bfbf148b20c8865f8bb37f Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 07:56:23 +0000 Subject: [PATCH 04/36] llama: add glm5next MoE feed-forward with clamped SwiGLU The routing is DeepSeek-V3 noaux_tc exactly as build_moe_ffn already implements it: sigmoid scores, exp_probs_b added for the top-k SELECTION only, weights gathered from the unbiased scores, normalised, then scaled by routed_scaling_factor. n_group and topk_group are both 1, so the group-limited stage is degenerate and build_moe_ffn's n_expert_groups > 1 guard skips it; no group keys are written and none are needed. The clamp is the one thing that needed a change outside this arch. glm5next clamps the gate max-only and the up symmetrically, both BEFORE the SiLU, which is what the branch behind the DEEPSEEK4/DFLASH arch gate already does; the else branch clamps after the SiLU and is a different function. Adding the arch to both gates reuses it rather than restating it. The two conditions are separate because the dense path and the MoE path read different hparams arrays. The leading dense layers clamp too. The reference builds them from the same Glm5NextTextMLP as the shared expert, so swiglu_limit is not MoE-only, and the converter already writes swiglu_clamp_shexp for every layer rather than only the sparse ones. The shared expert is added unscaled. --- src/llama-graph.cpp | 4 ++-- src/models/glm5next.cpp | 39 +++++++++++++++++++++++++++++++++++---- 2 files changed, 37 insertions(+), 6 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0e..9bdc7d0e7a1 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_GLM5NEXT || (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_GLM5NEXT || (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); diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 8a1594de8c5..539de24ac82 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -346,11 +346,42 @@ ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( const llama_model & model, ggml_tensor * cur, int il) const { - GGML_UNUSED(model); - GGML_UNUSED(cur); - GGML_UNUSED(il); + const auto & layer = model.layers[il]; + + // the leading dense layers clamp the same way the experts do: the reference + // routes both through one Glm5NextTextMLP, so swiglu_limit is not MoE-only + if (il < (int) hparams.n_layer_dense_lead) { + return 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); + } - throw std::runtime_error("glm5next: feed-forward not implemented yet"); + // noaux_tc: exp_probs_b biases top-k selection only, the weights are the + // unbiased sigmoid scores. n_group is 1, so the group mask is a no-op + 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, hparams.n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + + ggml_tensor * 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(shexp, "ffn_shexp", il); + + // shared expert unscaled: routed_scaling_factor is applied inside build_moe_ffn, + // after norm_topk_prob, to the routed weights only + return ggml_add(ctx0, moe_out, shexp); } llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_params & params) : From f320a28753ba95208de3bb76f0f9e8c9afc02e8d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 08:55:35 +0000 Subject: [PATCH 05/36] llama: add glm5next dense DSA attention nope-only MLA in the absorbed form, over every cached position. below index_topk + index_kpool - 1 resident tokens the indexer selects all of them, so this is exactly what the sparse path degenerates to, and it is a reference the sparse commit can be checked against. the attention half of the hybrid memory becomes the K-only variant: after absorption the cache holds the kv_lora_rank latent and V is a view of K. --- src/models/glm5next.cpp | 82 ++++++++++++++++++++++++++++++++++++----- src/models/models.h | 8 +++- 2 files changed, 80 insertions(+), 10 deletions(-) diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 539de24ac82..cc3df388ff0 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -327,19 +327,82 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( return cur; } -// DSA is stubbed until its own commit; the mHC wiring around both sublayers is final +// +// DSA layer, dense +// +// below index_topk + index_kpool - 1 resident positions the indexer selects every +// position, so sparse selection is exactly full attention. this builds that limit: +// correct MLA over the whole cache, no indexer, no kpool, no top-k +// +// the absorbed form is used, as in deepseek2/deepseek32/glm-dsa: q_nope is pushed +// through wk_b so that q.k is taken against the 512-wide latent directly, which is +// what the cache holds (is_mla() suppresses the V allocation and V becomes a view of +// K). the naive form would have to expand the latent back to n_head 256-wide keys and +// values on every step and would need a V cache this memory layout does not have +// +ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( + const llama_layer & layer, + llm_graph_input_attn_k * inp_attn, + ggml_tensor * cur, + int il) const { + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + + // nope-only. every other MLA port splits q and k into a nope and a rope half and + // ropes the second one; here that half is zero-width, so there is no split, no + // concat and no ggml_rope_ext anywhere in the text tower + GGML_ASSERT(hparams.n_rot() == 0); + + // the reference scales by qk_head_dim^-0.5, i.e. over the MLA head size, not over + // n_embd_head_k: after absorption q is kv_lora_rank wide and 1/sqrt(512) would be + // a different model + const float kq_scale = 1.0f/sqrtf(float(qk_head_dim)); + + ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur); + q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(q, "dsa_q_a_norm", il); + + q = ggml_mul_mat(ctx0, layer.wq_b, q); + q = ggml_reshape_3d(ctx0, q, qk_head_dim, n_head, n_tokens); + cb(q, "dsa_q_b", il); + + ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + cb(kv, "dsa_kv_a_norm", il); + + // {qk_head_dim, n_tokens, n_head} + q = ggml_permute(ctx0, q, 0, 2, 1, 3); + + // {qk_head_dim, kv_lora_rank, n_head} x {qk_head_dim, n_tokens, n_head} + q = ggml_mul_mat(ctx0, layer.wk_b, q); + + // {kv_lora_rank, n_head, n_tokens}. deepseek2 gets this contiguous for free out of + // the concat with the roped half, which does not exist here + q = ggml_cont(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3)); + cb(q, "dsa_q_absorbed", il); + + // absorbed MLA is MQA: one head of keys, and V is the same latent row as K + ggml_tensor * k = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens); + cb(k, "dsa_kv_latent", il); + + cur = build_attn(inp_attn, + layer.wo, nullptr, nullptr, + q, k, k, nullptr, nullptr, layer.wv_b, kq_scale, il); + cb(cur, "dsa_out", il); + + return cur; +} + ggml_tensor * llama_model_glm5next::graph::build_layer_attn( const llama_model & model, - llm_graph_input_rs * inp_rs, + llm_graph_input_mem_hybrid_k * inp_mem, ggml_tensor * cur, int il) { if (hparams.is_recr(il)) { - return build_kda_layer(model.layers[il], inp_rs, cur, il); + return build_kda_layer(model.layers[il], inp_mem->get_recr(), cur, il); } - GGML_UNUSED(cur); - - throw std::runtime_error("glm5next: DSA attention not implemented yet"); + return build_dsa_layer(model.layers[il], inp_mem->get_attn(), cur, il); } ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( @@ -391,8 +454,9 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp = build_inp_embd(model.tok_embd); ggml_tensor * inp_out_ids = build_inp_out_ids(); - // recurrent half of the hybrid memory; the attention half is unused until DSA - llm_graph_input_rs * inp_rs = build_inp_mem_hybrid()->get_recr(); + // MLA absorption leaves a K-only cache holding the kv_lora_rank latent, so the + // attention half of the hybrid memory is the _k variant, as in bailingmoe3 + llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); GGML_ASSERT(ubatch.n_seqs != 0); GGML_ASSERT(ubatch.equal_seqs()); @@ -426,7 +490,7 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); - cur = build_layer_attn(model, inp_rs, cur, il); + cur = build_layer_attn(model, inp_mem, cur, il); inpL = build_hc_post(cur, residual, post, comb, il); cb(inpL, "hc_attn_post", il); diff --git a/src/models/models.h b/src/models/models.h index 8670eda4179..8087aa7c482 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1354,7 +1354,7 @@ struct llama_model_glm5next : public llama_model_base { // not const: the delta-net helpers append to the graph through the base ggml_tensor * build_layer_attn( const llama_model & model, - llm_graph_input_rs * inp_rs, + llm_graph_input_mem_hybrid_k * inp_mem, ggml_tensor * cur, int il); @@ -1364,6 +1364,12 @@ struct llama_model_glm5next : public llama_model_base { ggml_tensor * cur, int il); + ggml_tensor * build_dsa_layer( + const llama_layer & layer, + llm_graph_input_attn_k * inp_attn, + ggml_tensor * cur, + int il) const; + ggml_tensor * build_layer_ffn( const llama_model & model, ggml_tensor * cur, From 7f2560e98aaed8fff96d1497b317e872a78b05c1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 08:55:41 +0000 Subject: [PATCH 06/36] llama: save the indexer kpool and hyper-connection keys both are required keys for glm5next, so a model saved without them cannot be loaded back. this is what stops test-llama-archs from round-tripping the arch. --- src/llama-model-saver.cpp | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 9adaa93f62e..a39a6a56dda 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -291,7 +291,11 @@ void llama_model_saver::add_kv_from_model() { 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_LOCAL_BLOCKS, hparams.indexer_local_blocks); + add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, hparams.indexer_kpool); add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true); + add_kv(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); + add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true); add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count); add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank); From cd69d60e5ce4f2e75cb5f3010e37355c2b00e01d Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 08:55:41 +0000 Subject: [PATCH 07/36] test-llama-archs: enable glm5next the DSA sublayer no longer throws, so the arch can construct and run. it needs the MLA head shape as well: with n_head_kv taken from the per-layer array it would size the K cache row n_head times wider than the latent the graph writes. --- 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 8b5d4f732c4..6f1855b0999 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -118,7 +118,10 @@ 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_GLM5NEXT || arch == LLM_ARCH_MISTRAL4) { + // MLA absorbs into MQA, so the K cache row is the latent: n_head_kv must be 1 + // or the per-layer head_count_kv array below sizes it n_head times too wide n_embd = 128; n_head = 1; n_ff = 192; From 6a58359c5e5b602e62365dbe9656118bd533c2e9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 09:10:53 +0000 Subject: [PATCH 08/36] llama: assert the glm5next indexer selection width index_topk + index_kpool - 1 is the number of positions the indexer keeps, and it is what makes the dense attention this branch builds exactly equal to the sparse path below that many cached tokens. an off-by-one in it is invisible to every output comparison measured so far, on both a dense and a sparse fixture, so it is checked against a second spelling of the same arithmetic instead. --- src/models/glm5next.cpp | 27 +++++++++++++++++++++++++++ 1 file changed, 27 insertions(+) diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index cc3df388ff0..7b37f804c10 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -11,6 +11,27 @@ // load time, so conversion/glm5next.py is the only place the sign is checked. // +// how many positions the indexer keeps: index_topk/index_kpool whole pools, plus the +// always-selected tail pool minus one. below this many cached tokens every position is +// selected and the sparse path is exactly the dense one built here. +// +// asserted rather than measured. an off-by-one here is invisible to every output +// comparison there is: the reference's own seeded off-by-one on this width is +// bit-identical on both the dense and the sparse fixtures. the second form below is the +// independent spelling the parity harness uses, so the two have to agree +static uint32_t glm5next_n_select(const llama_hparams & hparams) { + GGML_ASSERT(hparams.indexer_kpool > 0); + GGML_ASSERT(hparams.indexer_top_k >= hparams.indexer_kpool); + GGML_ASSERT(hparams.indexer_top_k % hparams.indexer_kpool == 0); + + const uint32_t n_select = hparams.indexer_top_k + hparams.indexer_kpool - 1; + + GGML_ASSERT(n_select > hparams.indexer_top_k); + GGML_ASSERT(n_select == (hparams.indexer_top_k/hparams.indexer_kpool + 1)*hparams.indexer_kpool - 1); + + return n_select; +} + void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); // indexer k_norm is a LayerNorm with bias; without this key it runs at eps 0 @@ -43,6 +64,12 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(hparams.indexer_kpool > 0); GGML_ASSERT(hparams.indexer_top_k % hparams.indexer_kpool == 0); + const uint32_t n_select = glm5next_n_select(hparams); + if (hparams.n_ctx_train > n_select) { + LLAMA_LOG_WARN("%s: attention is dense above %u cached tokens, but this checkpoint trains to %u. " + "the sparse selection is not implemented yet\n", __func__, n_select, hparams.n_ctx_train); + } + // 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); From 3841592347f99d48c0cca55ab2e1b3719b77dda1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 14:44:55 +0000 Subject: [PATCH 09/36] glm5next: trim DSA dense comments --- src/models/glm5next.cpp | 45 ++++++++++++++++---------------------- tests/test-llama-archs.cpp | 4 ++-- 2 files changed, 21 insertions(+), 28 deletions(-) diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 7b37f804c10..2128efdd3f9 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -11,14 +11,13 @@ // load time, so conversion/glm5next.py is the only place the sign is checked. // -// how many positions the indexer keeps: index_topk/index_kpool whole pools, plus the +// positions the indexer keeps: index_topk/index_kpool whole pools plus the // always-selected tail pool minus one. below this many cached tokens every position is -// selected and the sparse path is exactly the dense one built here. +// selected, so sparse selection is exactly the dense path built here. // -// asserted rather than measured. an off-by-one here is invisible to every output -// comparison there is: the reference's own seeded off-by-one on this width is -// bit-identical on both the dense and the sparse fixtures. the second form below is the -// independent spelling the parity harness uses, so the two have to agree +// asserted, not measured: an off-by-one here is invisible to output comparison (the +// reference's own off-by-one on this width is bit-identical on both fixtures). the +// second assert is the parity harness's independent spelling, so the two must agree static uint32_t glm5next_n_select(const llama_hparams & hparams) { GGML_ASSERT(hparams.indexer_kpool > 0); GGML_ASSERT(hparams.indexer_top_k >= hparams.indexer_kpool); @@ -355,17 +354,13 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( } // -// DSA layer, dense +// DSA layer, dense: full MLA over the whole cache, no indexer/kpool/top-k. this is the +// limit the sparse path collapses to below glm5next_n_select() resident positions // -// below index_topk + index_kpool - 1 resident positions the indexer selects every -// position, so sparse selection is exactly full attention. this builds that limit: -// correct MLA over the whole cache, no indexer, no kpool, no top-k -// -// the absorbed form is used, as in deepseek2/deepseek32/glm-dsa: q_nope is pushed -// through wk_b so that q.k is taken against the 512-wide latent directly, which is -// what the cache holds (is_mla() suppresses the V allocation and V becomes a view of -// K). the naive form would have to expand the latent back to n_head 256-wide keys and -// values on every step and would need a V cache this memory layout does not have +// absorbed form, as in deepseek2/deepseek32/glm-dsa: q_nope is pushed through wk_b so +// q.k is taken against the 512-wide latent the cache actually holds (is_mla() drops the +// V allocation and V becomes a view of K). the naive form would re-expand the latent to +// n_head 256-wide k/v every step and needs a V cache this layout does not have // ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( const llama_layer & layer, @@ -375,14 +370,12 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); const int64_t kv_lora_rank = hparams.n_lora_kv; - // nope-only. every other MLA port splits q and k into a nope and a rope half and - // ropes the second one; here that half is zero-width, so there is no split, no - // concat and no ggml_rope_ext anywhere in the text tower + // nope-only: the rope half is zero-width, so no split, no concat and no rope + // anywhere in the text tower GGML_ASSERT(hparams.n_rot() == 0); - // the reference scales by qk_head_dim^-0.5, i.e. over the MLA head size, not over - // n_embd_head_k: after absorption q is kv_lora_rank wide and 1/sqrt(512) would be - // a different model + // scale is over the MLA head size, as in the reference, not over the post-absorption + // width: 1/sqrt(kv_lora_rank) = 1/sqrt(512) would be a different model const float kq_scale = 1.0f/sqrtf(float(qk_head_dim)); ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur); @@ -403,8 +396,8 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( // {qk_head_dim, kv_lora_rank, n_head} x {qk_head_dim, n_tokens, n_head} q = ggml_mul_mat(ctx0, layer.wk_b, q); - // {kv_lora_rank, n_head, n_tokens}. deepseek2 gets this contiguous for free out of - // the concat with the roped half, which does not exist here + // {kv_lora_rank, n_head, n_tokens}. deepseek2 gets this contiguous for free from the + // concat with the roped half, which does not exist here q = ggml_cont(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3)); cb(q, "dsa_q_absorbed", il); @@ -481,8 +474,8 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp = build_inp_embd(model.tok_embd); ggml_tensor * inp_out_ids = build_inp_out_ids(); - // MLA absorption leaves a K-only cache holding the kv_lora_rank latent, so the - // attention half of the hybrid memory is the _k variant, as in bailingmoe3 + // MLA absorption leaves a K-only cache holding the latent, so the attention half of + // the hybrid memory is the _k variant, as in bailingmoe3 llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); GGML_ASSERT(ubatch.n_seqs != 0); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 6f1855b0999..523d57a3dfb 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -120,8 +120,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_GLM5NEXT || arch == LLM_ARCH_MISTRAL4) { - // MLA absorbs into MQA, so the K cache row is the latent: n_head_kv must be 1 - // or the per-layer head_count_kv array below sizes it n_head times too wide + // MLA absorbs into MQA, so n_head_kv must be 1: otherwise the per-layer + // head_count_kv array below sizes the latent K row n_head times too wide n_embd = 128; n_head = 1; n_ff = 192; From 2db8295eef9ab15e39bc36f94e0cbb096bf594dc Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 09:36:50 +0000 Subject: [PATCH 10/36] llama: third kv cache for the glm5next lightning indexer The DSA layers of this model score pools of index_kpool consecutive positions rather than single keys, and the pooled key cannot be rebuilt from the MLA latents. llama_memory_hybrid therefore gains an optional third cache holding one indexer key and one compressor gate per token, so the hybrid carries the KDA conv+recurrent state, the MLA latents and the indexer keys at once. Absent unless filter_idx is given, which defaults to null, so every existing architecture gets exactly what it got before, state file layout included. Two heads per cell, not one. GLM's compressor is not a mean pool: it is a per-channel softmax over the kpool slots with logits gate + ape, where the gate is a second projection of the hidden state of width indexer_head_size. Caching it beside the key is the only way a pool survives its member tokens leaving the batch. Architectures with indexer_kpool == 0 still get one head. The indexer cache is handed the attention cache's slot layout rather than finding its own, so the two agree cell for cell, and apply() asserts they do. It also keeps its own dtype: -ctk q8_0 would otherwise quantise the gates, which feed a softmax. llama-kv-cache-kpool.{h,cpp} builds the pool <-> cell map host side. Pools are defined on positions and cells are whatever find_slot handed out, so the correspondence cannot be derived in the graph. Nothing here emits a negative index: ggml_set_rows asserts i1 >= 0, so unpopulated entries are clamped into range and neutralised by an additive -INFINITY instead. Two things the map does that the qwen4exp shape it is ported from does not: - the top-k budget is indexer_top_k exactly, with the always-selected tail biased to -INFINITY so it spends none of it, and forced back in through a host-built base mask for the scatter. indexer_top_k is a whole number of pools, so the cut lands on a pool boundary; the reference's own output width of indexer_top_k + kpool - 1 does not, and ggml_top_k is unordered among equals on both CPU and CUDA. - one map per ubatch, shared by every indexer layer, since nothing in it depends on the layer. Measured on a 16 Ki cell cache with 512 tokens: ~4 ms once against ~4 ms x n_layers. A unified cache with more than one sequence would let two sequences at the same position pool each other's keys, so create_memory refuses it up front rather than aborting mid-run. tests/test-glm5next-memory.cpp: 74 checks, 0 failures, on both the full and the trunk-only fixture. test-llama-archs is byte identical to the same build without this commit at a fixed seed: 452 rows, 0 FAIL. Session state files for qwen3next, falcon-h1, minimax-01, qwen35moe and a real Falcon-H1-0.5B are byte identical too, across write, reload and rewrite. --- src/llama-kv-cache-kpool.cpp | 231 +++++++ src/llama-kv-cache-kpool.h | 134 ++++ src/llama-kv-cache.cpp | 8 + src/llama-kv-cache.h | 13 + src/llama-memory-hybrid.cpp | 101 ++- src/llama-memory-hybrid.h | 22 +- src/llama-model.cpp | 42 +- tests/CMakeLists.txt | 2 + tests/test-glm5next-memory.cpp | 1071 ++++++++++++++++++++++++++++++++ 9 files changed, 1617 insertions(+), 7 deletions(-) create mode 100644 src/llama-kv-cache-kpool.cpp create mode 100644 src/llama-kv-cache-kpool.h create mode 100644 tests/test-glm5next-memory.cpp diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp new file mode 100644 index 00000000000..cf9f01ae0d2 --- /dev/null +++ b/src/llama-kv-cache-kpool.cpp @@ -0,0 +1,231 @@ +#include "llama-kv-cache-kpool.h" + +#include "llama-batch.h" +#include "llama-kv-cache.h" +#include "llama-kv-cells.h" + +#include +#include +#include + +uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool) { + GGML_ASSERT(kpool > 0); + + return n_kv/kpool + 2; +} + +uint32_t llama_kpool_top_k_width(uint32_t n_kv, uint32_t indexer_top_k, uint32_t kpool) { + GGML_ASSERT(kpool > 0); + GGML_ASSERT(indexer_top_k % kpool == 0 && "indexer_top_k must be a whole number of pools"); + + return std::min(n_kv, indexer_top_k); +} + +void llama_kv_cache_set_input_kpool( + const llama_kv_cache * kv, + ggml_tensor * cell_pool, + ggml_tensor * pool_cells, + ggml_tensor * bias, + ggml_tensor * sel_mask, + const llama_ubatch * ubatch, + uint32_t kpool) { + GGML_ASSERT(kv != nullptr); + GGML_ASSERT(kpool > 0); + + GGML_ASSERT(ggml_backend_buffer_is_host(cell_pool ->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(bias ->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(sel_mask ->buffer)); + + // sel_mask is KQ-mask shaped, and every KQ mask in the tree is f16 under + // flash attention. writing floats into one would overrun the allocation by + // 2x, so refuse rather than trust the caller + GGML_ASSERT(cell_pool ->type == GGML_TYPE_I32); + GGML_ASSERT(pool_cells->type == GGML_TYPE_I32); + GGML_ASSERT(bias ->type == GGML_TYPE_F32); + GGML_ASSERT(sel_mask ->type == GGML_TYPE_F32 && "sel_mask must be f32 even when the KQ mask is f16"); + + // everything below is written through raw strides + GGML_ASSERT(ggml_is_contiguous(cell_pool)); + GGML_ASSERT(ggml_is_contiguous(pool_cells)); + GGML_ASSERT(ggml_is_contiguous(bias)); + GGML_ASSERT(ggml_is_contiguous(sel_mask)); + + const int64_t n_kv = cell_pool->ne[0]; + const int64_t n_ns = cell_pool->ne[1]; // streams in this ubatch + const int64_t r = kpool; + const int64_t n_pools = pool_cells->ne[0]/r; + const int64_t n_tokens = ubatch->n_tokens; + + GGML_ASSERT(pool_cells->ne[0] % r == 0); + GGML_ASSERT(pool_cells->ne[1] == n_ns); + GGML_ASSERT(bias->ne[0] == n_kv && bias->ne[2] == n_ns); + GGML_ASSERT(sel_mask->ne[0] == n_kv && sel_mask->ne[2] == 1 && sel_mask->ne[3] == n_ns); + GGML_ASSERT(n_tokens % n_ns == 0); + + const int64_t n_tps = n_tokens/n_ns; // tokens per stream + const int64_t n_padq = sel_mask->ne[1]; // KQ mask rows, >= n_tps + + GGML_ASSERT(bias->ne[1] == n_tps); + GGML_ASSERT(n_padq >= n_tps); + + int32_t * dst_cell_pool = (int32_t *) cell_pool ->data; + int32_t * dst_pool_cells = (int32_t *) pool_cells->data; + float * dst_bias = (float *) bias ->data; + float * dst_sel_mask = (float *) sel_mask ->data; + + // -1 marks a cell with no usable pool. Kept host side only: it is never + // copied into cell_pool, which ggml_get_rows would read as an index + std::vector pool_of(n_kv); + std::vector filled(n_pools); + std::vector pos_at; + + // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM] + // A pool is a set of cells grouped by position, and cell positions are only + // unambiguous within one sequence. Under a unified cache all sequences share + // one cells array, so two sequences holding the same position would collide + // in pool_cells and silently pool each other's keys. The pooled indexer + // therefore needs one sequence per stream: either a non-unified cache + // (n_stream == n_seq_max, the default) or a single sequence in flight. + // qwen4exp's set_input_qsa has the same requirement and no check. + GGML_ASSERT((int64_t) ubatch->n_seqs_unq == n_ns && + "the pooled indexer needs one sequence per stream; use a non-unified KV cache"); + + for (int64_t s = 0; s < n_ns; ++s) { + // the token at ubatch index s*n_tps belongs to this stream; ask the cache + // which cells array that sequence actually uses. same convention as + // llama_kv_cache::set_input_kq_mask + const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; + const auto & cells = kv->get_cells(seq_of_stream); + + int32_t * cur_cell_pool = dst_cell_pool + s*n_kv; + int32_t * cur_pool_cells = dst_pool_cells + s*(r*n_pools); + + std::fill(pool_of.begin(), pool_of.end(), -1); + std::fill(filled.begin(), filled.end(), 0); + std::fill(cur_pool_cells, cur_pool_cells + r*n_pools, 0); + + // hoist occupancy and sequence membership out of the O(n_kv * n_tokens) + // loop below: neither depends on the query, and this stream holds exactly + // one sequence. -1 means the cell holds nothing this stream may pool or + // attend to. under a unified cache `cells` is shared with the sequences + // that are not in this ubatch, and seq_has is what keeps their keys out + pos_at.resize(n_kv); + for (int64_t j = 0; j < n_kv; ++j) { + pos_at[j] = cells.is_empty(j) || !cells.seq_has(j, seq_of_stream) ? -1 : cells.pos_get(j); + } + + // Pools are cut out of the position line, so a pool ordinal is p/kpool - + // an absolute number that can be far larger than n_kv/kpool. The array + // slot is a separate thing, so rebase on the lowest resident pool of this + // stream. Grouping is untouched: every member of a pool shifts together. + // + // Not the reference's anchor. HF pools from the first *resident* key + // (valid_keys.argmax(-1)), which for a left-padded batch differs from + // p/kpool; vLLM and SGLang both anchor at p/kpool exactly as here, and + // that is the only choice that keeps a pool's identity stable between the + // prefill that built it and the decode steps that read it. + // + // The window is n_pools wide. Positions are contiguous in any batch the + // model actually sees, so the whole resident range fits and the base is + // just the lowest pool. seq_rm can leave a hole large enough that it does + // not, and then the newest pools are the ones worth keeping. + int64_t b_base = 0; + { + int64_t b_min = 0; + int64_t b_max = 0; + bool found = false; + + for (int64_t j = 0; j < n_kv; ++j) { + if (pos_at[j] < 0) { + continue; + } + const int64_t b = pos_at[j]/r; + b_min = found ? std::min(b_min, b) : b; + b_max = found ? std::max(b_max, b) : b; + found = true; + } + + b_base = std::max(b_min, b_max - (n_pools - 1)); + } + + for (int64_t j = 0; j < n_kv; ++j) { + if (pos_at[j] < 0) { + continue; + } + + const llama_pos p = pos_at[j]; + const int64_t bo = p/r - b_base; + + if (bo < 0 || bo >= n_pools) { + continue; + } + + pool_of[j] = (int32_t) bo; + cur_pool_cells[bo*r + (p%r)] = (int32_t) j; + filled[bo]++; + } + + // a pool that is not completely resident cannot be pooled: the learned + // compressor consumes all r member keys, and the reference demands + // pool_valid = grouped_valid_keys.all(-1). Those cells are the tail of + // the sequence, which sel_mask forces in below whatever score they carry, + // so they are pointed at pool slot 0 only to keep the gather in range + for (int64_t j = 0; j < n_kv; ++j) { + // != rather than <: two cells claiming one position would overwrite + // each other in pool_cells, so such a pool is not usable either + if (pool_of[j] >= 0 && filled[pool_of[j]] != (int32_t) r) { + pool_of[j] = -1; + } + cur_cell_pool[j] = pool_of[j] < 0 ? 0 : pool_of[j]; + } + + float * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv); + + // the rows below n_tps are written in full by the loop; only the KQ mask's + // padding rows need clearing + std::fill(cur_sel_mask + n_tps*n_kv, cur_sel_mask + n_padq*n_kv, -INFINITY); + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + const llama_pos q = ubatch->pos[i]; + + // q >= 0 is what makes the unsigned range test below a range test + GGML_ASSERT(q >= 0 && ubatch->seq_id[i][0] == seq_of_stream); + + // everything from here on is inside the query's own incomplete pool + // and is always attended to (index_kpool_always_select_tail), which + // is what makes the selection land on pool boundaries. (q + 1) % r + // cells, the query's own token included + const llama_pos tail_start = (q + 1)/r*r; + + // the reference tests visibility at a pool's LAST member, so a pool + // that straddles the query is dropped whole rather than partially + // masked. Pools are position-aligned here, so pool b's last member is + // position b*r + r - 1 and the test collapses to b*r < tail_start + const int64_t bo_vis = std::max(0, tail_start/r - b_base); + + float * cur_bias = dst_bias + i*n_kv; + float * cur_sel = cur_sel_mask + ii*n_kv; + + // the unsigned compares fold "empty or another sequence" (pos_at -1) + // and "no usable pool" (pool_of -1) into the same branch as the range + // test, which is what lets this vectorise + for (int64_t j = 0; j < n_kv; ++j) { + const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; + const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; + const bool tail = pos_at[j] >= tail_start; + + cur_bias[j] = vis && pooled ? 0.0f : -INFINITY; + cur_sel [j] = vis && tail ? 0.0f : -INFINITY; + } + } + } +} + +void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { + mctx_idx->set_input_k_idxs(k_idxs, ubatch); + + llama_kv_cache_set_input_kpool( + mctx_attn->get_kv(), cell_pool, pool_cells, bias, sel_mask, ubatch, kpool); +} diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h new file mode 100644 index 00000000000..974aa1a2c44 --- /dev/null +++ b/src/llama-kv-cache-kpool.h @@ -0,0 +1,134 @@ +#pragma once + +#include "ggml.h" +#include "llama.h" +#include "llama-graph.h" + +#include + +struct llama_ubatch; +class llama_kv_cache; +class llama_kv_cache_context; + +// +// GLM-5-Next indexer pooling +// +// GLM's lightning indexer scores pools of `kpool` consecutive *positions* rather +// than single tokens, and its top-k budget is counted in tokens (indexer_top_k), +// i.e. indexer_top_k/kpool whole pools. The reference requires +// indexer_top_k % kpool == 0, so that division is exact. +// +// A pool is defined on positions, but everything the graph indexes - the indexer +// key cache, the MLA KQ mask - is addressed by *cell*, and a cell index is +// whatever llama_kv_cache::find_slot happened to hand out. Cells of one pool are +// not adjacent, not ordered, and with a unified cache not even owned by the same +// sequence. The mapping between the two therefore cannot be derived in the +// graph; it is built here, host side, from the cache's own cells, and handed to +// the graph as plain input tensors. This mirrors what qwen4exp's QSA does for +// its compression blocks. +// +// Nothing here ever emits a negative index: ggml_set_rows asserts i1 >= 0 and +// ggml_get_rows has no sentinel either, so unpopulated entries are clamped into +// range and neutralised by the additive masks instead. +// + +// number of pool slots the graph has to allocate for `n_kv` cells. +// +// pool ordinals are position-derived and then rebased on the lowest resident +// pool of each stream, so a sequence whose positions do not start at 0 still +// lands inside the array. rebasing can cost one slot at each end, hence the +2. +uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool); + +// width to run ggml_top_k at. +// +// NOT indexer_top_k + kpool - 1, which is the width of the reference's *output* +// buffer, tail included. ggml_top_k is explicitly unordered among equals +// (ggml/src/ggml-cpu/ops.cpp uses std::partial_sort and then swaps the first two +// results to make the point; the CUDA op declares determinism::not_guaranteed), +// so a budget that does not end on a pool boundary picks an arbitrary 1..kpool-1 +// members out of the pool it cuts, and picks differently on CPU and on CUDA. +// With the tail biased to -INFINITY the budget is spent only on whole pools, and +// indexer_top_k is a multiple of kpool by construction, so the cut is exact. The +// tail is forced back in through `sel_mask` instead of through the budget. +uint32_t llama_kpool_top_k_width(uint32_t n_kv, uint32_t indexer_top_k, uint32_t kpool); + +// Fill the host-side inputs of the pooled indexer. +// +// cell_pool I32 [n_kv, n_stream] +// for cell j of stream s: the pool slot it belongs to, or 0 when it has no +// usable pool. Used to broadcast a pool's score back onto its member cells +// with ggml_get_rows, so that ggml_top_k over the replicated per-cell +// scores yields CELL indices directly and the cut still lands on a pool +// boundary (a pool's members tie bit-exactly). +// +// pool_cells I32 [kpool*n_pools, n_stream] +// for pool slot p, member m of stream s: the cell holding that member's +// position, or 0 when it is not resident. Used to gather a pool's member +// keys and gates before the learned compressor mixes them. +// +// bias F32 [n_kv, n_tokens/n_stream, n_stream] +// additive per-(cell, query) bias on the indexer SCORE: +// 0.0f the cell is in a complete pool whose last member the query +// can see, which is the reference's `pool_valid & +// pool_visible` +// -INFINITY everything else, the trailing incomplete pool included, so +// that no part of the top-k budget is spent on it +// +// sel_mask F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) +// the tensor the top-k scatter starts from, in place of the +// ggml_fill(kq_mask, -INFINITY) that opens the DSA mask build in +// llm_graph_context::build_attn. Forcing the tail in here rather than +// through the score is what lets the budget stay pool-aligned: +// 0.0f the cell is in the query's own incomplete trailing pool, +// which GLM always attends to +// (index_kpool_always_select_tail) +// -INFINITY everything else, including the padding rows +// +// `kv` must be the ATTENTION (MLA) cache, since the cells that define the pools +// are the ones the top-k indices are ultimately read against. The indexer cache +// is given the attention cache's slot layout by llama_memory_hybrid, so the two +// agree cell for cell. +void llama_kv_cache_set_input_kpool( + const llama_kv_cache * kv, + ggml_tensor * cell_pool, + ggml_tensor * pool_cells, + ggml_tensor * bias, + ggml_tensor * sel_mask, + const llama_ubatch * ubatch, + uint32_t kpool); + +// One pooling map per ubatch, shared by every indexer layer. +// +// All four tensors depend only on the cells and the ubatch, never on the layer, +// and so does the indexer cache's k_idxs. Rebuilding them per layer costs +// O(n_kv * n_tokens) host writes each time: at 128 Ki cells, 512 tokens and 11 +// DSA layers that is ~11 x 67M float stores per ubatch, which dominates prefill. +// The model graph creates one of these before the layer loop and reads the same +// tensors in every layer. +// +// Sharing `bias` and `sel_mask` too is only correct while every indexer layer +// sees the same candidate set. That holds for glm5next, whose indexer_types are +// all "full"; an architecture that mixed windowed indexers into the same model +// would need one bias per window. +class llm_graph_input_kpool : public llm_graph_input_i { +public: + llm_graph_input_kpool( + const llama_kv_cache_context * mctx_attn, + const llama_kv_cache_context * mctx_idx, + uint32_t kpool) : mctx_attn(mctx_attn), mctx_idx(mctx_idx), kpool(kpool) {} + + ~llm_graph_input_kpool() = default; + + void set_input(const llama_ubatch * ubatch) override; + + ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] + ggml_tensor * cell_pool = nullptr; // I32 [n_kv, n_stream] + ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] + ggml_tensor * bias = nullptr; // F32 [n_kv, n_tps, n_stream] + ggml_tensor * sel_mask = nullptr; // F32 [n_kv, n_batch, 1, n_stream] + + const llama_kv_cache_context * mctx_attn; + const llama_kv_cache_context * mctx_idx; + + const uint32_t kpool; +}; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ec0f5a75314..ba9129267a9 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2585,6 +2585,14 @@ uint32_t llama_kv_cache_context::get_n_kv() const { return n_kv; } +uint32_t llama_kv_cache_context::get_n_stream() const { + return sinfos[i_cur].s1 - sinfos[i_cur].s0 + 1; +} + +const llama_kv_cache * llama_kv_cache_context::get_kv() const { + return kv; +} + ggml_type llama_kv_cache_context::type_k() const { return kv->type_k(); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 6cb6dbd2f98..e9ee352d10a 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -366,6 +366,19 @@ class llama_kv_cache_context : public llama_memory_context_i { uint32_t get_n_kv() const; + // streams covered by the current slot info, matching the `ns` that get_k and + // get_v use for their stream dimension. 1 for a unified cache. + // + // note: this is the stream RANGE s1 - s0 + 1, not the number of sequences in + // the ubatch. They differ when the active sequences are not a contiguous run + // of slots, which is exactly when a per-cell input sized from this would stop + // agreeing with a KQ mask, sized from n_seqs_unq + uint32_t get_n_stream() const; + + // the cache this context is a view of, for host-side inputs that have to + // resolve cell -> position through the cells themselves + const llama_kv_cache * get_kv() const; + ggml_type type_k() const; ggml_type type_v() const; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6..328e1f846a3 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -4,6 +4,8 @@ #include "llama-model.h" #include "llama-context.h" +#include + // // llama_memory_hybrid // @@ -29,8 +31,11 @@ llama_memory_hybrid::llama_memory_hybrid( bool unified, /* layer filters */ const layer_filter_cb & filter_attn, - const layer_filter_cb & filter_recr) : + const layer_filter_cb & filter_recr, + const layer_filter_cb & filter_idx, + ggml_type type_idx) : hparams(model.hparams), + hparams_idx(model.hparams), mem_attn(new llama_kv_cache( model, model.hparams, @@ -62,7 +67,35 @@ llama_memory_hybrid::llama_memory_hybrid( filter_recr == nullptr ? [&](int32_t il) { return hparams.is_recr(il); } : filter_recr - )) {} + )), + mem_idx(filter_idx == nullptr ? nullptr : [&] { + // MQA with a single key head of indexer_head_size, the same shaping + // llama_kv_cache_dsa applies to its lightning-indexer cache. note that + // n_embd_head_k_full is the field n_embd_head_k(il) reads for a non-SWA + // layer, so this is unaffected by the model being MLA: an MLA model + // keeps is_mla() true here, which is what suppresses the V allocation + // the indexer does not need. + // + // A *pooling* indexer (indexer_kpool > 0, glm5next only) needs a second + // head: its compressor mixes the kpool member keys with a per-channel + // softmax over gate scores that are a projection of the same hidden + // state, so the gate has to be cached alongside the key or the pool + // cannot be rebuilt once the member tokens have left the batch. Every + // other architecture leaves indexer_kpool at 0 and gets one head, byte + // for byte what it got before. + const uint32_t n_head_idx = model.hparams.indexer_kpool > 0 ? 2 : 1; + + std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), n_head_idx); + hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; + + LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells, %u x %u per cell, type = %s\n", + __func__, kv_size, n_head_idx, hparams_idx.n_embd_head_k_full, ggml_type_name(type_idx)); + + return new llama_kv_cache( + model, hparams_idx, type_idx, type_idx, v_trans, offload, unified, + kv_size, n_seq_max, n_pad, n_swa, swa_type, + nullptr, filter_idx, nullptr, nullptr); + }()) {} llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { do { @@ -115,8 +148,18 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } + // The indexer cache is a side buffer addressed by the attention cache's + // cells, so it takes that same slot layout rather than finding its own. + // Allocating separately lets the two drift apart once the context is + // being rewritten between turns, and the top-k indices, which are read + // against the attention mask, then point at the wrong cells. + llama_kv_cache::slot_info_vec_t heads_idx; + if (mem_idx) { + heads_idx = heads_attn; + } + return std::make_unique( - this, std::move(heads_attn), std::move(ubatches)); + this, std::move(heads_attn), std::move(ubatches), std::move(heads_idx)); } while(false); return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); @@ -132,11 +175,15 @@ llama_memory_context_ptr llama_memory_hybrid::init_update(llama_context * lctx, bool llama_memory_hybrid::get_can_shift() const { // Shifting is trivially supported for recurrent + if (mem_idx && !mem_idx->get_can_shift()) { + return false; + } return mem_attn->get_can_shift(); } void llama_memory_hybrid::clear(bool data) { mem_attn->clear(data); + if (mem_idx) mem_idx->clear(data); mem_recr->clear(data); } @@ -146,26 +193,31 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 if (!mem_recr->seq_rm(seq_id, p0, p1)) { return false; } + if (mem_idx) mem_idx->seq_rm(seq_id, p0, p1); return mem_attn->seq_rm(seq_id, p0, p1); } void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1); + if (mem_idx) mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1); mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1); } void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { mem_attn->seq_keep(seq_id); + if (mem_idx) mem_idx->seq_keep(seq_id); mem_recr->seq_keep(seq_id); } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { mem_attn->seq_add(seq_id, p0, p1, shift); + if (mem_idx) mem_idx->seq_add(seq_id, p0, p1, shift); mem_recr->seq_add(seq_id, p0, p1, shift); } void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { mem_attn->seq_div(seq_id, p0, p1, d); + if (mem_idx) mem_idx->seq_div(seq_id, p0, p1, d); mem_recr->seq_div(seq_id, p0, p1, d); } @@ -184,12 +236,20 @@ std::map llama_memory_hybrid::memory_breakdo for (const auto & buft_size : mem_recr->memory_breakdown()) { mb[buft_size.first] += buft_size.second; } + if (mem_idx) { + for (const auto & buft_size : mem_idx->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + } return mb; } void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_write(io, seq_id, flags); + // the indexer keys are not recomputable from the attention cache, so a + // restored session that skipped them would select the wrong cells + if (mem_idx) mem_idx->state_write(io, seq_id, flags); } mem_recr->state_write(io, seq_id, flags); } @@ -197,6 +257,7 @@ void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_read(io, seq_id, flags); + if (mem_idx) mem_idx->state_read(io, seq_id, flags); } mem_recr->state_read(io, seq_id, flags); } @@ -209,11 +270,16 @@ llama_memory_recurrent * llama_memory_hybrid::get_mem_recr() const { return mem_recr.get(); } +llama_kv_cache * llama_memory_hybrid::get_mem_idx() const { + return mem_idx.get(); +} + llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_status status) : status(status) {} llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_hybrid * mem) : ctx_attn(mem->get_mem_attn()->init_full()), ctx_recr(mem->get_mem_recr()->init_full()), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : mem->get_mem_idx()->init_full()), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } @@ -223,17 +289,27 @@ llama_memory_hybrid_context::llama_memory_hybrid_context( bool optimize) : ctx_attn(mem->get_mem_attn()->init_update(lctx, optimize)), ctx_recr(mem->get_mem_recr()->init_update(lctx, optimize)), + // the indexer keys carry no positional encoding, so a shift has nothing to + // correct in them, but the pending per-cell delta still has to be cleared or + // the two caches disagree about whether a shift is outstanding. an indexer + // only exists for LLAMA_ROPE_TYPE_NONE architectures, which is exactly the + // case where llama_kv_cache::update skips the K-shift graph and does only + // that + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : mem->get_mem_idx()->init_update(lctx, optimize)), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } llama_memory_hybrid_context::llama_memory_hybrid_context( llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, - std::vector ubatches) : + std::vector ubatches, + slot_info_vec_t sinfos_idx) : ubatches(std::move(ubatches)), // note: here we copy the ubatches. not sure if this is ideal ctx_attn(new llama_kv_cache_context(mem->get_mem_attn(), std::move(sinfos_attn), this->ubatches)), ctx_recr(new llama_memory_recurrent_context(mem->get_mem_recr(), this->ubatches)), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), this->ubatches)), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } @@ -242,6 +318,7 @@ bool llama_memory_hybrid_context::next() { ctx_attn->next(); ctx_recr->next(); + if (ctx_idx) ctx_idx->next(); if (++i_next >= ubatches.size()) { return false; @@ -258,6 +335,18 @@ bool llama_memory_hybrid_context::apply() { res = res & ctx_attn->apply(); res = res & ctx_recr->apply(); + if (ctx_idx) { + res = res & ctx_idx->apply(); + + // the indexer is addressed by the attention cache's cells, so a top-k + // over indexer cells is only meaningful if the two cover the same window. + // only the batch context has slot infos to compare + if (!ubatches.empty()) { + GGML_ASSERT(get_idx()->get_n_kv() == get_attn()->get_n_kv()); + GGML_ASSERT(get_idx()->get_n_stream() == get_attn()->get_n_stream()); + } + } + return res; } @@ -277,3 +366,7 @@ const llama_kv_cache_context * llama_memory_hybrid_context::get_attn() const { const llama_memory_recurrent_context * llama_memory_hybrid_context::get_recr() const { return static_cast(ctx_recr.get()); } + +const llama_kv_cache_context * llama_memory_hybrid_context::get_idx() const { + return static_cast(ctx_idx.get()); +} diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb7499..d6c88f590ea 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -39,7 +39,13 @@ class llama_memory_hybrid : public llama_memory_i { bool unified, /* layer filters */ const layer_filter_cb & filter_attn = nullptr, - const layer_filter_cb & filter_recr = nullptr); + const layer_filter_cb & filter_recr = nullptr, + /* optional per-token indexer key cache, for hybrid + models whose attention layers are sparse. absent + unless filter_idx is given, so every existing + architecture is unaffected */ + const layer_filter_cb & filter_idx = nullptr, + ggml_type type_idx = GGML_TYPE_F16); ~llama_memory_hybrid() = default; @@ -82,12 +88,18 @@ class llama_memory_hybrid : public llama_memory_i { llama_kv_cache * get_mem_attn() const; llama_memory_recurrent * get_mem_recr() const; + llama_kv_cache * get_mem_idx() const; // nullptr when the model has no indexer private: const llama_hparams & hparams; + // geometry for the indexer cache: n_head_kv key heads of indexer_head_size, + // mirroring how llama_kv_cache_dsa builds its own + llama_hparams hparams_idx; + const std::unique_ptr mem_attn; const std::unique_ptr mem_recr; + const std::unique_ptr mem_idx; }; class llama_memory_hybrid_context : public llama_memory_context_i { @@ -110,7 +122,11 @@ class llama_memory_hybrid_context : public llama_memory_context_i { llama_memory_hybrid_context( llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, - std::vector ubatches); + std::vector ubatches, + // empty unless the model has an indexer cache. the + // indexer is a side buffer addressed by the attention + // cache's cells, so it is handed that cache's slots + slot_info_vec_t sinfos_idx = {}); ~llama_memory_hybrid_context() = default; @@ -126,6 +142,7 @@ class llama_memory_hybrid_context : public llama_memory_context_i { const llama_kv_cache_context * get_attn() const; const llama_memory_recurrent_context * get_recr() const; + const llama_kv_cache_context * get_idx() const; // nullptr without an indexer private: // the index of the next ubatch to process @@ -135,6 +152,7 @@ class llama_memory_hybrid_context : public llama_memory_context_i { const llama_memory_context_ptr ctx_attn; const llama_memory_context_ptr ctx_recr; + const llama_memory_context_ptr ctx_idx; // null unless the model has an indexer const llama_memory_status status; }; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 462982d1bac..d664e09249c 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2434,6 +2434,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; + // left null for every architecture but the sparse-attention + // ones, which is what keeps the indexer cache from existing + llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; + ggml_type type_idx = GGML_TYPE_F16; if (arch == LLM_ARCH_FALCON_H1) { filter_attn = [&](uint32_t) { return true; }; filter_recr = [&](uint32_t) { return true; }; @@ -2451,9 +2455,43 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_recr = [&](uint32_t il) { return il < hparams.n_layer() && hparams.is_recr(il); }; + + if (arch == LLM_ARCH_GLM5NEXT && hparams.indexer_head_size > 0) { + // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM] + // the indexer pools cells by position, and a unified + // cache gives every sequence the same cells array, so + // two sequences at the same position would pool each + // other's keys. Refuse here rather than aborting deep + // inside a set_input several thousand tokens in + if (cparams.kv_unified && cparams.n_seq_max > 1) { + throw std::runtime_error("glm5next: the pooled indexer needs one sequence per stream, so a unified KV cache is only supported with a single sequence"); + } + + // the DSA layers carry a lightning-indexer key cache; + // the KDA layers and the NextN block do not + filter_idx = [&](uint32_t il) { + return il < hparams.n_layer() && !hparams.is_recr(il); + }; + + // the pooling indexer caches the compressor gate score + // next to the key, and the gate feeds a softmax, so + // -ctk q8_0 would quantise something far more + // sensitive than a key. keep the indexer float + type_idx = params.type_k; + if (ggml_is_quantized(type_idx)) { + LLAMA_LOG_WARN("%s: indexer key cache stays %s rather than %s: it also holds the compressor gates\n", + __func__, ggml_type_name(GGML_TYPE_F16), ggml_type_name(type_idx)); + type_idx = GGML_TYPE_F16; + } + } } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { + // llama_memory_hybrid_iswa has no indexer cache. glm5next + // is swa_type NONE so it never lands here, but a sparse + // hybrid with SWA would silently lose its indexer + GGML_ASSERT(filter_idx == nullptr && "hybrid-iswa cannot carry an indexer cache"); + // Use hybrid-iswa for hybrid models with SWA res = new llama_memory_hybrid_iswa( /* model */ *this, @@ -2491,7 +2529,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* offload */ cparams.offload_kqv, /* unified */ cparams.kv_unified, /* filter_attn */ std::move(filter_attn), - /* filter_recr */ std::move(filter_recr)); + /* filter_recr */ std::move(filter_recr), + /* filter_idx */ std::move(filter_idx), + /* type_idx */ type_idx); } } else { llama_kv_cache::layer_filter_cb filter = nullptr; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index b9f9d4b78af..2c793e582e0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -196,6 +196,8 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # llama_build_and_test(test-double-float.cpp) # SLOW llama_build_and_test(test-llama-archs.cpp) + # needs a glm5next GGUF as argv[1], so it is built but not registered + llama_build(test-glm5next-memory.cpp) set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") diff --git a/tests/test-glm5next-memory.cpp b/tests/test-glm5next-memory.cpp new file mode 100644 index 00000000000..04f31c26ce5 --- /dev/null +++ b/tests/test-glm5next-memory.cpp @@ -0,0 +1,1071 @@ +// The GLM-5-Next hybrid memory. +// +// Bar: the memory object can be created for a real glm5next GGUF, it holds all +// three halves (KDA recurrent + conv state, MLA latent KV, pooled indexer key +// cache), the sizes are the ones the reference implies, a trivial ggml graph can +// reach every one of them, and the pooled top-k cuts on a pool boundary and +// agrees between CPU and CUDA. The indexer graph itself is not built here. +// +// Run as: test-glm5next-memory +// The GGUF is the one tests/glm5next_make_tiny_gguf.py writes. + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "llama.h" + +#include "../src/llama-model.h" +#include "../src/llama-hparams.h" +#include "../src/llama-cparams.h" +#include "../src/llama-memory-hybrid.h" +#include "../src/llama-memory-recurrent.h" +#include "../src/llama-kv-cache.h" +#include "../src/llama-kv-cache-kpool.h" +#include "../src/llama-batch.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static int n_fail = 0; + +#define CHECK(cond, ...) \ + do { \ + if (!(cond)) { \ + printf("FAIL %s:%d: ", __func__, __LINE__); \ + printf(__VA_ARGS__); \ + printf("\n"); \ + n_fail++; \ + } else { \ + printf("ok "); \ + printf(__VA_ARGS__); \ + printf("\n"); \ + } \ + } while (0) + +// +// numeric evidence for the top-k width, independent of any model +// +// The claim: running ggml_top_k at exactly indexer_top_k with the trailing +// incomplete pool biased to -INFINITY cuts on a pool boundary and gives the same +// selection on CPU and on CUDA, where the reference's own output width +// (indexer_top_k + kpool - 1) with the tail forced in by a +1e9 score does not. +// + +// run ggml_top_k(scores, width) on one backend and return the selected indices +static std::vector run_top_k( + ggml_backend_t backend, + const std::vector & scores, + int64_t n_kv, + int64_t n_rows, + int64_t width) { + ggml_init_params gparams = { + /* .mem_size */ ggml_tensor_overhead()*8 + ggml_graph_overhead(), + /* .mem_buffer */ nullptr, + /* .no_alloc */ true, + }; + + ggml_context * ctx = ggml_init(gparams); + + ggml_tensor * src = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_kv, n_rows); + ggml_set_input(src); + + ggml_tensor * dst = ggml_top_k(ctx, src, (int) width); + ggml_set_output(dst); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, dst); + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + std::vector out(width*n_rows); + + if (buf) { + ggml_backend_tensor_set(src, scores.data(), 0, scores.size()*sizeof(float)); + + if (ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS) { + ggml_backend_tensor_get(dst, out.data(), 0, out.size()*sizeof(int32_t)); + } else { + out.clear(); + } + + ggml_backend_buffer_free(buf); + } else { + out.clear(); + } + + ggml_free(ctx); + + return out; +} + +// how many pools of the selection are only partly present +static int64_t n_partial_pools(const std::vector & sel, int64_t off, int64_t width, int64_t kpool) { + std::map cnt; + for (int64_t i = 0; i < width; ++i) { + cnt[sel[off + i]/(int32_t) kpool]++; + } + + int64_t n = 0; + for (const auto & kv : cnt) { + n += kv.second != (int32_t) kpool; + } + + return n; +} + +static void test_top_k_boundary() { + printf("\n--- top-k boundary, standalone ---\n"); + + const int64_t kpool = 4; + const int64_t top_k = 2048; // the reference requires top_k %% kpool == 0 + const int64_t n_kv = 8192; + const int64_t n_rows = 4; // queries + + // pool p's score. all kpool members of a pool carry it bit-identically, which + // is what makes an intra-pool tie harmless: the pool is taken whole or not at + // all as long as the budget ends on a pool boundary + const int64_t n_pools = n_kv/kpool; + + std::vector pool_score(n_pools); + for (int64_t p = 0; p < n_pools; ++p) { + // deterministic, distinct, no exact ties between pools + pool_score[p] = std::sin(0.7f*(float) p)*1000.0f + 0.001f*(float) p; + } + + ggml_backend_t backend_cpu = ggml_backend_cpu_init(); + ggml_backend_t backend_gpu = nullptr; + { + ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); + if (dev) { + backend_gpu = ggml_backend_dev_init(dev, nullptr); + } + } + + printf("%s GPU backend for the top-k comparison: %s\n", + backend_gpu ? "ok " : "note", backend_gpu ? ggml_backend_name(backend_gpu) : "none, CPU only"); + + // t = (q + 1) %% kpool cells of trailing incomplete pool. t == 3 is the one + // residue for which the reference's own width happens to stay pool-aligned + for (int64_t t = 0; t < kpool; ++t) { + const int64_t n_tail = t; + const int64_t n_full = n_kv - n_tail; // cells belonging to complete pools + + // ---- what this change does: budget = top_k, tail out of the budget ---- + std::vector s_fix((size_t) n_kv*n_rows); + // ---- what the prototype did: budget = top_k + kpool - 1, tail at +1e9 -- + std::vector s_old((size_t) n_kv*n_rows); + + for (int64_t r = 0; r < n_rows; ++r) { + for (int64_t j = 0; j < n_kv; ++j) { + const bool tail = j >= n_full; + + s_fix[r*n_kv + j] = tail ? -INFINITY : pool_score[j/kpool]; + s_old[r*n_kv + j] = tail ? 1e9f : pool_score[j/kpool]; + } + } + + const int64_t w_fix = llama_kpool_top_k_width((uint32_t) n_kv, (uint32_t) top_k, (uint32_t) kpool); + const int64_t w_old = std::min(n_kv, top_k + kpool - 1); + + const auto sel_fix_cpu = run_top_k(backend_cpu, s_fix, n_kv, n_rows, w_fix); + const auto sel_old_cpu = run_top_k(backend_cpu, s_old, n_kv, n_rows, w_old); + + CHECK(w_fix == top_k, "t=%d: top-k runs at exactly indexer_top_k (%d)", (int) t, (int) w_fix); + + // 1. boundary exactness of the new scheme + int64_t partial_fix = 0; + int64_t tail_in_fix = 0; + for (int64_t r = 0; r < n_rows; ++r) { + partial_fix += n_partial_pools(sel_fix_cpu, r*w_fix, w_fix, kpool); + for (int64_t i = 0; i < w_fix; ++i) { + tail_in_fix += sel_fix_cpu[r*w_fix + i] >= n_full; + } + } + CHECK(partial_fix == 0, "t=%d: the %d-cell budget selects only whole pools (%d partial)", + (int) t, (int) w_fix, (int) partial_fix); + CHECK(tail_in_fix == 0, "t=%d: no tail cell consumes budget (%d did)", (int) t, (int) tail_in_fix); + + // 2. the same run on the prototype's width + int64_t partial_old = 0; + for (int64_t r = 0; r < n_rows; ++r) { + partial_old += n_partial_pools(sel_old_cpu, r*w_old, w_old, kpool); + } + // the tail itself is one partial pool by construction whenever t > 0 + const int64_t expect_old = t == 3 ? (t > 0 ? 1 : 0) : (t > 0 ? 2 : 1); + CHECK(partial_old/n_rows == expect_old, + "t=%d: width %d leaves %d partial pool(s) per query, expected %d", + (int) t, (int) w_old, (int) (partial_old/n_rows), (int) expect_old); + + // 3. CPU vs CUDA on the same input + if (backend_gpu) { + const auto sel_fix_gpu = run_top_k(backend_gpu, s_fix, n_kv, n_rows, w_fix); + const auto sel_old_gpu = run_top_k(backend_gpu, s_old, n_kv, n_rows, w_old); + + bool same_fix = sel_fix_gpu.size() == sel_fix_cpu.size(); + for (int64_t r = 0; same_fix && r < n_rows; ++r) { + std::set a(sel_fix_cpu.begin() + r*w_fix, sel_fix_cpu.begin() + (r + 1)*w_fix); + std::set b(sel_fix_gpu.begin() + r*w_fix, sel_fix_gpu.begin() + (r + 1)*w_fix); + same_fix = a == b; + } + CHECK(same_fix, "t=%d: CPU and %s select the same cell set at width %d", + (int) t, ggml_backend_name(backend_gpu), (int) w_fix); + + bool same_old = sel_old_gpu.size() == sel_old_cpu.size(); + for (int64_t r = 0; same_old && r < n_rows; ++r) { + std::set a(sel_old_cpu.begin() + r*w_old, sel_old_cpu.begin() + (r + 1)*w_old); + std::set b(sel_old_gpu.begin() + r*w_old, sel_old_gpu.begin() + (r + 1)*w_old); + same_old = a == b; + } + // reported, not asserted: whether the arbitrary intra-pool pick + // actually diverges depends on each backend's partial sort + printf("note t=%d: CPU and %s %s at width %d\n", + (int) t, ggml_backend_name(backend_gpu), + same_old ? "happen to agree" : "DISAGREE", (int) w_old); + } + } + + if (backend_gpu) { + ggml_backend_free(backend_gpu); + } + ggml_backend_free(backend_cpu); +} + +// +// the memory object itself +// + +struct kpool_tensors { + ggml_tensor * cell_pool = nullptr; + ggml_tensor * pool_cells = nullptr; + ggml_tensor * bias = nullptr; + ggml_tensor * sel_mask = nullptr; +}; + +static kpool_tensors alloc_kpool_tensors( + ggml_context * ctx, + int64_t n_kv, + int64_t n_tps, + int64_t n_padq, + int64_t n_stream, + int64_t kpool, + int64_t n_pools) { + kpool_tensors t; + + t.cell_pool = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_kv, n_stream); + t.pool_cells = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, kpool*n_pools, n_stream); + t.bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_kv, n_tps, n_stream); + t.sel_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_padq, 1, n_stream); + + ggml_set_input(t.cell_pool); + ggml_set_input(t.pool_cells); + ggml_set_input(t.bias); + ggml_set_input(t.sel_mask); + + return t; +} + +int main(int argc, char ** argv) { + if (argc < 2) { + printf("usage: %s \n", argv[0]); + return 1; + } + + setvbuf(stdout, nullptr, _IOLBF, 0); + + llama_backend_init(); + + test_top_k_boundary(); + + printf("\n--- model ---\n"); + + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = 0; + + llama_model * model = llama_model_load_from_file(argv[1], mparams); + if (model == nullptr) { + printf("FAIL: could not load %s\n", argv[1]); + return 1; + } + + const auto & hparams = model->hparams; + + printf("n_layer = %u (+%u nextn)\n", hparams.n_layer(), hparams.n_layer_nextn); + printf("n_head = %u\n", hparams.n_head()); + printf("n_embd_head_kda = %u\n", hparams.n_embd_head_kda); + printf("ssm_d_conv = %u\n", hparams.ssm_d_conv); + printf("n_lora_kv = %u\n", hparams.n_lora_kv); + printf("indexer_head_size= %u\n", hparams.indexer_head_size); + printf("indexer_kpool = %u\n", hparams.indexer_kpool); + printf("indexer_top_k = %u\n", hparams.indexer_top_k); + printf("n_embd_r() = %u\n", hparams.n_embd_r()); + printf("n_embd_s() = %u\n", hparams.n_embd_s()); + + // ---- sizing ------------------------------------------------------------ + { + const uint32_t d_inner = hparams.n_head()*hparams.n_embd_head_kda; + + CHECK(hparams.n_embd_r() == 3*(hparams.ssm_d_conv - 1)*d_inner, + "n_embd_r == 3 conv states of (d_conv-1)*n_head*head_dim (%u)", hparams.n_embd_r()); + CHECK(hparams.n_embd_s() == hparams.n_embd_head_kda*hparams.n_embd_head_kda*hparams.n_head(), + "n_embd_s == head_dim*head_dim per head (%u)", hparams.n_embd_s()); + CHECK(hparams.indexer_kpool > 0 && hparams.indexer_top_k % hparams.indexer_kpool == 0, + "indexer_top_k (%u) is a whole number of pools of %u", hparams.indexer_top_k, hparams.indexer_kpool); + } + + // ---- a multi-sequence unified cache is rejected at create time ---------- + // + // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM]. Pools group cells by position, and a + // unified cache gives every sequence the same cells array, so two sequences + // at the same position would pool each other's keys. -kvu with --parallel is + // reachable (llama-perplexity forces it for hellaswag / winogrande / + // multiple-choice), so this has to fail at startup rather than abort inside + // a set_input several thousand tokens in. + { + llama_cparams cp = {}; + cp.n_ctx = 256; + cp.n_ctx_seq = 256; + cp.n_batch = 32; + cp.n_ubatch = 32; + cp.n_seq_max = 2; + cp.kv_unified = true; + cp.causal_attn = true; + + llama_memory_params mp = {}; + mp.type_k = GGML_TYPE_F16; + mp.type_v = GGML_TYPE_F16; + mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; + + bool threw = false; + try { + delete model->create_memory(mp, cp); + } catch (const std::exception &) { + threw = true; + } + CHECK(threw, "-kvu with n_seq_max 2 is refused when the model has a pooling indexer"); + + // one sequence in flight is fine: the shared cells array holds only that + // sequence's keys, and the map filters on seq_has anyway + cp.n_seq_max = 1; + llama_memory_i * raw = nullptr; + try { + raw = model->create_memory(mp, cp); + } catch (const std::exception &) { + raw = nullptr; + } + CHECK(raw != nullptr, "-kvu with a single sequence is still allowed"); + delete raw; + } + + // ---- the indexer cache keeps its own dtype ------------------------------ + { + llama_cparams cp = {}; + cp.n_ctx = 256; + cp.n_ctx_seq = 256; + cp.n_batch = 32; + cp.n_ubatch = 32; + cp.n_seq_max = 1; + cp.kv_unified = false; + cp.causal_attn = true; + + llama_memory_params mp = {}; + mp.type_k = GGML_TYPE_Q8_0; + mp.type_v = GGML_TYPE_Q8_0; + mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; + + llama_memory_i * raw = model->create_memory(mp, cp); + auto * m = dynamic_cast(raw); + + CHECK(m != nullptr && m->get_mem_idx() != nullptr, "-ctk q8_0 still builds an indexer cache"); + if (m && m->get_mem_idx()) { + CHECK(!ggml_is_quantized(m->get_mem_idx()->type_k()), + "indexer cache stays %s under -ctk q8_0: it also holds the compressor gates", + ggml_type_name(m->get_mem_idx()->type_k())); + CHECK(ggml_is_quantized(m->get_mem_attn()->type_k()), + "the MLA cache still honours -ctk q8_0 (%s)", ggml_type_name(m->get_mem_attn()->type_k())); + } + + delete raw; + } + + // ---- create the memory ------------------------------------------------- + llama_cparams cparams = {}; + cparams.n_ctx = 512; + cparams.n_ctx_seq = 512; + cparams.n_batch = 32; + cparams.n_ubatch = 32; + cparams.n_seq_max = 2; + cparams.n_rs_seq = 0; + cparams.kv_unified = false; // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM] + cparams.offload_kqv = false; + cparams.flash_attn = false; + cparams.causal_attn = true; + + llama_memory_params mparams_mem = {}; + mparams_mem.type_k = GGML_TYPE_F16; + mparams_mem.type_v = GGML_TYPE_F16; + mparams_mem.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; + mparams_mem.swa_full = false; + + llama_memory_i * mem_raw = model->create_memory(mparams_mem, cparams); + CHECK(mem_raw != nullptr, "create_memory returned a memory object"); + if (!mem_raw) { + return 1; + } + + auto * mem = dynamic_cast(mem_raw); + CHECK(mem != nullptr, "the memory is a llama_memory_hybrid"); + if (!mem) { + return 1; + } + + llama_kv_cache * kv_attn = mem->get_mem_attn(); + llama_memory_recurrent * rs = mem->get_mem_recr(); + llama_kv_cache * kv_idx = mem->get_mem_idx(); + + CHECK(kv_attn != nullptr, "hybrid holds an attention (MLA) cache"); + CHECK(rs != nullptr, "hybrid holds a recurrent (KDA) cache"); + CHECK(kv_idx != nullptr, "hybrid holds an indexer key cache"); + if (!kv_attn || !rs || !kv_idx) { + return 1; + } + + // ---- layer partition --------------------------------------------------- + { + const auto ids_attn = kv_attn->get_layer_ids(); + const auto ids_idx = kv_idx ->get_layer_ids(); + + uint32_t n_dsa = 0; + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + n_dsa += !hparams.is_recr(il); + } + + CHECK(ids_attn.size() == n_dsa, "MLA cache holds the %u DSA trunk layers (%zu)", n_dsa, ids_attn.size()); + CHECK(ids_idx.size() == n_dsa, "indexer cache holds the same %u layers (%zu)", n_dsa, ids_idx.size()); + + bool same = ids_attn.size() == ids_idx.size(); + for (size_t i = 0; same && i < ids_attn.size(); ++i) { + same = ids_attn[i] == ids_idx[i]; + } + CHECK(same, "MLA and indexer caches cover exactly the same layers"); + + for (uint32_t il : ids_attn) { + CHECK(!hparams.is_recr(il), "cached attention layer %u is not recurrent", il); + } + } + + // ---- cache geometry ---------------------------------------------------- + { + const uint32_t il_dsa = kv_attn->get_layer_ids().front(); + + ggml_tensor * k_mla = kv_attn->get_k_storage(il_dsa); + ggml_tensor * k_idx = kv_idx ->get_k_storage(il_dsa); + + CHECK(k_mla != nullptr && k_idx != nullptr, "both caches expose K storage for layer %u", il_dsa); + + // nope-only MLA: the latent row is kv_lora_rank + qk_rope_head_dim + CHECK(k_mla->ne[0] == (int64_t) hparams.n_embd_head_k(il_dsa), + "MLA latent row = %d (n_embd_head_k = %u)", (int) k_mla->ne[0], hparams.n_embd_head_k(il_dsa)); + + // the pooling indexer caches the key AND the compressor gate score + CHECK(k_idx->ne[0] == (int64_t) (2*hparams.indexer_head_size), + "indexer row = %d (expected 2 x indexer_head_size = %u)", + (int) k_idx->ne[0], 2*hparams.indexer_head_size); + + CHECK(k_mla->ne[1] == k_idx->ne[1], "MLA and indexer caches have the same cell count (%d)", (int) k_mla->ne[1]); + CHECK(k_mla->ne[2] == k_idx->ne[2], "MLA and indexer caches have the same stream count (%d)", (int) k_mla->ne[2]); + + // is_mla() stays true for the indexer hparams copy, which is what keeps + // it from allocating a V tensor it would never read + { + size_t bytes = 0; + for (const auto & b : kv_idx->memory_breakdown()) { + bytes += b.second; + } + + const size_t k_only = (size_t) ggml_nbytes(k_idx)*kv_idx->get_layer_ids().size(); + + // is_mla() is what suppresses the V allocation, and the indexer + // hparams copy does not touch the fields it reads, so the byte count + // is the observable: K only, no V + CHECK(hparams.is_mla() && bytes == k_only, + "indexer cache allocates K and no V: %zu bytes for %zu layers", + bytes, kv_idx->get_layer_ids().size()); + } + } + + // ---- recurrent geometry ------------------------------------------------ + { + uint32_t n_kda = 0; + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + n_kda += hparams.is_recr(il); + } + + CHECK(rs->size >= cparams.n_seq_max, "recurrent cache has >= n_seq_max slots (%u)", rs->size); + CHECK(n_kda > 0, "the model has %u KDA layers", n_kda); + } + + // ---- drive a batch through it and reach all three halves in one graph --- + { + const int32_t n_tokens = 10; // pos 9 leaves (9+1) %% kpool = 2 tail cells + + std::vector tokens(n_tokens*cparams.n_seq_max, 0); + std::vector pos; + std::vector seqs; + for (uint32_t s = 0; s < cparams.n_seq_max; ++s) { + for (int32_t i = 0; i < n_tokens; ++i) { + pos.push_back(i); + seqs.push_back((llama_seq_id) s); + } + } + + llama_batch batch = {}; + batch.n_tokens = n_tokens*(int32_t) cparams.n_seq_max; + batch.token = tokens.data(); + batch.pos = pos.data(); + + std::vector seq_ptrs(batch.n_tokens); + for (int32_t i = 0; i < batch.n_tokens; ++i) { + seq_ptrs[i] = &seqs[i]; + } + std::vector n_seq_id(batch.n_tokens, 1); + std::vector logits(batch.n_tokens, 0); + logits.back() = 1; + + batch.seq_id = seq_ptrs.data(); + batch.n_seq_id = n_seq_id.data(); + batch.logits = logits.data(); + + llama_batch_allocr balloc(hparams.n_pos_per_embd()); + const bool ok = balloc.init(batch, model->vocab, nullptr, hparams.n_embd_inp(), + cparams.n_seq_max, true); + CHECK(ok, "batch allocr accepted a %d-token, %u-sequence batch", batch.n_tokens, cparams.n_seq_max); + + auto mctx_ptr = mem->init_batch(balloc, cparams.n_ubatch, false); + CHECK(mctx_ptr != nullptr && mctx_ptr->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, + "init_batch produced a usable memory context"); + + auto * mctx = dynamic_cast(mctx_ptr.get()); + CHECK(mctx != nullptr, "the context is a llama_memory_hybrid_context"); + + if (mctx && mctx->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { + const auto * ctx_attn = mctx->get_attn(); + const auto * ctx_recr = mctx->get_recr(); + const auto * ctx_idx = mctx->get_idx(); + + CHECK(ctx_attn != nullptr, "context exposes the attention half"); + CHECK(ctx_recr != nullptr, "context exposes the recurrent half"); + CHECK(ctx_idx != nullptr, "context exposes the indexer half"); + + // n_kv is only valid once the context has been applied. apply() also + // asserts the two caches agree + mctx->apply(); + + // apply() asserts both of these itself, so reaching this line at all + // is the result; printed rather than CHECKed so the count stays honest + printf("note indexer and attention caches agree on n_kv (%u) and n_stream (%u)\n", + ctx_idx->get_n_kv(), ctx_idx->get_n_stream()); + CHECK(ctx_idx && ctx_idx->get_kv() == kv_idx, "the indexer context is a view of the indexer cache"); + + ggml_init_params gparams = { + /* .mem_size */ ggml_tensor_overhead()*1024 + ggml_graph_overhead(), + /* .mem_buffer */ nullptr, + /* .no_alloc */ true, + }; + + ggml_context * ctx0 = ggml_init(gparams); + ggml_cgraph * gf = ggml_new_graph(ctx0); + + const uint32_t il_dsa = kv_attn->get_layer_ids().front(); + + uint32_t il_kda = 0; + while (il_kda < hparams.n_layer() && !hparams.is_recr(il_kda)) { + il_kda++; + } + + const int64_t n_kv = ctx_attn->get_n_kv(); + const int64_t n_stream = ctx_attn->get_n_stream(); + const int64_t n_tps = mctx->get_ubatch().n_tokens/n_stream; + + // 1. MLA latent K + ggml_tensor * k_mla = ctx_attn->get_k(ctx0, il_dsa); + CHECK(k_mla != nullptr, "graph reaches the MLA latent cache: [%d, %d, %d, %d]", + (int) k_mla->ne[0], (int) k_mla->ne[1], (int) k_mla->ne[2], (int) k_mla->ne[3]); + + // 2. indexer keys, split into the key half and the gate half + ggml_tensor * k_idx_all = ctx_idx->get_k(ctx0, il_dsa); + CHECK(k_idx_all->ne[1] == 2, "indexer cache view has 2 heads (key | compressor gate)"); + + const int64_t d_idx = hparams.indexer_head_size; + + ggml_tensor * k_idx_v = ggml_view_3d(ctx0, k_idx_all, d_idx, n_kv, n_stream, + k_idx_all->nb[2], k_idx_all->nb[3], 0); + ggml_tensor * g_idx_v = ggml_view_3d(ctx0, k_idx_all, d_idx, n_kv, n_stream, + k_idx_all->nb[2], k_idx_all->nb[3], k_idx_all->nb[1]); + + // 3. KDA recurrent + conv state + ggml_tensor * r_kda = ctx_recr->get_r_l(il_kda); + ggml_tensor * s_kda = ctx_recr->get_s_l(il_kda); + CHECK(r_kda != nullptr && s_kda != nullptr, + "graph reaches the KDA conv state [%d x %d] and recurrent state [%d x %d]", + (int) r_kda->ne[0], (int) r_kda->ne[1], (int) s_kda->ne[0], (int) s_kda->ne[1]); + + // 4. host-side pool map, the piece the indexer needs and nothing else has + const int64_t kpool = hparams.indexer_kpool; + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); + const int64_t n_padq = n_tps; // this tree does not pad the KQ mask + + kpool_tensors kt = alloc_kpool_tensors(ctx0, n_kv, n_tps, n_padq, n_stream, kpool, n_pools); + + // the shape of the real thing: gather the pool members, mix them, + // score, broadcast the pool score back onto its member cells, top-k. + ggml_tensor * members = ggml_get_rows(ctx0, k_idx_v, kt.pool_cells); + ggml_tensor * gates = ggml_get_rows(ctx0, g_idx_v, kt.pool_cells); + members = ggml_reshape_4d(ctx0, members, d_idx, kpool, n_pools, n_stream); + gates = ggml_reshape_4d(ctx0, gates, d_idx, kpool, n_pools, n_stream); + + // stand-in for softmax(gate + ape) * member, summed over kpool + ggml_tensor * pooled = ggml_mul(ctx0, members, gates); + pooled = ggml_reshape_3d(ctx0, pooled, d_idx*kpool, n_pools, n_stream); + // stand-in for the q . k_pool score: one number per (pool, query) + ggml_tensor * score = ggml_mul_mat(ctx0, + ggml_cont(ctx0, ggml_view_3d(ctx0, pooled, d_idx, n_pools, n_stream, + pooled->nb[1], pooled->nb[2], 0)), + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, d_idx, n_tps, n_stream)); + + // broadcast pool -> member cells, then top-k over CELL scores + ggml_tensor * expanded = ggml_get_rows(ctx0, + ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), kt.cell_pool); + expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); + expanded = ggml_add(ctx0, expanded, kt.bias); + + const int64_t width = llama_kpool_top_k_width((uint32_t) n_kv, hparams.indexer_top_k, (uint32_t) kpool); + ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, (int) width)); + + printf("note top-k over replicated per-cell scores yields %d I32 CELL indices\n", (int) width); + + // 5. the scatter the mask is built from, starting at sel_mask rather + // than at an all -INFINITY fill, which is what forces the tail in + ggml_tensor * top_k_4d = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream); + ggml_tensor * base = ggml_view_4d(ctx0, kt.sel_mask, 1, n_kv, n_padq, n_stream, + kt.sel_mask->nb[0], kt.sel_mask->nb[1], kt.sel_mask->nb[2], 0); + ggml_tensor * idxs = ggml_view_4d(ctx0, top_k_4d, width, n_tps, n_stream, 1, + top_k_4d->nb[1], top_k_4d->nb[2], n_stream*top_k_4d->nb[3], 0); + ggml_tensor * zeros = ggml_fill(ctx0, + ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, width, n_tps, n_stream), 0.0f); + ggml_tensor * mask = ggml_set_rows(ctx0, base, zeros, idxs); + + ggml_build_forward_expand(gf, k_mla); + ggml_build_forward_expand(gf, r_kda); + ggml_build_forward_expand(gf, s_kda); + ggml_build_forward_expand(gf, mask); + + printf("note one graph reaches all three halves in %d nodes\n", ggml_graph_n_nodes(gf)); + + // 6. the host-side pool map itself, on a real allocated buffer + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx0, backend); + CHECK(buf != nullptr, "allocated the graph's input tensors on the CPU backend"); + + if (buf) { + const llama_ubatch & ub = mctx->get_ubatch(); + llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.sel_mask, + &ub, (uint32_t) kpool); + + const int32_t * cp = (const int32_t *) kt.cell_pool->data; + const int32_t * pc = (const int32_t *) kt.pool_cells->data; + const float * bi = (const float *) kt.bias->data; + const float * sm = (const float *) kt.sel_mask->data; + + bool in_range = true; + for (int64_t i = 0; i < ggml_nelements(kt.cell_pool); ++i) { + in_range &= cp[i] >= 0 && cp[i] < n_pools; + } + for (int64_t i = 0; i < ggml_nelements(kt.pool_cells); ++i) { + in_range &= pc[i] >= 0 && pc[i] < n_kv; + } + CHECK(in_range, "every emitted index is non-negative and in range (ggml_set_rows asserts i1 >= 0)"); + + bool finite = true; + for (int64_t i = 0; i < ggml_nelements(kt.bias); ++i) { + finite &= bi[i] == 0.0f || bi[i] == -INFINITY; + } + for (int64_t i = 0; i < ggml_nelements(kt.sel_mask); ++i) { + finite &= sm[i] == 0.0f || sm[i] == -INFINITY; + } + CHECK(finite, "bias and sel_mask hold only 0 and -INFINITY: no +1e9 to meet a -inf"); + + // every query of every stream: (q+1) %% kpool cells sit in its own + // incomplete pool and must be forced in, the q+1-minus-that cells + // below it must be scored, and nothing else may be either. Both + // streams, so that the per-stream strides are covered and not + // just stream 0's, which is the one at offset 0 + { + const llama_ubatch & u = mctx->get_ubatch(); + + for (int64_t s = 0; s < n_stream; ++s) { + bool ok_tail = true; + bool ok_scored = true; + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const llama_pos q = u.pos[s*n_tps + ii]; + const int64_t t = (q + 1) % kpool; + + const float * row_s = sm + (s*n_padq + ii)*n_kv; + const float * row_b = bi + (s*n_tps + ii)*n_kv; + + int64_t n_forced = 0; + int64_t n_scored = 0; + for (int64_t j = 0; j < n_kv; ++j) { + n_forced += row_s[j] == 0.0f; + n_scored += row_b[j] == 0.0f; + } + + ok_tail &= n_forced == t; + ok_scored &= n_scored == (q + 1) - t; + } + + CHECK(ok_tail, "stream %d: every query forces in exactly its (q+1) %% %d tail cells", + (int) s, (int) kpool); + CHECK(ok_scored, "stream %d: and scores exactly the cells below that tail", + (int) s); + } + } + + // pool origin: while the window starts at position 0, pos/kpool + // and the reference's first-resident-key anchor are the same + // grouping, so the map must reproduce it exactly. Cells whose + // pool is not complete carry slot 0 and are masked instead + { + const llama_ubatch & u = mctx->get_ubatch(); + bool agree = true; + for (int64_t s = 0; s < n_stream; ++s) { + const llama_seq_id seq = u.seq_id[s*n_tps][0]; + const auto & cells = kv_attn->get_cells(seq); + + std::map members; + llama_pos p_min = -1; + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j)) { + continue; + } + members[cells.pos_get(j)/kpool]++; + if (p_min < 0 || cells.pos_get(j) < p_min) { + p_min = cells.pos_get(j); + } + } + agree &= p_min == 0; // HF's first_key would also be 0 + + for (int64_t j = 0; j < n_kv && agree; ++j) { + if (cells.is_empty(j)) { + continue; + } + const int64_t b = cells.pos_get(j)/kpool; + agree &= cp[s*n_kv + j] == (int32_t) (members[b] == kpool ? b : 0); + } + } + CHECK(agree, "pool ordinals match the reference anchor while the window starts at position 0"); + } + + // the graph actually runs, and nothing in it trips an assert + { + ggml_status st = ggml_backend_graph_compute(backend, gf); + CHECK(st == GGML_STATUS_SUCCESS, "the pooled-indexer-shaped graph computes (%d)", (int) st); + } + + ggml_backend_buffer_free(buf); + } + + ggml_backend_free(backend); + ggml_free(ctx0); + } + } + + // ---- pool origin across a front eviction -------------------------------- + // + // The reference anchors a pool at the first *resident* key + // (valid_keys.argmax(-1)); this port anchors at pos/kpool, as vLLM and SGLang + // do. They are the same grouping until the front of the window is dropped by + // a non-multiple of kpool, and then they differ: the reference regroups every + // surviving key, this port does not. Regrouping is what a cache cannot + // afford, since the pooled key a decode step scores was built during prefill. + { + printf("\n--- pool origin ---\n"); + + // seq 0 currently holds positions 0..9. drop 0..1, which is not a + // multiple of kpool = 4 + const llama_pos n_drop = 2; + mem->seq_rm(0, 0, n_drop); + + std::vector tok(1, 0); + std::vector pos(1, 10); + std::vector sid(1, 0); + std::vector sptr(1, sid.data()); + std::vector nsid(1, 1); + std::vector lg(1, 1); + + llama_batch b = {}; + b.n_tokens = 1; + b.token = tok.data(); + b.pos = pos.data(); + b.seq_id = sptr.data(); + b.n_seq_id = nsid.data(); + b.logits = lg.data(); + + llama_batch_allocr ba(hparams.n_pos_per_embd()); + if (ba.init(b, model->vocab, nullptr, hparams.n_embd_inp(), cparams.n_seq_max, true)) { + auto c = mem->init_batch(ba, cparams.n_ubatch, false); + auto * mc = dynamic_cast(c.get()); + + CHECK(mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, + "one more token fits after the eviction"); + + if (mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { + mc->apply(); + + const int64_t n_kv = mc->get_attn()->get_n_kv(); + const int64_t kpool = hparams.indexer_kpool; + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); + + ggml_init_params gp = { ggml_tensor_overhead()*16, nullptr, true }; + ggml_context * ctx = ggml_init(gp); + + kpool_tensors kt = alloc_kpool_tensors(ctx, n_kv, 1, 1, 1, kpool, n_pools); + + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + if (buf) { + llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.sel_mask, + &mc->get_ubatch(), (uint32_t) kpool); + + const int32_t * cp = (const int32_t *) kt.cell_pool->data; + const auto & cells = kv_attn->get_cells(0); + + // positions 4..7 were one pool before the eviction and must + // still be one pool, or the pooled key built during prefill no + // longer describes the cells it is scored against + std::map slot_of; + for (int64_t j = 0; j < n_kv; ++j) { + if (!cells.is_empty(j) && cells.seq_has(j, 0)) { + slot_of[cells.pos_get(j)] = cp[j]; + } + } + + const bool grouped = + slot_of.count(4) && slot_of.count(5) && slot_of.count(6) && slot_of.count(7) && + slot_of[4] == slot_of[5] && slot_of[5] == slot_of[6] && slot_of[6] == slot_of[7]; + CHECK(grouped, "positions 4..7 stay one pool after the eviction (slot %d)", + grouped ? (int) slot_of[4] : -1); + + // the reference's anchor would have made positions 2..5 the + // first pool instead. checked, not implemented: this port + // deliberately differs here + const bool differs = !slot_of.count(2) || slot_of[2] != slot_of[4]; + CHECK(differs, "positions 2 and 4 are NOT pooled together, where the reference anchor would"); + + // the leading remnant 2..3 is an incomplete pool, so it is not + // scored. it is also not in the query's tail, which is the one + // case where a cell visible to the query is neither + const float * bi = (const float *) kt.bias->data; + const float * sm = (const float *) kt.sel_mask->data; + int64_t n_orphan = 0; + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j) || !cells.seq_has(j, 0)) { + continue; + } + n_orphan += bi[j] == -INFINITY && sm[j] == -INFINITY; + } + CHECK(n_orphan == 2, "the %d-cell leading remnant is neither pooled nor tail (%d)", + (int) n_drop, (int) n_orphan); + + ggml_backend_buffer_free(buf); + } + + ggml_backend_free(backend); + ggml_free(ctx); + } + } + } + + delete mem_raw; + + // ---- the shared input: one fill per ubatch, not one per DSA layer ------- + { + printf("\n--- shared input ---\n"); + + llama_cparams cp = {}; + cp.n_ctx = 16384; + cp.n_ctx_seq = 16384; + cp.n_batch = 512; + cp.n_ubatch = 512; + cp.n_seq_max = 1; + cp.kv_unified = false; + cp.causal_attn = true; + + llama_memory_params mp = {}; + mp.type_k = GGML_TYPE_F16; + mp.type_v = GGML_TYPE_F16; + mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; + + llama_memory_i * raw = model->create_memory(mp, cp); + auto * m = dynamic_cast(raw); + + if (m) { + llama_kv_cache * kv = m->get_mem_attn(); + + const int64_t n_step = cp.n_ubatch; + const int64_t n_fill = cp.n_ctx_seq - n_step; + const uint32_t n_dsa = (uint32_t) kv->get_layer_ids().size(); + const int64_t kpool = hparams.indexer_kpool; + + std::vector tok(n_step, 0); + std::vector pos(n_step); + std::vector sid(n_step, 0); + std::vector sptr(n_step); + std::vector nsid(n_step, 1); + std::vector lg(n_step, 0); + + for (int64_t i = 0; i < n_step; ++i) { + sptr[i] = &sid[i]; + } + lg.back() = 1; + + llama_memory_context_ptr keep; + + for (int64_t base = 0; base <= n_fill; base += n_step) { + for (int64_t i = 0; i < n_step; ++i) { + pos[i] = (llama_pos) (base + i); + } + + llama_batch b = {}; + b.n_tokens = (int32_t) n_step; + b.token = tok.data(); + b.pos = pos.data(); + b.seq_id = sptr.data(); + b.n_seq_id = nsid.data(); + b.logits = lg.data(); + + llama_batch_allocr ba(hparams.n_pos_per_embd()); + if (!ba.init(b, model->vocab, nullptr, hparams.n_embd_inp(), cp.n_seq_max, true)) { + break; + } + + auto c = m->init_batch(ba, cp.n_ubatch, false); + if (!c || c->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { + break; + } + c->apply(); + keep = std::move(c); + } + + auto * mctx = dynamic_cast(keep.get()); + CHECK(mctx && mctx->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, + "filled a %d-cell cache in %d-token ubatches", (int) cp.n_ctx_seq, (int) cp.n_ubatch); + if (mctx && mctx->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { + const int64_t n_kv = mctx->get_attn()->get_n_kv(); + const int64_t n_tps = mctx->get_ubatch().n_tokens; + const int64_t n_padq = GGML_PAD(n_tps, 8) + 8; // exercise the padded-mask path + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); + + ggml_init_params gp = { + /* .mem_size */ ggml_tensor_overhead()*16, + /* .mem_buffer */ nullptr, + /* .no_alloc */ true, + }; + + ggml_context * ctx = ggml_init(gp); + kpool_tensors kt = alloc_kpool_tensors(ctx, n_kv, n_tps, n_padq, 1, kpool, n_pools); + + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + if (buf) { + const llama_ubatch & ub = mctx->get_ubatch(); + + // first pass faults in 60+ MiB of fresh pages; time the steady + // state, which is what a prefill actually pays per ubatch + double ms = 1e9; + for (int rep = 0; rep < 4; ++rep) { + const auto t0 = std::chrono::steady_clock::now(); + llama_kv_cache_set_input_kpool(kv, kt.cell_pool, kt.pool_cells, kt.bias, kt.sel_mask, + &ub, (uint32_t) kpool); + const auto t1 = std::chrono::steady_clock::now(); + + if (rep > 0) { + ms = std::min(ms, std::chrono::duration(t1 - t0).count()); + } + } + + printf("note n_kv = %d, n_tokens = %d, %u DSA layers\n", (int) n_kv, (int) n_tps, n_dsa); + printf("note one fill = %.2f ms; shared = %.2f ms/ubatch, per layer = %.2f ms/ubatch\n", + ms, ms, ms*n_dsa); + printf("note extrapolated to 128Ki x 512 x 11 layers: %.0f ms shared, %.0f ms per layer\n", + ms*(131072.0/(double) n_kv)*(512.0/(double) n_tps), + ms*(131072.0/(double) n_kv)*(512.0/(double) n_tps)*11.0); + + CHECK(n_dsa > 1, "the model has %u indexer layers, so sharing one fill saves %ux the host writes", + n_dsa, n_dsa); + + { + const float * sm = (const float *) kt.sel_mask->data; + + bool pad_masked = true; + for (int64_t ii = n_tps; ii < n_padq; ++ii) { + for (int64_t j = 0; j < n_kv; ++j) { + pad_masked &= sm[ii*n_kv + j] == -INFINITY; + } + } + CHECK(pad_masked, "the %d padding rows of a wider sel_mask stay -INFINITY", + (int) (n_padq - n_tps)); + } + + // the shared object refills the same tensors deterministically + std::vector first((size_t) ggml_nelements(kt.cell_pool)); + memcpy(first.data(), kt.cell_pool->data, first.size()*sizeof(int32_t)); + + llm_graph_input_kpool inp(mctx->get_attn(), mctx->get_idx(), (uint32_t) kpool); + inp.k_idxs = mctx->get_idx()->build_input_k_idxs(ctx, ub); + inp.cell_pool = kt.cell_pool; + inp.pool_cells = kt.pool_cells; + inp.bias = kt.bias; + inp.sel_mask = kt.sel_mask; + + ggml_backend_buffer_t buf2 = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (buf2) { + inp.set_input(&ub); + + CHECK(memcmp(first.data(), kt.cell_pool->data, first.size()*sizeof(int32_t)) == 0, + "llm_graph_input_kpool::set_input rebuilds the identical map, so one\n" + " object can back every indexer layer"); + + ggml_backend_buffer_free(buf2); + } + + ggml_backend_buffer_free(buf); + } + + ggml_backend_free(backend); + ggml_free(ctx); + } else { + printf("note could not fill a %d-cell cache; skipping the cost measurement\n", (int) cp.n_ctx_seq); + } + + keep.reset(); + } + + delete raw; + } + + llama_model_free(model); + llama_backend_free(); + + printf("\n%s: %d failure(s)\n", argv[0], n_fail); + return n_fail == 0 ? 0 : 1; +} From 839597c62c632ad8aa3cbb4281b54d66daf652a4 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 14:48:51 +0000 Subject: [PATCH 11/36] glm5next: trim pooled indexer comments --- src/llama-kv-cache-kpool.cpp | 89 ++++++++++------------- src/llama-kv-cache-kpool.h | 115 +++++++++++++---------------- src/llama-kv-cache.h | 15 ++-- src/llama-memory-hybrid.cpp | 50 ++++++------- src/llama-memory-hybrid.h | 15 ++-- src/llama-model.cpp | 29 ++++---- tests/test-glm5next-memory.cpp | 128 ++++++++++++++------------------- 7 files changed, 186 insertions(+), 255 deletions(-) diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index cf9f01ae0d2..c3f4cc56708 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -37,9 +37,8 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(ggml_backend_buffer_is_host(bias ->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(sel_mask ->buffer)); - // sel_mask is KQ-mask shaped, and every KQ mask in the tree is f16 under - // flash attention. writing floats into one would overrun the allocation by - // 2x, so refuse rather than trust the caller + // sel_mask is KQ-mask shaped, and KQ masks are f16 under flash attention; writing + // floats into one would overrun the allocation 2x, so check rather than trust GGML_ASSERT(cell_pool ->type == GGML_TYPE_I32); GGML_ASSERT(pool_cells->type == GGML_TYPE_I32); GGML_ASSERT(bias ->type == GGML_TYPE_F32); @@ -74,26 +73,23 @@ void llama_kv_cache_set_input_kpool( float * dst_bias = (float *) bias ->data; float * dst_sel_mask = (float *) sel_mask ->data; - // -1 marks a cell with no usable pool. Kept host side only: it is never - // copied into cell_pool, which ggml_get_rows would read as an index + // -1 marks a cell with no usable pool. host side only: never copied into cell_pool, + // where ggml_get_rows would read it as an index std::vector pool_of(n_kv); std::vector filled(n_pools); std::vector pos_at; // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM] - // A pool is a set of cells grouped by position, and cell positions are only - // unambiguous within one sequence. Under a unified cache all sequences share - // one cells array, so two sequences holding the same position would collide - // in pool_cells and silently pool each other's keys. The pooled indexer - // therefore needs one sequence per stream: either a non-unified cache - // (n_stream == n_seq_max, the default) or a single sequence in flight. - // qwen4exp's set_input_qsa has the same requirement and no check. + // positions are unambiguous only within one sequence, and a unified cache shares one + // cells array, so two sequences at the same position would collide in pool_cells and + // silently pool each other's keys. hence one sequence per stream: a non-unified cache + // (n_stream == n_seq_max, the default) or a single sequence in flight. qwen4exp's + // set_input_qsa has the same requirement and no check. GGML_ASSERT((int64_t) ubatch->n_seqs_unq == n_ns && "the pooled indexer needs one sequence per stream; use a non-unified KV cache"); for (int64_t s = 0; s < n_ns; ++s) { - // the token at ubatch index s*n_tps belongs to this stream; ask the cache - // which cells array that sequence actually uses. same convention as + // which cells array this stream's sequence uses; same convention as // llama_kv_cache::set_input_kq_mask const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; const auto & cells = kv->get_cells(seq_of_stream); @@ -105,31 +101,27 @@ void llama_kv_cache_set_input_kpool( std::fill(filled.begin(), filled.end(), 0); std::fill(cur_pool_cells, cur_pool_cells + r*n_pools, 0); - // hoist occupancy and sequence membership out of the O(n_kv * n_tokens) - // loop below: neither depends on the query, and this stream holds exactly - // one sequence. -1 means the cell holds nothing this stream may pool or - // attend to. under a unified cache `cells` is shared with the sequences - // that are not in this ubatch, and seq_has is what keeps their keys out + // hoist occupancy and sequence membership out of the O(n_kv * n_tokens) loop + // below; neither depends on the query. -1 means the cell holds nothing this + // stream may pool or attend to. under a unified cache `cells` is shared with + // sequences outside this ubatch, and seq_has is what keeps their keys out pos_at.resize(n_kv); for (int64_t j = 0; j < n_kv; ++j) { pos_at[j] = cells.is_empty(j) || !cells.seq_has(j, seq_of_stream) ? -1 : cells.pos_get(j); } - // Pools are cut out of the position line, so a pool ordinal is p/kpool - - // an absolute number that can be far larger than n_kv/kpool. The array - // slot is a separate thing, so rebase on the lowest resident pool of this - // stream. Grouping is untouched: every member of a pool shifts together. + // a pool ordinal is the absolute p/kpool, which can far exceed n_kv/kpool, so + // rebase on this stream's lowest resident pool. grouping is untouched, every + // member shifts together. // - // Not the reference's anchor. HF pools from the first *resident* key - // (valid_keys.argmax(-1)), which for a left-padded batch differs from - // p/kpool; vLLM and SGLang both anchor at p/kpool exactly as here, and - // that is the only choice that keeps a pool's identity stable between the - // prefill that built it and the decode steps that read it. + // anchoring at p/kpool follows vLLM and SGLang, not HF (which pools from the + // first *resident* key, valid_keys.argmax(-1), differing under left padding). + // it is the only anchor that keeps a pool's identity stable between the prefill + // that built it and the decodes that read it. // - // The window is n_pools wide. Positions are contiguous in any batch the - // model actually sees, so the whole resident range fits and the base is - // just the lowest pool. seq_rm can leave a hole large enough that it does - // not, and then the newest pools are the ones worth keeping. + // the window is n_pools wide; positions are contiguous in any real batch so the + // resident range fits. seq_rm can leave a hole large enough that it does not, + // and then the newest pools are the ones worth keeping. int64_t b_base = 0; { int64_t b_min = 0; @@ -166,14 +158,13 @@ void llama_kv_cache_set_input_kpool( filled[bo]++; } - // a pool that is not completely resident cannot be pooled: the learned - // compressor consumes all r member keys, and the reference demands - // pool_valid = grouped_valid_keys.all(-1). Those cells are the tail of - // the sequence, which sel_mask forces in below whatever score they carry, - // so they are pointed at pool slot 0 only to keep the gather in range + // an incompletely resident pool cannot be pooled: the compressor consumes all r + // member keys, and the reference demands pool_valid = grouped_valid_keys.all(-1). + // such cells are the sequence tail, which sel_mask forces in below regardless of + // score, so they point at pool slot 0 purely to keep the gather in range for (int64_t j = 0; j < n_kv; ++j) { - // != rather than <: two cells claiming one position would overwrite - // each other in pool_cells, so such a pool is not usable either + // != rather than <: two cells claiming one position overwrite each other in + // pool_cells, so an over-filled pool is not usable either if (pool_of[j] >= 0 && filled[pool_of[j]] != (int32_t) r) { pool_of[j] = -1; } @@ -182,8 +173,7 @@ void llama_kv_cache_set_input_kpool( float * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv); - // the rows below n_tps are written in full by the loop; only the KQ mask's - // padding rows need clearing + // the loop below writes rows < n_tps in full; only the padding rows need clearing std::fill(cur_sel_mask + n_tps*n_kv, cur_sel_mask + n_padq*n_kv, -INFINITY); for (int64_t ii = 0; ii < n_tps; ++ii) { @@ -193,24 +183,21 @@ void llama_kv_cache_set_input_kpool( // q >= 0 is what makes the unsigned range test below a range test GGML_ASSERT(q >= 0 && ubatch->seq_id[i][0] == seq_of_stream); - // everything from here on is inside the query's own incomplete pool - // and is always attended to (index_kpool_always_select_tail), which - // is what makes the selection land on pool boundaries. (q + 1) % r - // cells, the query's own token included + // the query's own incomplete pool ((q + 1) % r cells, its own token + // included) is always attended to (index_kpool_always_select_tail), which is + // what makes the selection land on pool boundaries const llama_pos tail_start = (q + 1)/r*r; // the reference tests visibility at a pool's LAST member, so a pool - // that straddles the query is dropped whole rather than partially - // masked. Pools are position-aligned here, so pool b's last member is - // position b*r + r - 1 and the test collapses to b*r < tail_start + // straddling the query is dropped whole. pools are position-aligned here, so + // that test collapses to b*r < tail_start const int64_t bo_vis = std::max(0, tail_start/r - b_base); float * cur_bias = dst_bias + i*n_kv; float * cur_sel = cur_sel_mask + ii*n_kv; - // the unsigned compares fold "empty or another sequence" (pos_at -1) - // and "no usable pool" (pool_of -1) into the same branch as the range - // test, which is what lets this vectorise + // the unsigned compares fold "empty or another sequence" (pos_at -1) and "no + // usable pool" (pool_of -1) into the range test, which lets this vectorise for (int64_t j = 0; j < n_kv; ++j) { const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index 974aa1a2c44..2385a6cfe8c 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -13,81 +13,69 @@ class llama_kv_cache_context; // // GLM-5-Next indexer pooling // -// GLM's lightning indexer scores pools of `kpool` consecutive *positions* rather -// than single tokens, and its top-k budget is counted in tokens (indexer_top_k), -// i.e. indexer_top_k/kpool whole pools. The reference requires -// indexer_top_k % kpool == 0, so that division is exact. +// GLM's lightning indexer scores pools of `kpool` consecutive *positions*, while its +// top-k budget is counted in tokens (indexer_top_k), i.e. indexer_top_k/kpool whole +// pools. The reference requires indexer_top_k % kpool == 0, so the division is exact. // -// A pool is defined on positions, but everything the graph indexes - the indexer -// key cache, the MLA KQ mask - is addressed by *cell*, and a cell index is -// whatever llama_kv_cache::find_slot happened to hand out. Cells of one pool are -// not adjacent, not ordered, and with a unified cache not even owned by the same -// sequence. The mapping between the two therefore cannot be derived in the -// graph; it is built here, host side, from the cache's own cells, and handed to -// the graph as plain input tensors. This mirrors what qwen4exp's QSA does for -// its compression blocks. +// Pools are defined on positions, but everything the graph indexes is addressed by +// *cell*, and a cell index is whatever llama_kv_cache::find_slot handed out: cells of +// one pool are not adjacent, not ordered, and under a unified cache not even owned by +// the same sequence. The mapping therefore cannot be derived in the graph. It is built +// here, host side, from the cache's cells, and passed in as plain input tensors, as +// qwen4exp's QSA does for its compression blocks. // -// Nothing here ever emits a negative index: ggml_set_rows asserts i1 >= 0 and -// ggml_get_rows has no sentinel either, so unpopulated entries are clamped into -// range and neutralised by the additive masks instead. +// Nothing here may emit a negative index: ggml_set_rows asserts i1 >= 0 and +// ggml_get_rows has no sentinel, so unpopulated entries are clamped into range and +// neutralised by the additive masks instead. // -// number of pool slots the graph has to allocate for `n_kv` cells. +// number of pool slots the graph must allocate for `n_kv` cells. // -// pool ordinals are position-derived and then rebased on the lowest resident -// pool of each stream, so a sequence whose positions do not start at 0 still -// lands inside the array. rebasing can cost one slot at each end, hence the +2. +// pool ordinals are position-derived, then rebased on each stream's lowest resident +// pool so sequences not starting at position 0 still land inside the array. rebasing +// can cost one slot at each end, hence the +2. uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool); // width to run ggml_top_k at. // -// NOT indexer_top_k + kpool - 1, which is the width of the reference's *output* -// buffer, tail included. ggml_top_k is explicitly unordered among equals -// (ggml/src/ggml-cpu/ops.cpp uses std::partial_sort and then swaps the first two -// results to make the point; the CUDA op declares determinism::not_guaranteed), -// so a budget that does not end on a pool boundary picks an arbitrary 1..kpool-1 -// members out of the pool it cuts, and picks differently on CPU and on CUDA. -// With the tail biased to -INFINITY the budget is spent only on whole pools, and -// indexer_top_k is a multiple of kpool by construction, so the cut is exact. The -// tail is forced back in through `sel_mask` instead of through the budget. +// NOT indexer_top_k + kpool - 1: that is the width of the reference's *output* buffer, +// not a top-k width. ggml_top_k is explicitly unordered among equals (the CPU op +// deliberately swaps the first two results; the CUDA op declares +// determinism::not_guaranteed), so a budget that does not end on a pool boundary would +// take an arbitrary, CPU/CUDA-divergent subset of the pool it cuts. Biasing the tail to +// -INFINITY keeps the budget on whole pools; the tail is forced back in via `sel_mask`. +// See the PR body for the full argument. uint32_t llama_kpool_top_k_width(uint32_t n_kv, uint32_t indexer_top_k, uint32_t kpool); // Fill the host-side inputs of the pooled indexer. // // cell_pool I32 [n_kv, n_stream] -// for cell j of stream s: the pool slot it belongs to, or 0 when it has no -// usable pool. Used to broadcast a pool's score back onto its member cells -// with ggml_get_rows, so that ggml_top_k over the replicated per-cell -// scores yields CELL indices directly and the cut still lands on a pool -// boundary (a pool's members tie bit-exactly). +// cell -> its pool slot, or 0 when it has no usable pool. ggml_get_rows +// broadcasts a pool's score onto its members, so ggml_top_k over the replicated +// scores yields CELL indices directly and still cuts on a pool boundary (members +// tie bit-exactly). // // pool_cells I32 [kpool*n_pools, n_stream] -// for pool slot p, member m of stream s: the cell holding that member's -// position, or 0 when it is not resident. Used to gather a pool's member -// keys and gates before the learned compressor mixes them. +// pool slot, member -> the cell holding that position, or 0 when not resident. +// gathers a pool's member keys and gates for the compressor. // // bias F32 [n_kv, n_tokens/n_stream, n_stream] -// additive per-(cell, query) bias on the indexer SCORE: -// 0.0f the cell is in a complete pool whose last member the query -// can see, which is the reference's `pool_valid & -// pool_visible` -// -INFINITY everything else, the trailing incomplete pool included, so -// that no part of the top-k budget is spent on it +// additive per-(cell, query) bias on the indexer SCORE: 0.0f in a complete pool +// whose last member the query can see (the reference's `pool_valid & +// pool_visible`), else -INFINITY, the trailing incomplete pool included so no +// budget is spent on it. // // sel_mask F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) -// the tensor the top-k scatter starts from, in place of the -// ggml_fill(kq_mask, -INFINITY) that opens the DSA mask build in -// llm_graph_context::build_attn. Forcing the tail in here rather than -// through the score is what lets the budget stay pool-aligned: -// 0.0f the cell is in the query's own incomplete trailing pool, -// which GLM always attends to -// (index_kpool_always_select_tail) -// -INFINITY everything else, including the padding rows +// what the top-k scatter starts from, replacing the ggml_fill(kq_mask, -INFINITY) +// opening the DSA mask build in llm_graph_context::build_attn: 0.0f for the +// query's own incomplete trailing pool, which GLM always attends to +// (index_kpool_always_select_tail), else -INFINITY, padding rows included. +// forcing the tail in here rather than through the score keeps the budget +// pool-aligned. // -// `kv` must be the ATTENTION (MLA) cache, since the cells that define the pools -// are the ones the top-k indices are ultimately read against. The indexer cache -// is given the attention cache's slot layout by llama_memory_hybrid, so the two -// agree cell for cell. +// `kv` must be the ATTENTION (MLA) cache: its cells define the pools and are what the +// top-k indices are ultimately read against. llama_memory_hybrid gives the indexer +// cache the attention cache's slot layout, so the two agree cell for cell. void llama_kv_cache_set_input_kpool( const llama_kv_cache * kv, ggml_tensor * cell_pool, @@ -97,19 +85,14 @@ void llama_kv_cache_set_input_kpool( const llama_ubatch * ubatch, uint32_t kpool); -// One pooling map per ubatch, shared by every indexer layer. +// One pooling map per ubatch, shared by every indexer layer: all four tensors, and +// k_idxs, depend only on the cells and the ubatch, never on the layer. Rebuilding them +// per layer costs O(n_kv * n_tokens) host writes each time (at 128 Ki cells, 512 tokens +// and 11 DSA layers, ~11 x 67M float stores per ubatch, which dominates prefill). // -// All four tensors depend only on the cells and the ubatch, never on the layer, -// and so does the indexer cache's k_idxs. Rebuilding them per layer costs -// O(n_kv * n_tokens) host writes each time: at 128 Ki cells, 512 tokens and 11 -// DSA layers that is ~11 x 67M float stores per ubatch, which dominates prefill. -// The model graph creates one of these before the layer loop and reads the same -// tensors in every layer. -// -// Sharing `bias` and `sel_mask` too is only correct while every indexer layer -// sees the same candidate set. That holds for glm5next, whose indexer_types are -// all "full"; an architecture that mixed windowed indexers into the same model -// would need one bias per window. +// Sharing `bias` and `sel_mask` is correct only while every indexer layer sees the same +// candidate set. That holds for glm5next, whose indexer_types are all "full"; a model +// mixing in windowed indexers would need one bias per window. class llm_graph_input_kpool : public llm_graph_input_i { public: llm_graph_input_kpool( diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index e9ee352d10a..983352d617c 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -366,17 +366,14 @@ class llama_kv_cache_context : public llama_memory_context_i { uint32_t get_n_kv() const; - // streams covered by the current slot info, matching the `ns` that get_k and - // get_v use for their stream dimension. 1 for a unified cache. - // - // note: this is the stream RANGE s1 - s0 + 1, not the number of sequences in - // the ubatch. They differ when the active sequences are not a contiguous run - // of slots, which is exactly when a per-cell input sized from this would stop - // agreeing with a KQ mask, sized from n_seqs_unq + // streams covered by the current slot info, matching the `ns` get_k/get_v use for + // their stream dimension. 1 for a unified cache. note: this is the stream RANGE + // s1 - s0 + 1, not n_seqs_unq; they differ when the active sequences are not a + // contiguous run of slots, i.e. exactly when a per-cell input sized from this would + // stop agreeing with a KQ mask uint32_t get_n_stream() const; - // the cache this context is a view of, for host-side inputs that have to - // resolve cell -> position through the cells themselves + // the cache this context views, for host-side inputs that resolve cell -> position const llama_kv_cache * get_kv() const; ggml_type type_k() const; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 328e1f846a3..3d842f99f84 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -69,20 +69,15 @@ llama_memory_hybrid::llama_memory_hybrid( : filter_recr )), mem_idx(filter_idx == nullptr ? nullptr : [&] { - // MQA with a single key head of indexer_head_size, the same shaping - // llama_kv_cache_dsa applies to its lightning-indexer cache. note that - // n_embd_head_k_full is the field n_embd_head_k(il) reads for a non-SWA - // layer, so this is unaffected by the model being MLA: an MLA model - // keeps is_mla() true here, which is what suppresses the V allocation - // the indexer does not need. + // MQA with one key head of indexer_head_size, as llama_kv_cache_dsa shapes its + // lightning-indexer cache. n_embd_head_k_full is what n_embd_head_k(il) reads + // for a non-SWA layer, so an MLA model still has is_mla() true here, which + // suppresses the V allocation the indexer does not need. // - // A *pooling* indexer (indexer_kpool > 0, glm5next only) needs a second - // head: its compressor mixes the kpool member keys with a per-channel - // softmax over gate scores that are a projection of the same hidden - // state, so the gate has to be cached alongside the key or the pool - // cannot be rebuilt once the member tokens have left the batch. Every - // other architecture leaves indexer_kpool at 0 and gets one head, byte - // for byte what it got before. + // a *pooling* indexer (indexer_kpool > 0, glm5next only) needs a second head: + // its compressor gate is a projection of the same hidden state, so it must be + // cached alongside the key or the pool cannot be rebuilt once the member tokens + // leave the batch. every other arch leaves indexer_kpool 0 and is unchanged. const uint32_t n_head_idx = model.hparams.indexer_kpool > 0 ? 2 : 1; std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), n_head_idx); @@ -148,11 +143,10 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } - // The indexer cache is a side buffer addressed by the attention cache's - // cells, so it takes that same slot layout rather than finding its own. - // Allocating separately lets the two drift apart once the context is - // being rewritten between turns, and the top-k indices, which are read - // against the attention mask, then point at the wrong cells. + // the indexer is a side buffer addressed by the attention cache's cells, so it + // takes that slot layout rather than finding its own: allocating separately lets + // the two drift apart when the context is rewritten between turns, and the top-k + // indices, read against the attention mask, would then point at the wrong cells llama_kv_cache::slot_info_vec_t heads_idx; if (mem_idx) { heads_idx = heads_attn; @@ -247,8 +241,8 @@ std::map llama_memory_hybrid::memory_breakdo void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_write(io, seq_id, flags); - // the indexer keys are not recomputable from the attention cache, so a - // restored session that skipped them would select the wrong cells + // indexer keys are not recomputable from the attention cache, so a restored + // session that skipped them would select the wrong cells if (mem_idx) mem_idx->state_write(io, seq_id, flags); } mem_recr->state_write(io, seq_id, flags); @@ -289,12 +283,11 @@ llama_memory_hybrid_context::llama_memory_hybrid_context( bool optimize) : ctx_attn(mem->get_mem_attn()->init_update(lctx, optimize)), ctx_recr(mem->get_mem_recr()->init_update(lctx, optimize)), - // the indexer keys carry no positional encoding, so a shift has nothing to - // correct in them, but the pending per-cell delta still has to be cleared or - // the two caches disagree about whether a shift is outstanding. an indexer - // only exists for LLAMA_ROPE_TYPE_NONE architectures, which is exactly the - // case where llama_kv_cache::update skips the K-shift graph and does only - // that + // indexer keys carry no positional encoding, so a shift has nothing to correct in + // them, but the pending per-cell delta must still be cleared or the two caches + // disagree about whether a shift is outstanding. an indexer only exists for + // LLAMA_ROPE_TYPE_NONE archs, which is exactly when llama_kv_cache::update skips the + // K-shift graph and does only that ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : mem->get_mem_idx()->init_update(lctx, optimize)), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } @@ -338,9 +331,8 @@ bool llama_memory_hybrid_context::apply() { if (ctx_idx) { res = res & ctx_idx->apply(); - // the indexer is addressed by the attention cache's cells, so a top-k - // over indexer cells is only meaningful if the two cover the same window. - // only the batch context has slot infos to compare + // a top-k over indexer cells is meaningful only if both caches cover the same + // window. only the batch context has slot infos to compare if (!ubatches.empty()) { GGML_ASSERT(get_idx()->get_n_kv() == get_attn()->get_n_kv()); GGML_ASSERT(get_idx()->get_n_stream() == get_attn()->get_n_stream()); diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index d6c88f590ea..92e576332e7 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -40,10 +40,8 @@ class llama_memory_hybrid : public llama_memory_i { /* layer filters */ const layer_filter_cb & filter_attn = nullptr, const layer_filter_cb & filter_recr = nullptr, - /* optional per-token indexer key cache, for hybrid - models whose attention layers are sparse. absent - unless filter_idx is given, so every existing - architecture is unaffected */ + /* optional per-token indexer key cache for sparse-attention + hybrids; absent unless filter_idx is given */ const layer_filter_cb & filter_idx = nullptr, ggml_type type_idx = GGML_TYPE_F16); @@ -93,8 +91,8 @@ class llama_memory_hybrid : public llama_memory_i { private: const llama_hparams & hparams; - // geometry for the indexer cache: n_head_kv key heads of indexer_head_size, - // mirroring how llama_kv_cache_dsa builds its own + // indexer cache geometry: n_head_kv key heads of indexer_head_size, as + // llama_kv_cache_dsa builds its own llama_hparams hparams_idx; const std::unique_ptr mem_attn; @@ -123,9 +121,8 @@ class llama_memory_hybrid_context : public llama_memory_context_i { llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, std::vector ubatches, - // empty unless the model has an indexer cache. the - // indexer is a side buffer addressed by the attention - // cache's cells, so it is handed that cache's slots + // empty without an indexer. the indexer is addressed by the + // attention cache's cells, so it gets that cache's slots slot_info_vec_t sinfos_idx = {}); ~llama_memory_hybrid_context() = default; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d664e09249c..15968434481 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2434,8 +2434,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; - // left null for every architecture but the sparse-attention - // ones, which is what keeps the indexer cache from existing + // null for every arch but the sparse-attention ones, which is what + // keeps the indexer cache from existing llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; ggml_type type_idx = GGML_TYPE_F16; if (arch == LLM_ARCH_FALCON_H1) { @@ -2458,25 +2458,22 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, if (arch == LLM_ARCH_GLM5NEXT && hparams.indexer_head_size > 0) { // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM] - // the indexer pools cells by position, and a unified - // cache gives every sequence the same cells array, so - // two sequences at the same position would pool each - // other's keys. Refuse here rather than aborting deep - // inside a set_input several thousand tokens in + // the indexer pools cells by position and a unified cache + // shares one cells array, so two sequences at the same + // position would pool each other's keys. refuse here rather + // than abort inside a set_input thousands of tokens in if (cparams.kv_unified && cparams.n_seq_max > 1) { throw std::runtime_error("glm5next: the pooled indexer needs one sequence per stream, so a unified KV cache is only supported with a single sequence"); } - // the DSA layers carry a lightning-indexer key cache; - // the KDA layers and the NextN block do not + // only the DSA layers carry an indexer key cache filter_idx = [&](uint32_t il) { return il < hparams.n_layer() && !hparams.is_recr(il); }; - // the pooling indexer caches the compressor gate score - // next to the key, and the gate feeds a softmax, so - // -ctk q8_0 would quantise something far more - // sensitive than a key. keep the indexer float + // the gate cached next to the key feeds a softmax, so -ctk + // q8_0 would quantise something far more sensitive than a + // key. keep the indexer float type_idx = params.type_k; if (ggml_is_quantized(type_idx)) { LLAMA_LOG_WARN("%s: indexer key cache stays %s rather than %s: it also holds the compressor gates\n", @@ -2487,9 +2484,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { - // llama_memory_hybrid_iswa has no indexer cache. glm5next - // is swa_type NONE so it never lands here, but a sparse - // hybrid with SWA would silently lose its indexer + // llama_memory_hybrid_iswa has no indexer cache; glm5next is + // swa_type NONE, but a sparse hybrid with SWA would silently + // lose its indexer GGML_ASSERT(filter_idx == nullptr && "hybrid-iswa cannot carry an indexer cache"); // Use hybrid-iswa for hybrid models with SWA diff --git a/tests/test-glm5next-memory.cpp b/tests/test-glm5next-memory.cpp index 04f31c26ce5..a6802a072bb 100644 --- a/tests/test-glm5next-memory.cpp +++ b/tests/test-glm5next-memory.cpp @@ -1,13 +1,11 @@ -// The GLM-5-Next hybrid memory. +// The GLM-5-Next hybrid memory: the memory object holds all three halves (KDA +// recurrent + conv state, MLA latent KV, pooled indexer key cache) at the sizes the +// reference implies, one ggml graph reaches every one of them, and the pooled top-k +// cuts on a pool boundary and agrees between CPU and CUDA. The indexer graph itself is +// not built here. // -// Bar: the memory object can be created for a real glm5next GGUF, it holds all -// three halves (KDA recurrent + conv state, MLA latent KV, pooled indexer key -// cache), the sizes are the ones the reference implies, a trivial ggml graph can -// reach every one of them, and the pooled top-k cuts on a pool boundary and -// agrees between CPU and CUDA. The indexer graph itself is not built here. -// -// Run as: test-glm5next-memory -// The GGUF is the one tests/glm5next_make_tiny_gguf.py writes. +// Run as: test-glm5next-memory , as written by +// tests/glm5next_make_tiny_gguf.py #include "ggml.h" #include "ggml-alloc.h" @@ -50,15 +48,12 @@ static int n_fail = 0; } while (0) // -// numeric evidence for the top-k width, independent of any model -// -// The claim: running ggml_top_k at exactly indexer_top_k with the trailing -// incomplete pool biased to -INFINITY cuts on a pool boundary and gives the same -// selection on CPU and on CUDA, where the reference's own output width -// (indexer_top_k + kpool - 1) with the tail forced in by a +1e9 score does not. +// numeric evidence for the top-k width, independent of any model: ggml_top_k at exactly +// indexer_top_k with the trailing incomplete pool biased to -INFINITY cuts on a pool +// boundary and agrees between CPU and CUDA, where the reference's output width +// (indexer_top_k + kpool - 1) with the tail forced in at +1e9 does not // -// run ggml_top_k(scores, width) on one backend and return the selected indices static std::vector run_top_k( ggml_backend_t backend, const std::vector & scores, @@ -105,7 +100,6 @@ static std::vector run_top_k( return out; } -// how many pools of the selection are only partly present static int64_t n_partial_pools(const std::vector & sel, int64_t off, int64_t width, int64_t kpool) { std::map cnt; for (int64_t i = 0; i < width; ++i) { @@ -128,9 +122,8 @@ static void test_top_k_boundary() { const int64_t n_kv = 8192; const int64_t n_rows = 4; // queries - // pool p's score. all kpool members of a pool carry it bit-identically, which - // is what makes an intra-pool tie harmless: the pool is taken whole or not at - // all as long as the budget ends on a pool boundary + // all kpool members of a pool carry the pool score bit-identically, so an + // intra-pool tie is harmless as long as the budget ends on a pool boundary const int64_t n_pools = n_kv/kpool; std::vector pool_score(n_pools); @@ -151,8 +144,8 @@ static void test_top_k_boundary() { printf("%s GPU backend for the top-k comparison: %s\n", backend_gpu ? "ok " : "note", backend_gpu ? ggml_backend_name(backend_gpu) : "none, CPU only"); - // t = (q + 1) %% kpool cells of trailing incomplete pool. t == 3 is the one - // residue for which the reference's own width happens to stay pool-aligned + // t = (q + 1) %% kpool cells of trailing incomplete pool. t == 3 is the one residue + // where the reference's own width happens to stay pool-aligned for (int64_t t = 0; t < kpool; ++t) { const int64_t n_tail = t; const int64_t n_full = n_kv - n_tail; // cells belonging to complete pools @@ -223,8 +216,8 @@ static void test_top_k_boundary() { std::set b(sel_old_gpu.begin() + r*w_old, sel_old_gpu.begin() + (r + 1)*w_old); same_old = a == b; } - // reported, not asserted: whether the arbitrary intra-pool pick - // actually diverges depends on each backend's partial sort + // reported, not asserted: whether the arbitrary pick actually diverges + // depends on each backend's partial sort printf("note t=%d: CPU and %s %s at width %d\n", (int) t, ggml_backend_name(backend_gpu), same_old ? "happen to agree" : "DISAGREE", (int) w_old); @@ -321,12 +314,10 @@ int main(int argc, char ** argv) { // ---- a multi-sequence unified cache is rejected at create time ---------- // - // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM]. Pools group cells by position, and a - // unified cache gives every sequence the same cells array, so two sequences - // at the same position would pool each other's keys. -kvu with --parallel is - // reachable (llama-perplexity forces it for hellaswag / winogrande / - // multiple-choice), so this has to fail at startup rather than abort inside - // a set_input several thousand tokens in. + // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM]. pools group cells by position and a unified + // cache shares one cells array, so two sequences at the same position would pool + // each other's keys. -kvu with --parallel is reachable (llama-perplexity forces it + // for hellaswag / winogrande / multiple-choice), so this must fail at startup. { llama_cparams cp = {}; cp.n_ctx = 256; @@ -350,8 +341,8 @@ int main(int argc, char ** argv) { } CHECK(threw, "-kvu with n_seq_max 2 is refused when the model has a pooling indexer"); - // one sequence in flight is fine: the shared cells array holds only that - // sequence's keys, and the map filters on seq_has anyway + // one sequence in flight is fine: the shared cells array holds only its keys, + // and the map filters on seq_has anyway cp.n_seq_max = 1; llama_memory_i * raw = nullptr; try { @@ -481,8 +472,7 @@ int main(int argc, char ** argv) { CHECK(k_mla->ne[1] == k_idx->ne[1], "MLA and indexer caches have the same cell count (%d)", (int) k_mla->ne[1]); CHECK(k_mla->ne[2] == k_idx->ne[2], "MLA and indexer caches have the same stream count (%d)", (int) k_mla->ne[2]); - // is_mla() stays true for the indexer hparams copy, which is what keeps - // it from allocating a V tensor it would never read + // is_mla() stays true for the indexer hparams copy, so no V is allocated { size_t bytes = 0; for (const auto & b : kv_idx->memory_breakdown()) { @@ -491,9 +481,7 @@ int main(int argc, char ** argv) { const size_t k_only = (size_t) ggml_nbytes(k_idx)*kv_idx->get_layer_ids().size(); - // is_mla() is what suppresses the V allocation, and the indexer - // hparams copy does not touch the fields it reads, so the byte count - // is the observable: K only, no V + // the byte count is the observable: K only, no V CHECK(hparams.is_mla() && bytes == k_only, "indexer cache allocates K and no V: %zu bytes for %zu layers", bytes, kv_idx->get_layer_ids().size()); @@ -563,12 +551,11 @@ int main(int argc, char ** argv) { CHECK(ctx_recr != nullptr, "context exposes the recurrent half"); CHECK(ctx_idx != nullptr, "context exposes the indexer half"); - // n_kv is only valid once the context has been applied. apply() also - // asserts the two caches agree + // n_kv is only valid once applied; apply() also asserts the caches agree mctx->apply(); - // apply() asserts both of these itself, so reaching this line at all - // is the result; printed rather than CHECKed so the count stays honest + // apply() asserts both itself, so reaching this line is the result; printed + // rather than CHECKed so the count stays honest printf("note indexer and attention caches agree on n_kv (%u) and n_stream (%u)\n", ctx_idx->get_n_kv(), ctx_idx->get_n_stream()); CHECK(ctx_idx && ctx_idx->get_kv() == kv_idx, "the indexer context is a view of the indexer cache"); @@ -616,15 +603,15 @@ int main(int argc, char ** argv) { "graph reaches the KDA conv state [%d x %d] and recurrent state [%d x %d]", (int) r_kda->ne[0], (int) r_kda->ne[1], (int) s_kda->ne[0], (int) s_kda->ne[1]); - // 4. host-side pool map, the piece the indexer needs and nothing else has + // 4. host-side pool map const int64_t kpool = hparams.indexer_kpool; const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); const int64_t n_padq = n_tps; // this tree does not pad the KQ mask kpool_tensors kt = alloc_kpool_tensors(ctx0, n_kv, n_tps, n_padq, n_stream, kpool, n_pools); - // the shape of the real thing: gather the pool members, mix them, - // score, broadcast the pool score back onto its member cells, top-k. + // shape of the real thing: gather pool members, mix, score, broadcast the + // pool score back onto its member cells, top-k ggml_tensor * members = ggml_get_rows(ctx0, k_idx_v, kt.pool_cells); ggml_tensor * gates = ggml_get_rows(ctx0, g_idx_v, kt.pool_cells); members = ggml_reshape_4d(ctx0, members, d_idx, kpool, n_pools, n_stream); @@ -650,8 +637,8 @@ int main(int argc, char ** argv) { printf("note top-k over replicated per-cell scores yields %d I32 CELL indices\n", (int) width); - // 5. the scatter the mask is built from, starting at sel_mask rather - // than at an all -INFINITY fill, which is what forces the tail in + // 5. the scatter the mask is built from, starting at sel_mask rather than + // an all -INFINITY fill, which is what forces the tail in ggml_tensor * top_k_4d = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream); ggml_tensor * base = ggml_view_4d(ctx0, kt.sel_mask, 1, n_kv, n_padq, n_stream, kt.sel_mask->nb[0], kt.sel_mask->nb[1], kt.sel_mask->nb[2], 0); @@ -668,7 +655,7 @@ int main(int argc, char ** argv) { printf("note one graph reaches all three halves in %d nodes\n", ggml_graph_n_nodes(gf)); - // 6. the host-side pool map itself, on a real allocated buffer + // 6. fill the pool map on a real allocated buffer ggml_backend_t backend = ggml_backend_cpu_init(); ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx0, backend); CHECK(buf != nullptr, "allocated the graph's input tensors on the CPU backend"); @@ -701,11 +688,9 @@ int main(int argc, char ** argv) { } CHECK(finite, "bias and sel_mask hold only 0 and -INFINITY: no +1e9 to meet a -inf"); - // every query of every stream: (q+1) %% kpool cells sit in its own - // incomplete pool and must be forced in, the q+1-minus-that cells - // below it must be scored, and nothing else may be either. Both - // streams, so that the per-stream strides are covered and not - // just stream 0's, which is the one at offset 0 + // per query: (q+1) %% kpool cells sit in its own incomplete pool and + // must be forced in, the q+1-minus-that below must be scored, and + // nothing else either. both streams, to cover the per-stream strides { const llama_ubatch & u = mctx->get_ubatch(); @@ -738,10 +723,9 @@ int main(int argc, char ** argv) { } } - // pool origin: while the window starts at position 0, pos/kpool - // and the reference's first-resident-key anchor are the same - // grouping, so the map must reproduce it exactly. Cells whose - // pool is not complete carry slot 0 and are masked instead + // while the window starts at position 0, pos/kpool and the reference's + // first-resident-key anchor are the same grouping, so the map must + // reproduce it. incomplete pools carry slot 0 and are masked instead { const llama_ubatch & u = mctx->get_ubatch(); bool agree = true; @@ -773,7 +757,6 @@ int main(int argc, char ** argv) { CHECK(agree, "pool ordinals match the reference anchor while the window starts at position 0"); } - // the graph actually runs, and nothing in it trips an assert { ggml_status st = ggml_backend_graph_compute(backend, gf); CHECK(st == GGML_STATUS_SUCCESS, "the pooled-indexer-shaped graph computes (%d)", (int) st); @@ -789,17 +772,15 @@ int main(int argc, char ** argv) { // ---- pool origin across a front eviction -------------------------------- // - // The reference anchors a pool at the first *resident* key - // (valid_keys.argmax(-1)); this port anchors at pos/kpool, as vLLM and SGLang - // do. They are the same grouping until the front of the window is dropped by - // a non-multiple of kpool, and then they differ: the reference regroups every - // surviving key, this port does not. Regrouping is what a cache cannot - // afford, since the pooled key a decode step scores was built during prefill. + // the reference anchors a pool at the first *resident* key (valid_keys.argmax(-1)); + // this port anchors at pos/kpool, as vLLM and SGLang do. same grouping until the + // window front is dropped by a non-multiple of kpool, when the reference regroups + // every surviving key and this port does not. a cache cannot afford regrouping: the + // pooled key a decode step scores was built during prefill. { printf("\n--- pool origin ---\n"); - // seq 0 currently holds positions 0..9. drop 0..1, which is not a - // multiple of kpool = 4 + // seq 0 holds positions 0..9. drop 0..1, not a multiple of kpool = 4 const llama_pos n_drop = 2; mem->seq_rm(0, 0, n_drop); @@ -848,9 +829,8 @@ int main(int argc, char ** argv) { const int32_t * cp = (const int32_t *) kt.cell_pool->data; const auto & cells = kv_attn->get_cells(0); - // positions 4..7 were one pool before the eviction and must - // still be one pool, or the pooled key built during prefill no - // longer describes the cells it is scored against + // positions 4..7 must stay one pool, or the pooled key built during + // prefill no longer describes the cells it is scored against std::map slot_of; for (int64_t j = 0; j < n_kv; ++j) { if (!cells.is_empty(j) && cells.seq_has(j, 0)) { @@ -864,15 +844,13 @@ int main(int argc, char ** argv) { CHECK(grouped, "positions 4..7 stay one pool after the eviction (slot %d)", grouped ? (int) slot_of[4] : -1); - // the reference's anchor would have made positions 2..5 the - // first pool instead. checked, not implemented: this port - // deliberately differs here + // the reference's anchor would make 2..5 the first pool; this port + // deliberately differs, so the difference is pinned rather than fixed const bool differs = !slot_of.count(2) || slot_of[2] != slot_of[4]; CHECK(differs, "positions 2 and 4 are NOT pooled together, where the reference anchor would"); - // the leading remnant 2..3 is an incomplete pool, so it is not - // scored. it is also not in the query's tail, which is the one - // case where a cell visible to the query is neither + // the leading remnant 2..3 is an incomplete pool and not in the + // query's tail, the one case where a visible cell is neither const float * bi = (const float *) kt.bias->data; const float * sm = (const float *) kt.sel_mask->data; int64_t n_orphan = 0; @@ -990,7 +968,7 @@ int main(int argc, char ** argv) { const llama_ubatch & ub = mctx->get_ubatch(); // first pass faults in 60+ MiB of fresh pages; time the steady - // state, which is what a prefill actually pays per ubatch + // state, which is what a prefill pays per ubatch double ms = 1e9; for (int rep = 0; rep < 4; ++rep) { const auto t0 = std::chrono::steady_clock::now(); From 8c498364077868edeccc6b7c0b9dddd302ee3c51 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 26 Aug 2026 14:47:02 +0000 Subject: [PATCH 12/36] glm5next: lightning indexer graph, pooled sparse selection Builds the pooled lightning indexer and gives the DSA layers a sparse attention path driven by it. Top-k runs over the POOL axis at select_k = index_topk/index_kpool, and the selected pools are expanded to their member cells through pool_cells. That is the reference's own two-step (modular_glm5_next.py, Glm5NextTextIndexer.forward: topk over the pool axis, then selected_indices = pool_indices[batch_idx, selected]), and it is not interchangeable with a single top-k of width index_topk over member cells. The argument for the cell-level form - a pool's members carry its score bit-exactly, so the cut must land on a pool boundary - assumes tie groups never span pools. They do: ReLU drives most pool scores to exactly 0.0, and ggml_top_k is explicitly unordered among equals, so the cut falls inside an inter-pool tie group and splits a pool. Measured on TinySparse at 512 tokens, the cell-level form leaves a partial pool on 7.51% of query rows at layer 3 and 5.93% at layer 7; this form leaves none. The indexer key and gate STORE is unconditional; only the SCORING is gated, on n_ctx > index_topk + index_kpool - 1. Gating the store the same way would leave every cell written below n_select with no indexer state, and the first ubatch to cross n_select would pool cells that were never written. Nothing here changes any other architecture: test-llama-archs produces a table byte-identical to the parent's, 300 rows over 143 archs, 0 FAIL. --- src/llama-graph.cpp | 175 +++++++++++++++++ src/llama-graph.h | 46 +++++ src/llama-kv-cache-kpool.cpp | 124 +++++++++--- src/llama-kv-cache-kpool.h | 120 +++++++++--- src/models/glm5next.cpp | 297 ++++++++++++++++++++++++++-- src/models/models.h | 21 ++ tests/test-glm5next-memory.cpp | 348 ++++++++++++++++++++++++--------- 7 files changed, 973 insertions(+), 158 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 9bdc7d0e7a1..b7dbe7ff710 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -12,6 +12,7 @@ #include "llama-kv-cache-dsa-iswa.h" #include "llama-kv-cache-msa.h" #include "llama-kv-cache-dsv4.h" +#include "llama-kv-cache-kpool.h" #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-recurrent.h" @@ -3552,6 +3553,180 @@ llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k() const return (llm_graph_input_mem_hybrid_k *) res->add_input(std::move(inp)); } +llm_graph_input_kpool * llm_graph_context::build_inp_kpool( + const llama_memory_hybrid_context * mctx_cur, + ggml_tensor * kq_mask, + bool scoring) const { + const auto * mctx_attn = mctx_cur->get_attn(); + const auto * mctx_idx = mctx_cur->get_idx(); + + GGML_ASSERT(mctx_idx != nullptr && "a pooled indexer needs the indexer KV cache"); + + const uint32_t kpool = hparams.indexer_kpool; + GGML_ASSERT(kpool > 0); + + auto inp = std::make_unique(mctx_attn, mctx_idx, kpool); + + inp->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + ggml_set_input(inp->k_idxs); + ggml_set_name(inp->k_idxs, "kpool_k_idxs"); + + if (scoring) { + const int64_t n_kv = mctx_attn->get_n_kv(); + + // spelled exactly as build_attn_inp_kq_mask spells it, so that sel_mask, + // cand_mask and the KQ mask are the same shape and add without a + // broadcast. get_n_stream() is the cache's stream RANGE and is a + // different number as soon as a server has non-contiguous slots busy + const int64_t n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq; + const int64_t n_tps = ubatch.n_tokens/n_stream; + + const int64_t n_pools = llama_kpool_n_pools(n_kv, kpool); + + GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[3] == n_stream); + + // this tree's build_attn_inp_kq_mask has no GGML_KQ_MASK_PAD, so the mask + // is exactly n_tps rows. the host side supports n_padq > n_tps, the graph + // below does not: the selection terms only exist for real queries + GGML_ASSERT(kq_mask->ne[1] == n_tps && "the pooled indexer needs an unpadded KQ mask"); + + inp->pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool*n_pools, n_stream); + ggml_set_input(inp->pool_cells); + ggml_set_name(inp->pool_cells, "kpool_pool_cells"); + + inp->pool_bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_pools, n_tps, n_stream); + ggml_set_input(inp->pool_bias); + ggml_set_name(inp->pool_bias, "kpool_pool_bias"); + + // f32 even under flash attention, where the KQ mask itself is f16: these + // are a scatter base and a mask addend, not the tensor the FA kernel sees + inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_kv, n_tps, 1, n_stream); + ggml_set_input(inp->sel_mask); + ggml_set_name(inp->sel_mask, "kpool_sel_mask"); + + inp->cand_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_kv, n_tps, 1, n_stream); + ggml_set_input(inp->cand_mask); + ggml_set_name(inp->cand_mask, "kpool_cand_mask"); + } + + return (llm_graph_input_kpool *) res->add_input(std::move(inp)); +} + +ggml_tensor * llm_graph_context::build_attn_sparse( + llm_graph_input_attn_k * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * kq_b, + ggml_tensor * sinks, + ggml_tensor * v_mla, + ggml_tensor * top_k, + ggml_tensor * sel_mask, + ggml_tensor * cand_mask, + float kq_scale, + int il) const { + ggml_build_forward_expand(gf, q_cur); + ggml_build_forward_expand(gf, v_cur); + ggml_build_forward_expand(gf, k_cur); + + const auto * mctx_cur = inp->mctx; + + // store to KV cache + { + const auto & k_idxs = inp->get_k_idxs(); + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + } + + const auto & kq_mask = inp->get_kq_mask(); + + GGML_ASSERT(sel_mask->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_are_same_shape(sel_mask, cand_mask)); + GGML_ASSERT(sel_mask->ne[0] == kq_mask->ne[0] && sel_mask->ne[1] == kq_mask->ne[1] && + sel_mask->ne[3] == kq_mask->ne[3]); + + // The dense DSA path (build_attn on llm_graph_input_attn_k_dsa) opens with + // ggml_fill(kq_mask, -INFINITY). Here the scatter starts from sel_mask + // instead, which already holds 0.0 on the query's own trailing incomplete + // pool: GLM always attends to that tail (index_kpool_always_select_tail), and + // keeping it out of the top-k budget is what lets the budget stay a whole + // number of pools. + // + // ggml_set_rows writes THROUGH to its destination and returns a view of it, + // and sel_mask is one shared per-ubatch input read by every indexer layer. + // Scattering into it directly makes each layer inherit the previous layer's + // unmasked cells: measured on TinySparse, layer 3 stayed inside its budget + // while layer 7, running second, reached 411 cells. The dense path never + // meets this because ggml_fill hands it a fresh tensor every layer. Take a + // private copy per layer. + ggml_tensor * mask_all = ggml_dup(ctx0, sel_mask); + + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + mask_all = ggml_view_4d(ctx0, mask_all, 1, mask_all->ne[0], mask_all->ne[1], mask_all->ne[3], + mask_all->nb[0], mask_all->nb[1], mask_all->nb[2], 0); + + // [n_select, n_tps, n_stream] -> [n_select, n_tps, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[2], 1, + top_k->nb[1], top_k->nb[2], top_k->ne[2]*top_k->nb[2], 0); + + // A constant 0, never the cell's own bias. The scatter must not be able to + // ERASE a zero sel_mask already granted: a tail cell can also be named by the + // top-k (through an over-budget pool whose unfilled slots point at it), and + // scattering that cell's -inf score bias would leave the query attending to + // nothing at all. Rejecting an over-budget selection is done additively, + // below, by cand_mask + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); + + ggml_tensor * mask_top_k = ggml_set_rows(ctx0, mask_all, zeros, top_k_3d); + + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + mask_top_k = ggml_view_4d(ctx0, mask_top_k, mask_top_k->ne[1], mask_top_k->ne[2], 1, mask_top_k->ne[3], + mask_top_k->nb[2], mask_top_k->nb[3], mask_top_k->nb[3], 0); + + // the reference's `selected_valid` gather, additively: a cell the top-k named + // that is not in the reference's candidate set goes back to -inf, and a tail + // cell that sel_mask granted stays at 0 because cand_mask is the UNION of the + // candidates and the tail + mask_top_k = ggml_add(ctx0, mask_top_k, cand_mask); + + // sel_mask and cand_mask are f32 by contract - llama_kv_cache_set_input_kpool + // writes floats through raw strides and refuses anything else - but the KQ + // mask is f16 whenever flash attention is on, which is the DEFAULT, and + // ggml_flash_attn_ext asserts its mask is f16. Cast rather than refuse: the + // only two values here are 0.0f and -INFINITY and both are exact in f16. + // Refusing instead aborts test-llama-archs, which runs with FA on + if (mask_top_k->type != kq_mask->type) { + mask_top_k = ggml_cast(ctx0, mask_top_k, kq_mask->type); + } + + // and finally re-apply causality, occupancy and padding. load bearing: it is + // what keeps an empty, future or foreign-sequence cell masked no matter what + // the top-k returned + mask_top_k = ggml_add(ctx0, mask_top_k, kq_mask); + cb(mask_top_k, "kpool_kq_mask", il); + + ggml_tensor * q = q_cur; + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); + + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, mask_top_k, sinks, v_mla, kq_scale, il); + cb(cur, "kqv_out", il); + + if (wo) { + cur = build_lora_mm(wo, cur, wo_s); + } + + if (wo_b) { + cur = ggml_add(ctx0, cur, wo_b); + } + + return cur; +} + llm_graph_input_mem_hybrid_iswa * llm_graph_context::build_inp_mem_hybrid_iswa() const { const auto * mctx_cur = static_cast(mctx); diff --git a/src/llama-graph.h b/src/llama-graph.h index b388e028cb5..a975eb1bb73 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -32,6 +32,10 @@ class llama_memory_recurrent_context; class llama_memory_hybrid_context; class llama_memory_hybrid_iswa_context; +// defined in llama-kv-cache-kpool.h, which includes this header, so it can only +// be forward declared here +class llm_graph_input_kpool; + // certain models (typically multi-modal) can produce different types of graphs enum llm_graph_type { LLM_GRAPH_TYPE_DEFAULT, @@ -1344,6 +1348,48 @@ struct llm_graph_context { llm_graph_input_mem_hybrid_iswa * build_inp_mem_hybrid_iswa() const; + // + // pooled (GLM-5-Next lightning) indexer + // + + // one pooling map per ubatch, shared by every indexer layer. see + // llama-kv-cache-kpool.h for what the tensors mean and why they are built + // host side rather than derived in the graph. + // + // `scoring` false allocates only k_idxs: the indexer key and gate store is + // unconditional, the selection is not. An input tensor with no consumer is + // never backed by the allocator, so the rest must not be created either + llm_graph_input_kpool * build_inp_kpool( + const llama_memory_hybrid_context * mctx_cur, + ggml_tensor * kq_mask, + bool scoring) const; + + // sparse (pooled top-k) variant of the llm_graph_input_attn_k build_attn. + // + // Identical to it except for the mask. `top_k` names the cells the pooled + // indexer selected, already expanded from whole pools; they are unmasked on + // top of `sel_mask`, which arrives already holding 0.0f on the query's own + // always-selected trailing pool. `cand_mask` is the reference's candidate + // set and is what makes an over-budget selection harmless: ggml_top_k + // returns a full budget of pool ordinals even when fewer pools carry a + // finite score, which during prefill is the normal state + ggml_tensor * build_attn_sparse( + llm_graph_input_attn_k * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens] + ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] + ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] + ggml_tensor * kq_b, + ggml_tensor * sinks, // [n_head_q] + ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] + ggml_tensor * top_k, // I32 [n_select, n_tokens/n_stream, n_stream] + ggml_tensor * sel_mask, // F32 [n_kv, n_batch, 1, n_stream] + ggml_tensor * cand_mask, // F32 [n_kv, n_batch, 1, n_stream] + float kq_scale, + int il) const; + // // pooling // diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index c3f4cc56708..68199bd2008 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -14,11 +14,13 @@ uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool) { return n_kv/kpool + 2; } -uint32_t llama_kpool_top_k_width(uint32_t n_kv, uint32_t indexer_top_k, uint32_t kpool) { +uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t kpool) { GGML_ASSERT(kpool > 0); + GGML_ASSERT(n_pools > 0); GGML_ASSERT(indexer_top_k % kpool == 0 && "indexer_top_k must be a whole number of pools"); - return std::min(n_kv, indexer_top_k); + // min(index_topk // index_kpool, n_pools), exactly the reference's select_k + return std::min(n_pools, indexer_top_k/kpool); } void llama_kv_cache_set_input_kpool( @@ -26,52 +28,74 @@ void llama_kv_cache_set_input_kpool( ggml_tensor * cell_pool, ggml_tensor * pool_cells, ggml_tensor * bias, + ggml_tensor * pool_bias, ggml_tensor * sel_mask, + ggml_tensor * cand_mask, const llama_ubatch * ubatch, uint32_t kpool) { GGML_ASSERT(kv != nullptr); GGML_ASSERT(kpool > 0); - GGML_ASSERT(ggml_backend_buffer_is_host(cell_pool ->buffer)); + // the per-CELL view is optional: the pooled graph does not consume it, and an + // input tensor with no consumer is never backed by the allocator GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer)); - GGML_ASSERT(ggml_backend_buffer_is_host(bias ->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(pool_bias ->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(sel_mask ->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(cand_mask ->buffer)); - // sel_mask is KQ-mask shaped, and KQ masks are f16 under flash attention; writing - // floats into one would overrun the allocation 2x, so check rather than trust - GGML_ASSERT(cell_pool ->type == GGML_TYPE_I32); + // sel_mask and cand_mask are KQ-mask shaped, and KQ masks are f16 under flash + // attention; writing floats into one would overrun the allocation 2x, so check + // rather than trust GGML_ASSERT(pool_cells->type == GGML_TYPE_I32); - GGML_ASSERT(bias ->type == GGML_TYPE_F32); + GGML_ASSERT(pool_bias ->type == GGML_TYPE_F32); GGML_ASSERT(sel_mask ->type == GGML_TYPE_F32 && "sel_mask must be f32 even when the KQ mask is f16"); + GGML_ASSERT(cand_mask ->type == GGML_TYPE_F32 && "cand_mask must be f32 even when the KQ mask is f16"); // everything below is written through raw strides - GGML_ASSERT(ggml_is_contiguous(cell_pool)); GGML_ASSERT(ggml_is_contiguous(pool_cells)); - GGML_ASSERT(ggml_is_contiguous(bias)); + GGML_ASSERT(ggml_is_contiguous(pool_bias)); GGML_ASSERT(ggml_is_contiguous(sel_mask)); + GGML_ASSERT(ggml_is_contiguous(cand_mask)); - const int64_t n_kv = cell_pool->ne[0]; - const int64_t n_ns = cell_pool->ne[1]; // streams in this ubatch + const int64_t n_kv = sel_mask->ne[0]; + const int64_t n_ns = sel_mask->ne[3]; // streams in this ubatch const int64_t r = kpool; const int64_t n_pools = pool_cells->ne[0]/r; const int64_t n_tokens = ubatch->n_tokens; GGML_ASSERT(pool_cells->ne[0] % r == 0); GGML_ASSERT(pool_cells->ne[1] == n_ns); - GGML_ASSERT(bias->ne[0] == n_kv && bias->ne[2] == n_ns); - GGML_ASSERT(sel_mask->ne[0] == n_kv && sel_mask->ne[2] == 1 && sel_mask->ne[3] == n_ns); + GGML_ASSERT(sel_mask->ne[2] == 1); + GGML_ASSERT(ggml_are_same_shape(cand_mask, sel_mask)); + GGML_ASSERT(pool_bias->ne[0] == n_pools && pool_bias->ne[2] == n_ns); GGML_ASSERT(n_tokens % n_ns == 0); const int64_t n_tps = n_tokens/n_ns; // tokens per stream const int64_t n_padq = sel_mask->ne[1]; // KQ mask rows, >= n_tps - GGML_ASSERT(bias->ne[1] == n_tps); + GGML_ASSERT(pool_bias->ne[1] == n_tps); GGML_ASSERT(n_padq >= n_tps); - int32_t * dst_cell_pool = (int32_t *) cell_pool ->data; + if (cell_pool) { + GGML_ASSERT(ggml_backend_buffer_is_host(cell_pool->buffer)); + GGML_ASSERT(cell_pool->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(cell_pool)); + GGML_ASSERT(cell_pool->ne[0] == n_kv && cell_pool->ne[1] == n_ns); + } + + if (bias) { + GGML_ASSERT(ggml_backend_buffer_is_host(bias->buffer)); + GGML_ASSERT(bias->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(bias)); + GGML_ASSERT(bias->ne[0] == n_kv && bias->ne[1] == n_tps && bias->ne[2] == n_ns); + } + + int32_t * dst_cell_pool = cell_pool ? (int32_t *) cell_pool->data : nullptr; int32_t * dst_pool_cells = (int32_t *) pool_cells->data; - float * dst_bias = (float *) bias ->data; + float * dst_bias = bias ? (float *) bias->data : nullptr; + float * dst_pool_bias = (float *) pool_bias ->data; float * dst_sel_mask = (float *) sel_mask ->data; + float * dst_cand_mask = (float *) cand_mask ->data; // -1 marks a cell with no usable pool. host side only: never copied into cell_pool, // where ggml_get_rows would read it as an index @@ -94,7 +118,7 @@ void llama_kv_cache_set_input_kpool( const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; const auto & cells = kv->get_cells(seq_of_stream); - int32_t * cur_cell_pool = dst_cell_pool + s*n_kv; + int32_t * cur_cell_pool = dst_cell_pool ? dst_cell_pool + s*n_kv : nullptr; int32_t * cur_pool_cells = dst_pool_cells + s*(r*n_pools); std::fill(pool_of.begin(), pool_of.end(), -1); @@ -168,13 +192,17 @@ void llama_kv_cache_set_input_kpool( if (pool_of[j] >= 0 && filled[pool_of[j]] != (int32_t) r) { pool_of[j] = -1; } - cur_cell_pool[j] = pool_of[j] < 0 ? 0 : pool_of[j]; + if (cur_cell_pool) { + cur_cell_pool[j] = pool_of[j] < 0 ? 0 : pool_of[j]; + } } - float * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv); + float * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv); + float * cur_cand_mask = dst_cand_mask + s*(n_padq*n_kv); // the loop below writes rows < n_tps in full; only the padding rows need clearing - std::fill(cur_sel_mask + n_tps*n_kv, cur_sel_mask + n_padq*n_kv, -INFINITY); + std::fill(cur_sel_mask + n_tps*n_kv, cur_sel_mask + n_padq*n_kv, -INFINITY); + std::fill(cur_cand_mask + n_tps*n_kv, cur_cand_mask + n_padq*n_kv, -INFINITY); for (int64_t ii = 0; ii < n_tps; ++ii) { const int64_t i = s*n_tps + ii; @@ -193,8 +221,9 @@ void llama_kv_cache_set_input_kpool( // that test collapses to b*r < tail_start const int64_t bo_vis = std::max(0, tail_start/r - b_base); - float * cur_bias = dst_bias + i*n_kv; - float * cur_sel = cur_sel_mask + ii*n_kv; + float * cur_bias = dst_bias ? dst_bias + i*n_kv : nullptr; + float * cur_sel = cur_sel_mask + ii*n_kv; + float * cur_cand = cur_cand_mask + ii*n_kv; // the unsigned compares fold "empty or another sequence" (pos_at -1) and "no // usable pool" (pool_of -1) into the range test, which lets this vectorise @@ -203,16 +232,61 @@ void llama_kv_cache_set_input_kpool( const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; const bool tail = pos_at[j] >= tail_start; - cur_bias[j] = vis && pooled ? 0.0f : -INFINITY; cur_sel [j] = vis && tail ? 0.0f : -INFINITY; + // max(bias, sel_mask): the reference's candidate set, which the + // top-k budget may overrun but must never escape + cur_cand[j] = vis && (pooled || tail) ? 0.0f : -INFINITY; + } + + if (cur_bias) { + for (int64_t j = 0; j < n_kv; ++j) { + const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; + const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; + + cur_bias[j] = vis && pooled ? 0.0f : -INFINITY; + } + } + + // The same predicate, per POOL, which is where the reference applies + // it: pool_valid (completely resident) & pool_visible (its LAST + // member is visible, so a pool the query straddles is dropped whole). + // Pools are position-aligned here, so pool bo's last member is at + // position (b_base + bo)*r + r - 1 and "last member visible" collapses + // to bo < bo_vis. + // + // NOT gathered from `bias` at the last member cell: an incomplete or + // absent pool has no resident last member, pool_cells points that slot + // at cell 0, and the pool would inherit cell 0's validity. + float * cur_pool_bias = dst_pool_bias + (s*n_tps + ii)*n_pools; + + for (int64_t p = 0; p < n_pools; ++p) { + const bool valid = filled[p] == (int32_t) r; + const bool visible = p < bo_vis; + + cur_pool_bias[p] = valid && visible ? 0.0f : -INFINITY; } } } } void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { + // unconditional: the indexer key and gate STORE runs on the dense path too, + // and k_idxs is what tells cpy_k where to put them. Gating the store the way + // the scoring is gated would leave every cell written below n_select - the + // first 2051 positions of every sequence on the real model - with no indexer + // state, and the first ubatch to cross n_select would pool cells that were + // never written mctx_idx->set_input_k_idxs(k_idxs, ubatch); + // the rest exists only when the graph scores. below n_select the indexer + // would select every visible position, so build_inp_kpool does not allocate + // these at all and there is nothing to fill + if (pool_cells == nullptr) { + return; + } + llama_kv_cache_set_input_kpool( - mctx_attn->get_kv(), cell_pool, pool_cells, bias, sel_mask, ubatch, kpool); + mctx_attn->get_kv(), + /* cell_pool */ nullptr, pool_cells, /* bias */ nullptr, pool_bias, + sel_mask, cand_mask, ubatch, kpool); } diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index 2385a6cfe8c..3a486cb0463 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -36,42 +36,99 @@ class llama_kv_cache_context; // can cost one slot at each end, hence the +2. uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool); -// width to run ggml_top_k at. -// -// NOT indexer_top_k + kpool - 1: that is the width of the reference's *output* buffer, -// not a top-k width. ggml_top_k is explicitly unordered among equals (the CPU op -// deliberately swaps the first two results; the CUDA op declares -// determinism::not_guaranteed), so a budget that does not end on a pool boundary would -// take an arbitrary, CPU/CUDA-divergent subset of the pool it cuts. Biasing the tail to -// -INFINITY keeps the budget on whole pools; the tail is forced back in via `sel_mask`. -// See the PR body for the full argument. -uint32_t llama_kpool_top_k_width(uint32_t n_kv, uint32_t indexer_top_k, uint32_t kpool); +// how many POOLS ggml_top_k selects. +// +// The reference (modular_glm5_next.py, Glm5NextTextIndexer.forward) scores the pool +// axis and takes select_k = min(index_topk // index_kpool, n_pools), then expands each +// selected pool to its members. This is that number, and top-k over pools rather than +// over cells is not an optimisation, it is the only correct spelling. +// +// The tempting alternative - broadcast a pool's score onto its kpool member cells and +// run one top-k of width indexer_top_k over cells - is WRONG, for a reason a +// set-similarity metric cannot see. Its argument is that a pool's members carry its +// score bit-exactly, so the cut must land on a pool boundary; that holds only if tie +// groups never span pools. They do: F.relu drives most pool scores to exactly 0.0 (on +// TinySparse, 92 cells tie at 0.0 for query 386 at layer 3), so the cut falls inside a +// tie group that straddles pools, and ggml_top_k - explicitly unordered among equals; +// the CPU op deliberately swaps the first two results, the CUDA op declares +// determinism::not_guaranteed - takes an arbitrary 1..kpool-1 members of the pool it +// cuts. Measured on TinySparse at 512 tokens the cell-level form leaves a partial pool +// on 7.51% of query rows at L3 and 5.93% at L7 (70 and 47 partial pools); the +// pool-level form leaves none. A pool-aligned top-k WIDTH does not help, because the +// ties are not aligned to anything. +// +// Selecting whole pools also removes the CPU/CUDA tie-break divergence structurally +// rather than making it unlikely: the backends may disagree about WHICH pools come out +// of a tie group, never about whether a pool is taken whole. +// +// NOT indexer_top_k + kpool - 1 either: that is the width of the reference's *output* +// buffer, tail included. The always-selected tail is forced in via `sel_mask`. +uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t kpool); // Fill the host-side inputs of the pooled indexer. // -// cell_pool I32 [n_kv, n_stream] -// cell -> its pool slot, or 0 when it has no usable pool. ggml_get_rows -// broadcasts a pool's score onto its members, so ggml_top_k over the replicated -// scores yields CELL indices directly and still cuts on a pool boundary (members -// tie bit-exactly). +// `cell_pool` and `bias` are the per-CELL view of the same information and may be +// nullptr. They are kept because they are the independent spelling +// tests/test-glm5next-memory.cpp checks `pool_bias` and `cand_mask` against - but the +// graph selects at POOL granularity (see llama_kpool_select_k), so it passes nullptr +// for both. An input tensor with no consumer is never backed by the allocator, so +// building them anyway would write through a null buffer. +// +// cell_pool I32 [n_kv, n_stream] OPTIONAL +// cell -> its pool slot, or 0 when it has no usable pool. // // pool_cells I32 [kpool*n_pools, n_stream] // pool slot, member -> the cell holding that position, or 0 when not resident. -// gathers a pool's member keys and gates for the compressor. +// Two consumers: the compressor gathers a pool's member keys and gates with it, +// and the top-k expands a selected POOL ordinal back into its kpool member CELLS +// with it - the reference's `pool_indices[batch_idx, selected]`. +// +// bias F32 [n_kv, n_tokens/n_stream, n_stream] OPTIONAL +// additive per-(cell, query) bias on a per-cell score: 0.0f in a complete pool +// whose last member the query can see, else -INFINITY, the trailing incomplete +// pool included. // -// bias F32 [n_kv, n_tokens/n_stream, n_stream] -// additive per-(cell, query) bias on the indexer SCORE: 0.0f in a complete pool -// whose last member the query can see (the reference's `pool_valid & -// pool_visible`), else -INFINITY, the trailing incomplete pool included so no -// budget is spent on it. +// pool_bias F32 [n_pools, n_tokens/n_stream, n_stream] +// the same predicate per POOL, which is where the reference applies it: 0.0f when +// pool p is completely resident and its LAST member is visible to query q +// (`pool_valid & pool_visible`), else -INFINITY, the query's own trailing pool +// included so no budget is spent on it. +// +// Derived here rather than in the graph on purpose. Gathering `bias` at +// each pool's last member looks equivalent and is not: an incomplete or +// entirely absent pool has no resident last member, `pool_cells` points +// that slot at cell 0 to keep the gather in range, and the pool would then +// inherit cell 0's validity and compete for budget with a finite score. // // sel_mask F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) // what the top-k scatter starts from, replacing the ggml_fill(kq_mask, -INFINITY) // opening the DSA mask build in llm_graph_context::build_attn: 0.0f for the // query's own incomplete trailing pool, which GLM always attends to // (index_kpool_always_select_tail), else -INFINITY, padding rows included. -// forcing the tail in here rather than through the score keeps the budget -// pool-aligned. +// Forcing the tail in here rather than through the score keeps the budget a whole +// number of pools. +// +// cand_mask F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) +// max(bias, sel_mask): the reference's candidate set, i.e. every cell it could +// return for this query. 0.0f in a complete visible pool OR in the query's own +// tail, else -INFINITY, padding rows included. +// +// ggml_top_k returns `select_k` pool ordinals even when fewer pools carry a +// finite score, and during prefill that is the NORMAL state, not a corner case: +// query q has only ~q/kpool complete visible pools against a budget of +// index_topk/kpool (on TinySparse, 20 of 20 query rows are under budget). The +// spilled ordinals are arbitrary among the -INFINITY ties and expanding them +// unmasks cells. Adding the causal KQ mask kills the spills that are empty, +// foreign or in the future. What it does not kill is a resident, causally visible +// cell in an INCOMPLETE pool below the tail - unreachable while positions are +// contiguous, reachable the moment a partial seq_rm leaves a hole. cand_mask +// kills exactly those, for one store per element in a loop that already computes +// both operands. +// +// (The other spill, an unfilled pool slot pointing at cell 0, is a provable +// no-op: the budget can only overflow once every finitely-scored pool is already +// selected, so cell 0 is either already selected through its own pool, or masked +// for the same reason it would have been anyway.) // // `kv` must be the ATTENTION (MLA) cache: its cells define the pools and are what the // top-k indices are ultimately read against. llama_memory_hybrid gives the indexer @@ -81,7 +138,9 @@ void llama_kv_cache_set_input_kpool( ggml_tensor * cell_pool, ggml_tensor * pool_cells, ggml_tensor * bias, + ggml_tensor * pool_bias, ggml_tensor * sel_mask, + ggml_tensor * cand_mask, const llama_ubatch * ubatch, uint32_t kpool); @@ -90,9 +149,14 @@ void llama_kv_cache_set_input_kpool( // per layer costs O(n_kv * n_tokens) host writes each time (at 128 Ki cells, 512 tokens // and 11 DSA layers, ~11 x 67M float stores per ubatch, which dominates prefill). // -// Sharing `bias` and `sel_mask` is correct only while every indexer layer sees the same -// candidate set. That holds for glm5next, whose indexer_types are all "full"; a model -// mixing in windowed indexers would need one bias per window. +// Sharing `pool_bias`, `sel_mask` and `cand_mask` is correct only while every indexer +// layer sees the same candidate set. That holds for glm5next, whose indexer_types are +// all "full"; a model mixing in windowed indexers would need one map per window. +// +// `k_idxs` is present whenever the model has an indexer cache; the rest only when the +// graph actually scores. The indexer key and gate STORE is not gated on the sparse path +// - below n_select the selection is a no-op but the cells still have to be written, or +// the first ubatch to cross n_select would pool cells that were never filled. class llm_graph_input_kpool : public llm_graph_input_i { public: llm_graph_input_kpool( @@ -105,10 +169,10 @@ class llm_graph_input_kpool : public llm_graph_input_i { void set_input(const llama_ubatch * ubatch) override; ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] - ggml_tensor * cell_pool = nullptr; // I32 [n_kv, n_stream] ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] - ggml_tensor * bias = nullptr; // F32 [n_kv, n_tps, n_stream] + ggml_tensor * pool_bias = nullptr; // F32 [n_pools, n_tps, n_stream] ggml_tensor * sel_mask = nullptr; // F32 [n_kv, n_batch, 1, n_stream] + ggml_tensor * cand_mask = nullptr; // F32 [n_kv, n_batch, 1, n_stream] const llama_kv_cache_context * mctx_attn; const llama_kv_cache_context * mctx_idx; diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 2128efdd3f9..e07b84e9472 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -1,6 +1,8 @@ #include "models.h" #include "llama-memory-recurrent.h" +#include "llama-memory-hybrid.h" +#include "llama-kv-cache-kpool.h" // // GLM-5.3-Flash: hybrid KDA (linear) + DSA (nope-only MLA) attention, mHC @@ -35,6 +37,18 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); // indexer k_norm is a LayerNorm with bias; without this key it runs at eps 0 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); + // The reference HARDCODES nn.LayerNorm(index_head_dim, eps=1e-6); it does not + // read it from the config, and rms_norm_eps is 1e-5. Warned about rather than + // asserted, for two reasons. No output comparison can see it: seeding eps=1e-5 + // into the reference leaves the logits BIT-IDENTICAL at 512 tokens and moves + // the index jaccard by 1.7e-5 at 2048, against a bf16 floor of 0.63. And an + // assert would abort test-llama-archs, whose synthetic models have no reason + // to carry a converter contract + if (hparams.f_norm_eps <= 0.0f || hparams.f_norm_eps > 2e-6f) { + LLAMA_LOG_WARN("%s: indexer k_norm eps is %g, but the reference hardcodes 1e-6. " + "this is invisible to every output comparison; check the converter\n", + __func__, (double) hparams.f_norm_eps); + } 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); @@ -63,11 +77,12 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(hparams.indexer_kpool > 0); GGML_ASSERT(hparams.indexer_top_k % hparams.indexer_kpool == 0); + // below n_select resident positions the indexer selects every visible one, so + // the dense path is not an approximation of the sparse one, it IS it const uint32_t n_select = glm5next_n_select(hparams); - if (hparams.n_ctx_train > n_select) { - LLAMA_LOG_WARN("%s: attention is dense above %u cached tokens, but this checkpoint trains to %u. " - "the sparse selection is not implemented yet\n", __func__, n_select, hparams.n_ctx_train); - } + LLAMA_LOG_INFO("%s: indexer selection width = %u cells (%u pools of %u, plus a %u-wide tail)\n", + __func__, n_select, hparams.indexer_top_k/hparams.indexer_kpool, + hparams.indexer_kpool, hparams.indexer_kpool - 1); // mHC ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); @@ -354,8 +369,213 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( } // -// DSA layer, dense: full MLA over the whole cache, no indexer/kpool/top-k. this is the -// limit the sparse path collapses to below glm5next_n_select() resident positions +// the lightning indexer, pooled +// +// Writes this layer's indexer key and compressor gate into the indexer cache, and +// then - when the cache is large enough for selection to bind - returns the I32 +// ATTENTION-cache cell indices this layer's queries select, ready for the scatter in +// build_attn_sparse. `qr` is the shared q LoRA residual, i.e. q_a_norm(q_a_proj(x)), +// which the indexer's own wq_b consumes; `cur` is the layer input. +// +// The store is NOT gated on the sparse path and the scoring is. Gating both the same +// way leaves every cell written below n_select - the first 2051 positions of every +// sequence on the real model - with no indexer state at all, and the first ubatch to +// cross n_select then pools cells that were never written. +// +// Three points where the reference implementations disagree, resolved 2-of-3 by +// reading transformers (modular_glm5_next.py Glm5NextTextIndexer), sglang +// (dsa_indexer_kpool.py) and vLLM (glm5next/nvidia/attention.py): +// +// * the weights_proj GEMM runs in fp32. sglang gives it params_dtype=fp32 and feeds +// x.float(); vLLM does the same and says why - bf16 head-gates move logits by +// ~1e-2, enough to flip near-tie pool rankings on long context. transformers is +// the outlier and rounds the activation to bf16. +// * k_norm is a LayerNorm WITH BIAS at eps 1e-6, hardcoded in transformers and in +// vLLM; sglang leaves torch's 1e-5 default. It is NOT f_norm_rms_eps (1e-5 here) +// and NOT ggml's 0 default. The value comes from the GGUF and load_arch_hparams +// warns when it is not 1e-6, because no output comparison can see it. +// * the ReLU between the QK dot and the head weighting is explicit in transformers +// and in sglang's non-pooled indexer; in the pooled path both engines hand it to +// the DeepGEMM MQA-logits kernel, so neither validates it. It is easy to drop and +// is written out here. +// +// No Hadamard rotation: sglang and vLLM rotate q and k by H128 before their fp8 +// kernel, but H is orthogonal so (Hq).(Hk) == q.k. It exists to spread magnitude ahead +// of fp8 quantisation; scoring in f32 here it would only cost accuracy. transformers, +// the semantic reference, has none. +// +ggml_tensor * llama_model_glm5next::graph::build_indexer( + const llama_layer & layer, + llm_graph_input_kpool * inp_kp, + ggml_tensor * cur, + ggml_tensor * qr, + bool scoring, + int il) const { + const int64_t d_idx = hparams.indexer_head_size; + const int64_t n_ihead = hparams.indexer_n_head; + const int64_t r = hparams.indexer_kpool; + + const auto * mctx_idx = inp_kp->mctx_idx; + + // a genuine LayerNorm: weight AND bias. glm-dsa runs this same norm at eps 0 today + GGML_ASSERT(layer.indexer_k_norm_b != nullptr && "the indexer k_norm is a LayerNorm with bias"); + + ggml_tensor * ik = build_norm(ggml_mul_mat(ctx0, layer.indexer_attn_k, cur), + layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il); + cb(ik, "indexer_k", il); + + // the pooling gate is a SECOND, INDEPENDENT projection of the hidden state, not a + // reuse of the indexer key. it has to be cached beside the key: the compressor + // mixes a pool's member keys with a softmax over these gates, and a pool is only + // rebuilt once its members have left the batch + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.indexer_comp_wgate, cur); + cb(gate, "indexer_gate", il); + + // {d_idx, 2, n_tokens}: head 0 is the key, head 1 the gate. llama_memory_hybrid + // allocates the second head exactly for this and only when indexer_kpool > 0 + ggml_tensor * packed = ggml_concat(ctx0, + ggml_reshape_3d(ctx0, ik, d_idx, 1, n_tokens), + ggml_reshape_3d(ctx0, gate, d_idx, 1, n_tokens), 1); + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, packed, inp_kp->k_idxs, il)); + + if (!scoring) { + return nullptr; + } + + // {d_idx, 2, n_kv, n_stream} + ggml_tensor * kbuf = mctx_idx->get_k(ctx0, il); + + const int64_t n_kv = kbuf->ne[2]; + const int64_t n_stream = kbuf->ne[3]; + const int64_t n_tps = n_tokens/n_stream; + const int64_t n_pools = inp_kp->pool_cells->ne[0]/r; + + GGML_ASSERT(kbuf->ne[0] == d_idx && kbuf->ne[1] == 2 && + "the pooled indexer cache needs a key head and a gate head"); + GGML_ASSERT(kbuf->nb[1] == (size_t) d_idx*kbuf->nb[0] && "key and gate must be adjacent in a cell"); + GGML_ASSERT(n_tokens == n_tps*n_stream); + + // one cell's key and gate as a single row, so that the members of a pool are + // gathered once rather than twice: {2*d_idx, n_kv, n_stream} + ggml_tensor * kg_rows = ggml_view_3d(ctx0, kbuf, 2*d_idx, n_kv, n_stream, + kbuf->nb[2], kbuf->nb[3], 0); + + // gather each pool's members. pool_cells names a cell per (pool, slot); slots that + // are not resident hold 0 rather than a negative sentinel, because ggml_get_rows + // has none. The resulting garbage pools are neutralised by pool_bias below, never + // by a NaN: unlike the reference, no -inf ever enters the compressor softmax. + // ggml_get_rows always yields F32, so the pooling runs in F32 even though the + // indexer cache is F16 + ggml_tensor * members = ggml_get_rows(ctx0, kg_rows, inp_kp->pool_cells); + cb(members, "indexer_pool_members", il); + + const size_t nb_mem = members->nb[1]; + + ggml_tensor * mem_k = ggml_view_4d(ctx0, members, d_idx, r, n_pools, n_stream, + nb_mem, nb_mem*r, members->nb[2], 0); + ggml_tensor * mem_g = ggml_view_4d(ctx0, members, d_idx, r, n_pools, n_stream, + nb_mem, nb_mem*r, members->nb[2], d_idx*members->nb[0]); + + // d_idx independent r-way softmaxes over the SLOT axis, so the slot axis has to be + // dim 0. ape is added PRE-softmax and is indexed by LOGICAL SLOT, so it broadcasts + // over pools and streams; pool_cells is built in position order, so slot m is + // position p % kpool and the two agree by construction + ggml_tensor * keys_t = ggml_cont(ctx0, ggml_permute(ctx0, mem_k, 1, 0, 2, 3)); + ggml_tensor * gate_t = ggml_cont(ctx0, ggml_permute(ctx0, mem_g, 1, 0, 2, 3)); + + ggml_tensor * ape = ggml_cont(ctx0, ggml_transpose(ctx0, layer.indexer_comp_ape)); + gate_t = ggml_add(ctx0, gate_t, ggml_reshape_4d(ctx0, ape, r, d_idx, 1, 1)); + + ggml_tensor * probs = ggml_soft_max(ctx0, gate_t); + cb(probs, "indexer_pool_probs", il); + + // per-channel weighted average over the pool's members -> {d_idx, n_pools, 1, n_stream} + ggml_tensor * pool_k = ggml_sum_rows(ctx0, ggml_mul(ctx0, keys_t, probs)); + pool_k = ggml_reshape_4d(ctx0, pool_k, d_idx, n_pools, 1, n_stream); + cb(pool_k, "indexer_pool_k", il); + + // {d_idx, n_tps, n_ihead, n_stream}. no rope: n_rot() is 0 for the whole text tower + ggml_tensor * iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr); + iq = ggml_reshape_4d(ctx0, iq, d_idx, n_ihead, n_tps, n_stream); + iq = ggml_permute(ctx0, iq, 0, 2, 1, 3); + cb(iq, "indexer_q", il); + + // {n_pools, n_tps, n_ihead, n_stream}: pool_k is MQA and broadcasts over the heads + ggml_tensor * kq = ggml_mul_mat(ctx0, pool_k, iq); + + // {n_ihead, n_tps, n_pools, n_stream}, contiguous for the relu and the head sum. + // the ReLU sits BETWEEN the per-head dot product and the head weighting: moving it + // to either side is a different function, because the head weights are sign-free + // and the sum is not a convex combination + kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); + ggml_tensor * score = ggml_relu(ctx0, kq); + cb(score, "indexer_score", il); + + // sign-unconstrained head weights: no softmax, no abs, no relu. Both scale + // constants - the reference's softmax_scale = d_idx^-0.5 and its n_heads^-0.5 head + // factor - are folded in here, on an {n_ihead, n_tokens} tensor rather than on the + // {n_pools, n_tps, n_ihead} score tensor. relu is positively homogeneous and both + // constants are positive, so this is exactly the same function, and it is what the + // engines and the in-tree glm-dsa both do + ggml_tensor * w = ggml_mul_mat(ctx0, layer.indexer_proj, cur); + ggml_mul_mat_set_prec(w, GGML_PREC_F32); + w = ggml_reshape_4d(ctx0, w, n_ihead, n_tps, 1, n_stream); + w = ggml_scale(ctx0, w, 1.0f/sqrtf(float(d_idx*n_ihead))); + cb(w, "indexer_weights", il); + + // {1, n_tps, n_pools, n_stream} -> {n_pools, n_tps, n_stream} + ggml_tensor * pool_score = ggml_sum_rows(ctx0, ggml_mul(ctx0, score, w)); + pool_score = ggml_cont(ctx0, ggml_permute(ctx0, pool_score, 2, 1, 0, 3)); + pool_score = ggml_reshape_3d(ctx0, pool_score, n_pools, n_tps, n_stream); + + // -INFINITY on every pool the reference's `pool_valid & pool_visible` rejects, the + // query's own trailing pool included, so that no budget is spent on it + pool_score = ggml_add(ctx0, pool_score, inp_kp->pool_bias); + cb(pool_score, "indexer_pool_score", il); + + // Top-k over POOLS at index_topk/index_kpool, then expand each selected pool to its + // members. This is the reference's own two-step (topk over the pool axis, then + // selected_indices = pool_indices[batch_idx, selected]) and it is NOT + // interchangeable with a single top-k of width index_topk over member cells. + // + // The cell-level form is the tempting one and the argument for it is false. It says + // a pool's members carry its score bit-exactly, so the cut must land on a pool + // boundary. But F.relu drives most pool scores to exactly 0.0, so tie groups SPAN + // pools, the cut falls inside one, and ggml_top_k - explicitly unordered among + // equals - takes an arbitrary 1..kpool-1 members of the pool it lands in. A + // pool-aligned top-k WIDTH does not save it: the ties are not aligned to anything. + // Measured on TinySparse at 512 tokens, the cell-level form leaves a partial pool + // on 7.51% of query rows at L3 and 5.93% at L7 (70 and 47 partial pools); this form + // leaves none. The index-set jaccard does not separate them - it scored the broken + // form at 0.9958/0.9764 against a bf16 noise floor of 0.9779/0.8385, i.e. ABOVE the + // floor - which is why scripts/glm5next_pool_integrity.py exists + const int64_t select_k = llama_kpool_select_k(n_pools, hparams.indexer_top_k, r); + GGML_ASSERT(select_k > 0 && select_k <= n_pools); + + // {select_k, n_tps, n_stream} of POOL ordinals + ggml_tensor * sel = ggml_cont(ctx0, ggml_top_k(ctx0, pool_score, (int) select_k)); + cb(sel, "indexer_top_k_pools", il); + + // Expand pools to members: gather whole rows of `kpool` cells out of pool_cells. + // The query axis folds into the gather's row axis, which is what lets ONE + // ggml_get_rows serve every query, while the stream axis stays where get_rows wants + // it (src0 dim 2 is indexed by the index tensor's dim 1) + ggml_tensor * pc3 = ggml_reshape_3d(ctx0, inp_kp->pool_cells, r, n_pools, n_stream); + ggml_tensor * sel_flat = ggml_reshape_2d(ctx0, sel, select_k*n_tps, n_stream); + + // {r, select_k*n_tps, n_stream} -> {r*select_k, n_tps, n_stream} + ggml_tensor * top_k = ggml_get_rows(ctx0, pc3, sel_flat); + GGML_ASSERT(top_k->type == GGML_TYPE_I32 && "pool_cells is I32, so the gather stays I32"); + top_k = ggml_reshape_3d(ctx0, top_k, r*select_k, n_tps, n_stream); + cb(top_k, "indexer_top_k", il); + + return top_k; +} + +// +// DSA layer. `scoring` false takes the dense limit: full MLA over the whole cache, no +// kpool, no top-k, which is what sparse selection collapses to below +// glm5next_n_select() resident positions // // absorbed form, as in deepseek2/deepseek32/glm-dsa: q_nope is pushed through wk_b so // q.k is taken against the 512-wide latent the cache actually holds (is_mla() drops the @@ -365,6 +585,8 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( const llama_layer & layer, llm_graph_input_attn_k * inp_attn, + llm_graph_input_kpool * inp_kp, + bool scoring, ggml_tensor * cur, int il) const { const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); @@ -378,11 +600,16 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( // width: 1/sqrt(kv_lora_rank) = 1/sqrt(512) would be a different model const float kq_scale = 1.0f/sqrtf(float(qk_head_dim)); - ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_a, cur); - q = build_norm(q, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); - cb(q, "dsa_q_a_norm", il); + 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, "dsa_q_a_norm", il); - q = ggml_mul_mat(ctx0, layer.wq_b, q); + // the indexer shares this q LoRA residual with the main MLA path and consumes it + // with its own wq_b, so it is built here rather than being handed the layer input + // twice. it also writes the indexer cache, which happens on the dense path too + ggml_tensor * top_k = inp_kp ? build_indexer(layer, inp_kp, cur, qr, scoring, il) : nullptr; + + ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr); q = ggml_reshape_3d(ctx0, q, qk_head_dim, n_head, n_tokens); cb(q, "dsa_q_b", il); @@ -405,9 +632,16 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( ggml_tensor * k = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens); cb(k, "dsa_kv_latent", il); - cur = build_attn(inp_attn, - layer.wo, nullptr, nullptr, - q, k, k, nullptr, nullptr, layer.wv_b, kq_scale, il); + if (top_k) { + cur = build_attn_sparse(inp_attn, + layer.wo, nullptr, nullptr, + q, k, k, nullptr, nullptr, layer.wv_b, + top_k, inp_kp->sel_mask, inp_kp->cand_mask, kq_scale, il); + } else { + cur = build_attn(inp_attn, + layer.wo, nullptr, nullptr, + q, k, k, nullptr, nullptr, layer.wv_b, kq_scale, il); + } cb(cur, "dsa_out", il); return cur; @@ -416,13 +650,15 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( ggml_tensor * llama_model_glm5next::graph::build_layer_attn( const llama_model & model, llm_graph_input_mem_hybrid_k * inp_mem, + llm_graph_input_kpool * inp_kp, + bool scoring, ggml_tensor * cur, int il) { if (hparams.is_recr(il)) { return build_kda_layer(model.layers[il], inp_mem->get_recr(), cur, il); } - return build_dsa_layer(model.layers[il], inp_mem->get_attn(), cur, il); + return build_dsa_layer(model.layers[il], inp_mem->get_attn(), inp_kp, scoring, cur, il); } ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( @@ -478,6 +714,37 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa // the hybrid memory is the _k variant, as in bailingmoe3 llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); + // One pooling map for the whole ubatch. pool_cells, pool_bias, sel_mask and + // cand_mask depend only on the cells and the ubatch, never on the layer, so + // rebuilding them per DSA layer would cost O(n_kv * n_tokens) host writes eleven + // times over - at 128 Ki cells and 512 tokens that dominates prefill on its own. + // + // The map is built whenever the model HAS an indexer cache, because the indexer + // key and gate store runs unconditionally. Only `scoring` is gated: below + // index_topk + index_kpool - 1 resident positions the reference selects every + // visible position, so the dense build_attn is not an approximation there, it is + // the same function. + // + // Gated on n_ctx and not on the ubatch's n_kv, even though n_kv is what actually + // decides whether the budget binds. n_kv grows as the cache fills, so gating on it + // would flip the graph's topology partway through a run. n_ctx is fixed for the + // lifetime of the context, so the topology is decided once. The cost is that a + // context configured larger than n_select runs the indexer even while the cache is + // still short, where it selects every visible pool: wasted work, never a wrong + // answer, and llama_kpool_select_k clamps the budget to the pools that exist. + llm_graph_input_kpool * inp_kp = nullptr; + bool indexer_scoring = false; + { + const auto * mctx_hyb = static_cast(mctx); + + if (mctx_hyb->get_idx() != nullptr) { + indexer_scoring = cparams.n_ctx > glm5next_n_select(hparams); + + inp_kp = build_inp_kpool(mctx_hyb, + inp_mem->get_attn()->get_kq_mask(), indexer_scoring); + } + } + GGML_ASSERT(ubatch.n_seqs != 0); GGML_ASSERT(ubatch.equal_seqs()); GGML_ASSERT(ubatch.n_tokens == ubatch.n_seq_tokens * ubatch.n_seqs); @@ -510,7 +777,7 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); cb(cur, "attn_norm", il); - cur = build_layer_attn(model, inp_mem, cur, il); + cur = build_layer_attn(model, inp_mem, inp_kp, indexer_scoring, cur, il); inpL = build_hc_post(cur, residual, post, comb, il); cb(inpL, "hc_attn_post", il); diff --git a/src/models/models.h b/src/models/models.h index 8087aa7c482..1182fd7207b 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1355,6 +1355,8 @@ struct llama_model_glm5next : public llama_model_base { ggml_tensor * build_layer_attn( const llama_model & model, llm_graph_input_mem_hybrid_k * inp_mem, + llm_graph_input_kpool * inp_kp, + bool scoring, ggml_tensor * cur, int il); @@ -1364,12 +1366,31 @@ struct llama_model_glm5next : public llama_model_base { ggml_tensor * cur, int il); + // `scoring` false keeps the dense path: at or below + // index_topk + index_kpool - 1 resident positions the indexer selects + // every visible position, so dense is not an approximation there, it is + // the same function without the pooling work. The indexer STORE still + // runs; only the selection is skipped ggml_tensor * build_dsa_layer( const llama_layer & layer, llm_graph_input_attn_k * inp_attn, + llm_graph_input_kpool * inp_kp, + bool scoring, ggml_tensor * cur, int il) const; + // writes this layer's indexer key and compressor gate into the indexer + // cache - always - and then, when `scoring`, returns the I32 attention + // cache CELL indices this layer's queries select, already expanded from + // whole pools: [kpool*select_k, n_tps, n_stream] + ggml_tensor * build_indexer( + const llama_layer & layer, + llm_graph_input_kpool * inp_kp, + ggml_tensor * cur, + ggml_tensor * qr, + bool scoring, + int il) const; + ggml_tensor * build_layer_ffn( const llama_model & model, ggml_tensor * cur, diff --git a/tests/test-glm5next-memory.cpp b/tests/test-glm5next-memory.cpp index a6802a072bb..b30ffbfeb74 100644 --- a/tests/test-glm5next-memory.cpp +++ b/tests/test-glm5next-memory.cpp @@ -48,10 +48,20 @@ static int n_fail = 0; } while (0) // -// numeric evidence for the top-k width, independent of any model: ggml_top_k at exactly -// indexer_top_k with the trailing incomplete pool biased to -INFINITY cuts on a pool -// boundary and agrees between CPU and CUDA, where the reference's output width -// (indexer_top_k + kpool - 1) with the tail forced in at +1e9 does not +// numeric evidence for WHERE the top-k runs, independent of any model +// +// The claim under test is the one that decides the whole design: a top-k over +// CELLS cannot be made pool-aligned by choosing its width, and a top-k over +// POOLS is pool-aligned by construction. +// +// The tempting argument for the cell-level form is that a pool's members carry +// its score bit-exactly, so an intra-pool tie is harmless and a budget that is a +// whole multiple of kpool must cut on a pool boundary. That argument silently +// assumes tie groups never SPAN pools. F.relu drives most pool scores to exactly +// 0.0, so they do, and ggml_top_k - explicitly unordered among equals - then +// takes an arbitrary 1..kpool-1 members of the pool it lands in. This +// reproduces that with no model at all: most pools at exactly 0.0, a minority +// positive, fewer positive pools than the budget. // static std::vector run_top_k( @@ -115,23 +125,28 @@ static int64_t n_partial_pools(const std::vector & sel, int64_t off, in } static void test_top_k_boundary() { - printf("\n--- top-k boundary, standalone ---\n"); + printf("\n--- top-k granularity, standalone ---\n"); - const int64_t kpool = 4; - const int64_t top_k = 2048; // the reference requires top_k %% kpool == 0 - const int64_t n_kv = 8192; - const int64_t n_rows = 4; // queries - - // all kpool members of a pool carry the pool score bit-identically, so an - // intra-pool tie is harmless as long as the budget ends on a pool boundary + const int64_t kpool = 4; + const int64_t top_k = 2048; // the reference requires top_k %% kpool == 0 + const int64_t n_kv = 8192; + const int64_t n_rows = 4; // queries const int64_t n_pools = n_kv/kpool; - std::vector pool_score(n_pools); - for (int64_t p = 0; p < n_pools; ++p) { - // deterministic, distinct, no exact ties between pools - pool_score[p] = std::sin(0.7f*(float) p)*1000.0f + 0.001f*(float) p; + // the shape ReLU actually produces: a minority of pools with a positive score, + // every other pool at EXACTLY 0.0. Fewer positive pools than the budget, which + // during prefill is the normal state and not a corner case + const int64_t n_pos = 300; + + std::vector pool_score(n_pools, 0.0f); + for (int64_t p = 0; p < n_pos; ++p) { + // deterministic, distinct, strictly positive + pool_score[(p*7919) % n_pools] = 1.0f + std::fabs(std::sin(0.7f*(float) p))*1000.0f; } + const int64_t select_k = llama_kpool_select_k((uint32_t) n_pools, (uint32_t) top_k, (uint32_t) kpool); + CHECK(select_k == top_k/kpool, "llama_kpool_select_k is index_topk/index_kpool (%d pools)", (int) select_k); + ggml_backend_t backend_cpu = ggml_backend_cpu_init(); ggml_backend_t backend_gpu = nullptr; { @@ -144,83 +159,105 @@ static void test_top_k_boundary() { printf("%s GPU backend for the top-k comparison: %s\n", backend_gpu ? "ok " : "note", backend_gpu ? ggml_backend_name(backend_gpu) : "none, CPU only"); - // t = (q + 1) %% kpool cells of trailing incomplete pool. t == 3 is the one residue - // where the reference's own width happens to stay pool-aligned + // t = (q + 1) %% kpool cells of trailing incomplete pool, biased out of the + // budget. swept so that the per-stream tail residue is covered, not just t == 0 for (int64_t t = 0; t < kpool; ++t) { const int64_t n_tail = t; const int64_t n_full = n_kv - n_tail; // cells belonging to complete pools - // ---- what this change does: budget = top_k, tail out of the budget ---- - std::vector s_fix((size_t) n_kv*n_rows); - // ---- what the prototype did: budget = top_k + kpool - 1, tail at +1e9 -- - std::vector s_old((size_t) n_kv*n_rows); + // ---- A: the WRONG design. one score per cell, top-k of width index_topk -- + std::vector s_cell((size_t) n_kv*n_rows); + // ---- B: what ships. one score per pool, top-k of width index_topk/kpool -- + std::vector s_pool((size_t) n_pools*n_rows); for (int64_t r = 0; r < n_rows; ++r) { for (int64_t j = 0; j < n_kv; ++j) { - const bool tail = j >= n_full; - - s_fix[r*n_kv + j] = tail ? -INFINITY : pool_score[j/kpool]; - s_old[r*n_kv + j] = tail ? 1e9f : pool_score[j/kpool]; + s_cell[r*n_kv + j] = j >= n_full ? -INFINITY : pool_score[j/kpool]; + } + for (int64_t p = 0; p < n_pools; ++p) { + // a pool that the tail bites into is not a candidate at all + s_pool[r*n_pools + p] = (p + 1)*kpool > n_full ? -INFINITY : pool_score[p]; } } - const int64_t w_fix = llama_kpool_top_k_width((uint32_t) n_kv, (uint32_t) top_k, (uint32_t) kpool); - const int64_t w_old = std::min(n_kv, top_k + kpool - 1); + const auto sel_cell = run_top_k(backend_cpu, s_cell, n_kv, n_rows, top_k); + const auto sel_pool = run_top_k(backend_cpu, s_pool, n_pools, n_rows, select_k); - const auto sel_fix_cpu = run_top_k(backend_cpu, s_fix, n_kv, n_rows, w_fix); - const auto sel_old_cpu = run_top_k(backend_cpu, s_old, n_kv, n_rows, w_old); - - CHECK(w_fix == top_k, "t=%d: top-k runs at exactly indexer_top_k (%d)", (int) t, (int) w_fix); - - // 1. boundary exactness of the new scheme - int64_t partial_fix = 0; - int64_t tail_in_fix = 0; + // 1. the cell-level form, at a budget that IS a whole number of pools, + // still splits pools. this is the measurement the design turns on + int64_t partial_cell = 0; + for (int64_t r = 0; r < n_rows; ++r) { + partial_cell += n_partial_pools(sel_cell, r*top_k, top_k, kpool); + } + CHECK(partial_cell > 0, + "t=%d: a CELL-level top-k at the pool-aligned width %d still leaves %d partial pool(s); " + "a pool-aligned WIDTH does not make a pool-aligned CUT", + (int) t, (int) top_k, (int) partial_cell); + + // 2. the pool-level form: expand each selected pool ordinal to its kpool + // member cells, exactly as the graph does through pool_cells + std::vector expanded((size_t) select_k*kpool*n_rows); for (int64_t r = 0; r < n_rows; ++r) { - partial_fix += n_partial_pools(sel_fix_cpu, r*w_fix, w_fix, kpool); - for (int64_t i = 0; i < w_fix; ++i) { - tail_in_fix += sel_fix_cpu[r*w_fix + i] >= n_full; + for (int64_t i = 0; i < select_k; ++i) { + const int32_t p = sel_pool[r*select_k + i]; + for (int64_t m = 0; m < kpool; ++m) { + expanded[r*select_k*kpool + i*kpool + m] = (int32_t) (p*kpool + m); + } } } - CHECK(partial_fix == 0, "t=%d: the %d-cell budget selects only whole pools (%d partial)", - (int) t, (int) w_fix, (int) partial_fix); - CHECK(tail_in_fix == 0, "t=%d: no tail cell consumes budget (%d did)", (int) t, (int) tail_in_fix); - // 2. the same run on the prototype's width - int64_t partial_old = 0; + int64_t partial_pool = 0; + int64_t tail_in_pool = 0; for (int64_t r = 0; r < n_rows; ++r) { - partial_old += n_partial_pools(sel_old_cpu, r*w_old, w_old, kpool); + partial_pool += n_partial_pools(expanded, r*select_k*kpool, select_k*kpool, kpool); + for (int64_t i = 0; i < select_k*kpool; ++i) { + tail_in_pool += expanded[r*select_k*kpool + i] >= n_full; + } } - // the tail itself is one partial pool by construction whenever t > 0 - const int64_t expect_old = t == 3 ? (t > 0 ? 1 : 0) : (t > 0 ? 2 : 1); - CHECK(partial_old/n_rows == expect_old, - "t=%d: width %d leaves %d partial pool(s) per query, expected %d", - (int) t, (int) w_old, (int) (partial_old/n_rows), (int) expect_old); + CHECK(partial_pool == 0, + "t=%d: a POOL-level top-k of %d pools expands to only whole pools (%d partial)", + (int) t, (int) select_k, (int) partial_pool); + CHECK(tail_in_pool == 0, "t=%d: no tail cell consumes budget (%d did)", (int) t, (int) tail_in_pool); // 3. CPU vs CUDA on the same input if (backend_gpu) { - const auto sel_fix_gpu = run_top_k(backend_gpu, s_fix, n_kv, n_rows, w_fix); - const auto sel_old_gpu = run_top_k(backend_gpu, s_old, n_kv, n_rows, w_old); - - bool same_fix = sel_fix_gpu.size() == sel_fix_cpu.size(); - for (int64_t r = 0; same_fix && r < n_rows; ++r) { - std::set a(sel_fix_cpu.begin() + r*w_fix, sel_fix_cpu.begin() + (r + 1)*w_fix); - std::set b(sel_fix_gpu.begin() + r*w_fix, sel_fix_gpu.begin() + (r + 1)*w_fix); - same_fix = a == b; + const auto gpu_cell = run_top_k(backend_gpu, s_cell, n_kv, n_rows, top_k); + const auto gpu_pool = run_top_k(backend_gpu, s_pool, n_pools, n_rows, select_k); + + // the pool SET may legitimately differ between backends when the cut + // falls inside a tie group. What may not differ is pool integrity, + // and that is what is asserted + int64_t partial_gpu = 0; + bool ok_gpu = gpu_pool.size() == sel_pool.size(); + for (int64_t r = 0; ok_gpu && r < n_rows; ++r) { + for (int64_t i = 0; i < select_k; ++i) { + const int32_t p = gpu_pool[r*select_k + i]; + for (int64_t m = 0; m < kpool; ++m) { + expanded[r*select_k*kpool + i*kpool + m] = (int32_t) (p*kpool + m); + } + } + partial_gpu += n_partial_pools(expanded, r*select_k*kpool, select_k*kpool, kpool); + } + CHECK(ok_gpu && partial_gpu == 0, + "t=%d: %s also expands to only whole pools (%d partial)", + (int) t, ggml_backend_name(backend_gpu), (int) partial_gpu); + + bool same_pool = gpu_pool.size() == sel_pool.size(); + for (int64_t r = 0; same_pool && r < n_rows; ++r) { + std::set a(sel_pool.begin() + r*select_k, sel_pool.begin() + (r + 1)*select_k); + std::set b(gpu_pool.begin() + r*select_k, gpu_pool.begin() + (r + 1)*select_k); + same_pool = a == b; } - CHECK(same_fix, "t=%d: CPU and %s select the same cell set at width %d", - (int) t, ggml_backend_name(backend_gpu), (int) w_fix); - - bool same_old = sel_old_gpu.size() == sel_old_cpu.size(); - for (int64_t r = 0; same_old && r < n_rows; ++r) { - std::set a(sel_old_cpu.begin() + r*w_old, sel_old_cpu.begin() + (r + 1)*w_old); - std::set b(sel_old_gpu.begin() + r*w_old, sel_old_gpu.begin() + (r + 1)*w_old); - same_old = a == b; + int64_t partial_gpu_cell = 0; + for (int64_t r = 0; r < n_rows; ++r) { + partial_gpu_cell += n_partial_pools(gpu_cell, r*top_k, top_k, kpool); } - // reported, not asserted: whether the arbitrary pick actually diverges - // depends on each backend's partial sort - printf("note t=%d: CPU and %s %s at width %d\n", + // reported, not asserted: which members of a cut pool each backend's + // partial sort happens to keep is not contractual + printf("note t=%d: CPU and %s %s on the selected POOL set; the cell-level form leaves " + "%d partial pool(s) there too\n", (int) t, ggml_backend_name(backend_gpu), - same_old ? "happen to agree" : "DISAGREE", (int) w_old); + same_pool ? "agree" : "DISAGREE", (int) partial_gpu_cell); } } @@ -238,9 +275,14 @@ struct kpool_tensors { ggml_tensor * cell_pool = nullptr; ggml_tensor * pool_cells = nullptr; ggml_tensor * bias = nullptr; + ggml_tensor * pool_bias = nullptr; ggml_tensor * sel_mask = nullptr; + ggml_tensor * cand_mask = nullptr; }; +// The test asks for all six, including the two the pooled graph does not consume. +// cell_pool and bias are the per-CELL spelling of the same predicate, computed by a +// different loop, and they are what pool_bias and cand_mask are checked against here static kpool_tensors alloc_kpool_tensors( ggml_context * ctx, int64_t n_kv, @@ -254,12 +296,16 @@ static kpool_tensors alloc_kpool_tensors( t.cell_pool = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_kv, n_stream); t.pool_cells = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, kpool*n_pools, n_stream); t.bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_kv, n_tps, n_stream); + t.pool_bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_pools, n_tps, n_stream); t.sel_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_padq, 1, n_stream); + t.cand_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_padq, 1, n_stream); ggml_set_input(t.cell_pool); ggml_set_input(t.pool_cells); ggml_set_input(t.bias); + ggml_set_input(t.pool_bias); ggml_set_input(t.sel_mask); + ggml_set_input(t.cand_mask); return t; } @@ -626,16 +672,26 @@ int main(int argc, char ** argv) { pooled->nb[1], pooled->nb[2], 0)), ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, d_idx, n_tps, n_stream)); - // broadcast pool -> member cells, then top-k over CELL scores - ggml_tensor * expanded = ggml_get_rows(ctx0, - ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), kt.cell_pool); - expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); - expanded = ggml_add(ctx0, expanded, kt.bias); + // mask the pools the reference rejects, then top-k over POOLS + score = ggml_add(ctx0, ggml_cont(ctx0, score), kt.pool_bias); + + const int64_t select_k = llama_kpool_select_k((uint32_t) n_pools, hparams.indexer_top_k, (uint32_t) kpool); + ggml_tensor * sel = ggml_cont(ctx0, ggml_top_k(ctx0, score, (int) select_k)); + + // expand each selected POOL ordinal into its kpool member CELLS, which is + // the reference's selected_indices = pool_indices[batch_idx, selected]. + // the query axis folds into the gather's row axis so that one get_rows + // serves every query, while the stream axis stays where get_rows wants it + ggml_tensor * pc3 = ggml_reshape_3d(ctx0, kt.pool_cells, kpool, n_pools, n_stream); + ggml_tensor * sel_flat = ggml_reshape_2d(ctx0, sel, select_k*n_tps, n_stream); - const int64_t width = llama_kpool_top_k_width((uint32_t) n_kv, hparams.indexer_top_k, (uint32_t) kpool); - ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, (int) width)); + const int64_t width = kpool*select_k; + ggml_tensor * top_k = ggml_reshape_3d(ctx0, + ggml_get_rows(ctx0, pc3, sel_flat), width, n_tps, n_stream); - printf("note top-k over replicated per-cell scores yields %d I32 CELL indices\n", (int) width); + CHECK(top_k->type == GGML_TYPE_I32, "the pool -> cell expansion stays I32 (pool_cells is I32)"); + printf("note top-k over %d POOL scores expands to %d I32 CELL indices per query\n", + (int) select_k, (int) width); // 5. the scatter the mask is built from, starting at sel_mask rather than // an all -INFINITY fill, which is what forces the tail in @@ -662,7 +718,7 @@ int main(int argc, char ** argv) { if (buf) { const llama_ubatch & ub = mctx->get_ubatch(); - llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.sel_mask, + llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, kt.sel_mask, kt.cand_mask, &ub, (uint32_t) kpool); const int32_t * cp = (const int32_t *) kt.cell_pool->data; @@ -686,7 +742,114 @@ int main(int argc, char ** argv) { for (int64_t i = 0; i < ggml_nelements(kt.sel_mask); ++i) { finite &= sm[i] == 0.0f || sm[i] == -INFINITY; } - CHECK(finite, "bias and sel_mask hold only 0 and -INFINITY: no +1e9 to meet a -inf"); + for (int64_t i = 0; i < ggml_nelements(kt.pool_bias); ++i) { + const float * pb = (const float *) kt.pool_bias->data; + finite &= pb[i] == 0.0f || pb[i] == -INFINITY; + } + for (int64_t i = 0; i < ggml_nelements(kt.cand_mask); ++i) { + const float * cm = (const float *) kt.cand_mask->data; + finite &= cm[i] == 0.0f || cm[i] == -INFINITY; + } + CHECK(finite, "bias, pool_bias, sel_mask and cand_mask hold only 0 and -INFINITY: " + "no +1e9 to meet a -inf"); + + // cand_mask is exactly max(bias, sel_mask) lifted to KQ shape, and it + // is what stops an over-budget top-k from escaping the reference's + // candidate set. + // + // ggml_top_k always returns select_k pool ordinals even when fewer + // than select_k pools are finite. During prefill the query at + // position q has only ~q/kpool complete visible pools against a + // budget of index_topk/kpool, so the budget spills into -INFINITY + // pools and picks among them arbitrarily. Adding the causal mask + // kills the expansions that are empty or in the future. What it does + // not kill is a resident, causally visible cell that sits in an + // INCOMPLETE pool below the tail: the reference never selects it + // (pool_valid = grouped_valid_keys.all(-1)), the graph would. + // Unreachable while positions are contiguous, reachable the moment a + // partial seq_rm leaves a hole. + { + const float * cm = (const float *) kt.cand_mask->data; + + bool is_union = true; + int64_t n_spill = 0; // candidate pools strictly fewer than the budget + + const int64_t select_k = + llama_kpool_select_k((uint32_t) n_pools, hparams.indexer_top_k, (uint32_t) kpool); + + for (int64_t s = 0; s < n_stream; ++s) { + for (int64_t ii = 0; ii < n_tps; ++ii) { + const float * rb = bi + (s*n_tps + ii)*n_kv; + const float * rs = sm + (s*n_padq + ii)*n_kv; + const float * rc = cm + (s*n_padq + ii)*n_kv; + + for (int64_t j = 0; j < n_kv; ++j) { + is_union &= rc[j] == std::max(rb[j], rs[j]); + } + + const float * rp = (const float *) kt.pool_bias->data + (s*n_tps + ii)*n_pools; + int64_t n_cand = 0; + for (int64_t p = 0; p < n_pools; ++p) { + n_cand += rp[p] == 0.0f; + } + n_spill += n_cand < select_k; + } + } + + CHECK(is_union, "cand_mask == max(bias, sel_mask): the reference's candidate set"); + + // not a failure: it is the normal prefill state, and the whole + // reason the gate has to exist. Printed so that a future change + // making it zero cannot quietly turn the check above vacuous + printf("note %d of %d (query, stream) rows have fewer candidate pools than the\n" + " top-k budget of %d, so the budget spills and cand_mask is load bearing\n", + (int) n_spill, (int) (n_stream*n_tps), (int) select_k); + } + + // pool_bias is the same predicate as bias, evaluated where the + // reference evaluates it. For a COMPLETE pool the two must agree + // cell for cell; for an incomplete or absent pool bias has no cell + // to speak for it, which is exactly why pool_bias is computed here + // rather than gathered from bias at each pool's last member + { + const float * pb = (const float *) kt.pool_bias->data; + + bool agree = true; + bool exact = true; + + for (int64_t s = 0; s < n_stream; ++s) { + for (int64_t ii = 0; ii < n_tps; ++ii) { + const float * rb = bi + (s*n_tps + ii)*n_kv; + const float * rp = pb + (s*n_tps + ii)*n_pools; + + std::set pools_of_scored_cells; + + for (int64_t j = 0; j < n_kv; ++j) { + if (rb[j] == 0.0f) { + // a scored cell's pool must be a candidate pool + agree &= rp[cp[s*n_kv + j]] == 0.0f; + pools_of_scored_cells.insert(cp[s*n_kv + j]); + } + } + + int64_t n_cand = 0; + for (int64_t p = 0; p < n_pools; ++p) { + n_cand += rp[p] == 0.0f; + } + + // and nothing else may be one. an entirely absent pool + // would pass the first check vacuously; this is what + // catches it, and it is exactly the failure mode of + // gathering pool_bias from bias at each pool's last + // member instead of computing it + exact &= n_cand == (int64_t) pools_of_scored_cells.size(); + } + } + + CHECK(agree, "every cell bias scores lies in a pool pool_bias also accepts"); + CHECK(exact, "and pool_bias accepts no pool that has no scored cell " + "(an absent pool must not become a candidate)"); + } // per query: (q+1) %% kpool cells sit in its own incomplete pool and // must be forced in, the q+1-minus-that below must be scored, and @@ -823,7 +986,7 @@ int main(int argc, char ** argv) { ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); if (buf) { - llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.sel_mask, + llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, kt.sel_mask, kt.cand_mask, &mc->get_ubatch(), (uint32_t) kpool); const int32_t * cp = (const int32_t *) kt.cell_pool->data; @@ -972,7 +1135,7 @@ int main(int argc, char ** argv) { double ms = 1e9; for (int rep = 0; rep < 4; ++rep) { const auto t0 = std::chrono::steady_clock::now(); - llama_kv_cache_set_input_kpool(kv, kt.cell_pool, kt.pool_cells, kt.bias, kt.sel_mask, + llama_kv_cache_set_input_kpool(kv, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, kt.sel_mask, kt.cand_mask, &ub, (uint32_t) kpool); const auto t1 = std::chrono::steady_clock::now(); @@ -992,34 +1155,39 @@ int main(int argc, char ** argv) { n_dsa, n_dsa); { - const float * sm = (const float *) kt.sel_mask->data; + const float * sm = (const float *) kt.sel_mask ->data; + const float * cm = (const float *) kt.cand_mask->data; bool pad_masked = true; for (int64_t ii = n_tps; ii < n_padq; ++ii) { for (int64_t j = 0; j < n_kv; ++j) { pad_masked &= sm[ii*n_kv + j] == -INFINITY; + pad_masked &= cm[ii*n_kv + j] == -INFINITY; } } - CHECK(pad_masked, "the %d padding rows of a wider sel_mask stay -INFINITY", + CHECK(pad_masked, "the %d padding rows of a wider sel_mask/cand_mask stay -INFINITY", (int) (n_padq - n_tps)); } - // the shared object refills the same tensors deterministically - std::vector first((size_t) ggml_nelements(kt.cell_pool)); - memcpy(first.data(), kt.cell_pool->data, first.size()*sizeof(int32_t)); + // the shared object refills the same tensors deterministically. + // it fills only what the pooled graph consumes - cell_pool and + // bias have no consumer there, so it passes nullptr for both and + // pool_cells is what gets compared + std::vector first((size_t) ggml_nelements(kt.pool_cells)); + memcpy(first.data(), kt.pool_cells->data, first.size()*sizeof(int32_t)); llm_graph_input_kpool inp(mctx->get_attn(), mctx->get_idx(), (uint32_t) kpool); inp.k_idxs = mctx->get_idx()->build_input_k_idxs(ctx, ub); - inp.cell_pool = kt.cell_pool; inp.pool_cells = kt.pool_cells; - inp.bias = kt.bias; + inp.pool_bias = kt.pool_bias; inp.sel_mask = kt.sel_mask; + inp.cand_mask = kt.cand_mask; ggml_backend_buffer_t buf2 = ggml_backend_alloc_ctx_tensors(ctx, backend); if (buf2) { inp.set_input(&ub); - CHECK(memcmp(first.data(), kt.cell_pool->data, first.size()*sizeof(int32_t)) == 0, + CHECK(memcmp(first.data(), kt.pool_cells->data, first.size()*sizeof(int32_t)) == 0, "llm_graph_input_kpool::set_input rebuilds the identical map, so one\n" " object can back every indexer layer"); From fe95953cc087bfeedbf85b8f4cd97a8815a16587 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 07:50:52 +0000 Subject: [PATCH 13/36] mtmd: don't normalize patch embeddings when norm_embd is absent --- tools/mtmd/models/glm4v.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/tools/mtmd/models/glm4v.cpp b/tools/mtmd/models/glm4v.cpp index 0e1d596b41b..ad36e7a483d 100644 --- a/tools/mtmd/models/glm4v.cpp +++ b/tools/mtmd/models/glm4v.cpp @@ -42,7 +42,11 @@ ggml_cgraph * clip_graph_glm4v::build() { cb(inp, "patch_bias", -1); // pos-conv norm - inp = build_norm(inp, model.norm_embd_w, model.norm_embd_b, norm_t, eps, -1); + // Note: GLM-OCR does not have a post-conv norm, and build_norm still normalizes when the + // weight is null, so the whole call must be skipped rather than just the affine scale + if (model.norm_embd_w != nullptr) { + 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 From 582fe816e5fd57e7a52d842c4f2081a942869773 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 07:57:04 +0000 Subject: [PATCH 14/36] convert: refuse to write a GGUF with zero tensors --- conversion/base.py | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/conversion/base.py b/conversion/base.py index 56547ace009..3ff5634407c 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -2537,6 +2537,20 @@ def set_gguf_parameters(self): if not self.has_vision_encoder and not self.has_audio_encoder: raise ValueError("MmprojModel must have either vision or audio encoder") + def prepare_tensors(self): + super().prepare_tensors() + + # an mmproj has no vocab-only mode, so an empty one is always a silent mapping failure + # rather than a supported output; count and size are both checked because a tensor map + # that produces only zero-sized entries is just as broken as one that produces none + n_tensors = sum(len(t) for t in self.gguf_writer.tensors) + n_bytes = sum(ti.nbytes for t in self.gguf_writer.tensors for ti in t.values()) + if n_tensors == 0 or n_bytes == 0: + raise ValueError( + f"refusing to write an mmproj with no tensor data (n_tensors = {n_tensors}, n_bytes = {n_bytes}); " + "check that the model's vision/audio tensors are named as the tensor map expects" + ) + def write_vocab(self): raise ValueError("MmprojModel does not support vocab writing") From 9901ab4bdb700bd747f50af731a1670efa0bc6eb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 08:32:44 +0000 Subject: [PATCH 15/36] mtmd: add glm5next (GLM-5.3-Flash) vision tower The tower is the GLM-OCR ViT with a clamped SwiGLU: the gate is bounded above only, the up projection on both sides, and both before the SiLU. ggml_swiglu_oai clamps the same way but then adds one to the up branch, which is a gpt-oss detail this model does not share, so this adds an FFN_SILU_CLAMP op rather than reusing it. The clamp sits at the per-block MLP and again at the merger. Both read hparams.ffn_op, so the graph body stays the GLM-4V one and the pair is covered together. It gets its own projector type rather than a flag on glm4v because the image token limits differ (16/8000 against 8/4096, per the GLM-5.3-Flash preprocessor) and those are hardcoded per projector, and because the clamp must stay off for GLM-4V and GLM-OCR. Also writes clip.vision.spatial_merge_size. No GLM4V-family mmproj has ever carried it: Glm4VVisionModel skips the Qwen3VL parameters, which is where it is written, so clip.cpp's hardcoded 2 has been carrying it. Images only. glm5next spells video with its own token pair and distinct start/end spans, and that is not handled here. --- conversion/glm5next.py | 23 +++++++++++++--- conversion/qwen3vl.py | 5 +++- gguf-py/gguf/constants.py | 1 + tools/mtmd/CMakeLists.txt | 1 + tools/mtmd/clip-impl.h | 3 +++ tools/mtmd/clip-model.h | 4 +++ tools/mtmd/clip.cpp | 39 +++++++++++++++++++++++++++ tools/mtmd/models/glm5next-vision.cpp | 16 +++++++++++ tools/mtmd/models/models.h | 5 ++++ tools/mtmd/mtmd.cpp | 2 ++ 10 files changed, 94 insertions(+), 5 deletions(-) create mode 100644 tools/mtmd/models/glm5next-vision.cpp diff --git a/conversion/glm5next.py b/conversion/glm5next.py index 24290041e6d..f2db5867765 100644 --- a/conversion/glm5next.py +++ b/conversion/glm5next.py @@ -251,8 +251,23 @@ def prepare_tensors(self): @ModelBase.register("Glm5NextForConditionalGeneration") # [TAG_HF_EXAMPLE_MISSING] class Glm5NextVisionModel(Glm4VVisionModel): - """The GLM-4.5V ViT under a `model.visual.` prefix, registered so the mmproj - can be produced. The clip graph is not ported yet: glm4v.cpp normalises the - patch embeddings unconditionally but this tower has no post_conv_layernorm, - and vision_config.swiglu_limit has no key. + """The vision tower is the GLM-OCR ViT under a `model.visual.` prefix. + + Every tensor already maps through the GLM-4V entries. The one structural + difference is a clamp on the SwiGLU gate and up projections, applied in the + per-block MLP and again in the merger, so it gets its own projector type + rather than a flag on glm4v. """ + + clip_projector_type = gguf.VisionProjectorType.GLM5NEXT + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_vision is not None + + # Glm4VVisionModel bypasses Qwen3VLVisionModel entirely, which is also where + # the merge size is written, so no GLM4V-family mmproj carries this key and + # clip.cpp falls back to a hardcoded 2. Write it rather than rely on that + self.gguf_writer.add_vision_spatial_merge_size(int(self.hparams_vision["spatial_merge_size"])) + + self.gguf_writer.add_vision_swiglu_limit(float(self.hparams_vision["swiglu_limit"])) diff --git a/conversion/qwen3vl.py b/conversion/qwen3vl.py index 4fec708c9ff..385f47149fc 100644 --- a/conversion/qwen3vl.py +++ b/conversion/qwen3vl.py @@ -228,10 +228,13 @@ 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): + # subclasses that share this tower but need their own clip graph override this + clip_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.clip_projector_type) hidden_act = str(self.hparams_vision.get("hidden_act", "")).lower() if hidden_act == "gelu": diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index bc38ba87a85..6d740cca197 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -5588,6 +5588,7 @@ class VisionProjectorType: LFM2A = "lfm2a" # audio MUSIC_FLAMINGO = "musicflamingo" # audio GLM4V = "glm4v" + GLM5NEXT = "glm5next" YOUTUVL = "youtuvl" NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index e60c9c8787a..4ff8db9e6be 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -38,6 +38,7 @@ add_library(mtmd models/gemma4ua.cpp models/gemma4uv.cpp models/glm4v.cpp + models/glm5next-vision.cpp models/granite-speech.cpp models/granite4-vision.cpp models/hunyuanvl.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index f6045093c63..5c4821f4027 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_VISION_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_GLM5NEXT, 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_GLM5NEXT, "glm5next"}, { 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 060938d86e3..47d28f9018a 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -17,6 +17,7 @@ enum ffn_op_type { FFN_SILU, FFN_GELU_QUICK, FFN_RELU_SQR, + FFN_SILU_CLAMP, }; enum norm_type { @@ -89,6 +90,9 @@ struct clip_hparams { ffn_op_type ffn_op = FFN_GELU; + // clamp applied to the SwiGLU gate/up before the activation (FFN_SILU_CLAMP) + float swiglu_limit = 0.0f; + patch_merge_type mm_patch_merge_type = PATCH_MERGE_FLAT; float eps = 1e-6; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 90de1957586..9a7350a13c3 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -684,6 +684,21 @@ ggml_tensor * clip_graph::build_ffn( cur = ggml_sqr(ctx0, cur); cb(cur, "ffn_relu_sqr", il); } break; + case FFN_SILU_CLAMP: + { + // the gate is bounded above only, the up projection on both sides, and both + // before the activation. ggml_swiglu_oai clamps the same way but then adds + // one to the up branch, which is a gpt-oss detail this model does not share + GGML_ASSERT(gate && "FFN_SILU_CLAMP is a gated activation"); + const float limit = hparams.swiglu_limit; + GGML_ASSERT(limit > 0.0f); + tmp = ggml_clamp(ctx0, tmp, -limit, limit); + cb(tmp, "ffn_up_clamped", il); + cur = ggml_clamp(ctx0, cur, -INFINITY, limit); + cb(cur, "ffn_gate_clamped", il); + cur = ggml_swiglu_split(ctx0, cur, tmp); + cb(cur, "ffn_swiglu_limited", il); + } break; } if (down) { @@ -1081,6 +1096,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_GLM5NEXT: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_QWEN3A: { builder = std::make_unique(ctx, img); @@ -1726,6 +1745,20 @@ 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_GLM5NEXT: + { + hparams.rope_theta = 10000.0f; + hparams.n_merge = 2; + hparams.image_resize_algo = RESIZE_ALGO_BILINEAR; + get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); + // drives the clamp in both the per-block MLP and the merger + get_f32(KEY_VISION_SWIGLU_LIMIT, hparams.swiglu_limit); + hparams.ffn_op = FFN_SILU_CLAMP; + log_ffn_op = "silu_clamp"; + // min_pixels/max_pixels of the GLM-5.3-Flash preprocessor, in tokens + hparams.set_limit_image_tokens(16, 8000); + hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup + } break; case PROJECTOR_TYPE_LLAMA4: { hparams.rope_theta = 10000.0f; @@ -2543,6 +2576,7 @@ struct clip_model_loader { } } break; case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5NEXT: { 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 +4035,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_GLM5NEXT: case PROJECTOR_TYPE_PADDLEOCR: case PROJECTOR_TYPE_HUNYUANVL: case PROJECTOR_TYPE_YOUTUVL: @@ -4027,6 +4062,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_GLM5NEXT: case PROJECTOR_TYPE_PADDLEOCR: case PROJECTOR_TYPE_HUNYUANVL: case PROJECTOR_TYPE_YOUTUVL: @@ -4108,6 +4144,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_GLM5NEXT: case PROJECTOR_TYPE_YOUTUVL: case PROJECTOR_TYPE_MUSE_GLIMMER: { @@ -4733,6 +4770,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_GLM5NEXT: { const int merge_ratio = hparams.n_merge; const int pw = image_size_width / patch_size; @@ -5881,6 +5919,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_GLM5NEXT: 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/glm5next-vision.cpp b/tools/mtmd/models/glm5next-vision.cpp new file mode 100644 index 00000000000..e1130f1cc5d --- /dev/null +++ b/tools/mtmd/models/glm5next-vision.cpp @@ -0,0 +1,16 @@ +#include "models.h" + +// GLM-5.3-Flash reuses the GLM-OCR ViT unchanged: no learned position embeddings, no +// post-conv norm, q/k norms and biases throughout, and an RMSNorm called post_layernorm. +// ref: https://huggingface.co/zai-org/GLM-5.3-Flash/blob/main/modeling_glm5_next.py +// +// The one difference is the clamp on the SwiGLU gate and up projections. It applies to +// the per-block MLP and to the merger alike, and both read hparams.ffn_op, so selecting +// FFN_SILU_CLAMP for this projector type is enough to cover the pair. +ggml_cgraph * clip_graph_glm5next::build() { + GGML_ASSERT(hparams.ffn_op == FFN_SILU_CLAMP); + GGML_ASSERT(model.norm_embd_w == nullptr); + GGML_ASSERT(model.position_embeddings == nullptr); + + return clip_graph_glm4v::build(); +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 10546fa5dc7..7a7e9e69870 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -178,6 +178,11 @@ struct clip_graph_glm4v : clip_graph { ggml_cgraph * build() override; }; +struct clip_graph_glm5next : clip_graph_glm4v { + clip_graph_glm5next(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph_glm4v(ctx, img) {} + ggml_cgraph * build() override; +}; + struct clip_graph_hunyuanvl : clip_graph { clip_graph_hunyuanvl(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 5b306180d62..0975419eb7c 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -854,8 +854,10 @@ struct mtmd_context { image_preproc = std::make_unique(ctx_v); } break; case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5NEXT: { // <|begin_of_image|> ... (image embeddings) ... <|end_of_image|> + // glm5next spells video with its own token pair, but video is not supported here img_beg = "<|begin_of_image|>"; img_end = "<|end_of_image|>"; image_preproc = std::make_unique(ctx_v); From 41cd9235fbc0f4dfcf03fe9e8403254a345cb9fe Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 09:52:04 +0000 Subject: [PATCH 16/36] mtmd: implement the glm5next 0826 image preprocessor the vision tower shipped with the shared dynamic-size preprocessor, which is a qwen-style smart_resize. the 2026-08-26 GLM-5-Next adaptation resizes differently: both edges are aligned up by ceil rather than round, an over-budget image is fitted by binary searching the content height for the largest aligned canvas still within max_pixels, and the resized content is pasted into the top-left of that canvas rather than centred and stretched to fill it. an image already at or above min_pixels is never upscaled. min_pixels/max_pixels stay in tokens. the reference scales them by temporal_factor * factor**2 and compares against aligned_frames * area, and aligned_frames equals temporal_factor for a still image, so the two cancel and hparams.image_min_pixels / image_max_pixels (16 and 8000 tokens, 12544 and 6272000 pixels) are used directly. glm4v and glm-ocr keep the dynamic-size preprocessor. images only. video has its own token pair (154855, distinct from the image token 154854) with its own start/end spans, and is out of scope here. the resize arithmetic is covered in test-mtmd-impl against values taken from the reference processor, including the 16- and 8000-token boundaries, extreme aspect ratios, and inputs where the binary search and smart_resize disagree. --- tools/mtmd/mtmd-image.cpp | 105 ++++++++++++++++++++++++++++++++++++++ tools/mtmd/mtmd.cpp | 9 +++- 2 files changed, 112 insertions(+), 2 deletions(-) diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 0dda8770f29..5dc80faae4e 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1583,3 +1583,108 @@ mtmd_image_preproc_out mtmd_image_preprocessor_muse_glimmer::preprocess(const cl output.append(hparams, resized_image, true); return output; } + +// +// mtmd_image_preprocessor_glm5next +// + +// for a still image the reference's temporal_factor cancels on both sides of its budget comparison, +// leaving a plain pixel area against image_min_pixels / image_max_pixels (n_tokens * factor**2) +clip_image_size mtmd_image_preprocessor_glm5next::smart_resize(const clip_hparams & hparams, const clip_image_size & size) { + const int factor = hparams.patch_size * hparams.n_merge; + GGML_ASSERT(factor > 0); + GGML_ASSERT(hparams.image_min_pixels > 0 && hparams.image_max_pixels > 0); + + const int height = size.height; + const int width = size.width; + if (height <= 0 || width <= 0) { + // an empty bitmap is reachable through the public API, same tolerance as calc_size_preserved_ratio + return { 0, 0 }; + } + + const int64_t min_pixels = hparams.image_min_pixels; + const int64_t max_pixels = hparams.image_max_pixels; + + auto align = [factor](int64_t value) { + return (int) (((value + factor - 1) / factor) * factor); + }; + + int aligned_height = align(height); + int aligned_width = align(width); + + // upscale an image that is too small to spend the minimum token budget + if ((int64_t) aligned_height * aligned_width < min_pixels) { + const double scale = std::sqrt((double) min_pixels / ((double) height * (double) width)); + aligned_height = align(std::max(1, (int64_t) std::ceil(height * scale))); + aligned_width = align(std::max(1, (int64_t) std::ceil(width * scale))); + } + + if ((int64_t) aligned_height * aligned_width > max_pixels) { + // binary search the tallest content height whose aligned canvas still fits the budget. the Qwen + // sqrt(area / max_pixels) scale leaves budget unspent: aligning both edges is not monotone in it + int low = 1; + int high = height; + aligned_height = factor; + aligned_width = factor; + while (low <= high) { + const int content_height = (low + high) / 2; + // the reference divides in float64 before flooring, keep it bit-identical + const int content_width = std::max(1, (int) std::floor((double) width * content_height / (double) height)); + const int cand_height = align(content_height); + const int cand_width = align(content_width); + + if ((int64_t) cand_height * cand_width <= max_pixels) { + aligned_height = cand_height; + aligned_width = cand_width; + low = content_height + 1; + } else { + high = content_height - 1; + } + } + } + + return { aligned_width, aligned_height }; +} + +mtmd_image_preprocessor_glm5next::geometry mtmd_image_preprocessor_glm5next::get_geometry(const clip_hparams & hparams, const clip_image_size & size) { + const clip_image_size canvas = smart_resize(hparams, size); + + const int height = size.height; + const int width = size.width; + if (canvas.width == 0 || canvas.height == 0) { + return { canvas, canvas }; + } + + double scale = std::min((double) canvas.height / height, (double) canvas.width / width); + if ((int64_t) height * (int64_t) width >= hparams.image_min_pixels) { + // an image already spending the minimum budget is only shrunk, never upscaled to fill the canvas + scale = std::min(1.0, scale); + } + + geometry geo; + geo.canvas = canvas; + geo.content = { + std::max(1, std::min(canvas.width, (int) std::floor(width * scale))), + std::max(1, std::min(canvas.height, (int) std::floor(height * scale))), + }; + return geo; +} + +mtmd_image_preproc_out mtmd_image_preprocessor_glm5next::preprocess(const clip_image_u8 & img) { + const geometry geo = get_geometry(hparams, img.get_size()); + + clip_image_u8 content; + img_tool::resize(img, content, geo.content, hparams.image_resize_algo, PAD_NONE); + + // the reference pads bottom/right, not centred like img_tool::resize would + clip_image_u8 canvas; + canvas.set_size(geo.canvas, img.is_placeholder()); + if (!img.is_placeholder()) { + img_tool::fill(canvas, hparams.image_pad_color); + img_tool::composite(canvas, content, 0, 0); + } + + mtmd_image_preproc_out output; + output.append(hparams, canvas, true); + return output; +} diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 0975419eb7c..1753bd59174 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -854,14 +854,19 @@ struct mtmd_context { image_preproc = std::make_unique(ctx_v); } break; case PROJECTOR_TYPE_GLM4V: - case PROJECTOR_TYPE_GLM5NEXT: { // <|begin_of_image|> ... (image embeddings) ... <|end_of_image|> - // glm5next spells video with its own token pair, but video is not supported here img_beg = "<|begin_of_image|>"; img_end = "<|end_of_image|>"; image_preproc = std::make_unique(ctx_v); } break; + case PROJECTOR_TYPE_GLM5NEXT: + { + // glm5next spells video with its own token pair, but video is not supported here + 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 02fa4a43a2cd9f6b4476084fe07d807ff32ddb42 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 09:55:18 +0000 Subject: [PATCH 17/36] mtmd: resample glm5next images with bicubic, matching the reference --- tools/mtmd/clip.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 9a7350a13c3..913f50e59b1 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1749,7 +1749,10 @@ struct clip_model_loader { { hparams.rope_theta = 10000.0f; hparams.n_merge = 2; - hparams.image_resize_algo = RESIZE_ALGO_BILINEAR; + // the reference asks for BICUBIC (resample = PILImageResampling.BICUBIC); + // the bilinear here came from the GLM-4V case this was copied from. neither + // filter matches torchvision's exactly, so this is closer in kind, not exact + hparams.image_resize_algo = RESIZE_ALGO_BICUBIC; get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); // drives the clamp in both the per-block MLP and the merger get_f32(KEY_VISION_SWIGLU_LIMIT, hparams.swiglu_limit); From 81e3e6716f6de6877364568cfd6b71941eacf52d Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 26 Aug 2026 11:47:30 +0000 Subject: [PATCH 18/36] llama: enable ignore_merges for the glm4 pre-tokenizer MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit glm4 / chatglm-bpe tokenizer.json files set "ignore_merges": true, meaning a pre-token that is already a vocab entry is emitted directly and the merge loop never runs. llama.cpp implements this (llama-vocab.cpp, the get_ignore_merges() short-circuit) but only enables it for a hardcoded list of pre-tokenizer names, and glm4 was never added. Without it the merges are applied - correctly - and reach a different answer, because greedy BPE cannot always reconstruct a vocab entry from its bytes. " 王" (Ġçİĭ, id 102322) is the case that exposed it: from Ġ ç İ ĭ the only merges available are (Ġ,ç)=27944, (ç,İ)=76417 and (çİ,ĭ)=239209, so the lowest rank wins first and yields Ġç İ ĭ, at which point neither (Ġç,İ) nor (İ,ĭ) exists and it stops three tokens short. Reaching Ġçİĭ needs (Ġ,çİĭ) at 242943, which requires never taking (Ġ,ç) at 27944. The trigger is whitespace immediately before a CJK character, so pure Chinese prose is unaffected and mixed Chinese-English is not: pure Chinese prose 620 vs 620 tokens, already identical mixed Chinese-English 680 -> 600 tokens, now identical to HF (-13.3%) wikitext-2 (289569 tok) one divergence -> byte-identical Found while comparing GLM-5.3-Flash perplexity against transformers, vLLM and SGLang: the mismatch bounded how many scoring windows could be compared at long context, and reads exactly like a model-port defect rather than a tokenizer one. --- src/llama-vocab.cpp | 24 ++++++++++++++++++++++++ 1 file changed, 24 insertions(+) diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index ff926ceecd1..0a2b4538ac1 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2257,6 +2257,30 @@ 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; + // glm4 tokenizer.json sets "ignore_merges": true, i.e. a pre-token that is + // already a vocab entry is emitted directly and the merge loop never runs. + // without it the merges are applied - correctly - and reach a different + // answer, because greedy BPE cannot always reconstruct a vocab entry from + // its bytes. + // + // " 王" (Ġçİĭ, id 102322) is the case that exposed this. From Ġ ç İ ĭ the + // only merges that exist are (Ġ,ç)=27944, (ç,İ)=76417 and (çİ,ĭ)=239209; + // the lowest rank wins first, giving Ġç İ ĭ, and then neither (Ġç,İ) nor + // (İ,ĭ) exists, so it stops three tokens short. Reaching Ġçİĭ needs + // (Ġ,çİĭ) at 242943, which requires never taking (Ġ,ç) at 27944. + // + // Trigger is whitespace before a CJK character, so pure Chinese is + // unaffected and mixed Chinese-English inflates ~13%. + // + // Deliberately NOT applied to chatglm-bpe, which shares this pre_type but + // is a different tokenizer (ChatGLM3). I have no ChatGLM3 checkpoint here + // to confirm it declares the flag, and turning it on where the tokenizer + // does not set it would silently change that model's tokenization in the + // same invisible way this bug did. Someone with the checkpoint should + // check tokenizer.json and widen this if it applies. + if (tokenizer_pre == "glm4") { + ignore_merges = true; + } } else if ( tokenizer_pre == "viking") { pre_type = LLAMA_VOCAB_PRE_TYPE_VIKING; From 2d9570d2c7f8c3aaa62c6672d7719ad5381fbfeb Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 14:20:31 +0000 Subject: [PATCH 19/36] llama: shorten comment --- src/llama-vocab.cpp | 27 ++++++--------------------- 1 file changed, 6 insertions(+), 21 deletions(-) diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index 0a2b4538ac1..d37f49a5398 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2257,27 +2257,12 @@ 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; - // glm4 tokenizer.json sets "ignore_merges": true, i.e. a pre-token that is - // already a vocab entry is emitted directly and the merge loop never runs. - // without it the merges are applied - correctly - and reach a different - // answer, because greedy BPE cannot always reconstruct a vocab entry from - // its bytes. - // - // " 王" (Ġçİĭ, id 102322) is the case that exposed this. From Ġ ç İ ĭ the - // only merges that exist are (Ġ,ç)=27944, (ç,İ)=76417 and (çİ,ĭ)=239209; - // the lowest rank wins first, giving Ġç İ ĭ, and then neither (Ġç,İ) nor - // (İ,ĭ) exists, so it stops three tokens short. Reaching Ġçİĭ needs - // (Ġ,çİĭ) at 242943, which requires never taking (Ġ,ç) at 27944. - // - // Trigger is whitespace before a CJK character, so pure Chinese is - // unaffected and mixed Chinese-English inflates ~13%. - // - // Deliberately NOT applied to chatglm-bpe, which shares this pre_type but - // is a different tokenizer (ChatGLM3). I have no ChatGLM3 checkpoint here - // to confirm it declares the flag, and turning it on where the tokenizer - // does not set it would silently change that model's tokenization in the - // same invisible way this bug did. Someone with the checkpoint should - // check tokenizer.json and widen this if it applies. + // glm4 tokenizer.json sets "ignore_merges": true. without it greedy BPE + // cannot reach some vocab entries: " 王" (Ġçİĭ, 102322) needs (Ġ,çİĭ)=242943 + // but (Ġ,ç)=27944 wins first, so it stops three tokens short. triggered by + // whitespace before CJK, inflating mixed Chinese-English ~13%. + // NOT applied to chatglm-bpe, which shares this pre_type but is a different + // tokenizer (ChatGLM3) not confirmed to declare the flag. if (tokenizer_pre == "glm4") { ignore_merges = true; } From 6c59f9b8e41eac519866b22779c2e49e2b0960b5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Wed, 26 Aug 2026 16:53:03 +0000 Subject: [PATCH 20/36] glm5next: repair conflict resolution from the rebase onto master The scripted resolution used for the rebase mangled four files: it spliced a condition into the middle of graph_max_nodes' multi-line else-if, dropped the mtmd_image_preprocessor_glm5next declaration, dropped llama-kv-cache-kpool.cpp from src/CMakeLists.txt (undefined llama_kpool_* and the llm_graph_input_kpool vtable at link time), and left an "} else {" immediately followed by an "} else if" in test-llama-archs. These files are byte-identical between this base and the tree the glm5next work was verified on, so each is taken from there verbatim. --- gguf-py/gguf/constants.py | 1 + gguf-py/gguf/gguf_writer.py | 2 ++ src/CMakeLists.txt | 1 + src/llama-arch.cpp | 6 +++--- src/llama-context.cpp | 7 ++++--- src/llama-model.cpp | 2 +- tests/test-llama-archs.cpp | 19 ++++++++++++++++--- tools/mtmd/mtmd-image.h | 19 +++++++++++++++++++ 8 files changed, 47 insertions(+), 10 deletions(-) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 6d740cca197..cf28513db38 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -365,6 +365,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" # glm5next: clamp on the SwiGLU gate/up 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" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 276432623a4..420b40b639e 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1330,6 +1330,8 @@ def add_vision_image_std(self, values: Sequence[float]) -> None: def add_vision_spatial_merge_size(self, value: int) -> None: self.add_uint32(Keys.ClipVision.SPATIAL_MERGE_SIZE, value) + def add_vision_swiglu_limit(self, value: float) -> None: + self.add_float32(Keys.ClipVision.SWIGLU_LIMIT, value) def add_vision_expert_count_per_layer(self, value: Sequence[int]) -> None: self.add_array(Keys.ClipVision.EXPERT_COUNT_PER_LAYER, value) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c6df19f2ecf..fceddfdb479 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -26,6 +26,7 @@ add_library(llama llama-kv-cache-iswa.cpp llama-kv-cache-dsa.cpp llama-kv-cache-dsa-iswa.cpp + llama-kv-cache-kpool.cpp llama-kv-cache-msa.cpp llama-kv-cache-dsv4.cpp llama-memory.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5a145f9c1ac..076dbfc2c89 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -1012,8 +1012,8 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: - case LLM_ARCH_MINIMAX_01: case LLM_ARCH_GLM5NEXT: + case LLM_ARCH_MINIMAX_01: return true; default: return false; @@ -1037,12 +1037,12 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_LFM2: case LLM_ARCH_LFM2MOE: case LLM_ARCH_BAILINGMOE3: - case LLM_ARCH_GLM5NEXT: return true; default: return false; @@ -1076,9 +1076,9 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_KIMI_K3: - case LLM_ARCH_GLM5NEXT: case LLM_ARCH_QWEN3TTS: return false; default: diff --git a/src/llama-context.cpp b/src/llama-context.cpp index df10311ba70..4542cd9c0a7 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2293,11 +2293,12 @@ 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) { - // the n_tokens*40 budget below is exhausted at ubatch 3840 + if (model.arch == LLM_ARCH_KIMI_K3 || model.arch == LLM_ARCH_GLM5NEXT) { + // the n_tokens*40 budget below is exhausted at ubatch 3840 for kimi-k3, and + // earlier for glm5next: each KDA layer costs 182 nodes + ~16/token, so its 34 + // KDA layers alone need 6.2k + 31.9*n_tokens before DSA or the MoE res = std::max(n_tokens * 160, 64u * model.n_tensors()); } else if (model.arch == LLM_ARCH_QWEN3NEXT || - if (model.arch == LLM_ARCH_GLM5NEXT) { model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_BAILINGMOE3 || model.arch == LLM_ARCH_QWEN35 || diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 15968434481..db26cf5dd69 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2797,8 +2797,8 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_KIMI_LINEAR: - case LLM_ARCH_KIMI_K3: case LLM_ARCH_GLM5NEXT: + case LLM_ARCH_KIMI_K3: return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 523d57a3dfb..ce2febd4209 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -116,9 +116,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_KIMI_LINEAR + || arch == LLM_ARCH_GLM5NEXT || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 - || arch == LLM_ARCH_GLM5NEXT || arch == LLM_ARCH_MISTRAL4) { // MLA absorbs into MQA, so n_head_kv must be 1: otherwise the per-layer // head_count_kv array below sizes the latent K row n_head times too wide @@ -217,6 +217,20 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types); } } else if (arch == LLM_ARCH_GLM5NEXT) { + // nope-only MLA: rope dimension count must be written as 0, not omitted, or + // the generic loader defaults it non-zero and load_arch_hparams rejects it + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(512)); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512)); + ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(0)); + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192)); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, uint32_t(4)); + // build_hc_pre hard-codes a 4-wide residual + 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, 1e-6f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } else if (arch == LLM_ARCH_MINIMAX_M3) { // partial rotary: n_rot must not exceed the indexer key length (64) ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); @@ -448,8 +462,8 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_PADDLEOCR: case LLM_ARCH_MIMO2: case LLM_ARCH_KIMI_LINEAR: - case LLM_ARCH_KIMI_K3: case LLM_ARCH_GLM5NEXT: + case LLM_ARCH_KIMI_K3: case LLM_ARCH_STEP35: case LLM_ARCH_MISTRAL4: case LLM_ARCH_MELLUM: @@ -510,7 +524,6 @@ static bool arch_supported(const llm_arch arch) { if (arch == LLM_ARCH_DEEPSEEK2OCR) { return false; } - if (arch == LLM_ARCH_GLM5NEXT) { // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE) { diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index 732e27379d2..4e0e707b2b8 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -239,3 +239,22 @@ struct mtmd_image_preprocessor_muse_glimmer : mtmd_image_preprocessor { mtmd_image_preprocessor_muse_glimmer(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; }; + +// GLM-5-Next / GLM-5.3-Flash dynamic resize +// ref: Glm5NextImageProcessor.{smart_resize,resize} of the 2026-08-26 adaptation +// +// unlike the Qwen-style smart_resize in mtmd_image_preprocessor_dyn_size, an over-budget image is fitted +// by binary search on the content height, then pasted top-left rather than centred and stretched to fill +struct mtmd_image_preprocessor_glm5next : mtmd_image_preprocessor { + mtmd_image_preprocessor_glm5next(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} + mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; + + struct geometry { + clip_image_size canvas; // aligned, padded output size + clip_image_size content; // resized image, placed at the top-left of the canvas + }; + + // pure arithmetic, exposed so tests can reach it without a clip_ctx + static clip_image_size smart_resize(const clip_hparams & hparams, const clip_image_size & size); + static geometry get_geometry(const clip_hparams & hparams, const clip_image_size & size); +}; From eab9ee93252fca9beb0893084945bc5fca03f82c Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 27 Aug 2026 01:44:15 +0000 Subject: [PATCH 21/36] glm5next: do not inherit deepseek4's n_embd_out deepseek4 sets n_embd_out_impl to hc_mult*n_embd to size its MTP h input. glm5next inherited that, but our t_embd is build_norm(build_hc_mean(...)), which is [n_embd, n_tokens]. n_embd_out() therefore reported 4*n_embd while the tensor held n_embd, and llama-context read n_outputs*n_embd_out floats out of it, four times what is there. The assert at that site sizes the destination buffer, so nothing catches the short source. Only --embeddings and llama_get_embeddings* reach the path, which is why plain generation never showed it. Note for when the NextN graph starts consuming h: give MTP its own width rather than widening n_embd_out again. --- src/models/glm5next.cpp | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index e07b84e9472..92cef19a48e 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -91,8 +91,19 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(hparams.dsv4_hc_mult > 0); // trunk residual is hc_mult streams wide (deepseek4); lm_head still sees - // n_embd, the streams are averaged first - hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd; + // n_embd, the streams are averaged first -- so n_embd_out stays n_embd. + // + // deepseek4 sets this to hc_mult*n_embd, and inheriting that here was a bug: + // our t_embd is build_norm(build_hc_mean(...)), which is [n_embd, n_tokens], + // but n_embd_out() would report 4*n_embd, so llama-context.cpp reads + // n_outputs*n_embd_out floats out of a tensor holding a quarter of that. + // The assert there sizes the DESTINATION, so nothing catches the short SOURCE. + // Only --embeddings / llama_get_embeddings* reach it, which is why plain + // generation never showed it. + // + // deepseek4 needs the wide value to size its MTP `h` input; when our NextN + // graph starts consuming it, give MTP its own width rather than widening this. + hparams.n_embd_out_impl = 0; // MoE ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); From 869e878798cd32fc90b04f0a7b4db82989c84ed7 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 27 Aug 2026 01:44:15 +0000 Subject: [PATCH 22/36] glm5next: keep precision-sensitive tensors at source precision The mHC residual mixers, the lightning indexer (selection gate, learned k-pool position table, and the three indexer projections) and the KDA recurrence gates are about 1 GiB in total on GLM-5.3-Flash, so the size cost is noise against a 100-240 GB quant. Quantizing them perturbs which pools the indexer selects and how much state each KDA step retains, and those errors compound along a sequence rather than averaging out. Both spellings are required. The compressor tensors arrived with the DeepSeek-V4 merge and use an underscore (indexer_compressor_ape / _gate), while the projections use a dot (indexer.proj / .attn_k / .attn_q_b), so a single "indexer." prefix test silently misses the compressor pair. attn_q_a, attn_kv_a_mqa, attn_k_b and attn_v_b are deliberately not listed. They are precision sensitive too, but the release recipe pins them to q8_0 via --tensor-type, and that is the configuration the shipped quants were measured in. Verified with llama-quantize --dry-run q4_k_m on the BF16: all 12 pinned families report 0 quantized (45 mHC, 12 indexer, 34 KDA each), while ffn_gate_exps 43/43, attn_q_a 12/12 and attn_output 46/46 still quantize. --- src/llama-quant.cpp | 41 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 41 insertions(+) diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index f93f95fddab..97b8f4d22ed 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -350,6 +350,47 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param // do not quantize relative position bias (T5) quantize &= name.find("attn_rel_b.weight") == std::string::npos; + // glm5next: small, precision-sensitive tensors that must stay at source + // precision -- the mHC residual mixers, the lightning indexer (selection gate, + // learned k-pool position table, and the three indexer projections) and the KDA + // recurrence gates. ~1 GiB total on GLM-5.3-Flash, so the size cost is noise + // against a 100-240 GB quant, while quantizing them perturbs *which* pools the + // indexer selects and *how much* state each KDA step retains -- errors that + // compound over a sequence instead of averaging out. + // + // Note both spellings are required. The compressor tensors came in with the + // DeepSeek-V4 merge and are named with an UNDERSCORE (indexer_compressor_ape / + // _gate), while the projections use a DOT (indexer.proj / .attn_k / .attn_q_b). + // A single "indexer." prefix test silently misses the compressor pair. + // + // Deliberately NOT listed: attn_q_a / attn_kv_a_mqa / attn_k_b / attn_v_b. They + // are precision-sensitive too, but our release recipe pins them to q8_0 via + // --tensor-type, and that is the configuration the shipped quants were measured + // in (KLD 0.027 at Q5_K_XL). Forcing them full precision here would change quant + // sizes and invalidate those measurements. + // + // indexer.k_norm.weight needs no entry -- the generic "_norm.weight" rule above + // already excludes it. + if (arch == LLM_ARCH_GLM5NEXT) { + static const char * const glm5next_full_precision[] = { + "hc_attn_fn", + "hc_ffn_fn", + "indexer_compressor_ape", + "indexer_compressor_gate", + "indexer.proj", + "indexer.attn_k", + "indexer.attn_q_b", + "ssm_f_a", + "ssm_f_b", + "ssm_g_a", + "ssm_g_b", + "ssm_beta", + }; + for (const char * pin : glm5next_full_precision) { + quantize &= name.find(pin) == std::string::npos; + } + } + // do not quantize specific multimodal tensors quantize &= name.find(".position_embd") == std::string::npos; quantize &= name.find("sam.pos_embd") == std::string::npos; From 282ef610a4f1a948448305e7be1512f646c24d27 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 27 Aug 2026 01:44:15 +0000 Subject: [PATCH 23/36] glm5next: restore the vision preprocessor tests The 101-line glm5next block was dropped from test-mtmd-impl.cpp when the vision work was rebased, even though the commit message still claimed the resize arithmetic was covered there. It holds the 36-case table over the 16- and 8000-token budget boundaries, including six cases annotated as ones where a naive smart_resize disagrees, so it is the guard against sliding back to stretch-to-fill instead of ceil-align plus zero pad. Restored from 29c096371. test-mtmd-impl now runs 216 assertions, of which glm5next_resize contributes 185. --- tests/test-mtmd-impl.cpp | 101 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 101 insertions(+) diff --git a/tests/test-mtmd-impl.cpp b/tests/test-mtmd-impl.cpp index df18b0a42e2..f6b908344b1 100644 --- a/tests/test-mtmd-impl.cpp +++ b/tests/test-mtmd-impl.cpp @@ -70,6 +70,107 @@ MAKE_TEST(test_image_preprocessor_lfm2) { } } +// GLM-5-Next / GLM-5.3-Flash, hparams as loaded by clip.cpp for PROJECTOR_TYPE_GLM5NEXT +static clip_hparams glm5next_hparams() { + clip_hparams hparams; + hparams.patch_size = 14; + hparams.n_merge = 2; + hparams.set_limit_image_tokens(16, 8000); + return hparams; +} + +// set_limit_image_tokens takes token counts; for a still image the reference's temporal factors cancel, +// leaving factor**2 == 28*28 +MAKE_TEST(test_image_preprocessor_glm5next_budget) { + const clip_hparams hparams = glm5next_hparams(); + + t.assert_equal("image_min_pixels", 16 * 28 * 28, hparams.image_min_pixels); + t.assert_equal("image_max_pixels", 8000 * 28 * 28, hparams.image_max_pixels); +} + +MAKE_TEST(test_image_preprocessor_glm5next_resize) { + const clip_hparams hparams = glm5next_hparams(); + + struct test_case { + clip_image_size input; + clip_image_size canvas; // padded output + clip_image_size content; // resized image inside the canvas + int n_tokens; + }; + + // expected values come from Glm5NextImageProcessor.resize in + // adapt_zips/extract_0826/image_processing_glm5_next.py + const std::vector cases = { + // inside the budget + { { 224, 224 }, { 224, 224 }, { 224, 224 }, 64 }, + { { 448, 448 }, { 448, 448 }, { 448, 448 }, 256 }, + { { 1024, 1024 }, { 1036, 1036 }, { 1024, 1024 }, 1369 }, + // the 16-token floor: 112x112 is exactly 16 tokens + { { 112, 112 }, { 112, 112 }, { 112, 112 }, 16 }, + { { 111, 111 }, { 112, 112 }, { 111, 111 }, 16 }, + { { 113, 113 }, { 140, 140 }, { 113, 113 }, 25 }, + { { 56, 56 }, { 112, 112 }, { 112, 112 }, 16 }, + { { 50, 50 }, { 112, 112 }, { 112, 112 }, 16 }, + { { 28, 28 }, { 112, 112 }, { 112, 112 }, 16 }, + { { 1, 1 }, { 112, 112 }, { 112, 112 }, 16 }, + // the 8000-token ceiling: 2492x2492 still fits, 2520x2520 does not + { { 2492, 2492 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + { { 2493, 2493 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + { { 2504, 2504 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + { { 2520, 2520 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + { { 4000, 4000 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + // wide and tall + { { 800, 600 }, { 812, 616 }, { 800, 600 }, 638 }, + { { 600, 800 }, { 616, 812 }, { 600, 800 }, 638 }, + { { 1920, 480 }, { 1932, 504 }, { 1920, 480 }, 1242 }, + { { 480, 1920 }, { 504, 1932 }, { 480, 1920 }, 1242 }, + { { 1920, 1080 }, { 1932, 1092 }, { 1920, 1080 }, 2691 }, + { { 1080, 1920 }, { 1092, 1932 }, { 1080, 1920 }, 2691 }, + // extreme aspect ratios, single-pixel edges survive the resize + { { 4000, 16 }, { 4004, 28 }, { 4000, 16 }, 143 }, + { { 16, 4000 }, { 28, 4004 }, { 16, 4000 }, 143 }, + { { 5000, 1 }, { 5012, 28 }, { 5012, 1 }, 179 }, + { { 1, 5000 }, { 28, 5012 }, { 1, 5012 }, 179 }, + { { 12000, 100 }, { 12012, 112 }, { 12000, 100 }, 1716 }, + { { 100, 12000 }, { 112, 12012 }, { 100, 12000 }, 1716 }, + // over budget, neither edge a multiple of 28 + { { 4007, 3001 }, { 2884, 2156 }, { 2878, 2156 }, 7931 }, + { { 3001, 4007 }, { 2156, 2884 }, { 2156, 2878 }, 7931 }, + { { 3333, 5000 }, { 2044, 3052 }, { 2034, 3052 }, 7957 }, + { { 2729, 2731 }, { 2492, 2492 }, { 2490, 2492 }, 7921 }, + // the 0826 binary search and a Qwen-style smart_resize disagree on all of these, so a regression + // to the old dynamic-size preprocessor cannot pass + { { 4618, 2282 }, { 3528, 1764 }, { 3528, 1743 }, 7938 }, // smart_resize: 3556x1736 + { { 2794, 6096 }, { 1708, 3668 }, { 1681, 3668 }, 7991 }, // smart_resize: 1680x3696 + { { 8858, 1315 }, { 6412, 952 }, { 6412, 951 }, 7786 }, // smart_resize: 6496x952 + { { 1350, 5856 }, { 1204, 5208 }, { 1200, 5208 }, 7998 }, // smart_resize: 1176x5208 + { { 2134, 1472 }, { 2156, 1484 }, { 2134, 1472 }, 4081 }, // smart_resize: 2128x1484 + { { 1021, 4268 }, { 1036, 4284 }, { 1021, 4268 }, 5661 }, // smart_resize: 1008x4256 + }; + + auto fmt = [](const clip_image_size & s) { + return std::to_string(s.width) + "x" + std::to_string(s.height); + }; + + for (const auto & tc : cases) { + const auto geo = mtmd_image_preprocessor_glm5next::get_geometry(hparams, tc.input); + const auto name = " for " + fmt(tc.input); + + t.assert_equal("canvas" + name, fmt(tc.canvas), fmt(geo.canvas)); + t.assert_equal("content" + name, fmt(tc.content), fmt(geo.content)); + + // mirrors clip_n_output_tokens for PROJECTOR_TYPE_GLM5NEXT + const int n_tokens = (geo.canvas.width / (hparams.patch_size * hparams.n_merge)) + * (geo.canvas.height / (hparams.patch_size * hparams.n_merge)); + t.assert_equal("n_tokens" + name, tc.n_tokens, n_tokens); + + t.assert_true("content fits the canvas" + name, + geo.content.width <= geo.canvas.width && geo.content.height <= geo.canvas.height); + t.assert_true("canvas is within the token budget" + name, + geo.canvas.width * geo.canvas.height <= hparams.image_max_pixels); + } +} + // // mtmd temporal merge // From cadbe97b7ed5601fcfecb02c4d46b43ca83c93b0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 27 Aug 2026 02:33:57 +0000 Subject: [PATCH 24/36] glm5next: fix E301 lint in gguf_writer add_vision_swiglu_limit was inserted directly above the next method with no blank line between them, which flake8 flags as E301. Caught by ggml-org CI. --- gguf-py/gguf/gguf_writer.py | 1 + 1 file changed, 1 insertion(+) diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 420b40b639e..b60806071fb 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1332,6 +1332,7 @@ def add_vision_spatial_merge_size(self, value: int) -> None: def add_vision_swiglu_limit(self, value: float) -> None: self.add_float32(Keys.ClipVision.SWIGLU_LIMIT, value) + def add_vision_expert_count_per_layer(self, value: Sequence[int]) -> None: self.add_array(Keys.ClipVision.EXPERT_COUNT_PER_LAYER, value) From 204fa7003c3371677f30abff5f370e8f91239b00 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 27 Aug 2026 03:19:04 +0000 Subject: [PATCH 25/36] glm5next: pool the indexer per sequence so a unified KV cache works llama-embedding turns -np 1 into kv_unified with n_seq_max 256, so every --embeddings run hit the refusal in create_memory and llama-embedding then dereferenced the null context. Two fixes. The pool map is now per SEQUENCE rather than per stream. A non-unified cache already gives one sequence per stream, so nothing changes there. A unified cache puts every sequence of the ubatch in stream 0, and the stream's pool table is cut into one contiguous run per sequence, each rebased on its own lowest resident pool. pool_bias is -INFINITY outside the query's own run, so a query never spends budget on a foreign pool, and cand_mask already kept foreign cells out of the attention mask. Packed runs, not one full-width table per sequence: the indexer scores every pool slot against every query, so a full-width table per sequence multiplies the score tensor by the sequence count and graph_reserve asked for 286 GB at n_seq_max 256. The table is n_kv/kpool shared plus 2 slots per sequence for rebasing, which is exact while the sequences' cells are disjoint. A prefix shared through seq_cp can oversubscribe it, and then a sequence keeps its newest pools -- the same cut a large hole in the cache already forces. The top-k stays over POOLS and pool_cells still holds whole pools, so pool integrity is untouched. examples/embedding also checked only the model for null, not the context. --- examples/embedding/embedding.cpp | 6 + src/llama-graph.cpp | 12 +- src/llama-kv-cache-kpool.cpp | 348 ++++++++++++++++++++----------- src/llama-kv-cache-kpool.h | 30 ++- src/llama-model.cpp | 10 +- 5 files changed, 268 insertions(+), 138 deletions(-) diff --git a/examples/embedding/embedding.cpp b/examples/embedding/embedding.cpp index f6a20ef9d07..a30d271253f 100644 --- a/examples/embedding/embedding.cpp +++ b/examples/embedding/embedding.cpp @@ -144,6 +144,12 @@ int main(int argc, char ** argv) { return 1; } + // a context can fail on its own, and llama_n_ctx below dereferences it + if (ctx == NULL) { + LOG_ERR("%s: unable to create context\n", __func__); + return 1; + } + const llama_vocab * vocab = llama_model_get_vocab(model); const int n_ctx_train = llama_model_n_ctx_train(model); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index b7dbe7ff710..934d76aaf53 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3581,7 +3581,17 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( const int64_t n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq; const int64_t n_tps = ubatch.n_tokens/n_stream; - const int64_t n_pools = llama_kpool_n_pools(n_kv, kpool); + // one pool map per SEQUENCE. a non-unified cache has one sequence per stream, a + // unified cache puts every sequence of the ubatch in stream 0 and cuts the + // stream's pool table into one run per sequence, so the table needs the rebasing + // slack once per sequence. sized on the ubatch and not on n_seq_max, which + // llama-embedding sets to 256; n_seqs_unq is already part of + // llm_graph_params::allow_reuse, so the shape holds while a graph is reused + const int64_t n_ps = (int64_t) ubatch.n_seqs_unq/n_stream; + + GGML_ASSERT(n_ps >= 1 && (int64_t) ubatch.n_seqs_unq == n_ps*n_stream); + + const int64_t n_pools = llama_kpool_n_pools(n_kv, kpool, n_ps); GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[3] == n_stream); diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index 68199bd2008..54efd67e5e9 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -8,10 +8,11 @@ #include #include -uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool) { +uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool, uint32_t n_seqs) { GGML_ASSERT(kpool > 0); + GGML_ASSERT(n_seqs > 0); - return n_kv/kpool + 2; + return n_kv/kpool + 2*n_seqs; } uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t kpool) { @@ -60,10 +61,21 @@ void llama_kv_cache_set_input_kpool( const int64_t n_kv = sel_mask->ne[0]; const int64_t n_ns = sel_mask->ne[3]; // streams in this ubatch const int64_t r = kpool; - const int64_t n_pools = pool_cells->ne[0]/r; const int64_t n_tokens = ubatch->n_tokens; + // [TAG_KPOOL_SEQ_PARTITION] + // positions are unambiguous only within one sequence, so one pool map per SEQUENCE, + // not per stream. a non-unified cache gives each stream its own cells array and one + // sequence (n_ps == 1, the layout this file had before), a unified cache gives one + // stream carrying every sequence in the ubatch + GGML_ASSERT(n_ns == 1 || (int64_t) ubatch->n_seqs_unq == n_ns); + + const int64_t n_ps = (int64_t) ubatch->n_seqs_unq/n_ns; // sequences per stream + const int64_t n_pools = pool_cells->ne[0]/r; // pool slots per stream + + GGML_ASSERT(n_ps > 0 && (int64_t) ubatch->n_seqs_unq == n_ns*n_ps); GGML_ASSERT(pool_cells->ne[0] % r == 0); + GGML_ASSERT(n_pools >= 2*n_ps); GGML_ASSERT(pool_cells->ne[1] == n_ns); GGML_ASSERT(sel_mask->ne[2] == 1); GGML_ASSERT(ggml_are_same_shape(cand_mask, sel_mask)); @@ -81,6 +93,10 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(cell_pool->type == GGML_TYPE_I32); GGML_ASSERT(ggml_is_contiguous(cell_pool)); GGML_ASSERT(cell_pool->ne[0] == n_kv && cell_pool->ne[1] == n_ns); + + // one row per stream, so a cell that two sequences of one stream share has + // nowhere to put its second pool. the graph never asks for this view + GGML_ASSERT(n_ps == 1 && "the per-cell pool view needs one sequence per stream"); } if (bias) { @@ -103,169 +119,253 @@ void llama_kv_cache_set_input_kpool( std::vector filled(n_pools); std::vector pos_at; - // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM] - // positions are unambiguous only within one sequence, and a unified cache shares one - // cells array, so two sequences at the same position would collide in pool_cells and - // silently pool each other's keys. hence one sequence per stream: a non-unified cache - // (n_stream == n_seq_max, the default) or a single sequence in flight. qwen4exp's - // set_input_qsa has the same requirement and no check. - GGML_ASSERT((int64_t) ubatch->n_seqs_unq == n_ns && - "the pooled indexer needs one sequence per stream; use a non-unified KV cache"); + // one contiguous run of pool slots per sequence of the stream + std::vector run_off(n_ps); + std::vector run_len(n_ps); - for (int64_t s = 0; s < n_ns; ++s) { - // which cells array this stream's sequence uses; same convention as - // llama_kv_cache::set_input_kq_mask - const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; - const auto & cells = kv->get_cells(seq_of_stream); + // which cells array a sequence of this stream uses; same convention as + // llama_kv_cache::set_input_kq_mask. with one sequence per stream the ubatch's + // unique list and the stream's own sequence are the same thing + auto seq_of = [&](int64_t s, int64_t ps) { + return n_ps == 1 ? ubatch->seq_id[s*n_tps][0] : ubatch->seq_id_unq[ps]; + }; - int32_t * cur_cell_pool = dst_cell_pool ? dst_cell_pool + s*n_kv : nullptr; + for (int64_t s = 0; s < n_ns; ++s) { int32_t * cur_pool_cells = dst_pool_cells + s*(r*n_pools); + float * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv); + float * cur_cand_mask = dst_cand_mask + s*(n_padq*n_kv); + float * cur_pool_bias = dst_pool_bias + s*(n_tps*n_pools); - std::fill(pool_of.begin(), pool_of.end(), -1); - std::fill(filled.begin(), filled.end(), 0); + // slots of a pool that is not resident, and slots outside the query's own + // sequence run, are cleared once per stream here. the per-sequence pass below + // only writes what it owns std::fill(cur_pool_cells, cur_pool_cells + r*n_pools, 0); + std::fill(cur_pool_bias, cur_pool_bias + n_tps*n_pools, -INFINITY); - // hoist occupancy and sequence membership out of the O(n_kv * n_tokens) loop - // below; neither depends on the query. -1 means the cell holds nothing this - // stream may pool or attend to. under a unified cache `cells` is shared with - // sequences outside this ubatch, and seq_has is what keeps their keys out - pos_at.resize(n_kv); - for (int64_t j = 0; j < n_kv; ++j) { - pos_at[j] = cells.is_empty(j) || !cells.seq_has(j, seq_of_stream) ? -1 : cells.pos_get(j); - } + // the token loop writes rows < n_tps in full; only the padding rows need clearing + std::fill(cur_sel_mask + n_tps*n_kv, cur_sel_mask + n_padq*n_kv, -INFINITY); + std::fill(cur_cand_mask + n_tps*n_kv, cur_cand_mask + n_padq*n_kv, -INFINITY); - // a pool ordinal is the absolute p/kpool, which can far exceed n_kv/kpool, so - // rebase on this stream's lowest resident pool. grouping is untouched, every - // member shifts together. + // [TAG_KPOOL_PACK] + // cut the stream's pool table into one run per sequence, sized on the pool range + // that sequence actually holds. the table is n_kv/kpool shared plus 2 per sequence + // for rebasing, which covers it whenever the sequences' cells are disjoint - every + // case but a prefix shared through llama_memory_seq_cp. that one can ask for more + // slots than exist, and then a sequence keeps its newest pools, the same cut a + // large hole already forces. // - // anchoring at p/kpool follows vLLM and SGLang, not HF (which pools from the - // first *resident* key, valid_keys.argmax(-1), differing under left padding). - // it is the only anchor that keeps a pool's identity stable between the prefill - // that built it and the decodes that read it. - // - // the window is n_pools wide; positions are contiguous in any real batch so the - // resident range fits. seq_rm can leave a hole large enough that it does not, - // and then the newest pools are the ones worth keeping. - int64_t b_base = 0; + // NOT one full-width table per sequence: the indexer scores every slot against + // every query, so a full-width table would multiply the score tensor by the + // sequence count, and llama-embedding asks for n_seq_max 256 { - int64_t b_min = 0; - int64_t b_max = 0; - bool found = false; + int64_t n_want = 0; - for (int64_t j = 0; j < n_kv; ++j) { - if (pos_at[j] < 0) { - continue; + for (int64_t ps = 0; ps < n_ps; ++ps) { + const llama_seq_id seq = seq_of(s, ps); + const auto & cells = kv->get_cells(seq); + + int64_t b_min = 0; + int64_t b_max = 0; + bool found = false; + + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j) || !cells.seq_has(j, seq)) { + continue; + } + const int64_t b = cells.pos_get(j)/r; + b_min = found ? std::min(b_min, b) : b; + b_max = found ? std::max(b_max, b) : b; + found = true; } - const int64_t b = pos_at[j]/r; - b_min = found ? std::min(b_min, b) : b; - b_max = found ? std::max(b_max, b) : b; - found = true; + + run_len[ps] = found ? b_max - b_min + 1 : 0; + n_want += run_len[ps]; } - b_base = std::max(b_min, b_max - (n_pools - 1)); - } + if (n_want > n_pools) { + int64_t rem = n_pools; - for (int64_t j = 0; j < n_kv; ++j) { - if (pos_at[j] < 0) { - continue; + for (int64_t ps = 0; ps < n_ps; ++ps) { + run_len[ps] = std::min(run_len[ps], rem/(n_ps - ps)); + rem -= run_len[ps]; + } } - const llama_pos p = pos_at[j]; - const int64_t bo = p/r - b_base; - - if (bo < 0 || bo >= n_pools) { - continue; + int64_t off = 0; + for (int64_t ps = 0; ps < n_ps; ++ps) { + run_off[ps] = off; + off += run_len[ps]; } - pool_of[j] = (int32_t) bo; - cur_pool_cells[bo*r + (p%r)] = (int32_t) j; - filled[bo]++; + GGML_ASSERT(off <= n_pools); } - // an incompletely resident pool cannot be pooled: the compressor consumes all r - // member keys, and the reference demands pool_valid = grouped_valid_keys.all(-1). - // such cells are the sequence tail, which sel_mask forces in below regardless of - // score, so they point at pool slot 0 purely to keep the gather in range - for (int64_t j = 0; j < n_kv; ++j) { - // != rather than <: two cells claiming one position overwrite each other in - // pool_cells, so an over-filled pool is not usable either - if (pool_of[j] >= 0 && filled[pool_of[j]] != (int32_t) r) { - pool_of[j] = -1; - } - if (cur_cell_pool) { - cur_cell_pool[j] = pool_of[j] < 0 ? 0 : pool_of[j]; + int64_t n_done = 0; + + for (int64_t ps = 0; ps < n_ps; ++ps) { + const llama_seq_id seq_of_pool = seq_of(s, ps); + const auto & cells = kv->get_cells(seq_of_pool); + + const int64_t n_run = run_len[ps]; + + int32_t * cur_cell_pool = dst_cell_pool ? dst_cell_pool + s*n_kv : nullptr; + int32_t * part_pool_cells = cur_pool_cells + run_off[ps]*r; + + std::fill(pool_of.begin(), pool_of.end(), -1); + std::fill(filled.begin(), filled.end(), 0); + + // hoist occupancy and sequence membership out of the O(n_kv * n_tokens) loop + // below; neither depends on the query. -1 means the cell holds nothing this + // sequence may pool or attend to. under a unified cache `cells` is shared with + // every other sequence, and seq_has is what keeps their keys out + pos_at.resize(n_kv); + for (int64_t j = 0; j < n_kv; ++j) { + pos_at[j] = cells.is_empty(j) || !cells.seq_has(j, seq_of_pool) ? -1 : cells.pos_get(j); } - } - float * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv); - float * cur_cand_mask = dst_cand_mask + s*(n_padq*n_kv); + // a pool ordinal is the absolute p/kpool, which can far exceed n_kv/kpool, so + // rebase on this sequence's lowest resident pool. grouping is untouched, every + // member shifts together. + // + // anchoring at p/kpool follows vLLM and SGLang, not HF (which pools from the + // first *resident* key, valid_keys.argmax(-1), differing under left padding). + // it is the only anchor that keeps a pool's identity stable between the prefill + // that built it and the decodes that read it. + // + // the window is this sequence's run; positions are contiguous in any real + // batch so the resident range fits. seq_rm can leave a hole large enough that + // it does not, and then the newest pools are the ones worth keeping. + int64_t b_base = 0; + { + int64_t b_min = 0; + int64_t b_max = 0; + bool found = false; - // the loop below writes rows < n_tps in full; only the padding rows need clearing - std::fill(cur_sel_mask + n_tps*n_kv, cur_sel_mask + n_padq*n_kv, -INFINITY); - std::fill(cur_cand_mask + n_tps*n_kv, cur_cand_mask + n_padq*n_kv, -INFINITY); + for (int64_t j = 0; j < n_kv; ++j) { + if (pos_at[j] < 0) { + continue; + } + const int64_t b = pos_at[j]/r; + b_min = found ? std::min(b_min, b) : b; + b_max = found ? std::max(b_max, b) : b; + found = true; + } - for (int64_t ii = 0; ii < n_tps; ++ii) { - const int64_t i = s*n_tps + ii; - const llama_pos q = ubatch->pos[i]; + b_base = std::max(b_min, b_max - (n_run - 1)); + } - // q >= 0 is what makes the unsigned range test below a range test - GGML_ASSERT(q >= 0 && ubatch->seq_id[i][0] == seq_of_stream); + for (int64_t j = 0; j < n_kv; ++j) { + if (pos_at[j] < 0) { + continue; + } - // the query's own incomplete pool ((q + 1) % r cells, its own token - // included) is always attended to (index_kpool_always_select_tail), which is - // what makes the selection land on pool boundaries - const llama_pos tail_start = (q + 1)/r*r; + const llama_pos p = pos_at[j]; + const int64_t bo = p/r - b_base; - // the reference tests visibility at a pool's LAST member, so a pool - // straddling the query is dropped whole. pools are position-aligned here, so - // that test collapses to b*r < tail_start - const int64_t bo_vis = std::max(0, tail_start/r - b_base); + if (bo < 0 || bo >= n_run) { + continue; + } - float * cur_bias = dst_bias ? dst_bias + i*n_kv : nullptr; - float * cur_sel = cur_sel_mask + ii*n_kv; - float * cur_cand = cur_cand_mask + ii*n_kv; + pool_of[j] = (int32_t) bo; + part_pool_cells[bo*r + (p%r)] = (int32_t) j; + filled[bo]++; + } - // the unsigned compares fold "empty or another sequence" (pos_at -1) and "no - // usable pool" (pool_of -1) into the range test, which lets this vectorise + // an incompletely resident pool cannot be pooled: the compressor consumes all r + // member keys, and the reference demands pool_valid = grouped_valid_keys.all(-1). + // such cells are the sequence tail, which sel_mask forces in below regardless of + // score, so they point at pool slot 0 purely to keep the gather in range for (int64_t j = 0; j < n_kv; ++j) { - const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; - const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; - const bool tail = pos_at[j] >= tail_start; - - cur_sel [j] = vis && tail ? 0.0f : -INFINITY; - // max(bias, sel_mask): the reference's candidate set, which the - // top-k budget may overrun but must never escape - cur_cand[j] = vis && (pooled || tail) ? 0.0f : -INFINITY; + // != rather than <: two cells claiming one position overwrite each other in + // pool_cells, so an over-filled pool is not usable either + if (pool_of[j] >= 0 && filled[pool_of[j]] != (int32_t) r) { + pool_of[j] = -1; + } + if (cur_cell_pool) { + cur_cell_pool[j] = pool_of[j] < 0 ? 0 : pool_of[j]; + } } - if (cur_bias) { + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + + // a query is pooled by the sequence it is being written to. with several + // sequences in one stream the other partitions own the rest of the rows + if (ubatch->seq_id[i][0] != seq_of_pool) { + continue; + } + + const llama_pos q = ubatch->pos[i]; + + // q >= 0 is what makes the unsigned range test below a range test + GGML_ASSERT(q >= 0); + + n_done++; + + // the query's own incomplete pool ((q + 1) % r cells, its own token + // included) is always attended to (index_kpool_always_select_tail), which is + // what makes the selection land on pool boundaries + const llama_pos tail_start = (q + 1)/r*r; + + // the reference tests visibility at a pool's LAST member, so a pool + // straddling the query is dropped whole. pools are position-aligned here, so + // that test collapses to b*r < tail_start + const int64_t bo_vis = std::max(0, tail_start/r - b_base); + + float * cur_bias = dst_bias ? dst_bias + i*n_kv : nullptr; + float * cur_sel = cur_sel_mask + ii*n_kv; + float * cur_cand = cur_cand_mask + ii*n_kv; + + // the unsigned compares fold "empty or another sequence" (pos_at -1) and "no + // usable pool" (pool_of -1) into the range test, which lets this vectorise for (int64_t j = 0; j < n_kv; ++j) { const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; + const bool tail = pos_at[j] >= tail_start; - cur_bias[j] = vis && pooled ? 0.0f : -INFINITY; + cur_sel [j] = vis && tail ? 0.0f : -INFINITY; + // max(bias, sel_mask): the reference's candidate set, which the + // top-k budget may overrun but must never escape + cur_cand[j] = vis && (pooled || tail) ? 0.0f : -INFINITY; } - } - // The same predicate, per POOL, which is where the reference applies - // it: pool_valid (completely resident) & pool_visible (its LAST - // member is visible, so a pool the query straddles is dropped whole). - // Pools are position-aligned here, so pool bo's last member is at - // position (b_base + bo)*r + r - 1 and "last member visible" collapses - // to bo < bo_vis. - // - // NOT gathered from `bias` at the last member cell: an incomplete or - // absent pool has no resident last member, pool_cells points that slot - // at cell 0, and the pool would inherit cell 0's validity. - float * cur_pool_bias = dst_pool_bias + (s*n_tps + ii)*n_pools; + if (cur_bias) { + for (int64_t j = 0; j < n_kv; ++j) { + const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; + const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; - for (int64_t p = 0; p < n_pools; ++p) { - const bool valid = filled[p] == (int32_t) r; - const bool visible = p < bo_vis; + cur_bias[j] = vis && pooled ? 0.0f : -INFINITY; + } + } - cur_pool_bias[p] = valid && visible ? 0.0f : -INFINITY; + // The same predicate, per POOL, which is where the reference applies + // it: pool_valid (completely resident) & pool_visible (its LAST + // member is visible, so a pool the query straddles is dropped whole). + // Pools are position-aligned here, so pool bo's last member is at + // position (b_base + bo)*r + r - 1 and "last member visible" collapses + // to bo < bo_vis. + // + // NOT gathered from `bias` at the last member cell: an incomplete or + // absent pool has no resident last member, pool_cells points that slot + // at cell 0, and the pool would inherit cell 0's validity. + // + // the query's own sequence run only. every other slot stays at the + // -INFINITY the per-stream fill above left, which is what keeps a foreign + // pool out of the budget + float * q_pool_bias = cur_pool_bias + ii*n_pools + run_off[ps]; + + for (int64_t p = 0; p < n_run; ++p) { + const bool valid = filled[p] == (int32_t) r; + const bool visible = p < bo_vis; + + q_pool_bias[p] = valid && visible ? 0.0f : -INFINITY; + } } } + + // every row of sel_mask, cand_mask and pool_bias must have been written by + // exactly one partition, or a query is left reading another sequence's pools + GGML_ASSERT(n_done == n_tps && "every query must belong to a sequence of the ubatch"); } } diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index 3a486cb0463..d3cd86c546c 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -24,17 +24,35 @@ class llama_kv_cache_context; // here, host side, from the cache's cells, and passed in as plain input tensors, as // qwen4exp's QSA does for its compression blocks. // +// A position only names a pool inside ONE sequence, so the map is per SEQUENCE. A +// non-unified cache gives every sequence its own stream and its own cells array, and +// the two are the same thing. A unified cache puts every sequence of the ubatch in one +// stream and one cells array, so the stream's pool table is PARTITIONED: each sequence +// gets a contiguous run of slots and rebases inside it. `pool_bias` is -INFINITY outside +// the query's own run, which is what stops a query from selecting a foreign pool. +// +// The runs are packed, not one full-width table per sequence. A full-width table per +// sequence would multiply the pool axis by the sequence count, and the indexer scores +// every pool against every query, so llama-embedding (which asks for n_seq_max 256 with +// a unified cache) reserved 286 GB of compute buffer that way. +// // Nothing here may emit a negative index: ggml_set_rows asserts i1 >= 0 and // ggml_get_rows has no sentinel, so unpopulated entries are clamped into range and // neutralised by the additive masks instead. // -// number of pool slots the graph must allocate for `n_kv` cells. +// number of pool slots the graph must allocate for `n_kv` cells shared by `n_seqs` +// sequences. // -// pool ordinals are position-derived, then rebased on each stream's lowest resident +// pool ordinals are position-derived, then rebased on each sequence's lowest resident // pool so sequences not starting at position 0 still land inside the array. rebasing -// can cost one slot at each end, hence the +2. -uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool); +// can cost one slot at each end, hence 2 per sequence. +// +// n_kv/kpool is the budget the sequences share, which is exact while their cells are +// disjoint: a cell belongs to one pool of one sequence. llama_memory_seq_cp breaks that +// - a shared prefix cell is pooled by every sequence holding it - and then the runs are +// cut to the newest pools rather than the table being grown, see [TAG_KPOOL_PACK]. +uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool, uint32_t n_seqs = 1); // how many POOLS ggml_top_k selects. // @@ -79,6 +97,7 @@ uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t // // pool_cells I32 [kpool*n_pools, n_stream] // pool slot, member -> the cell holding that position, or 0 when not resident. +// One stream's slots are cut into one contiguous run per sequence. // Two consumers: the compressor gathers a pool's member keys and gates with it, // and the top-k expands a selected POOL ordinal back into its kpool member CELLS // with it - the reference's `pool_indices[batch_idx, selected]`. @@ -92,7 +111,8 @@ uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t // the same predicate per POOL, which is where the reference applies it: 0.0f when // pool p is completely resident and its LAST member is visible to query q // (`pool_valid & pool_visible`), else -INFINITY, the query's own trailing pool -// included so no budget is spent on it. +// included so no budget is spent on it. Every slot outside the query's own +// sequence run is -INFINITY, so a query never selects another sequence's pool. // // Derived here rather than in the graph on purpose. Gathering `bias` at // each pool's last member looks equivalent and is not: an incomplete or diff --git a/src/llama-model.cpp b/src/llama-model.cpp index db26cf5dd69..0d1bfee0dca 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2457,14 +2457,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, }; if (arch == LLM_ARCH_GLM5NEXT && hparams.indexer_head_size > 0) { - // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM] - // the indexer pools cells by position and a unified cache - // shares one cells array, so two sequences at the same - // position would pool each other's keys. refuse here rather - // than abort inside a set_input thousands of tokens in - if (cparams.kv_unified && cparams.n_seq_max > 1) { - throw std::runtime_error("glm5next: the pooled indexer needs one sequence per stream, so a unified KV cache is only supported with a single sequence"); - } + // a unified cache is fine here: the pool map is per SEQUENCE, + // not per stream. see [TAG_KPOOL_SEQ_PARTITION] // only the DSA layers carry an indexer key cache filter_idx = [&](uint32_t il) { From ef531f827e435bdd7fe75d03f3c4eb3486265bed Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Thu, 27 Aug 2026 03:28:26 +0000 Subject: [PATCH 26/36] glm5next: test the per-sequence pool runs instead of the old refusal test-glm5next-memory asserted that create_memory REFUSES -kvu with n_seq_max 2, which was the contract before the pool map became per sequence. Assert the new one: the cache is built, and one ubatch holding both sequences is driven through llama_kv_cache_set_input_kpool. Three checks replace the guard. llama_kpool_n_pools is n_kv/kpool plus 2 slots per sequence, so the table is a shared budget and not one full-width table each. Every pool a query may spend budget on holds only that query's own visible cells -- the invariant a shared cells array breaks if the map is keyed per stream. The two sequences get disjoint runs and neither run is empty. Both cell-level checks fail if the runs are made to overlap, so they are not tautologies. cell_pool is not requested here: it has one row per stream and a cell that two sequences share has nowhere to put its second pool. --- tests/test-glm5next-memory.cpp | 180 +++++++++++++++++++++++++++++---- 1 file changed, 162 insertions(+), 18 deletions(-) diff --git a/tests/test-glm5next-memory.cpp b/tests/test-glm5next-memory.cpp index b30ffbfeb74..b64e71cee5c 100644 --- a/tests/test-glm5next-memory.cpp +++ b/tests/test-glm5next-memory.cpp @@ -358,12 +358,14 @@ int main(int argc, char ** argv) { "indexer_top_k (%u) is a whole number of pools of %u", hparams.indexer_top_k, hparams.indexer_kpool); } - // ---- a multi-sequence unified cache is rejected at create time ---------- + // ---- a multi-sequence unified cache pools per SEQUENCE ------------------ // - // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM]. pools group cells by position and a unified - // cache shares one cells array, so two sequences at the same position would pool - // each other's keys. -kvu with --parallel is reachable (llama-perplexity forces it - // for hellaswag / winogrande / multiple-choice), so this must fail at startup. + // [TAG_KPOOL_SEQ_PARTITION]. pools group cells by position, and a position only + // names a pool inside one sequence. a unified cache shares one cells array, so the + // stream's pool table is cut into one run per sequence instead of the cache being + // refused. -kvu with --parallel is reachable (llama-perplexity forces it for + // hellaswag / winogrande / multiple-choice, and llama-embedding turns -np 1 into + // -kvu with n_seq_max 256), so this has to work rather than fail at startup. { llama_cparams cp = {}; cp.n_ctx = 256; @@ -379,25 +381,167 @@ int main(int argc, char ** argv) { mp.type_v = GGML_TYPE_F16; mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; - bool threw = false; + llama_memory_i * raw = nullptr; try { - delete model->create_memory(mp, cp); - } catch (const std::exception &) { - threw = true; + raw = model->create_memory(mp, cp); + } catch (const std::exception & e) { + printf("note create_memory threw: %s\n", e.what()); + raw = nullptr; } - CHECK(threw, "-kvu with n_seq_max 2 is refused when the model has a pooling indexer"); + CHECK(raw != nullptr, "-kvu with n_seq_max 2 is accepted: the pool map is per sequence"); + + const uint32_t kpool = hparams.indexer_kpool; + + // one shared n_kv/kpool budget plus the rebasing slack per sequence, NOT one + // full-width table each: the indexer scores every slot against every query, so a + // full-width table per sequence multiplies the score tensor by the sequence count + CHECK(llama_kpool_n_pools(256, kpool, 1) == 256/kpool + 2 && + llama_kpool_n_pools(256, kpool, 2) == 256/kpool + 4, + "llama_kpool_n_pools is n_kv/kpool + 2 per sequence (%u slots for 2 sequences)", + llama_kpool_n_pools(256, kpool, 2)); + + auto * m2 = dynamic_cast(raw); + CHECK(m2 != nullptr && m2->get_mem_idx() != nullptr, "-kvu builds an indexer cache too"); + + // drive one ubatch holding BOTH sequences through the map. 16 positions each, so + // every sequence owns several complete pools and the runs have to be told apart + if (m2) { + const int64_t n_tok = 16; + + std::vector tok; + std::vector pos; + std::vector sid; + std::vector sptr; + std::vector nsid; + std::vector lg; + + for (llama_seq_id s = 0; s < 2; ++s) { + for (int64_t i = 0; i < n_tok; ++i) { + tok.push_back(0); + pos.push_back((llama_pos) i); + sid.push_back(s); + nsid.push_back(1); + lg.push_back(1); + } + } + for (size_t i = 0; i < sid.size(); ++i) { + sptr.push_back(&sid[i]); + } + + llama_batch b = {}; + b.n_tokens = (int32_t) tok.size(); + b.token = tok.data(); + b.pos = pos.data(); + b.seq_id = sptr.data(); + b.n_seq_id = nsid.data(); + b.logits = lg.data(); + + llama_batch_allocr ba(hparams.n_pos_per_embd()); + if (ba.init(b, model->vocab, nullptr, hparams.n_embd_inp(), cp.n_seq_max, true)) { + auto c = m2->init_batch(ba, cp.n_ubatch, false); + auto * mc = dynamic_cast(c.get()); + + CHECK(mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, + "both sequences fit one unified ubatch"); + + if (mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { + mc->apply(); + + const llama_ubatch & ub = mc->get_ubatch(); - // one sequence in flight is fine: the shared cells array holds only its keys, - // and the map filters on seq_has anyway + const int64_t n_kv = mc->get_attn()->get_n_kv(); + const int64_t n_stream = mc->get_attn()->get_n_stream(); + const int64_t n_tps = ub.n_tokens/n_stream; + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, kpool, ub.n_seqs_unq); + + CHECK(n_stream == 1 && ub.n_seqs_unq == 2, + "a unified cache carries both sequences in one stream (n_seqs_unq %u)", + ub.n_seqs_unq); + + ggml_init_params gp = { ggml_tensor_overhead()*16, nullptr, true }; + ggml_context * ctx = ggml_init(gp); + + kpool_tensors kt = alloc_kpool_tensors(ctx, n_kv, n_tps, n_tps, n_stream, kpool, n_pools); + + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + if (buf) { + // no cell_pool: the per-cell view has one row per stream, and a + // cell two sequences share has nowhere to put its second pool + llama_kv_cache_set_input_kpool(m2->get_mem_attn(), + /* cell_pool */ nullptr, kt.pool_cells, kt.bias, kt.pool_bias, + kt.sel_mask, kt.cand_mask, &ub, kpool); + + const int32_t * pc = (const int32_t *) kt.pool_cells->data; + const float * pb = (const float *) kt.pool_bias->data; + + const auto & cells = m2->get_mem_attn()->get_cells(0); + + // the invariant the partitioning exists for: a pool the query may + // spend budget on holds only that query's own visible cells. this + // is what a shared cells array breaks if the map is per stream + bool own = true; + int64_t n_finite = 0; + std::map seq_mask_of_slot; + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const llama_seq_id sq = ub.seq_id[ii][0]; + const llama_pos q = ub.pos[ii]; + + for (int64_t p = 0; p < n_pools; ++p) { + if (pb[ii*n_pools + p] != 0.0f) { + continue; + } + + n_finite++; + seq_mask_of_slot[p] |= 1 << sq; + + for (int64_t m = 0; m < (int64_t) kpool; ++m) { + const int32_t c = pc[p*kpool + m]; + own &= c >= 0 && c < n_kv && !cells.is_empty(c) && + cells.seq_has(c, sq) && cells.pos_get(c) <= q; + } + } + } + + CHECK(own && n_finite > 0, + "every pool a query may select holds only that query's own visible cells (%d selectable (query, pool) pairs)", + (int) n_finite); + + int mask_all = 0; + bool disjoint = true; + for (const auto & kv : seq_mask_of_slot) { + mask_all |= kv.second; + disjoint &= (kv.second & (kv.second - 1)) == 0; + } + + CHECK(disjoint, "the two sequences get disjoint pool runs (%d slots in use of %d)", + (int) seq_mask_of_slot.size(), (int) n_pools); + CHECK(mask_all == 0x3, "both sequences own selectable pools, so neither run is empty"); + + ggml_backend_buffer_free(buf); + } + + ggml_backend_free(backend); + ggml_free(ctx); + } + } + } + + delete raw; + + // one sequence in flight stays the simple case: the shared cells array holds + // only its keys, and the map filters on seq_has anyway cp.n_seq_max = 1; - llama_memory_i * raw = nullptr; + llama_memory_i * raw1 = nullptr; try { - raw = model->create_memory(mp, cp); + raw1 = model->create_memory(mp, cp); } catch (const std::exception &) { - raw = nullptr; + raw1 = nullptr; } - CHECK(raw != nullptr, "-kvu with a single sequence is still allowed"); - delete raw; + CHECK(raw1 != nullptr, "-kvu with a single sequence is still allowed"); + delete raw1; } // ---- the indexer cache keeps its own dtype ------------------------------ @@ -439,7 +583,7 @@ int main(int argc, char ** argv) { cparams.n_ubatch = 32; cparams.n_seq_max = 2; cparams.n_rs_seq = 0; - cparams.kv_unified = false; // [TAG_KPOOL_NEEDS_ONE_SEQ_PER_STREAM] + cparams.kv_unified = false; // one sequence per stream, the n_ps == 1 layout cparams.offload_kqv = false; cparams.flash_attn = false; cparams.causal_attn = true; From e88c92d02d01037753eff358873af5b66a5a99f9 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 27 Aug 2026 03:06:45 +0000 Subject: [PATCH 27/36] glm5next: store the indexer selection masks in f16 sel_mask and cand_mask are KQ-mask shaped and hold only 0.0f and -INFINITY, both exact in f16, so storing them in half the bytes is lossless. At n_ctx = 1 Mi, n_ubatch = 512 that is 2 GiB saved per mask plus 1 GiB on the per-layer ggml_dup. ggml_add gives its result src0's type and f16 + f32 -> f16 is a supported bin_bcast on CUDA and on the CPU, so the f16 selection mask absorbs the f32 KQ mask that flash-attention-off builds, and ggml_soft_max_ext takes an f16 mask as readily as an f32 one. Under flash attention the KQ mask is already f16 and the per-layer ggml_cast disappears. llama_kv_cache_set_input_kpool now writes either width and asserts the two masks share a type instead of asserting f32. --- src/llama-graph.cpp | 39 +++++++++----- src/llama-graph.h | 11 ++-- src/llama-kv-cache-kpool.cpp | 99 +++++++++++++++++++++++++--------- src/llama-kv-cache-kpool.h | 8 +-- tests/test-glm5next-memory.cpp | 56 +++++++++++++++++++ 5 files changed, 167 insertions(+), 46 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 934d76aaf53..06fdcfd69e7 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3608,13 +3608,22 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( ggml_set_input(inp->pool_bias); ggml_set_name(inp->pool_bias, "kpool_pool_bias"); - // f32 even under flash attention, where the KQ mask itself is f16: these - // are a scatter base and a mask addend, not the tensor the FA kernel sees - inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_kv, n_tps, 1, n_stream); + // f16, not f32. The only two values either mask holds are 0.0f and + // -INFINITY, both exact in f16, so the narrower type is lossless and halves + // two [n_kv, n_tps, n_stream] inputs that live for the whole ubatch: 2 GiB + // each at n_ctx = 1 Mi, n_ubatch = 512. + // + // Every consumer takes f16. Under flash attention this is the KQ mask's own + // type, so build_attn_sparse adds the two with no conversion at all. With + // flash attention off - which GLM-5.3-Flash requires, so it is the path that + // matters here - the KQ mask is f32, and ggml_add gives its result src0's + // type: f16 + f32 -> f16 is a supported bin_bcast on CUDA and on the CPU, + // and ggml_soft_max_ext takes an f16 mask as readily as an f32 one + inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, n_kv, n_tps, 1, n_stream); ggml_set_input(inp->sel_mask); ggml_set_name(inp->sel_mask, "kpool_sel_mask"); - inp->cand_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_kv, n_tps, 1, n_stream); + inp->cand_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, n_kv, n_tps, 1, n_stream); ggml_set_input(inp->cand_mask); ggml_set_name(inp->cand_mask, "kpool_cand_mask"); } @@ -3653,7 +3662,8 @@ ggml_tensor * llm_graph_context::build_attn_sparse( const auto & kq_mask = inp->get_kq_mask(); - GGML_ASSERT(sel_mask->type == GGML_TYPE_F32); + GGML_ASSERT(sel_mask->type == GGML_TYPE_F16 || sel_mask->type == GGML_TYPE_F32); + GGML_ASSERT(sel_mask->type == cand_mask->type); GGML_ASSERT(ggml_are_same_shape(sel_mask, cand_mask)); GGML_ASSERT(sel_mask->ne[0] == kq_mask->ne[0] && sel_mask->ne[1] == kq_mask->ne[1] && sel_mask->ne[3] == kq_mask->ne[3]); @@ -3687,7 +3697,10 @@ ggml_tensor * llm_graph_context::build_attn_sparse( // top-k (through an over-budget pool whose unfilled slots point at it), and // scattering that cell's -inf score bias would leave the query attending to // nothing at all. Rejecting an over-budget selection is done additively, - // below, by cand_mask + // below, by cand_mask. + // + // f32 whatever mask_all is: ggml_set_rows converts its values into the + // destination, and the CUDA backend only advertises SET_ROWS for f32 values ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); zeros = ggml_fill(ctx0, zeros, 0.0f); @@ -3703,14 +3716,12 @@ ggml_tensor * llm_graph_context::build_attn_sparse( // candidates and the tail mask_top_k = ggml_add(ctx0, mask_top_k, cand_mask); - // sel_mask and cand_mask are f32 by contract - llama_kv_cache_set_input_kpool - // writes floats through raw strides and refuses anything else - but the KQ - // mask is f16 whenever flash attention is on, which is the DEFAULT, and - // ggml_flash_attn_ext asserts its mask is f16. Cast rather than refuse: the - // only two values here are 0.0f and -INFINITY and both are exact in f16. - // Refusing instead aborts test-llama-archs, which runs with FA on - if (mask_top_k->type != kq_mask->type) { - mask_top_k = ggml_cast(ctx0, mask_top_k, kq_mask->type); + // ggml_add gives its result src0's type, so an f16 selection mask absorbs an + // f32 KQ mask and no cast is needed. The one direction that still needs one is + // an f32 mask meeting an f16 KQ mask: the add would yield f32 and + // ggml_flash_attn_ext asserts its mask is f16 + if (mask_top_k->type == GGML_TYPE_F32 && kq_mask->type == GGML_TYPE_F16) { + mask_top_k = ggml_cast(ctx0, mask_top_k, GGML_TYPE_F16); } // and finally re-apply causality, occupancy and padding. load bearing: it is diff --git a/src/llama-graph.h b/src/llama-graph.h index a975eb1bb73..c994a1774d5 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1372,7 +1372,12 @@ struct llm_graph_context { // always-selected trailing pool. `cand_mask` is the reference's candidate // set and is what makes an over-budget selection harmless: ggml_top_k // returns a full budget of pool ordinals even when fewer pools carry a - // finite score, which during prefill is the normal state + // finite score, which during prefill is the normal state. + // + // Both masks may be f16 or f32, and build_inp_kpool allocates f16: the only + // values either holds are 0.0f and -INFINITY, both exact in f16. The combined + // mask inherits their type, which the KQ mask then adds into whatever its own + // type is ggml_tensor * build_attn_sparse( llm_graph_input_attn_k * inp, ggml_tensor * wo, @@ -1385,8 +1390,8 @@ struct llm_graph_context { ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] ggml_tensor * top_k, // I32 [n_select, n_tokens/n_stream, n_stream] - ggml_tensor * sel_mask, // F32 [n_kv, n_batch, 1, n_stream] - ggml_tensor * cand_mask, // F32 [n_kv, n_batch, 1, n_stream] + ggml_tensor * sel_mask, // F16/F32 [n_kv, n_batch, 1, n_stream] + ggml_tensor * cand_mask, // F16/F32 [n_kv, n_batch, 1, n_stream] float kq_scale, int il) const; diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index 54efd67e5e9..9f214dd16cd 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -24,6 +24,50 @@ uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t return std::min(n_pools, indexer_top_k/kpool); } +// sel_mask and cand_mask hold only 0.0f and -INFINITY, so they can be written in the +// KQ mask's f16 exactly as in f32 +template struct kpool_mask_of; + +template <> struct kpool_mask_of { + static float from(float v) { return v; } +}; + +template <> struct kpool_mask_of { + static ggml_fp16_t from(float v) { return ggml_fp32_to_fp16(v); } +}; + +template +static void kpool_mask_fill(T * dst, int64_t n) { + std::fill(dst, dst + n, kpool_mask_of::from(-INFINITY)); +} + +// one query's row of both masks; the two predicates share their operands, and the +// unsigned compares are what let this vectorise +template +static void kpool_mask_row( + T * cur_sel, + T * cur_cand, + const llama_pos * pos_at, + const int32_t * pool_of, + int64_t n_kv, + llama_pos q, + llama_pos tail_start, + int64_t bo_vis) { + const T v_sel = kpool_mask_of::from(0.0f); + const T v_mask = kpool_mask_of::from(-INFINITY); + + for (int64_t j = 0; j < n_kv; ++j) { + const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; + const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; + const bool tail = pos_at[j] >= tail_start; + + cur_sel [j] = vis && tail ? v_sel : v_mask; + // max(bias, sel_mask): the reference's candidate set, which the + // top-k budget may overrun but must never escape + cur_cand[j] = vis && (pooled || tail) ? v_sel : v_mask; + } +} + void llama_kv_cache_set_input_kpool( const llama_kv_cache * kv, ggml_tensor * cell_pool, @@ -44,13 +88,13 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(ggml_backend_buffer_is_host(sel_mask ->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(cand_mask ->buffer)); - // sel_mask and cand_mask are KQ-mask shaped, and KQ masks are f16 under flash - // attention; writing floats into one would overrun the allocation 2x, so check - // rather than trust + // both masks are written through raw strides below, so writing the wrong width + // would overrun the allocation 2x. check rather than trust GGML_ASSERT(pool_cells->type == GGML_TYPE_I32); GGML_ASSERT(pool_bias ->type == GGML_TYPE_F32); - GGML_ASSERT(sel_mask ->type == GGML_TYPE_F32 && "sel_mask must be f32 even when the KQ mask is f16"); - GGML_ASSERT(cand_mask ->type == GGML_TYPE_F32 && "cand_mask must be f32 even when the KQ mask is f16"); + GGML_ASSERT((sel_mask->type == GGML_TYPE_F16 || sel_mask->type == GGML_TYPE_F32) && + "sel_mask must be f16 or f32"); + GGML_ASSERT(cand_mask->type == sel_mask->type && "both masks must have the KQ mask's type"); // everything below is written through raw strides GGML_ASSERT(ggml_is_contiguous(pool_cells)); @@ -110,8 +154,11 @@ void llama_kv_cache_set_input_kpool( int32_t * dst_pool_cells = (int32_t *) pool_cells->data; float * dst_bias = bias ? (float *) bias->data : nullptr; float * dst_pool_bias = (float *) pool_bias ->data; - float * dst_sel_mask = (float *) sel_mask ->data; - float * dst_cand_mask = (float *) cand_mask ->data; + char * dst_sel_mask = (char *) sel_mask ->data; + char * dst_cand_mask = (char *) cand_mask ->data; + + const bool mask_f16 = sel_mask->type == GGML_TYPE_F16; + const size_t mask_ts = ggml_type_size(sel_mask->type); // -1 marks a cell with no usable pool. host side only: never copied into cell_pool, // where ggml_get_rows would read it as an index @@ -132,8 +179,8 @@ void llama_kv_cache_set_input_kpool( for (int64_t s = 0; s < n_ns; ++s) { int32_t * cur_pool_cells = dst_pool_cells + s*(r*n_pools); - float * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv); - float * cur_cand_mask = dst_cand_mask + s*(n_padq*n_kv); + char * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv)*mask_ts; + char * cur_cand_mask = dst_cand_mask + s*(n_padq*n_kv)*mask_ts; float * cur_pool_bias = dst_pool_bias + s*(n_tps*n_pools); // slots of a pool that is not resident, and slots outside the query's own @@ -143,8 +190,13 @@ void llama_kv_cache_set_input_kpool( std::fill(cur_pool_bias, cur_pool_bias + n_tps*n_pools, -INFINITY); // the token loop writes rows < n_tps in full; only the padding rows need clearing - std::fill(cur_sel_mask + n_tps*n_kv, cur_sel_mask + n_padq*n_kv, -INFINITY); - std::fill(cur_cand_mask + n_tps*n_kv, cur_cand_mask + n_padq*n_kv, -INFINITY); + if (mask_f16) { + kpool_mask_fill((ggml_fp16_t *) (cur_sel_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); + kpool_mask_fill((ggml_fp16_t *) (cur_cand_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); + } else { + kpool_mask_fill((float *) (cur_sel_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); + kpool_mask_fill((float *) (cur_cand_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); + } // [TAG_KPOOL_PACK] // cut the stream's pool table into one run per sequence, sized on the pool range @@ -313,20 +365,17 @@ void llama_kv_cache_set_input_kpool( const int64_t bo_vis = std::max(0, tail_start/r - b_base); float * cur_bias = dst_bias ? dst_bias + i*n_kv : nullptr; - float * cur_sel = cur_sel_mask + ii*n_kv; - float * cur_cand = cur_cand_mask + ii*n_kv; - - // the unsigned compares fold "empty or another sequence" (pos_at -1) and "no - // usable pool" (pool_of -1) into the range test, which lets this vectorise - for (int64_t j = 0; j < n_kv; ++j) { - const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; - const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; - const bool tail = pos_at[j] >= tail_start; - - cur_sel [j] = vis && tail ? 0.0f : -INFINITY; - // max(bias, sel_mask): the reference's candidate set, which the - // top-k budget may overrun but must never escape - cur_cand[j] = vis && (pooled || tail) ? 0.0f : -INFINITY; + char * cur_sel = cur_sel_mask + ii*n_kv*mask_ts; + char * cur_cand = cur_cand_mask + ii*n_kv*mask_ts; + + // the unsigned compares inside fold "empty or another sequence" (pos_at -1) + // and "no usable pool" (pool_of -1) into the range test + if (mask_f16) { + kpool_mask_row((ggml_fp16_t *) cur_sel, (ggml_fp16_t *) cur_cand, + pos_at.data(), pool_of.data(), n_kv, q, tail_start, bo_vis); + } else { + kpool_mask_row((float *) cur_sel, (float *) cur_cand, + pos_at.data(), pool_of.data(), n_kv, q, tail_start, bo_vis); } if (cur_bias) { diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index d3cd86c546c..294843a1210 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -120,7 +120,7 @@ uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t // that slot at cell 0 to keep the gather in range, and the pool would then // inherit cell 0's validity and compete for budget with a finite score. // -// sel_mask F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) +// sel_mask F16/F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) // what the top-k scatter starts from, replacing the ggml_fill(kq_mask, -INFINITY) // opening the DSA mask build in llm_graph_context::build_attn: 0.0f for the // query's own incomplete trailing pool, which GLM always attends to @@ -128,7 +128,7 @@ uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t // Forcing the tail in here rather than through the score keeps the budget a whole // number of pools. // -// cand_mask F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) +// cand_mask F16/F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) // max(bias, sel_mask): the reference's candidate set, i.e. every cell it could // return for this query. 0.0f in a complete visible pool OR in the query's own // tail, else -INFINITY, padding rows included. @@ -191,8 +191,8 @@ class llm_graph_input_kpool : public llm_graph_input_i { ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] ggml_tensor * pool_bias = nullptr; // F32 [n_pools, n_tps, n_stream] - ggml_tensor * sel_mask = nullptr; // F32 [n_kv, n_batch, 1, n_stream] - ggml_tensor * cand_mask = nullptr; // F32 [n_kv, n_batch, 1, n_stream] + ggml_tensor * sel_mask = nullptr; // F16 [n_kv, n_batch, 1, n_stream] + ggml_tensor * cand_mask = nullptr; // F16 [n_kv, n_batch, 1, n_stream] const llama_kv_cache_context * mctx_attn; const llama_kv_cache_context * mctx_idx; diff --git a/tests/test-glm5next-memory.cpp b/tests/test-glm5next-memory.cpp index b64e71cee5c..7c4aa16071b 100644 --- a/tests/test-glm5next-memory.cpp +++ b/tests/test-glm5next-memory.cpp @@ -278,8 +278,26 @@ struct kpool_tensors { ggml_tensor * pool_bias = nullptr; ggml_tensor * sel_mask = nullptr; ggml_tensor * cand_mask = nullptr; + + // the same two masks in the type the graph actually allocates. the values are + // only ever 0.0f and -INFINITY, so f16 must reproduce f32 bit for bit + ggml_tensor * sel_mask_f16 = nullptr; + ggml_tensor * cand_mask_f16 = nullptr; }; +static bool kpool_mask_f16_matches(const ggml_tensor * f32, const ggml_tensor * f16) { + const float * a = (const float *) f32->data; + const ggml_fp16_t * b = (const ggml_fp16_t *) f16->data; + + for (int64_t i = 0; i < ggml_nelements(f32); ++i) { + if (ggml_fp16_to_fp32(b[i]) != a[i]) { + return false; + } + } + + return true; +} + // The test asks for all six, including the two the pooled graph does not consume. // cell_pool and bias are the per-CELL spelling of the same predicate, computed by a // different loop, and they are what pool_bias and cand_mask are checked against here @@ -300,6 +318,12 @@ static kpool_tensors alloc_kpool_tensors( t.sel_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_padq, 1, n_stream); t.cand_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_padq, 1, n_stream); + t.sel_mask_f16 = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_padq, 1, n_stream); + t.cand_mask_f16 = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_padq, 1, n_stream); + + ggml_set_input(t.sel_mask_f16); + ggml_set_input(t.cand_mask_f16); + ggml_set_input(t.cell_pool); ggml_set_input(t.pool_cells); ggml_set_input(t.bias); @@ -520,6 +544,17 @@ int main(int argc, char ** argv) { (int) seq_mask_of_slot.size(), (int) n_pools); CHECK(mask_all == 0x3, "both sequences own selectable pools, so neither run is empty"); + // the masks are filled by one partition per sequence, each writing + // the rows it owns into the shared stream buffer, so this is where + // a wrong element width would cross partitions + llama_kv_cache_set_input_kpool(m2->get_mem_attn(), + /* cell_pool */ nullptr, kt.pool_cells, kt.bias, kt.pool_bias, + kt.sel_mask_f16, kt.cand_mask_f16, &ub, kpool); + + CHECK(kpool_mask_f16_matches(kt.sel_mask, kt.sel_mask_f16) && + kpool_mask_f16_matches(kt.cand_mask, kt.cand_mask_f16), + "the f16 masks match the f32 ones across two sequences in one stream"); + ggml_backend_buffer_free(buf); } @@ -897,6 +932,17 @@ int main(int argc, char ** argv) { CHECK(finite, "bias, pool_bias, sel_mask and cand_mask hold only 0 and -INFINITY: " "no +1e9 to meet a -inf"); + // which is why the graph stores them in the KQ mask's f16 and halves + // two [n_kv, n_batch, n_stream] inputs. same fill, half the bytes + { + llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, + kt.sel_mask_f16, kt.cand_mask_f16, &ub, (uint32_t) kpool); + + CHECK(kpool_mask_f16_matches(kt.sel_mask, kt.sel_mask_f16) && + kpool_mask_f16_matches(kt.cand_mask, kt.cand_mask_f16), + "an f16 sel_mask/cand_mask is the exact image of the f32 one"); + } + // cand_mask is exactly max(bias, sel_mask) lifted to KQ shape, and it // is what stops an over-budget top-k from escaping the reference's // candidate set. @@ -1313,6 +1359,16 @@ int main(int argc, char ** argv) { (int) (n_padq - n_tps)); } + // the padding fill has its own code path per mask type + { + llama_kv_cache_set_input_kpool(kv, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, + kt.sel_mask_f16, kt.cand_mask_f16, &ub, (uint32_t) kpool); + + CHECK(kpool_mask_f16_matches(kt.sel_mask, kt.sel_mask_f16) && + kpool_mask_f16_matches(kt.cand_mask, kt.cand_mask_f16), + "the f16 masks match the f32 ones on a padded ubatch too"); + } + // the shared object refills the same tensors deterministically. // it fills only what the pooled graph consumes - cell_pool and // bias have no consumer there, so it passes nullptr for both and From 2e0e57f1008053bae4902a772da85e3eb99d4aff Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Thu, 27 Aug 2026 03:51:29 +0000 Subject: [PATCH 28/36] glm5next : use the fused lightning indexer for the pool scores Replaces the 7-node score chain (mul_mat, cont/permute, relu, mul, sum_rows, cont/permute, add) with one ggml_lightning_indexer, as glm-dsa, deepseek4, deepseek32 and dots3note already do. The op needs an f16 mask, so pool_bias is cast once per graph in build_inp_kpool rather than once per DSA layer. pool_k is left in f32 so the CUDA op takes its f32 vector path, not the f16 wmma path, which would undo the GGML_PREC_F32 on the head weights. The unfused chain stays behind cparams.fused_lid, plus a LLAMA_FUSED_LID_DISABLE escape hatch. --- src/llama-context.cpp | 10 ++++++ src/llama-graph.cpp | 9 ++++++ src/llama-kv-cache-kpool.h | 7 +++++ src/models/glm5next.cpp | 63 +++++++++++++++++++++++++------------- 4 files changed, 67 insertions(+), 22 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 4542cd9c0a7..6ed803bc58a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -236,6 +236,16 @@ llama_context::llama_context( cparams.fused_lid = true; cparams.auto_flid = true; + { + // escape hatch: the fused kernel sums the heads in a different order than the + // unfused graph, so a near-tied indexer top-k can come out different + const char * LLAMA_FUSED_LID_DISABLE = getenv("LLAMA_FUSED_LID_DISABLE"); + if (LLAMA_FUSED_LID_DISABLE && atoi(LLAMA_FUSED_LID_DISABLE) != 0) { + cparams.fused_lid = false; + cparams.auto_flid = false; + } + } + cparams.fused_dsv4_hc_pre = true; cparams.fused_dsv4_hc_comb = true; cparams.fused_dsv4_hc_post = true; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 06fdcfd69e7..1bb06702a8a 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3608,6 +3608,15 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( ggml_set_input(inp->pool_bias); ggml_set_name(inp->pool_bias, "kpool_pool_bias"); + // the fused lightning indexer wants an f16 mask. built here, once, because + // every indexer layer shares it + if (cparams.fused_lid) { + inp->pool_bias_f16 = ggml_cast(ctx0, + ggml_reshape_4d(ctx0, inp->pool_bias, n_pools, n_tps, 1, n_stream), + GGML_TYPE_F16); + ggml_set_name(inp->pool_bias_f16, "kpool_pool_bias_f16"); + } + // f16, not f32. The only two values either mask holds are 0.0f and // -INFINITY, both exact in f16, so the narrower type is lossless and halves // two [n_kv, n_tps, n_stream] inputs that live for the whole ubatch: 2 GiB diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index 294843a1210..2ad259c8dcf 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -191,6 +191,13 @@ class llm_graph_input_kpool : public llm_graph_input_i { ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] ggml_tensor * pool_bias = nullptr; // F32 [n_pools, n_tps, n_stream] + + // pool_bias in the shape and type the fused lightning indexer wants for its mask. + // A ggml_cast node, not an input: built once per graph rather than once per DSA + // layer, since every indexer layer shares the same mask. The cast is exact - + // pool_bias only ever holds 0.0f or -INFINITY. nullptr when the fused path is off + ggml_tensor * pool_bias_f16 = nullptr; // F16 [n_pools, n_tps, 1, n_stream] + ggml_tensor * sel_mask = nullptr; // F16 [n_kv, n_batch, 1, n_stream] ggml_tensor * cand_mask = nullptr; // F16 [n_kv, n_batch, 1, n_stream] diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 92cef19a48e..2cb8d4c20b3 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -505,44 +505,63 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( pool_k = ggml_reshape_4d(ctx0, pool_k, d_idx, n_pools, 1, n_stream); cb(pool_k, "indexer_pool_k", il); - // {d_idx, n_tps, n_ihead, n_stream}. no rope: n_rot() is 0 for the whole text tower + // {d_idx, n_ihead, n_tps, n_stream}. no rope: n_rot() is 0 for the whole text tower ggml_tensor * iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr); iq = ggml_reshape_4d(ctx0, iq, d_idx, n_ihead, n_tps, n_stream); - iq = ggml_permute(ctx0, iq, 0, 2, 1, 3); cb(iq, "indexer_q", il); - // {n_pools, n_tps, n_ihead, n_stream}: pool_k is MQA and broadcasts over the heads - ggml_tensor * kq = ggml_mul_mat(ctx0, pool_k, iq); - - // {n_ihead, n_tps, n_pools, n_stream}, contiguous for the relu and the head sum. - // the ReLU sits BETWEEN the per-head dot product and the head weighting: moving it - // to either side is a different function, because the head weights are sign-free - // and the sum is not a convex combination - kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); - ggml_tensor * score = ggml_relu(ctx0, kq); - cb(score, "indexer_score", il); - // sign-unconstrained head weights: no softmax, no abs, no relu. Both scale // constants - the reference's softmax_scale = d_idx^-0.5 and its n_heads^-0.5 head // factor - are folded in here, on an {n_ihead, n_tokens} tensor rather than on the // {n_pools, n_tps, n_ihead} score tensor. relu is positively homogeneous and both // constants are positive, so this is exactly the same function, and it is what the - // engines and the in-tree glm-dsa both do + // engines and the in-tree glm-dsa both do. + // + // GGML_PREC_F32 is not cosmetic: on vLLM a bf16 head gate moves a logit by ~1e-2, + // which is enough to swap two near-tied pools, and the top-k below is a hard cut ggml_tensor * w = ggml_mul_mat(ctx0, layer.indexer_proj, cur); ggml_mul_mat_set_prec(w, GGML_PREC_F32); w = ggml_reshape_4d(ctx0, w, n_ihead, n_tps, 1, n_stream); w = ggml_scale(ctx0, w, 1.0f/sqrtf(float(d_idx*n_ihead))); cb(w, "indexer_weights", il); - // {1, n_tps, n_pools, n_stream} -> {n_pools, n_tps, n_stream} - ggml_tensor * pool_score = ggml_sum_rows(ctx0, ggml_mul(ctx0, score, w)); - pool_score = ggml_cont(ctx0, ggml_permute(ctx0, pool_score, 2, 1, 0, 3)); - pool_score = ggml_reshape_3d(ctx0, pool_score, n_pools, n_tps, n_stream); + // both paths end with pool_bias added: -INFINITY on every pool the reference's + // `pool_valid & pool_visible` rejects, the query's own trailing pool included, so + // that no budget is spent on it + ggml_tensor * pool_score = nullptr; + + if (cparams.fused_lid) { + // one node for the whole dot product -> relu -> head sum -> mask chain, and the + // mask add comes for free. pool_k stays f32, so the kernel takes its f32 vector + // path rather than the f16 wmma path, which would undo the GGML_PREC_F32 above + ggml_tensor * pool_kf = ggml_reshape_4d(ctx0, pool_k, d_idx, 1, n_pools, n_stream); - // -INFINITY on every pool the reference's `pool_valid & pool_visible` rejects, the - // query's own trailing pool included, so that no budget is spent on it - pool_score = ggml_add(ctx0, pool_score, inp_kp->pool_bias); - cb(pool_score, "indexer_pool_score", il); + pool_score = ggml_lightning_indexer(ctx0, iq, pool_kf, w, inp_kp->pool_bias_f16); + cb(pool_score, "indexer_pool_score", il); + + res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, pool_score, il}); + + pool_score = ggml_reshape_3d(ctx0, pool_score, n_pools, n_tps, n_stream); + } else { + // {n_pools, n_tps, n_ihead, n_stream}: pool_k is MQA and broadcasts over the heads + ggml_tensor * kq = ggml_mul_mat(ctx0, pool_k, ggml_permute(ctx0, iq, 0, 2, 1, 3)); + + // {n_ihead, n_tps, n_pools, n_stream}, contiguous for the relu and the head sum. + // the ReLU sits BETWEEN the per-head dot product and the head weighting: moving it + // to either side is a different function, because the head weights are sign-free + // and the sum is not a convex combination + kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); + ggml_tensor * score = ggml_relu(ctx0, kq); + cb(score, "indexer_score", il); + + // {1, n_tps, n_pools, n_stream} -> {n_pools, n_tps, n_stream} + pool_score = ggml_sum_rows(ctx0, ggml_mul(ctx0, score, w)); + pool_score = ggml_cont(ctx0, ggml_permute(ctx0, pool_score, 2, 1, 0, 3)); + pool_score = ggml_reshape_3d(ctx0, pool_score, n_pools, n_tps, n_stream); + + pool_score = ggml_add(ctx0, pool_score, inp_kp->pool_bias); + cb(pool_score, "indexer_pool_score", il); + } // Top-k over POOLS at index_topk/index_kpool, then expand each selected pool to its // members. This is the reference's own two-step (topk over the pool axis, then From cfa63e83d4e2de2b7013c9e2ab9449adf92c5384 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 28 Aug 2026 01:31:53 +0000 Subject: [PATCH 29/36] glm5next: drop the test suite for the slim variant --- tests/CMakeLists.txt | 2 - tests/test-glm5next-memory.cpp | 1417 -------------------------------- tests/test-llama-archs.cpp | 21 +- tests/test-mtmd-impl.cpp | 101 --- 4 files changed, 1 insertion(+), 1540 deletions(-) delete mode 100644 tests/test-glm5next-memory.cpp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2c793e582e0..b9f9d4b78af 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -196,8 +196,6 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # llama_build_and_test(test-double-float.cpp) # SLOW llama_build_and_test(test-llama-archs.cpp) - # needs a glm5next GGUF as argv[1], so it is built but not registered - llama_build(test-glm5next-memory.cpp) set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") diff --git a/tests/test-glm5next-memory.cpp b/tests/test-glm5next-memory.cpp deleted file mode 100644 index 7c4aa16071b..00000000000 --- a/tests/test-glm5next-memory.cpp +++ /dev/null @@ -1,1417 +0,0 @@ -// The GLM-5-Next hybrid memory: the memory object holds all three halves (KDA -// recurrent + conv state, MLA latent KV, pooled indexer key cache) at the sizes the -// reference implies, one ggml graph reaches every one of them, and the pooled top-k -// cuts on a pool boundary and agrees between CPU and CUDA. The indexer graph itself is -// not built here. -// -// Run as: test-glm5next-memory , as written by -// tests/glm5next_make_tiny_gguf.py - -#include "ggml.h" -#include "ggml-alloc.h" -#include "ggml-backend.h" -#include "llama.h" - -#include "../src/llama-model.h" -#include "../src/llama-hparams.h" -#include "../src/llama-cparams.h" -#include "../src/llama-memory-hybrid.h" -#include "../src/llama-memory-recurrent.h" -#include "../src/llama-kv-cache.h" -#include "../src/llama-kv-cache-kpool.h" -#include "../src/llama-batch.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include - -static int n_fail = 0; - -#define CHECK(cond, ...) \ - do { \ - if (!(cond)) { \ - printf("FAIL %s:%d: ", __func__, __LINE__); \ - printf(__VA_ARGS__); \ - printf("\n"); \ - n_fail++; \ - } else { \ - printf("ok "); \ - printf(__VA_ARGS__); \ - printf("\n"); \ - } \ - } while (0) - -// -// numeric evidence for WHERE the top-k runs, independent of any model -// -// The claim under test is the one that decides the whole design: a top-k over -// CELLS cannot be made pool-aligned by choosing its width, and a top-k over -// POOLS is pool-aligned by construction. -// -// The tempting argument for the cell-level form is that a pool's members carry -// its score bit-exactly, so an intra-pool tie is harmless and a budget that is a -// whole multiple of kpool must cut on a pool boundary. That argument silently -// assumes tie groups never SPAN pools. F.relu drives most pool scores to exactly -// 0.0, so they do, and ggml_top_k - explicitly unordered among equals - then -// takes an arbitrary 1..kpool-1 members of the pool it lands in. This -// reproduces that with no model at all: most pools at exactly 0.0, a minority -// positive, fewer positive pools than the budget. -// - -static std::vector run_top_k( - ggml_backend_t backend, - const std::vector & scores, - int64_t n_kv, - int64_t n_rows, - int64_t width) { - ggml_init_params gparams = { - /* .mem_size */ ggml_tensor_overhead()*8 + ggml_graph_overhead(), - /* .mem_buffer */ nullptr, - /* .no_alloc */ true, - }; - - ggml_context * ctx = ggml_init(gparams); - - ggml_tensor * src = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_kv, n_rows); - ggml_set_input(src); - - ggml_tensor * dst = ggml_top_k(ctx, src, (int) width); - ggml_set_output(dst); - - ggml_cgraph * gf = ggml_new_graph(ctx); - ggml_build_forward_expand(gf, dst); - - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); - - std::vector out(width*n_rows); - - if (buf) { - ggml_backend_tensor_set(src, scores.data(), 0, scores.size()*sizeof(float)); - - if (ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS) { - ggml_backend_tensor_get(dst, out.data(), 0, out.size()*sizeof(int32_t)); - } else { - out.clear(); - } - - ggml_backend_buffer_free(buf); - } else { - out.clear(); - } - - ggml_free(ctx); - - return out; -} - -static int64_t n_partial_pools(const std::vector & sel, int64_t off, int64_t width, int64_t kpool) { - std::map cnt; - for (int64_t i = 0; i < width; ++i) { - cnt[sel[off + i]/(int32_t) kpool]++; - } - - int64_t n = 0; - for (const auto & kv : cnt) { - n += kv.second != (int32_t) kpool; - } - - return n; -} - -static void test_top_k_boundary() { - printf("\n--- top-k granularity, standalone ---\n"); - - const int64_t kpool = 4; - const int64_t top_k = 2048; // the reference requires top_k %% kpool == 0 - const int64_t n_kv = 8192; - const int64_t n_rows = 4; // queries - const int64_t n_pools = n_kv/kpool; - - // the shape ReLU actually produces: a minority of pools with a positive score, - // every other pool at EXACTLY 0.0. Fewer positive pools than the budget, which - // during prefill is the normal state and not a corner case - const int64_t n_pos = 300; - - std::vector pool_score(n_pools, 0.0f); - for (int64_t p = 0; p < n_pos; ++p) { - // deterministic, distinct, strictly positive - pool_score[(p*7919) % n_pools] = 1.0f + std::fabs(std::sin(0.7f*(float) p))*1000.0f; - } - - const int64_t select_k = llama_kpool_select_k((uint32_t) n_pools, (uint32_t) top_k, (uint32_t) kpool); - CHECK(select_k == top_k/kpool, "llama_kpool_select_k is index_topk/index_kpool (%d pools)", (int) select_k); - - ggml_backend_t backend_cpu = ggml_backend_cpu_init(); - ggml_backend_t backend_gpu = nullptr; - { - ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); - if (dev) { - backend_gpu = ggml_backend_dev_init(dev, nullptr); - } - } - - printf("%s GPU backend for the top-k comparison: %s\n", - backend_gpu ? "ok " : "note", backend_gpu ? ggml_backend_name(backend_gpu) : "none, CPU only"); - - // t = (q + 1) %% kpool cells of trailing incomplete pool, biased out of the - // budget. swept so that the per-stream tail residue is covered, not just t == 0 - for (int64_t t = 0; t < kpool; ++t) { - const int64_t n_tail = t; - const int64_t n_full = n_kv - n_tail; // cells belonging to complete pools - - // ---- A: the WRONG design. one score per cell, top-k of width index_topk -- - std::vector s_cell((size_t) n_kv*n_rows); - // ---- B: what ships. one score per pool, top-k of width index_topk/kpool -- - std::vector s_pool((size_t) n_pools*n_rows); - - for (int64_t r = 0; r < n_rows; ++r) { - for (int64_t j = 0; j < n_kv; ++j) { - s_cell[r*n_kv + j] = j >= n_full ? -INFINITY : pool_score[j/kpool]; - } - for (int64_t p = 0; p < n_pools; ++p) { - // a pool that the tail bites into is not a candidate at all - s_pool[r*n_pools + p] = (p + 1)*kpool > n_full ? -INFINITY : pool_score[p]; - } - } - - const auto sel_cell = run_top_k(backend_cpu, s_cell, n_kv, n_rows, top_k); - const auto sel_pool = run_top_k(backend_cpu, s_pool, n_pools, n_rows, select_k); - - // 1. the cell-level form, at a budget that IS a whole number of pools, - // still splits pools. this is the measurement the design turns on - int64_t partial_cell = 0; - for (int64_t r = 0; r < n_rows; ++r) { - partial_cell += n_partial_pools(sel_cell, r*top_k, top_k, kpool); - } - CHECK(partial_cell > 0, - "t=%d: a CELL-level top-k at the pool-aligned width %d still leaves %d partial pool(s); " - "a pool-aligned WIDTH does not make a pool-aligned CUT", - (int) t, (int) top_k, (int) partial_cell); - - // 2. the pool-level form: expand each selected pool ordinal to its kpool - // member cells, exactly as the graph does through pool_cells - std::vector expanded((size_t) select_k*kpool*n_rows); - for (int64_t r = 0; r < n_rows; ++r) { - for (int64_t i = 0; i < select_k; ++i) { - const int32_t p = sel_pool[r*select_k + i]; - for (int64_t m = 0; m < kpool; ++m) { - expanded[r*select_k*kpool + i*kpool + m] = (int32_t) (p*kpool + m); - } - } - } - - int64_t partial_pool = 0; - int64_t tail_in_pool = 0; - for (int64_t r = 0; r < n_rows; ++r) { - partial_pool += n_partial_pools(expanded, r*select_k*kpool, select_k*kpool, kpool); - for (int64_t i = 0; i < select_k*kpool; ++i) { - tail_in_pool += expanded[r*select_k*kpool + i] >= n_full; - } - } - CHECK(partial_pool == 0, - "t=%d: a POOL-level top-k of %d pools expands to only whole pools (%d partial)", - (int) t, (int) select_k, (int) partial_pool); - CHECK(tail_in_pool == 0, "t=%d: no tail cell consumes budget (%d did)", (int) t, (int) tail_in_pool); - - // 3. CPU vs CUDA on the same input - if (backend_gpu) { - const auto gpu_cell = run_top_k(backend_gpu, s_cell, n_kv, n_rows, top_k); - const auto gpu_pool = run_top_k(backend_gpu, s_pool, n_pools, n_rows, select_k); - - // the pool SET may legitimately differ between backends when the cut - // falls inside a tie group. What may not differ is pool integrity, - // and that is what is asserted - int64_t partial_gpu = 0; - bool ok_gpu = gpu_pool.size() == sel_pool.size(); - for (int64_t r = 0; ok_gpu && r < n_rows; ++r) { - for (int64_t i = 0; i < select_k; ++i) { - const int32_t p = gpu_pool[r*select_k + i]; - for (int64_t m = 0; m < kpool; ++m) { - expanded[r*select_k*kpool + i*kpool + m] = (int32_t) (p*kpool + m); - } - } - partial_gpu += n_partial_pools(expanded, r*select_k*kpool, select_k*kpool, kpool); - } - CHECK(ok_gpu && partial_gpu == 0, - "t=%d: %s also expands to only whole pools (%d partial)", - (int) t, ggml_backend_name(backend_gpu), (int) partial_gpu); - - bool same_pool = gpu_pool.size() == sel_pool.size(); - for (int64_t r = 0; same_pool && r < n_rows; ++r) { - std::set a(sel_pool.begin() + r*select_k, sel_pool.begin() + (r + 1)*select_k); - std::set b(gpu_pool.begin() + r*select_k, gpu_pool.begin() + (r + 1)*select_k); - same_pool = a == b; - } - int64_t partial_gpu_cell = 0; - for (int64_t r = 0; r < n_rows; ++r) { - partial_gpu_cell += n_partial_pools(gpu_cell, r*top_k, top_k, kpool); - } - // reported, not asserted: which members of a cut pool each backend's - // partial sort happens to keep is not contractual - printf("note t=%d: CPU and %s %s on the selected POOL set; the cell-level form leaves " - "%d partial pool(s) there too\n", - (int) t, ggml_backend_name(backend_gpu), - same_pool ? "agree" : "DISAGREE", (int) partial_gpu_cell); - } - } - - if (backend_gpu) { - ggml_backend_free(backend_gpu); - } - ggml_backend_free(backend_cpu); -} - -// -// the memory object itself -// - -struct kpool_tensors { - ggml_tensor * cell_pool = nullptr; - ggml_tensor * pool_cells = nullptr; - ggml_tensor * bias = nullptr; - ggml_tensor * pool_bias = nullptr; - ggml_tensor * sel_mask = nullptr; - ggml_tensor * cand_mask = nullptr; - - // the same two masks in the type the graph actually allocates. the values are - // only ever 0.0f and -INFINITY, so f16 must reproduce f32 bit for bit - ggml_tensor * sel_mask_f16 = nullptr; - ggml_tensor * cand_mask_f16 = nullptr; -}; - -static bool kpool_mask_f16_matches(const ggml_tensor * f32, const ggml_tensor * f16) { - const float * a = (const float *) f32->data; - const ggml_fp16_t * b = (const ggml_fp16_t *) f16->data; - - for (int64_t i = 0; i < ggml_nelements(f32); ++i) { - if (ggml_fp16_to_fp32(b[i]) != a[i]) { - return false; - } - } - - return true; -} - -// The test asks for all six, including the two the pooled graph does not consume. -// cell_pool and bias are the per-CELL spelling of the same predicate, computed by a -// different loop, and they are what pool_bias and cand_mask are checked against here -static kpool_tensors alloc_kpool_tensors( - ggml_context * ctx, - int64_t n_kv, - int64_t n_tps, - int64_t n_padq, - int64_t n_stream, - int64_t kpool, - int64_t n_pools) { - kpool_tensors t; - - t.cell_pool = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_kv, n_stream); - t.pool_cells = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, kpool*n_pools, n_stream); - t.bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_kv, n_tps, n_stream); - t.pool_bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_pools, n_tps, n_stream); - t.sel_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_padq, 1, n_stream); - t.cand_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_padq, 1, n_stream); - - t.sel_mask_f16 = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_padq, 1, n_stream); - t.cand_mask_f16 = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_padq, 1, n_stream); - - ggml_set_input(t.sel_mask_f16); - ggml_set_input(t.cand_mask_f16); - - ggml_set_input(t.cell_pool); - ggml_set_input(t.pool_cells); - ggml_set_input(t.bias); - ggml_set_input(t.pool_bias); - ggml_set_input(t.sel_mask); - ggml_set_input(t.cand_mask); - - return t; -} - -int main(int argc, char ** argv) { - if (argc < 2) { - printf("usage: %s \n", argv[0]); - return 1; - } - - setvbuf(stdout, nullptr, _IOLBF, 0); - - llama_backend_init(); - - test_top_k_boundary(); - - printf("\n--- model ---\n"); - - llama_model_params mparams = llama_model_default_params(); - mparams.n_gpu_layers = 0; - - llama_model * model = llama_model_load_from_file(argv[1], mparams); - if (model == nullptr) { - printf("FAIL: could not load %s\n", argv[1]); - return 1; - } - - const auto & hparams = model->hparams; - - printf("n_layer = %u (+%u nextn)\n", hparams.n_layer(), hparams.n_layer_nextn); - printf("n_head = %u\n", hparams.n_head()); - printf("n_embd_head_kda = %u\n", hparams.n_embd_head_kda); - printf("ssm_d_conv = %u\n", hparams.ssm_d_conv); - printf("n_lora_kv = %u\n", hparams.n_lora_kv); - printf("indexer_head_size= %u\n", hparams.indexer_head_size); - printf("indexer_kpool = %u\n", hparams.indexer_kpool); - printf("indexer_top_k = %u\n", hparams.indexer_top_k); - printf("n_embd_r() = %u\n", hparams.n_embd_r()); - printf("n_embd_s() = %u\n", hparams.n_embd_s()); - - // ---- sizing ------------------------------------------------------------ - { - const uint32_t d_inner = hparams.n_head()*hparams.n_embd_head_kda; - - CHECK(hparams.n_embd_r() == 3*(hparams.ssm_d_conv - 1)*d_inner, - "n_embd_r == 3 conv states of (d_conv-1)*n_head*head_dim (%u)", hparams.n_embd_r()); - CHECK(hparams.n_embd_s() == hparams.n_embd_head_kda*hparams.n_embd_head_kda*hparams.n_head(), - "n_embd_s == head_dim*head_dim per head (%u)", hparams.n_embd_s()); - CHECK(hparams.indexer_kpool > 0 && hparams.indexer_top_k % hparams.indexer_kpool == 0, - "indexer_top_k (%u) is a whole number of pools of %u", hparams.indexer_top_k, hparams.indexer_kpool); - } - - // ---- a multi-sequence unified cache pools per SEQUENCE ------------------ - // - // [TAG_KPOOL_SEQ_PARTITION]. pools group cells by position, and a position only - // names a pool inside one sequence. a unified cache shares one cells array, so the - // stream's pool table is cut into one run per sequence instead of the cache being - // refused. -kvu with --parallel is reachable (llama-perplexity forces it for - // hellaswag / winogrande / multiple-choice, and llama-embedding turns -np 1 into - // -kvu with n_seq_max 256), so this has to work rather than fail at startup. - { - llama_cparams cp = {}; - cp.n_ctx = 256; - cp.n_ctx_seq = 256; - cp.n_batch = 32; - cp.n_ubatch = 32; - cp.n_seq_max = 2; - cp.kv_unified = true; - cp.causal_attn = true; - - llama_memory_params mp = {}; - mp.type_k = GGML_TYPE_F16; - mp.type_v = GGML_TYPE_F16; - mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; - - llama_memory_i * raw = nullptr; - try { - raw = model->create_memory(mp, cp); - } catch (const std::exception & e) { - printf("note create_memory threw: %s\n", e.what()); - raw = nullptr; - } - CHECK(raw != nullptr, "-kvu with n_seq_max 2 is accepted: the pool map is per sequence"); - - const uint32_t kpool = hparams.indexer_kpool; - - // one shared n_kv/kpool budget plus the rebasing slack per sequence, NOT one - // full-width table each: the indexer scores every slot against every query, so a - // full-width table per sequence multiplies the score tensor by the sequence count - CHECK(llama_kpool_n_pools(256, kpool, 1) == 256/kpool + 2 && - llama_kpool_n_pools(256, kpool, 2) == 256/kpool + 4, - "llama_kpool_n_pools is n_kv/kpool + 2 per sequence (%u slots for 2 sequences)", - llama_kpool_n_pools(256, kpool, 2)); - - auto * m2 = dynamic_cast(raw); - CHECK(m2 != nullptr && m2->get_mem_idx() != nullptr, "-kvu builds an indexer cache too"); - - // drive one ubatch holding BOTH sequences through the map. 16 positions each, so - // every sequence owns several complete pools and the runs have to be told apart - if (m2) { - const int64_t n_tok = 16; - - std::vector tok; - std::vector pos; - std::vector sid; - std::vector sptr; - std::vector nsid; - std::vector lg; - - for (llama_seq_id s = 0; s < 2; ++s) { - for (int64_t i = 0; i < n_tok; ++i) { - tok.push_back(0); - pos.push_back((llama_pos) i); - sid.push_back(s); - nsid.push_back(1); - lg.push_back(1); - } - } - for (size_t i = 0; i < sid.size(); ++i) { - sptr.push_back(&sid[i]); - } - - llama_batch b = {}; - b.n_tokens = (int32_t) tok.size(); - b.token = tok.data(); - b.pos = pos.data(); - b.seq_id = sptr.data(); - b.n_seq_id = nsid.data(); - b.logits = lg.data(); - - llama_batch_allocr ba(hparams.n_pos_per_embd()); - if (ba.init(b, model->vocab, nullptr, hparams.n_embd_inp(), cp.n_seq_max, true)) { - auto c = m2->init_batch(ba, cp.n_ubatch, false); - auto * mc = dynamic_cast(c.get()); - - CHECK(mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, - "both sequences fit one unified ubatch"); - - if (mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { - mc->apply(); - - const llama_ubatch & ub = mc->get_ubatch(); - - const int64_t n_kv = mc->get_attn()->get_n_kv(); - const int64_t n_stream = mc->get_attn()->get_n_stream(); - const int64_t n_tps = ub.n_tokens/n_stream; - const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, kpool, ub.n_seqs_unq); - - CHECK(n_stream == 1 && ub.n_seqs_unq == 2, - "a unified cache carries both sequences in one stream (n_seqs_unq %u)", - ub.n_seqs_unq); - - ggml_init_params gp = { ggml_tensor_overhead()*16, nullptr, true }; - ggml_context * ctx = ggml_init(gp); - - kpool_tensors kt = alloc_kpool_tensors(ctx, n_kv, n_tps, n_tps, n_stream, kpool, n_pools); - - ggml_backend_t backend = ggml_backend_cpu_init(); - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); - - if (buf) { - // no cell_pool: the per-cell view has one row per stream, and a - // cell two sequences share has nowhere to put its second pool - llama_kv_cache_set_input_kpool(m2->get_mem_attn(), - /* cell_pool */ nullptr, kt.pool_cells, kt.bias, kt.pool_bias, - kt.sel_mask, kt.cand_mask, &ub, kpool); - - const int32_t * pc = (const int32_t *) kt.pool_cells->data; - const float * pb = (const float *) kt.pool_bias->data; - - const auto & cells = m2->get_mem_attn()->get_cells(0); - - // the invariant the partitioning exists for: a pool the query may - // spend budget on holds only that query's own visible cells. this - // is what a shared cells array breaks if the map is per stream - bool own = true; - int64_t n_finite = 0; - std::map seq_mask_of_slot; - - for (int64_t ii = 0; ii < n_tps; ++ii) { - const llama_seq_id sq = ub.seq_id[ii][0]; - const llama_pos q = ub.pos[ii]; - - for (int64_t p = 0; p < n_pools; ++p) { - if (pb[ii*n_pools + p] != 0.0f) { - continue; - } - - n_finite++; - seq_mask_of_slot[p] |= 1 << sq; - - for (int64_t m = 0; m < (int64_t) kpool; ++m) { - const int32_t c = pc[p*kpool + m]; - own &= c >= 0 && c < n_kv && !cells.is_empty(c) && - cells.seq_has(c, sq) && cells.pos_get(c) <= q; - } - } - } - - CHECK(own && n_finite > 0, - "every pool a query may select holds only that query's own visible cells (%d selectable (query, pool) pairs)", - (int) n_finite); - - int mask_all = 0; - bool disjoint = true; - for (const auto & kv : seq_mask_of_slot) { - mask_all |= kv.second; - disjoint &= (kv.second & (kv.second - 1)) == 0; - } - - CHECK(disjoint, "the two sequences get disjoint pool runs (%d slots in use of %d)", - (int) seq_mask_of_slot.size(), (int) n_pools); - CHECK(mask_all == 0x3, "both sequences own selectable pools, so neither run is empty"); - - // the masks are filled by one partition per sequence, each writing - // the rows it owns into the shared stream buffer, so this is where - // a wrong element width would cross partitions - llama_kv_cache_set_input_kpool(m2->get_mem_attn(), - /* cell_pool */ nullptr, kt.pool_cells, kt.bias, kt.pool_bias, - kt.sel_mask_f16, kt.cand_mask_f16, &ub, kpool); - - CHECK(kpool_mask_f16_matches(kt.sel_mask, kt.sel_mask_f16) && - kpool_mask_f16_matches(kt.cand_mask, kt.cand_mask_f16), - "the f16 masks match the f32 ones across two sequences in one stream"); - - ggml_backend_buffer_free(buf); - } - - ggml_backend_free(backend); - ggml_free(ctx); - } - } - } - - delete raw; - - // one sequence in flight stays the simple case: the shared cells array holds - // only its keys, and the map filters on seq_has anyway - cp.n_seq_max = 1; - llama_memory_i * raw1 = nullptr; - try { - raw1 = model->create_memory(mp, cp); - } catch (const std::exception &) { - raw1 = nullptr; - } - CHECK(raw1 != nullptr, "-kvu with a single sequence is still allowed"); - delete raw1; - } - - // ---- the indexer cache keeps its own dtype ------------------------------ - { - llama_cparams cp = {}; - cp.n_ctx = 256; - cp.n_ctx_seq = 256; - cp.n_batch = 32; - cp.n_ubatch = 32; - cp.n_seq_max = 1; - cp.kv_unified = false; - cp.causal_attn = true; - - llama_memory_params mp = {}; - mp.type_k = GGML_TYPE_Q8_0; - mp.type_v = GGML_TYPE_Q8_0; - mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; - - llama_memory_i * raw = model->create_memory(mp, cp); - auto * m = dynamic_cast(raw); - - CHECK(m != nullptr && m->get_mem_idx() != nullptr, "-ctk q8_0 still builds an indexer cache"); - if (m && m->get_mem_idx()) { - CHECK(!ggml_is_quantized(m->get_mem_idx()->type_k()), - "indexer cache stays %s under -ctk q8_0: it also holds the compressor gates", - ggml_type_name(m->get_mem_idx()->type_k())); - CHECK(ggml_is_quantized(m->get_mem_attn()->type_k()), - "the MLA cache still honours -ctk q8_0 (%s)", ggml_type_name(m->get_mem_attn()->type_k())); - } - - delete raw; - } - - // ---- create the memory ------------------------------------------------- - llama_cparams cparams = {}; - cparams.n_ctx = 512; - cparams.n_ctx_seq = 512; - cparams.n_batch = 32; - cparams.n_ubatch = 32; - cparams.n_seq_max = 2; - cparams.n_rs_seq = 0; - cparams.kv_unified = false; // one sequence per stream, the n_ps == 1 layout - cparams.offload_kqv = false; - cparams.flash_attn = false; - cparams.causal_attn = true; - - llama_memory_params mparams_mem = {}; - mparams_mem.type_k = GGML_TYPE_F16; - mparams_mem.type_v = GGML_TYPE_F16; - mparams_mem.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; - mparams_mem.swa_full = false; - - llama_memory_i * mem_raw = model->create_memory(mparams_mem, cparams); - CHECK(mem_raw != nullptr, "create_memory returned a memory object"); - if (!mem_raw) { - return 1; - } - - auto * mem = dynamic_cast(mem_raw); - CHECK(mem != nullptr, "the memory is a llama_memory_hybrid"); - if (!mem) { - return 1; - } - - llama_kv_cache * kv_attn = mem->get_mem_attn(); - llama_memory_recurrent * rs = mem->get_mem_recr(); - llama_kv_cache * kv_idx = mem->get_mem_idx(); - - CHECK(kv_attn != nullptr, "hybrid holds an attention (MLA) cache"); - CHECK(rs != nullptr, "hybrid holds a recurrent (KDA) cache"); - CHECK(kv_idx != nullptr, "hybrid holds an indexer key cache"); - if (!kv_attn || !rs || !kv_idx) { - return 1; - } - - // ---- layer partition --------------------------------------------------- - { - const auto ids_attn = kv_attn->get_layer_ids(); - const auto ids_idx = kv_idx ->get_layer_ids(); - - uint32_t n_dsa = 0; - for (uint32_t il = 0; il < hparams.n_layer(); ++il) { - n_dsa += !hparams.is_recr(il); - } - - CHECK(ids_attn.size() == n_dsa, "MLA cache holds the %u DSA trunk layers (%zu)", n_dsa, ids_attn.size()); - CHECK(ids_idx.size() == n_dsa, "indexer cache holds the same %u layers (%zu)", n_dsa, ids_idx.size()); - - bool same = ids_attn.size() == ids_idx.size(); - for (size_t i = 0; same && i < ids_attn.size(); ++i) { - same = ids_attn[i] == ids_idx[i]; - } - CHECK(same, "MLA and indexer caches cover exactly the same layers"); - - for (uint32_t il : ids_attn) { - CHECK(!hparams.is_recr(il), "cached attention layer %u is not recurrent", il); - } - } - - // ---- cache geometry ---------------------------------------------------- - { - const uint32_t il_dsa = kv_attn->get_layer_ids().front(); - - ggml_tensor * k_mla = kv_attn->get_k_storage(il_dsa); - ggml_tensor * k_idx = kv_idx ->get_k_storage(il_dsa); - - CHECK(k_mla != nullptr && k_idx != nullptr, "both caches expose K storage for layer %u", il_dsa); - - // nope-only MLA: the latent row is kv_lora_rank + qk_rope_head_dim - CHECK(k_mla->ne[0] == (int64_t) hparams.n_embd_head_k(il_dsa), - "MLA latent row = %d (n_embd_head_k = %u)", (int) k_mla->ne[0], hparams.n_embd_head_k(il_dsa)); - - // the pooling indexer caches the key AND the compressor gate score - CHECK(k_idx->ne[0] == (int64_t) (2*hparams.indexer_head_size), - "indexer row = %d (expected 2 x indexer_head_size = %u)", - (int) k_idx->ne[0], 2*hparams.indexer_head_size); - - CHECK(k_mla->ne[1] == k_idx->ne[1], "MLA and indexer caches have the same cell count (%d)", (int) k_mla->ne[1]); - CHECK(k_mla->ne[2] == k_idx->ne[2], "MLA and indexer caches have the same stream count (%d)", (int) k_mla->ne[2]); - - // is_mla() stays true for the indexer hparams copy, so no V is allocated - { - size_t bytes = 0; - for (const auto & b : kv_idx->memory_breakdown()) { - bytes += b.second; - } - - const size_t k_only = (size_t) ggml_nbytes(k_idx)*kv_idx->get_layer_ids().size(); - - // the byte count is the observable: K only, no V - CHECK(hparams.is_mla() && bytes == k_only, - "indexer cache allocates K and no V: %zu bytes for %zu layers", - bytes, kv_idx->get_layer_ids().size()); - } - } - - // ---- recurrent geometry ------------------------------------------------ - { - uint32_t n_kda = 0; - for (uint32_t il = 0; il < hparams.n_layer(); ++il) { - n_kda += hparams.is_recr(il); - } - - CHECK(rs->size >= cparams.n_seq_max, "recurrent cache has >= n_seq_max slots (%u)", rs->size); - CHECK(n_kda > 0, "the model has %u KDA layers", n_kda); - } - - // ---- drive a batch through it and reach all three halves in one graph --- - { - const int32_t n_tokens = 10; // pos 9 leaves (9+1) %% kpool = 2 tail cells - - std::vector tokens(n_tokens*cparams.n_seq_max, 0); - std::vector pos; - std::vector seqs; - for (uint32_t s = 0; s < cparams.n_seq_max; ++s) { - for (int32_t i = 0; i < n_tokens; ++i) { - pos.push_back(i); - seqs.push_back((llama_seq_id) s); - } - } - - llama_batch batch = {}; - batch.n_tokens = n_tokens*(int32_t) cparams.n_seq_max; - batch.token = tokens.data(); - batch.pos = pos.data(); - - std::vector seq_ptrs(batch.n_tokens); - for (int32_t i = 0; i < batch.n_tokens; ++i) { - seq_ptrs[i] = &seqs[i]; - } - std::vector n_seq_id(batch.n_tokens, 1); - std::vector logits(batch.n_tokens, 0); - logits.back() = 1; - - batch.seq_id = seq_ptrs.data(); - batch.n_seq_id = n_seq_id.data(); - batch.logits = logits.data(); - - llama_batch_allocr balloc(hparams.n_pos_per_embd()); - const bool ok = balloc.init(batch, model->vocab, nullptr, hparams.n_embd_inp(), - cparams.n_seq_max, true); - CHECK(ok, "batch allocr accepted a %d-token, %u-sequence batch", batch.n_tokens, cparams.n_seq_max); - - auto mctx_ptr = mem->init_batch(balloc, cparams.n_ubatch, false); - CHECK(mctx_ptr != nullptr && mctx_ptr->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, - "init_batch produced a usable memory context"); - - auto * mctx = dynamic_cast(mctx_ptr.get()); - CHECK(mctx != nullptr, "the context is a llama_memory_hybrid_context"); - - if (mctx && mctx->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { - const auto * ctx_attn = mctx->get_attn(); - const auto * ctx_recr = mctx->get_recr(); - const auto * ctx_idx = mctx->get_idx(); - - CHECK(ctx_attn != nullptr, "context exposes the attention half"); - CHECK(ctx_recr != nullptr, "context exposes the recurrent half"); - CHECK(ctx_idx != nullptr, "context exposes the indexer half"); - - // n_kv is only valid once applied; apply() also asserts the caches agree - mctx->apply(); - - // apply() asserts both itself, so reaching this line is the result; printed - // rather than CHECKed so the count stays honest - printf("note indexer and attention caches agree on n_kv (%u) and n_stream (%u)\n", - ctx_idx->get_n_kv(), ctx_idx->get_n_stream()); - CHECK(ctx_idx && ctx_idx->get_kv() == kv_idx, "the indexer context is a view of the indexer cache"); - - ggml_init_params gparams = { - /* .mem_size */ ggml_tensor_overhead()*1024 + ggml_graph_overhead(), - /* .mem_buffer */ nullptr, - /* .no_alloc */ true, - }; - - ggml_context * ctx0 = ggml_init(gparams); - ggml_cgraph * gf = ggml_new_graph(ctx0); - - const uint32_t il_dsa = kv_attn->get_layer_ids().front(); - - uint32_t il_kda = 0; - while (il_kda < hparams.n_layer() && !hparams.is_recr(il_kda)) { - il_kda++; - } - - const int64_t n_kv = ctx_attn->get_n_kv(); - const int64_t n_stream = ctx_attn->get_n_stream(); - const int64_t n_tps = mctx->get_ubatch().n_tokens/n_stream; - - // 1. MLA latent K - ggml_tensor * k_mla = ctx_attn->get_k(ctx0, il_dsa); - CHECK(k_mla != nullptr, "graph reaches the MLA latent cache: [%d, %d, %d, %d]", - (int) k_mla->ne[0], (int) k_mla->ne[1], (int) k_mla->ne[2], (int) k_mla->ne[3]); - - // 2. indexer keys, split into the key half and the gate half - ggml_tensor * k_idx_all = ctx_idx->get_k(ctx0, il_dsa); - CHECK(k_idx_all->ne[1] == 2, "indexer cache view has 2 heads (key | compressor gate)"); - - const int64_t d_idx = hparams.indexer_head_size; - - ggml_tensor * k_idx_v = ggml_view_3d(ctx0, k_idx_all, d_idx, n_kv, n_stream, - k_idx_all->nb[2], k_idx_all->nb[3], 0); - ggml_tensor * g_idx_v = ggml_view_3d(ctx0, k_idx_all, d_idx, n_kv, n_stream, - k_idx_all->nb[2], k_idx_all->nb[3], k_idx_all->nb[1]); - - // 3. KDA recurrent + conv state - ggml_tensor * r_kda = ctx_recr->get_r_l(il_kda); - ggml_tensor * s_kda = ctx_recr->get_s_l(il_kda); - CHECK(r_kda != nullptr && s_kda != nullptr, - "graph reaches the KDA conv state [%d x %d] and recurrent state [%d x %d]", - (int) r_kda->ne[0], (int) r_kda->ne[1], (int) s_kda->ne[0], (int) s_kda->ne[1]); - - // 4. host-side pool map - const int64_t kpool = hparams.indexer_kpool; - const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); - const int64_t n_padq = n_tps; // this tree does not pad the KQ mask - - kpool_tensors kt = alloc_kpool_tensors(ctx0, n_kv, n_tps, n_padq, n_stream, kpool, n_pools); - - // shape of the real thing: gather pool members, mix, score, broadcast the - // pool score back onto its member cells, top-k - ggml_tensor * members = ggml_get_rows(ctx0, k_idx_v, kt.pool_cells); - ggml_tensor * gates = ggml_get_rows(ctx0, g_idx_v, kt.pool_cells); - members = ggml_reshape_4d(ctx0, members, d_idx, kpool, n_pools, n_stream); - gates = ggml_reshape_4d(ctx0, gates, d_idx, kpool, n_pools, n_stream); - - // stand-in for softmax(gate + ape) * member, summed over kpool - ggml_tensor * pooled = ggml_mul(ctx0, members, gates); - pooled = ggml_reshape_3d(ctx0, pooled, d_idx*kpool, n_pools, n_stream); - // stand-in for the q . k_pool score: one number per (pool, query) - ggml_tensor * score = ggml_mul_mat(ctx0, - ggml_cont(ctx0, ggml_view_3d(ctx0, pooled, d_idx, n_pools, n_stream, - pooled->nb[1], pooled->nb[2], 0)), - ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, d_idx, n_tps, n_stream)); - - // mask the pools the reference rejects, then top-k over POOLS - score = ggml_add(ctx0, ggml_cont(ctx0, score), kt.pool_bias); - - const int64_t select_k = llama_kpool_select_k((uint32_t) n_pools, hparams.indexer_top_k, (uint32_t) kpool); - ggml_tensor * sel = ggml_cont(ctx0, ggml_top_k(ctx0, score, (int) select_k)); - - // expand each selected POOL ordinal into its kpool member CELLS, which is - // the reference's selected_indices = pool_indices[batch_idx, selected]. - // the query axis folds into the gather's row axis so that one get_rows - // serves every query, while the stream axis stays where get_rows wants it - ggml_tensor * pc3 = ggml_reshape_3d(ctx0, kt.pool_cells, kpool, n_pools, n_stream); - ggml_tensor * sel_flat = ggml_reshape_2d(ctx0, sel, select_k*n_tps, n_stream); - - const int64_t width = kpool*select_k; - ggml_tensor * top_k = ggml_reshape_3d(ctx0, - ggml_get_rows(ctx0, pc3, sel_flat), width, n_tps, n_stream); - - CHECK(top_k->type == GGML_TYPE_I32, "the pool -> cell expansion stays I32 (pool_cells is I32)"); - printf("note top-k over %d POOL scores expands to %d I32 CELL indices per query\n", - (int) select_k, (int) width); - - // 5. the scatter the mask is built from, starting at sel_mask rather than - // an all -INFINITY fill, which is what forces the tail in - ggml_tensor * top_k_4d = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream); - ggml_tensor * base = ggml_view_4d(ctx0, kt.sel_mask, 1, n_kv, n_padq, n_stream, - kt.sel_mask->nb[0], kt.sel_mask->nb[1], kt.sel_mask->nb[2], 0); - ggml_tensor * idxs = ggml_view_4d(ctx0, top_k_4d, width, n_tps, n_stream, 1, - top_k_4d->nb[1], top_k_4d->nb[2], n_stream*top_k_4d->nb[3], 0); - ggml_tensor * zeros = ggml_fill(ctx0, - ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, width, n_tps, n_stream), 0.0f); - ggml_tensor * mask = ggml_set_rows(ctx0, base, zeros, idxs); - - ggml_build_forward_expand(gf, k_mla); - ggml_build_forward_expand(gf, r_kda); - ggml_build_forward_expand(gf, s_kda); - ggml_build_forward_expand(gf, mask); - - printf("note one graph reaches all three halves in %d nodes\n", ggml_graph_n_nodes(gf)); - - // 6. fill the pool map on a real allocated buffer - ggml_backend_t backend = ggml_backend_cpu_init(); - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx0, backend); - CHECK(buf != nullptr, "allocated the graph's input tensors on the CPU backend"); - - if (buf) { - const llama_ubatch & ub = mctx->get_ubatch(); - llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, kt.sel_mask, kt.cand_mask, - &ub, (uint32_t) kpool); - - const int32_t * cp = (const int32_t *) kt.cell_pool->data; - const int32_t * pc = (const int32_t *) kt.pool_cells->data; - const float * bi = (const float *) kt.bias->data; - const float * sm = (const float *) kt.sel_mask->data; - - bool in_range = true; - for (int64_t i = 0; i < ggml_nelements(kt.cell_pool); ++i) { - in_range &= cp[i] >= 0 && cp[i] < n_pools; - } - for (int64_t i = 0; i < ggml_nelements(kt.pool_cells); ++i) { - in_range &= pc[i] >= 0 && pc[i] < n_kv; - } - CHECK(in_range, "every emitted index is non-negative and in range (ggml_set_rows asserts i1 >= 0)"); - - bool finite = true; - for (int64_t i = 0; i < ggml_nelements(kt.bias); ++i) { - finite &= bi[i] == 0.0f || bi[i] == -INFINITY; - } - for (int64_t i = 0; i < ggml_nelements(kt.sel_mask); ++i) { - finite &= sm[i] == 0.0f || sm[i] == -INFINITY; - } - for (int64_t i = 0; i < ggml_nelements(kt.pool_bias); ++i) { - const float * pb = (const float *) kt.pool_bias->data; - finite &= pb[i] == 0.0f || pb[i] == -INFINITY; - } - for (int64_t i = 0; i < ggml_nelements(kt.cand_mask); ++i) { - const float * cm = (const float *) kt.cand_mask->data; - finite &= cm[i] == 0.0f || cm[i] == -INFINITY; - } - CHECK(finite, "bias, pool_bias, sel_mask and cand_mask hold only 0 and -INFINITY: " - "no +1e9 to meet a -inf"); - - // which is why the graph stores them in the KQ mask's f16 and halves - // two [n_kv, n_batch, n_stream] inputs. same fill, half the bytes - { - llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, - kt.sel_mask_f16, kt.cand_mask_f16, &ub, (uint32_t) kpool); - - CHECK(kpool_mask_f16_matches(kt.sel_mask, kt.sel_mask_f16) && - kpool_mask_f16_matches(kt.cand_mask, kt.cand_mask_f16), - "an f16 sel_mask/cand_mask is the exact image of the f32 one"); - } - - // cand_mask is exactly max(bias, sel_mask) lifted to KQ shape, and it - // is what stops an over-budget top-k from escaping the reference's - // candidate set. - // - // ggml_top_k always returns select_k pool ordinals even when fewer - // than select_k pools are finite. During prefill the query at - // position q has only ~q/kpool complete visible pools against a - // budget of index_topk/kpool, so the budget spills into -INFINITY - // pools and picks among them arbitrarily. Adding the causal mask - // kills the expansions that are empty or in the future. What it does - // not kill is a resident, causally visible cell that sits in an - // INCOMPLETE pool below the tail: the reference never selects it - // (pool_valid = grouped_valid_keys.all(-1)), the graph would. - // Unreachable while positions are contiguous, reachable the moment a - // partial seq_rm leaves a hole. - { - const float * cm = (const float *) kt.cand_mask->data; - - bool is_union = true; - int64_t n_spill = 0; // candidate pools strictly fewer than the budget - - const int64_t select_k = - llama_kpool_select_k((uint32_t) n_pools, hparams.indexer_top_k, (uint32_t) kpool); - - for (int64_t s = 0; s < n_stream; ++s) { - for (int64_t ii = 0; ii < n_tps; ++ii) { - const float * rb = bi + (s*n_tps + ii)*n_kv; - const float * rs = sm + (s*n_padq + ii)*n_kv; - const float * rc = cm + (s*n_padq + ii)*n_kv; - - for (int64_t j = 0; j < n_kv; ++j) { - is_union &= rc[j] == std::max(rb[j], rs[j]); - } - - const float * rp = (const float *) kt.pool_bias->data + (s*n_tps + ii)*n_pools; - int64_t n_cand = 0; - for (int64_t p = 0; p < n_pools; ++p) { - n_cand += rp[p] == 0.0f; - } - n_spill += n_cand < select_k; - } - } - - CHECK(is_union, "cand_mask == max(bias, sel_mask): the reference's candidate set"); - - // not a failure: it is the normal prefill state, and the whole - // reason the gate has to exist. Printed so that a future change - // making it zero cannot quietly turn the check above vacuous - printf("note %d of %d (query, stream) rows have fewer candidate pools than the\n" - " top-k budget of %d, so the budget spills and cand_mask is load bearing\n", - (int) n_spill, (int) (n_stream*n_tps), (int) select_k); - } - - // pool_bias is the same predicate as bias, evaluated where the - // reference evaluates it. For a COMPLETE pool the two must agree - // cell for cell; for an incomplete or absent pool bias has no cell - // to speak for it, which is exactly why pool_bias is computed here - // rather than gathered from bias at each pool's last member - { - const float * pb = (const float *) kt.pool_bias->data; - - bool agree = true; - bool exact = true; - - for (int64_t s = 0; s < n_stream; ++s) { - for (int64_t ii = 0; ii < n_tps; ++ii) { - const float * rb = bi + (s*n_tps + ii)*n_kv; - const float * rp = pb + (s*n_tps + ii)*n_pools; - - std::set pools_of_scored_cells; - - for (int64_t j = 0; j < n_kv; ++j) { - if (rb[j] == 0.0f) { - // a scored cell's pool must be a candidate pool - agree &= rp[cp[s*n_kv + j]] == 0.0f; - pools_of_scored_cells.insert(cp[s*n_kv + j]); - } - } - - int64_t n_cand = 0; - for (int64_t p = 0; p < n_pools; ++p) { - n_cand += rp[p] == 0.0f; - } - - // and nothing else may be one. an entirely absent pool - // would pass the first check vacuously; this is what - // catches it, and it is exactly the failure mode of - // gathering pool_bias from bias at each pool's last - // member instead of computing it - exact &= n_cand == (int64_t) pools_of_scored_cells.size(); - } - } - - CHECK(agree, "every cell bias scores lies in a pool pool_bias also accepts"); - CHECK(exact, "and pool_bias accepts no pool that has no scored cell " - "(an absent pool must not become a candidate)"); - } - - // per query: (q+1) %% kpool cells sit in its own incomplete pool and - // must be forced in, the q+1-minus-that below must be scored, and - // nothing else either. both streams, to cover the per-stream strides - { - const llama_ubatch & u = mctx->get_ubatch(); - - for (int64_t s = 0; s < n_stream; ++s) { - bool ok_tail = true; - bool ok_scored = true; - - for (int64_t ii = 0; ii < n_tps; ++ii) { - const llama_pos q = u.pos[s*n_tps + ii]; - const int64_t t = (q + 1) % kpool; - - const float * row_s = sm + (s*n_padq + ii)*n_kv; - const float * row_b = bi + (s*n_tps + ii)*n_kv; - - int64_t n_forced = 0; - int64_t n_scored = 0; - for (int64_t j = 0; j < n_kv; ++j) { - n_forced += row_s[j] == 0.0f; - n_scored += row_b[j] == 0.0f; - } - - ok_tail &= n_forced == t; - ok_scored &= n_scored == (q + 1) - t; - } - - CHECK(ok_tail, "stream %d: every query forces in exactly its (q+1) %% %d tail cells", - (int) s, (int) kpool); - CHECK(ok_scored, "stream %d: and scores exactly the cells below that tail", - (int) s); - } - } - - // while the window starts at position 0, pos/kpool and the reference's - // first-resident-key anchor are the same grouping, so the map must - // reproduce it. incomplete pools carry slot 0 and are masked instead - { - const llama_ubatch & u = mctx->get_ubatch(); - bool agree = true; - for (int64_t s = 0; s < n_stream; ++s) { - const llama_seq_id seq = u.seq_id[s*n_tps][0]; - const auto & cells = kv_attn->get_cells(seq); - - std::map members; - llama_pos p_min = -1; - for (int64_t j = 0; j < n_kv; ++j) { - if (cells.is_empty(j)) { - continue; - } - members[cells.pos_get(j)/kpool]++; - if (p_min < 0 || cells.pos_get(j) < p_min) { - p_min = cells.pos_get(j); - } - } - agree &= p_min == 0; // HF's first_key would also be 0 - - for (int64_t j = 0; j < n_kv && agree; ++j) { - if (cells.is_empty(j)) { - continue; - } - const int64_t b = cells.pos_get(j)/kpool; - agree &= cp[s*n_kv + j] == (int32_t) (members[b] == kpool ? b : 0); - } - } - CHECK(agree, "pool ordinals match the reference anchor while the window starts at position 0"); - } - - { - ggml_status st = ggml_backend_graph_compute(backend, gf); - CHECK(st == GGML_STATUS_SUCCESS, "the pooled-indexer-shaped graph computes (%d)", (int) st); - } - - ggml_backend_buffer_free(buf); - } - - ggml_backend_free(backend); - ggml_free(ctx0); - } - } - - // ---- pool origin across a front eviction -------------------------------- - // - // the reference anchors a pool at the first *resident* key (valid_keys.argmax(-1)); - // this port anchors at pos/kpool, as vLLM and SGLang do. same grouping until the - // window front is dropped by a non-multiple of kpool, when the reference regroups - // every surviving key and this port does not. a cache cannot afford regrouping: the - // pooled key a decode step scores was built during prefill. - { - printf("\n--- pool origin ---\n"); - - // seq 0 holds positions 0..9. drop 0..1, not a multiple of kpool = 4 - const llama_pos n_drop = 2; - mem->seq_rm(0, 0, n_drop); - - std::vector tok(1, 0); - std::vector pos(1, 10); - std::vector sid(1, 0); - std::vector sptr(1, sid.data()); - std::vector nsid(1, 1); - std::vector lg(1, 1); - - llama_batch b = {}; - b.n_tokens = 1; - b.token = tok.data(); - b.pos = pos.data(); - b.seq_id = sptr.data(); - b.n_seq_id = nsid.data(); - b.logits = lg.data(); - - llama_batch_allocr ba(hparams.n_pos_per_embd()); - if (ba.init(b, model->vocab, nullptr, hparams.n_embd_inp(), cparams.n_seq_max, true)) { - auto c = mem->init_batch(ba, cparams.n_ubatch, false); - auto * mc = dynamic_cast(c.get()); - - CHECK(mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, - "one more token fits after the eviction"); - - if (mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { - mc->apply(); - - const int64_t n_kv = mc->get_attn()->get_n_kv(); - const int64_t kpool = hparams.indexer_kpool; - const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); - - ggml_init_params gp = { ggml_tensor_overhead()*16, nullptr, true }; - ggml_context * ctx = ggml_init(gp); - - kpool_tensors kt = alloc_kpool_tensors(ctx, n_kv, 1, 1, 1, kpool, n_pools); - - ggml_backend_t backend = ggml_backend_cpu_init(); - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); - - if (buf) { - llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, kt.sel_mask, kt.cand_mask, - &mc->get_ubatch(), (uint32_t) kpool); - - const int32_t * cp = (const int32_t *) kt.cell_pool->data; - const auto & cells = kv_attn->get_cells(0); - - // positions 4..7 must stay one pool, or the pooled key built during - // prefill no longer describes the cells it is scored against - std::map slot_of; - for (int64_t j = 0; j < n_kv; ++j) { - if (!cells.is_empty(j) && cells.seq_has(j, 0)) { - slot_of[cells.pos_get(j)] = cp[j]; - } - } - - const bool grouped = - slot_of.count(4) && slot_of.count(5) && slot_of.count(6) && slot_of.count(7) && - slot_of[4] == slot_of[5] && slot_of[5] == slot_of[6] && slot_of[6] == slot_of[7]; - CHECK(grouped, "positions 4..7 stay one pool after the eviction (slot %d)", - grouped ? (int) slot_of[4] : -1); - - // the reference's anchor would make 2..5 the first pool; this port - // deliberately differs, so the difference is pinned rather than fixed - const bool differs = !slot_of.count(2) || slot_of[2] != slot_of[4]; - CHECK(differs, "positions 2 and 4 are NOT pooled together, where the reference anchor would"); - - // the leading remnant 2..3 is an incomplete pool and not in the - // query's tail, the one case where a visible cell is neither - const float * bi = (const float *) kt.bias->data; - const float * sm = (const float *) kt.sel_mask->data; - int64_t n_orphan = 0; - for (int64_t j = 0; j < n_kv; ++j) { - if (cells.is_empty(j) || !cells.seq_has(j, 0)) { - continue; - } - n_orphan += bi[j] == -INFINITY && sm[j] == -INFINITY; - } - CHECK(n_orphan == 2, "the %d-cell leading remnant is neither pooled nor tail (%d)", - (int) n_drop, (int) n_orphan); - - ggml_backend_buffer_free(buf); - } - - ggml_backend_free(backend); - ggml_free(ctx); - } - } - } - - delete mem_raw; - - // ---- the shared input: one fill per ubatch, not one per DSA layer ------- - { - printf("\n--- shared input ---\n"); - - llama_cparams cp = {}; - cp.n_ctx = 16384; - cp.n_ctx_seq = 16384; - cp.n_batch = 512; - cp.n_ubatch = 512; - cp.n_seq_max = 1; - cp.kv_unified = false; - cp.causal_attn = true; - - llama_memory_params mp = {}; - mp.type_k = GGML_TYPE_F16; - mp.type_v = GGML_TYPE_F16; - mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; - - llama_memory_i * raw = model->create_memory(mp, cp); - auto * m = dynamic_cast(raw); - - if (m) { - llama_kv_cache * kv = m->get_mem_attn(); - - const int64_t n_step = cp.n_ubatch; - const int64_t n_fill = cp.n_ctx_seq - n_step; - const uint32_t n_dsa = (uint32_t) kv->get_layer_ids().size(); - const int64_t kpool = hparams.indexer_kpool; - - std::vector tok(n_step, 0); - std::vector pos(n_step); - std::vector sid(n_step, 0); - std::vector sptr(n_step); - std::vector nsid(n_step, 1); - std::vector lg(n_step, 0); - - for (int64_t i = 0; i < n_step; ++i) { - sptr[i] = &sid[i]; - } - lg.back() = 1; - - llama_memory_context_ptr keep; - - for (int64_t base = 0; base <= n_fill; base += n_step) { - for (int64_t i = 0; i < n_step; ++i) { - pos[i] = (llama_pos) (base + i); - } - - llama_batch b = {}; - b.n_tokens = (int32_t) n_step; - b.token = tok.data(); - b.pos = pos.data(); - b.seq_id = sptr.data(); - b.n_seq_id = nsid.data(); - b.logits = lg.data(); - - llama_batch_allocr ba(hparams.n_pos_per_embd()); - if (!ba.init(b, model->vocab, nullptr, hparams.n_embd_inp(), cp.n_seq_max, true)) { - break; - } - - auto c = m->init_batch(ba, cp.n_ubatch, false); - if (!c || c->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { - break; - } - c->apply(); - keep = std::move(c); - } - - auto * mctx = dynamic_cast(keep.get()); - CHECK(mctx && mctx->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, - "filled a %d-cell cache in %d-token ubatches", (int) cp.n_ctx_seq, (int) cp.n_ubatch); - if (mctx && mctx->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { - const int64_t n_kv = mctx->get_attn()->get_n_kv(); - const int64_t n_tps = mctx->get_ubatch().n_tokens; - const int64_t n_padq = GGML_PAD(n_tps, 8) + 8; // exercise the padded-mask path - const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); - - ggml_init_params gp = { - /* .mem_size */ ggml_tensor_overhead()*16, - /* .mem_buffer */ nullptr, - /* .no_alloc */ true, - }; - - ggml_context * ctx = ggml_init(gp); - kpool_tensors kt = alloc_kpool_tensors(ctx, n_kv, n_tps, n_padq, 1, kpool, n_pools); - - ggml_backend_t backend = ggml_backend_cpu_init(); - ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); - - if (buf) { - const llama_ubatch & ub = mctx->get_ubatch(); - - // first pass faults in 60+ MiB of fresh pages; time the steady - // state, which is what a prefill pays per ubatch - double ms = 1e9; - for (int rep = 0; rep < 4; ++rep) { - const auto t0 = std::chrono::steady_clock::now(); - llama_kv_cache_set_input_kpool(kv, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, kt.sel_mask, kt.cand_mask, - &ub, (uint32_t) kpool); - const auto t1 = std::chrono::steady_clock::now(); - - if (rep > 0) { - ms = std::min(ms, std::chrono::duration(t1 - t0).count()); - } - } - - printf("note n_kv = %d, n_tokens = %d, %u DSA layers\n", (int) n_kv, (int) n_tps, n_dsa); - printf("note one fill = %.2f ms; shared = %.2f ms/ubatch, per layer = %.2f ms/ubatch\n", - ms, ms, ms*n_dsa); - printf("note extrapolated to 128Ki x 512 x 11 layers: %.0f ms shared, %.0f ms per layer\n", - ms*(131072.0/(double) n_kv)*(512.0/(double) n_tps), - ms*(131072.0/(double) n_kv)*(512.0/(double) n_tps)*11.0); - - CHECK(n_dsa > 1, "the model has %u indexer layers, so sharing one fill saves %ux the host writes", - n_dsa, n_dsa); - - { - const float * sm = (const float *) kt.sel_mask ->data; - const float * cm = (const float *) kt.cand_mask->data; - - bool pad_masked = true; - for (int64_t ii = n_tps; ii < n_padq; ++ii) { - for (int64_t j = 0; j < n_kv; ++j) { - pad_masked &= sm[ii*n_kv + j] == -INFINITY; - pad_masked &= cm[ii*n_kv + j] == -INFINITY; - } - } - CHECK(pad_masked, "the %d padding rows of a wider sel_mask/cand_mask stay -INFINITY", - (int) (n_padq - n_tps)); - } - - // the padding fill has its own code path per mask type - { - llama_kv_cache_set_input_kpool(kv, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, - kt.sel_mask_f16, kt.cand_mask_f16, &ub, (uint32_t) kpool); - - CHECK(kpool_mask_f16_matches(kt.sel_mask, kt.sel_mask_f16) && - kpool_mask_f16_matches(kt.cand_mask, kt.cand_mask_f16), - "the f16 masks match the f32 ones on a padded ubatch too"); - } - - // the shared object refills the same tensors deterministically. - // it fills only what the pooled graph consumes - cell_pool and - // bias have no consumer there, so it passes nullptr for both and - // pool_cells is what gets compared - std::vector first((size_t) ggml_nelements(kt.pool_cells)); - memcpy(first.data(), kt.pool_cells->data, first.size()*sizeof(int32_t)); - - llm_graph_input_kpool inp(mctx->get_attn(), mctx->get_idx(), (uint32_t) kpool); - inp.k_idxs = mctx->get_idx()->build_input_k_idxs(ctx, ub); - inp.pool_cells = kt.pool_cells; - inp.pool_bias = kt.pool_bias; - inp.sel_mask = kt.sel_mask; - inp.cand_mask = kt.cand_mask; - - ggml_backend_buffer_t buf2 = ggml_backend_alloc_ctx_tensors(ctx, backend); - if (buf2) { - inp.set_input(&ub); - - CHECK(memcmp(first.data(), kt.pool_cells->data, first.size()*sizeof(int32_t)) == 0, - "llm_graph_input_kpool::set_input rebuilds the identical map, so one\n" - " object can back every indexer layer"); - - ggml_backend_buffer_free(buf2); - } - - ggml_backend_buffer_free(buf); - } - - ggml_backend_free(backend); - ggml_free(ctx); - } else { - printf("note could not fill a %d-cell cache; skipping the cost measurement\n", (int) cp.n_ctx_seq); - } - - keep.reset(); - } - - delete raw; - } - - llama_model_free(model); - llama_backend_free(); - - printf("\n%s: %d failure(s)\n", argv[0], n_fail); - return n_fail == 0 ? 0 : 1; -} diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index ce2febd4209..b8fd66ccae5 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -116,12 +116,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_KIMI_LINEAR - || arch == LLM_ARCH_GLM5NEXT || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_MISTRAL4) { - // MLA absorbs into MQA, so n_head_kv must be 1: otherwise the per-layer - // head_count_kv array below sizes the latent K row n_head times too wide n_embd = 128; n_head = 1; n_ff = 192; @@ -168,7 +165,7 @@ 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_GLM5NEXT) { + arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3) { GGML_ASSERT(n_layer >= 2); std::vector n_head_per_layer; n_head_per_layer.reserve(n_layer); @@ -216,21 +213,6 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types); } - } else if (arch == LLM_ARCH_GLM5NEXT) { - // nope-only MLA: rope dimension count must be written as 0, not omitted, or - // the generic loader defaults it non-zero and load_arch_hparams rejects it - ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(512)); - ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512)); - ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(0)); - ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192)); - ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128)); - ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, uint32_t(4)); - // build_hc_pre hard-codes a 4-wide residual - 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, 1e-6f); - ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); - ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } else if (arch == LLM_ARCH_MINIMAX_M3) { // partial rotary: n_rot must not exceed the indexer key length (64) ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); @@ -462,7 +444,6 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_PADDLEOCR: case LLM_ARCH_MIMO2: case LLM_ARCH_KIMI_LINEAR: - case LLM_ARCH_GLM5NEXT: case LLM_ARCH_KIMI_K3: case LLM_ARCH_STEP35: case LLM_ARCH_MISTRAL4: diff --git a/tests/test-mtmd-impl.cpp b/tests/test-mtmd-impl.cpp index f6b908344b1..df18b0a42e2 100644 --- a/tests/test-mtmd-impl.cpp +++ b/tests/test-mtmd-impl.cpp @@ -70,107 +70,6 @@ MAKE_TEST(test_image_preprocessor_lfm2) { } } -// GLM-5-Next / GLM-5.3-Flash, hparams as loaded by clip.cpp for PROJECTOR_TYPE_GLM5NEXT -static clip_hparams glm5next_hparams() { - clip_hparams hparams; - hparams.patch_size = 14; - hparams.n_merge = 2; - hparams.set_limit_image_tokens(16, 8000); - return hparams; -} - -// set_limit_image_tokens takes token counts; for a still image the reference's temporal factors cancel, -// leaving factor**2 == 28*28 -MAKE_TEST(test_image_preprocessor_glm5next_budget) { - const clip_hparams hparams = glm5next_hparams(); - - t.assert_equal("image_min_pixels", 16 * 28 * 28, hparams.image_min_pixels); - t.assert_equal("image_max_pixels", 8000 * 28 * 28, hparams.image_max_pixels); -} - -MAKE_TEST(test_image_preprocessor_glm5next_resize) { - const clip_hparams hparams = glm5next_hparams(); - - struct test_case { - clip_image_size input; - clip_image_size canvas; // padded output - clip_image_size content; // resized image inside the canvas - int n_tokens; - }; - - // expected values come from Glm5NextImageProcessor.resize in - // adapt_zips/extract_0826/image_processing_glm5_next.py - const std::vector cases = { - // inside the budget - { { 224, 224 }, { 224, 224 }, { 224, 224 }, 64 }, - { { 448, 448 }, { 448, 448 }, { 448, 448 }, 256 }, - { { 1024, 1024 }, { 1036, 1036 }, { 1024, 1024 }, 1369 }, - // the 16-token floor: 112x112 is exactly 16 tokens - { { 112, 112 }, { 112, 112 }, { 112, 112 }, 16 }, - { { 111, 111 }, { 112, 112 }, { 111, 111 }, 16 }, - { { 113, 113 }, { 140, 140 }, { 113, 113 }, 25 }, - { { 56, 56 }, { 112, 112 }, { 112, 112 }, 16 }, - { { 50, 50 }, { 112, 112 }, { 112, 112 }, 16 }, - { { 28, 28 }, { 112, 112 }, { 112, 112 }, 16 }, - { { 1, 1 }, { 112, 112 }, { 112, 112 }, 16 }, - // the 8000-token ceiling: 2492x2492 still fits, 2520x2520 does not - { { 2492, 2492 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, - { { 2493, 2493 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, - { { 2504, 2504 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, - { { 2520, 2520 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, - { { 4000, 4000 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, - // wide and tall - { { 800, 600 }, { 812, 616 }, { 800, 600 }, 638 }, - { { 600, 800 }, { 616, 812 }, { 600, 800 }, 638 }, - { { 1920, 480 }, { 1932, 504 }, { 1920, 480 }, 1242 }, - { { 480, 1920 }, { 504, 1932 }, { 480, 1920 }, 1242 }, - { { 1920, 1080 }, { 1932, 1092 }, { 1920, 1080 }, 2691 }, - { { 1080, 1920 }, { 1092, 1932 }, { 1080, 1920 }, 2691 }, - // extreme aspect ratios, single-pixel edges survive the resize - { { 4000, 16 }, { 4004, 28 }, { 4000, 16 }, 143 }, - { { 16, 4000 }, { 28, 4004 }, { 16, 4000 }, 143 }, - { { 5000, 1 }, { 5012, 28 }, { 5012, 1 }, 179 }, - { { 1, 5000 }, { 28, 5012 }, { 1, 5012 }, 179 }, - { { 12000, 100 }, { 12012, 112 }, { 12000, 100 }, 1716 }, - { { 100, 12000 }, { 112, 12012 }, { 100, 12000 }, 1716 }, - // over budget, neither edge a multiple of 28 - { { 4007, 3001 }, { 2884, 2156 }, { 2878, 2156 }, 7931 }, - { { 3001, 4007 }, { 2156, 2884 }, { 2156, 2878 }, 7931 }, - { { 3333, 5000 }, { 2044, 3052 }, { 2034, 3052 }, 7957 }, - { { 2729, 2731 }, { 2492, 2492 }, { 2490, 2492 }, 7921 }, - // the 0826 binary search and a Qwen-style smart_resize disagree on all of these, so a regression - // to the old dynamic-size preprocessor cannot pass - { { 4618, 2282 }, { 3528, 1764 }, { 3528, 1743 }, 7938 }, // smart_resize: 3556x1736 - { { 2794, 6096 }, { 1708, 3668 }, { 1681, 3668 }, 7991 }, // smart_resize: 1680x3696 - { { 8858, 1315 }, { 6412, 952 }, { 6412, 951 }, 7786 }, // smart_resize: 6496x952 - { { 1350, 5856 }, { 1204, 5208 }, { 1200, 5208 }, 7998 }, // smart_resize: 1176x5208 - { { 2134, 1472 }, { 2156, 1484 }, { 2134, 1472 }, 4081 }, // smart_resize: 2128x1484 - { { 1021, 4268 }, { 1036, 4284 }, { 1021, 4268 }, 5661 }, // smart_resize: 1008x4256 - }; - - auto fmt = [](const clip_image_size & s) { - return std::to_string(s.width) + "x" + std::to_string(s.height); - }; - - for (const auto & tc : cases) { - const auto geo = mtmd_image_preprocessor_glm5next::get_geometry(hparams, tc.input); - const auto name = " for " + fmt(tc.input); - - t.assert_equal("canvas" + name, fmt(tc.canvas), fmt(geo.canvas)); - t.assert_equal("content" + name, fmt(tc.content), fmt(geo.content)); - - // mirrors clip_n_output_tokens for PROJECTOR_TYPE_GLM5NEXT - const int n_tokens = (geo.canvas.width / (hparams.patch_size * hparams.n_merge)) - * (geo.canvas.height / (hparams.patch_size * hparams.n_merge)); - t.assert_equal("n_tokens" + name, tc.n_tokens, n_tokens); - - t.assert_true("content fits the canvas" + name, - geo.content.width <= geo.canvas.width && geo.content.height <= geo.canvas.height); - t.assert_true("canvas is within the token budget" + name, - geo.canvas.width * geo.canvas.height <= hparams.image_max_pixels); - } -} - // // mtmd temporal merge // From 6dd91866f44c6e8d88b1d725fa89ddd535da3499 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Fri, 28 Aug 2026 01:49:51 +0000 Subject: [PATCH 30/36] glm5next: cut comment volume by 72% Deletes comments that restate the code, section banners, paragraph spacers, pointers, development narration and measured numbers that belong in the PR description. What survives is limited to correctness constraints, reference implementation citations, warnings that a tempting alternative is wrong, and explanations of real bugs - each stated in one or two lines. Comments only: verified by stripping every comment from each file and comparing the normalised source against the pre-pass baseline. --- examples/embedding/embedding.cpp | 1 - src/llama-context.cpp | 8 +- src/llama-graph.cpp | 71 ++----- src/llama-graph.h | 30 +-- src/llama-kv-cache-kpool.cpp | 129 +++---------- src/llama-kv-cache-kpool.h | 177 +++-------------- src/llama-kv-cache.h | 7 +- src/llama-memory-hybrid.cpp | 33 ++-- src/llama-memory-hybrid.h | 7 +- src/llama-model.cpp | 14 +- src/llama-quant.cpp | 24 +-- src/llama-vocab.cpp | 9 +- src/models/delta-net-base.cpp | 7 +- src/models/glm5next.cpp | 268 +++++--------------------- src/models/models.h | 22 +-- tools/mtmd/clip.cpp | 11 +- tools/mtmd/models/glm4v.cpp | 3 +- tools/mtmd/models/glm5next-vision.cpp | 8 +- tools/mtmd/mtmd-image.cpp | 11 +- tools/mtmd/mtmd-image.h | 9 +- 20 files changed, 156 insertions(+), 693 deletions(-) diff --git a/examples/embedding/embedding.cpp b/examples/embedding/embedding.cpp index a30d271253f..f1b4df32e43 100644 --- a/examples/embedding/embedding.cpp +++ b/examples/embedding/embedding.cpp @@ -144,7 +144,6 @@ int main(int argc, char ** argv) { return 1; } - // a context can fail on its own, and llama_n_ctx below dereferences it if (ctx == NULL) { LOG_ERR("%s: unable to create context\n", __func__); return 1; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 6ed803bc58a..0c65305810d 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -237,8 +237,7 @@ llama_context::llama_context( cparams.auto_flid = true; { - // escape hatch: the fused kernel sums the heads in a different order than the - // unfused graph, so a near-tied indexer top-k can come out different + // the fused kernel sums heads in a different order, so a near-tied top-k can differ const char * LLAMA_FUSED_LID_DISABLE = getenv("LLAMA_FUSED_LID_DISABLE"); if (LLAMA_FUSED_LID_DISABLE && atoi(LLAMA_FUSED_LID_DISABLE) != 0) { cparams.fused_lid = false; @@ -2304,9 +2303,8 @@ 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 || model.arch == LLM_ARCH_GLM5NEXT) { - // the n_tokens*40 budget below is exhausted at ubatch 3840 for kimi-k3, and - // earlier for glm5next: each KDA layer costs 182 nodes + ~16/token, so its 34 - // KDA layers alone need 6.2k + 31.9*n_tokens before DSA or the MoE + // the n_tokens*40 budget below runs out by ubatch 3840: KDA costs 182 nodes + ~16/token + // per layer, so 34 KDA layers alone need 6.2k + 31.9*n_tokens before DSA or the MoE res = std::max(n_tokens * 160, 64u * model.n_tensors()); } else if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 1bb06702a8a..e7ab7fc7bec 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3574,19 +3574,11 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( if (scoring) { const int64_t n_kv = mctx_attn->get_n_kv(); - // spelled exactly as build_attn_inp_kq_mask spells it, so that sel_mask, - // cand_mask and the KQ mask are the same shape and add without a - // broadcast. get_n_stream() is the cache's stream RANGE and is a - // different number as soon as a server has non-contiguous slots busy + // must match build_attn_inp_kq_mask; get_n_stream() is the stream RANGE and is wrong const int64_t n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq; const int64_t n_tps = ubatch.n_tokens/n_stream; - // one pool map per SEQUENCE. a non-unified cache has one sequence per stream, a - // unified cache puts every sequence of the ubatch in stream 0 and cuts the - // stream's pool table into one run per sequence, so the table needs the rebasing - // slack once per sequence. sized on the ubatch and not on n_seq_max, which - // llama-embedding sets to 256; n_seqs_unq is already part of - // llm_graph_params::allow_reuse, so the shape holds while a graph is reused + // pool maps are per SEQUENCE; sized on the ubatch, not n_seq_max (256 in llama-embedding) const int64_t n_ps = (int64_t) ubatch.n_seqs_unq/n_stream; GGML_ASSERT(n_ps >= 1 && (int64_t) ubatch.n_seqs_unq == n_ps*n_stream); @@ -3595,9 +3587,7 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[3] == n_stream); - // this tree's build_attn_inp_kq_mask has no GGML_KQ_MASK_PAD, so the mask - // is exactly n_tps rows. the host side supports n_padq > n_tps, the graph - // below does not: the selection terms only exist for real queries + // the selection terms below exist only for real queries GGML_ASSERT(kq_mask->ne[1] == n_tps && "the pooled indexer needs an unpadded KQ mask"); inp->pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool*n_pools, n_stream); @@ -3608,8 +3598,7 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( ggml_set_input(inp->pool_bias); ggml_set_name(inp->pool_bias, "kpool_pool_bias"); - // the fused lightning indexer wants an f16 mask. built here, once, because - // every indexer layer shares it + // the fused indexer wants f16; built once, shared by every indexer layer if (cparams.fused_lid) { inp->pool_bias_f16 = ggml_cast(ctx0, ggml_reshape_4d(ctx0, inp->pool_bias, n_pools, n_tps, 1, n_stream), @@ -3617,17 +3606,7 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( ggml_set_name(inp->pool_bias_f16, "kpool_pool_bias_f16"); } - // f16, not f32. The only two values either mask holds are 0.0f and - // -INFINITY, both exact in f16, so the narrower type is lossless and halves - // two [n_kv, n_tps, n_stream] inputs that live for the whole ubatch: 2 GiB - // each at n_ctx = 1 Mi, n_ubatch = 512. - // - // Every consumer takes f16. Under flash attention this is the KQ mask's own - // type, so build_attn_sparse adds the two with no conversion at all. With - // flash attention off - which GLM-5.3-Flash requires, so it is the path that - // matters here - the KQ mask is f32, and ggml_add gives its result src0's - // type: f16 + f32 -> f16 is a supported bin_bcast on CUDA and on the CPU, - // and ggml_soft_max_ext takes an f16 mask as readily as an f32 one + // lossless in f16 (only 0.0f and -INFINITY), and f16 + f32 -> f16 adds the KQ mask uncast inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, n_kv, n_tps, 1, n_stream); ggml_set_input(inp->sel_mask); ggml_set_name(inp->sel_mask, "kpool_sel_mask"); @@ -3677,20 +3656,7 @@ ggml_tensor * llm_graph_context::build_attn_sparse( GGML_ASSERT(sel_mask->ne[0] == kq_mask->ne[0] && sel_mask->ne[1] == kq_mask->ne[1] && sel_mask->ne[3] == kq_mask->ne[3]); - // The dense DSA path (build_attn on llm_graph_input_attn_k_dsa) opens with - // ggml_fill(kq_mask, -INFINITY). Here the scatter starts from sel_mask - // instead, which already holds 0.0 on the query's own trailing incomplete - // pool: GLM always attends to that tail (index_kpool_always_select_tail), and - // keeping it out of the top-k budget is what lets the budget stay a whole - // number of pools. - // - // ggml_set_rows writes THROUGH to its destination and returns a view of it, - // and sel_mask is one shared per-ubatch input read by every indexer layer. - // Scattering into it directly makes each layer inherit the previous layer's - // unmasked cells: measured on TinySparse, layer 3 stayed inside its budget - // while layer 7, running second, reached 411 cells. The dense path never - // meets this because ggml_fill hands it a fresh tensor every layer. Take a - // private copy per layer. + // ggml_set_rows writes THROUGH, and sel_mask is shared per ubatch: scatter into a copy ggml_tensor * mask_all = ggml_dup(ctx0, sel_mask); // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] @@ -3701,15 +3667,8 @@ ggml_tensor * llm_graph_context::build_attn_sparse( ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[2], 1, top_k->nb[1], top_k->nb[2], top_k->ne[2]*top_k->nb[2], 0); - // A constant 0, never the cell's own bias. The scatter must not be able to - // ERASE a zero sel_mask already granted: a tail cell can also be named by the - // top-k (through an over-budget pool whose unfilled slots point at it), and - // scattering that cell's -inf score bias would leave the query attending to - // nothing at all. Rejecting an over-budget selection is done additively, - // below, by cand_mask. - // - // f32 whatever mask_all is: ggml_set_rows converts its values into the - // destination, and the CUDA backend only advertises SET_ROWS for f32 values + // a constant 0, never the cell's bias: scattering -inf would ERASE a zero granted to the + // tail (cand_mask rejects over-budget picks below). f32: CUDA only does SET_ROWS for f32 ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); zeros = ggml_fill(ctx0, zeros, 0.0f); @@ -3719,23 +3678,15 @@ ggml_tensor * llm_graph_context::build_attn_sparse( mask_top_k = ggml_view_4d(ctx0, mask_top_k, mask_top_k->ne[1], mask_top_k->ne[2], 1, mask_top_k->ne[3], mask_top_k->nb[2], mask_top_k->nb[3], mask_top_k->nb[3], 0); - // the reference's `selected_valid` gather, additively: a cell the top-k named - // that is not in the reference's candidate set goes back to -inf, and a tail - // cell that sel_mask granted stays at 0 because cand_mask is the UNION of the - // candidates and the tail + // the reference's `selected_valid` gather, additively; cand_mask is candidates UNION tail mask_top_k = ggml_add(ctx0, mask_top_k, cand_mask); - // ggml_add gives its result src0's type, so an f16 selection mask absorbs an - // f32 KQ mask and no cast is needed. The one direction that still needs one is - // an f32 mask meeting an f16 KQ mask: the add would yield f32 and - // ggml_flash_attn_ext asserts its mask is f16 + // ggml_flash_attn_ext asserts an f16 mask, and ggml_add would yield src0's f32 if (mask_top_k->type == GGML_TYPE_F32 && kq_mask->type == GGML_TYPE_F16) { mask_top_k = ggml_cast(ctx0, mask_top_k, GGML_TYPE_F16); } - // and finally re-apply causality, occupancy and padding. load bearing: it is - // what keeps an empty, future or foreign-sequence cell masked no matter what - // the top-k returned + // load bearing: keeps an empty, future or foreign-sequence cell masked whatever top-k said mask_top_k = ggml_add(ctx0, mask_top_k, kq_mask); cb(mask_top_k, "kpool_kq_mask", il); diff --git a/src/llama-graph.h b/src/llama-graph.h index c994a1774d5..258a66dd53e 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -32,8 +32,7 @@ class llama_memory_recurrent_context; class llama_memory_hybrid_context; class llama_memory_hybrid_iswa_context; -// defined in llama-kv-cache-kpool.h, which includes this header, so it can only -// be forward declared here +// defined in llama-kv-cache-kpool.h, which includes this header, so forward declared only class llm_graph_input_kpool; // certain models (typically multi-modal) can produce different types of graphs @@ -1348,36 +1347,13 @@ struct llm_graph_context { llm_graph_input_mem_hybrid_iswa * build_inp_mem_hybrid_iswa() const; - // - // pooled (GLM-5-Next lightning) indexer - // - - // one pooling map per ubatch, shared by every indexer layer. see - // llama-kv-cache-kpool.h for what the tensors mean and why they are built - // host side rather than derived in the graph. - // - // `scoring` false allocates only k_idxs: the indexer key and gate store is - // unconditional, the selection is not. An input tensor with no consumer is - // never backed by the allocator, so the rest must not be created either + // one pooling map per ubatch (see llama-kv-cache-kpool.h); `scoring` false gives only k_idxs llm_graph_input_kpool * build_inp_kpool( const llama_memory_hybrid_context * mctx_cur, ggml_tensor * kq_mask, bool scoring) const; - // sparse (pooled top-k) variant of the llm_graph_input_attn_k build_attn. - // - // Identical to it except for the mask. `top_k` names the cells the pooled - // indexer selected, already expanded from whole pools; they are unmasked on - // top of `sel_mask`, which arrives already holding 0.0f on the query's own - // always-selected trailing pool. `cand_mask` is the reference's candidate - // set and is what makes an over-budget selection harmless: ggml_top_k - // returns a full budget of pool ordinals even when fewer pools carry a - // finite score, which during prefill is the normal state. - // - // Both masks may be f16 or f32, and build_inp_kpool allocates f16: the only - // values either holds are 0.0f and -INFINITY, both exact in f16. The combined - // mask inherits their type, which the KQ mask then adds into whatever its own - // type is + // build_attn, but masking with `top_k` over `sel_mask`; `cand_mask` drops over-budget picks ggml_tensor * build_attn_sparse( llm_graph_input_attn_k * inp, ggml_tensor * wo, diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index 9f214dd16cd..3fc1eaca616 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -20,12 +20,10 @@ uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t GGML_ASSERT(n_pools > 0); GGML_ASSERT(indexer_top_k % kpool == 0 && "indexer_top_k must be a whole number of pools"); - // min(index_topk // index_kpool, n_pools), exactly the reference's select_k return std::min(n_pools, indexer_top_k/kpool); } -// sel_mask and cand_mask hold only 0.0f and -INFINITY, so they can be written in the -// KQ mask's f16 exactly as in f32 +// sel_mask and cand_mask hold only 0.0f and -INFINITY, so f16 is exact here template struct kpool_mask_of; template <> struct kpool_mask_of { @@ -41,8 +39,6 @@ static void kpool_mask_fill(T * dst, int64_t n) { std::fill(dst, dst + n, kpool_mask_of::from(-INFINITY)); } -// one query's row of both masks; the two predicates share their operands, and the -// unsigned compares are what let this vectorise template static void kpool_mask_row( T * cur_sel, @@ -62,8 +58,7 @@ static void kpool_mask_row( const bool tail = pos_at[j] >= tail_start; cur_sel [j] = vis && tail ? v_sel : v_mask; - // max(bias, sel_mask): the reference's candidate set, which the - // top-k budget may overrun but must never escape + // the candidate set, which the top-k budget may overrun but must never escape cur_cand[j] = vis && (pooled || tail) ? v_sel : v_mask; } } @@ -81,41 +76,33 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(kv != nullptr); GGML_ASSERT(kpool > 0); - // the per-CELL view is optional: the pooled graph does not consume it, and an - // input tensor with no consumer is never backed by the allocator GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(pool_bias ->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(sel_mask ->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(cand_mask ->buffer)); - // both masks are written through raw strides below, so writing the wrong width - // would overrun the allocation 2x. check rather than trust GGML_ASSERT(pool_cells->type == GGML_TYPE_I32); GGML_ASSERT(pool_bias ->type == GGML_TYPE_F32); GGML_ASSERT((sel_mask->type == GGML_TYPE_F16 || sel_mask->type == GGML_TYPE_F32) && "sel_mask must be f16 or f32"); GGML_ASSERT(cand_mask->type == sel_mask->type && "both masks must have the KQ mask's type"); - // everything below is written through raw strides GGML_ASSERT(ggml_is_contiguous(pool_cells)); GGML_ASSERT(ggml_is_contiguous(pool_bias)); GGML_ASSERT(ggml_is_contiguous(sel_mask)); GGML_ASSERT(ggml_is_contiguous(cand_mask)); const int64_t n_kv = sel_mask->ne[0]; - const int64_t n_ns = sel_mask->ne[3]; // streams in this ubatch + const int64_t n_ns = sel_mask->ne[3]; const int64_t r = kpool; const int64_t n_tokens = ubatch->n_tokens; - // [TAG_KPOOL_SEQ_PARTITION] - // positions are unambiguous only within one sequence, so one pool map per SEQUENCE, - // not per stream. a non-unified cache gives each stream its own cells array and one - // sequence (n_ps == 1, the layout this file had before), a unified cache gives one - // stream carrying every sequence in the ubatch + // [TAG_KPOOL_SEQ_PARTITION] positions are unambiguous only within one sequence, so + // one pool map per SEQUENCE, not per stream GGML_ASSERT(n_ns == 1 || (int64_t) ubatch->n_seqs_unq == n_ns); - const int64_t n_ps = (int64_t) ubatch->n_seqs_unq/n_ns; // sequences per stream - const int64_t n_pools = pool_cells->ne[0]/r; // pool slots per stream + const int64_t n_ps = (int64_t) ubatch->n_seqs_unq/n_ns; + const int64_t n_pools = pool_cells->ne[0]/r; GGML_ASSERT(n_ps > 0 && (int64_t) ubatch->n_seqs_unq == n_ns*n_ps); GGML_ASSERT(pool_cells->ne[0] % r == 0); @@ -126,8 +113,8 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(pool_bias->ne[0] == n_pools && pool_bias->ne[2] == n_ns); GGML_ASSERT(n_tokens % n_ns == 0); - const int64_t n_tps = n_tokens/n_ns; // tokens per stream - const int64_t n_padq = sel_mask->ne[1]; // KQ mask rows, >= n_tps + const int64_t n_tps = n_tokens/n_ns; + const int64_t n_padq = sel_mask->ne[1]; GGML_ASSERT(pool_bias->ne[1] == n_tps); GGML_ASSERT(n_padq >= n_tps); @@ -138,8 +125,7 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(ggml_is_contiguous(cell_pool)); GGML_ASSERT(cell_pool->ne[0] == n_kv && cell_pool->ne[1] == n_ns); - // one row per stream, so a cell that two sequences of one stream share has - // nowhere to put its second pool. the graph never asks for this view + // one row per stream, so a shared cell has nowhere to put its second pool GGML_ASSERT(n_ps == 1 && "the per-cell pool view needs one sequence per stream"); } @@ -160,19 +146,14 @@ void llama_kv_cache_set_input_kpool( const bool mask_f16 = sel_mask->type == GGML_TYPE_F16; const size_t mask_ts = ggml_type_size(sel_mask->type); - // -1 marks a cell with no usable pool. host side only: never copied into cell_pool, - // where ggml_get_rows would read it as an index + // -1 marks a cell with no usable pool; host side only, never copied into cell_pool std::vector pool_of(n_kv); std::vector filled(n_pools); std::vector pos_at; - // one contiguous run of pool slots per sequence of the stream std::vector run_off(n_ps); std::vector run_len(n_ps); - // which cells array a sequence of this stream uses; same convention as - // llama_kv_cache::set_input_kq_mask. with one sequence per stream the ubatch's - // unique list and the stream's own sequence are the same thing auto seq_of = [&](int64_t s, int64_t ps) { return n_ps == 1 ? ubatch->seq_id[s*n_tps][0] : ubatch->seq_id_unq[ps]; }; @@ -183,9 +164,6 @@ void llama_kv_cache_set_input_kpool( char * cur_cand_mask = dst_cand_mask + s*(n_padq*n_kv)*mask_ts; float * cur_pool_bias = dst_pool_bias + s*(n_tps*n_pools); - // slots of a pool that is not resident, and slots outside the query's own - // sequence run, are cleared once per stream here. the per-sequence pass below - // only writes what it owns std::fill(cur_pool_cells, cur_pool_cells + r*n_pools, 0); std::fill(cur_pool_bias, cur_pool_bias + n_tps*n_pools, -INFINITY); @@ -198,17 +176,11 @@ void llama_kv_cache_set_input_kpool( kpool_mask_fill((float *) (cur_cand_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); } - // [TAG_KPOOL_PACK] - // cut the stream's pool table into one run per sequence, sized on the pool range - // that sequence actually holds. the table is n_kv/kpool shared plus 2 per sequence - // for rebasing, which covers it whenever the sequences' cells are disjoint - every - // case but a prefix shared through llama_memory_seq_cp. that one can ask for more - // slots than exist, and then a sequence keeps its newest pools, the same cut a - // large hole already forces. - // + // [TAG_KPOOL_PACK] one packed run per sequence, sized on the pool range it holds. // NOT one full-width table per sequence: the indexer scores every slot against - // every query, so a full-width table would multiply the score tensor by the - // sequence count, and llama-embedding asks for n_seq_max 256 + // every query, so that multiplies the score tensor by n_seq_max. + // llama_memory_seq_cp can ask for more slots than exist; then a sequence keeps its + // newest pools, the same cut a large hole already forces. { int64_t n_want = 0; @@ -266,27 +238,14 @@ void llama_kv_cache_set_input_kpool( std::fill(pool_of.begin(), pool_of.end(), -1); std::fill(filled.begin(), filled.end(), 0); - // hoist occupancy and sequence membership out of the O(n_kv * n_tokens) loop - // below; neither depends on the query. -1 means the cell holds nothing this - // sequence may pool or attend to. under a unified cache `cells` is shared with - // every other sequence, and seq_has is what keeps their keys out pos_at.resize(n_kv); for (int64_t j = 0; j < n_kv; ++j) { pos_at[j] = cells.is_empty(j) || !cells.seq_has(j, seq_of_pool) ? -1 : cells.pos_get(j); } - // a pool ordinal is the absolute p/kpool, which can far exceed n_kv/kpool, so - // rebase on this sequence's lowest resident pool. grouping is untouched, every - // member shifts together. - // - // anchoring at p/kpool follows vLLM and SGLang, not HF (which pools from the - // first *resident* key, valid_keys.argmax(-1), differing under left padding). - // it is the only anchor that keeps a pool's identity stable between the prefill - // that built it and the decodes that read it. - // - // the window is this sequence's run; positions are contiguous in any real - // batch so the resident range fits. seq_rm can leave a hole large enough that - // it does not, and then the newest pools are the ones worth keeping. + // anchoring at the absolute p/kpool follows vLLM and SGLang, not HF + // (valid_keys.argmax(-1)): it is the only anchor that keeps a pool's identity + // stable from the prefill that built it to the decodes that read it. int64_t b_base = 0; { int64_t b_min = 0; @@ -323,13 +282,9 @@ void llama_kv_cache_set_input_kpool( filled[bo]++; } - // an incompletely resident pool cannot be pooled: the compressor consumes all r - // member keys, and the reference demands pool_valid = grouped_valid_keys.all(-1). - // such cells are the sequence tail, which sel_mask forces in below regardless of - // score, so they point at pool slot 0 purely to keep the gather in range + // pool_valid = grouped_valid_keys.all(-1): the compressor consumes all r keys for (int64_t j = 0; j < n_kv; ++j) { - // != rather than <: two cells claiming one position overwrite each other in - // pool_cells, so an over-filled pool is not usable either + // != rather than <: two cells claiming one position overwrite each other if (pool_of[j] >= 0 && filled[pool_of[j]] != (int32_t) r) { pool_of[j] = -1; } @@ -341,8 +296,6 @@ void llama_kv_cache_set_input_kpool( for (int64_t ii = 0; ii < n_tps; ++ii) { const int64_t i = s*n_tps + ii; - // a query is pooled by the sequence it is being written to. with several - // sequences in one stream the other partitions own the rest of the rows if (ubatch->seq_id[i][0] != seq_of_pool) { continue; } @@ -354,22 +307,17 @@ void llama_kv_cache_set_input_kpool( n_done++; - // the query's own incomplete pool ((q + 1) % r cells, its own token - // included) is always attended to (index_kpool_always_select_tail), which is - // what makes the selection land on pool boundaries + // index_kpool_always_select_tail, which lands selection on pool boundaries const llama_pos tail_start = (q + 1)/r*r; - // the reference tests visibility at a pool's LAST member, so a pool - // straddling the query is dropped whole. pools are position-aligned here, so - // that test collapses to b*r < tail_start + // the reference tests visibility at a pool's LAST member, so a pool the + // query straddles is dropped whole const int64_t bo_vis = std::max(0, tail_start/r - b_base); float * cur_bias = dst_bias ? dst_bias + i*n_kv : nullptr; char * cur_sel = cur_sel_mask + ii*n_kv*mask_ts; char * cur_cand = cur_cand_mask + ii*n_kv*mask_ts; - // the unsigned compares inside fold "empty or another sequence" (pos_at -1) - // and "no usable pool" (pool_of -1) into the range test if (mask_f16) { kpool_mask_row((ggml_fp16_t *) cur_sel, (ggml_fp16_t *) cur_cand, pos_at.data(), pool_of.data(), n_kv, q, tail_start, bo_vis); @@ -387,20 +335,8 @@ void llama_kv_cache_set_input_kpool( } } - // The same predicate, per POOL, which is where the reference applies - // it: pool_valid (completely resident) & pool_visible (its LAST - // member is visible, so a pool the query straddles is dropped whole). - // Pools are position-aligned here, so pool bo's last member is at - // position (b_base + bo)*r + r - 1 and "last member visible" collapses - // to bo < bo_vis. - // - // NOT gathered from `bias` at the last member cell: an incomplete or - // absent pool has no resident last member, pool_cells points that slot - // at cell 0, and the pool would inherit cell 0's validity. - // - // the query's own sequence run only. every other slot stays at the - // -INFINITY the per-stream fill above left, which is what keeps a foreign - // pool out of the budget + // the query's own sequence run only; every other slot keeps the -INFINITY + // of the fill above, which is what keeps a foreign pool out of the budget float * q_pool_bias = cur_pool_bias + ii*n_pools + run_off[ps]; for (int64_t p = 0; p < n_run; ++p) { @@ -412,24 +348,17 @@ void llama_kv_cache_set_input_kpool( } } - // every row of sel_mask, cand_mask and pool_bias must have been written by - // exactly one partition, or a query is left reading another sequence's pools + // exactly one partition per row, or a query reads another sequence's pools GGML_ASSERT(n_done == n_tps && "every query must belong to a sequence of the ubatch"); } } void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { - // unconditional: the indexer key and gate STORE runs on the dense path too, - // and k_idxs is what tells cpy_k where to put them. Gating the store the way - // the scoring is gated would leave every cell written below n_select - the - // first 2051 positions of every sequence on the real model - with no indexer - // state, and the first ubatch to cross n_select would pool cells that were - // never written + // unconditional: the key and gate STORE runs on the dense path too. gating it the + // way the scoring is gated would leave every cell below n_select with no indexer + // state, and the first ubatch to cross n_select would pool cells never written mctx_idx->set_input_k_idxs(k_idxs, ubatch); - // the rest exists only when the graph scores. below n_select the indexer - // would select every visible position, so build_inp_kpool does not allocate - // these at all and there is nothing to fill if (pool_cells == nullptr) { return; } diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index 2ad259c8dcf..7bd9d862a29 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -10,149 +10,29 @@ struct llama_ubatch; class llama_kv_cache; class llama_kv_cache_context; -// -// GLM-5-Next indexer pooling -// -// GLM's lightning indexer scores pools of `kpool` consecutive *positions*, while its -// top-k budget is counted in tokens (indexer_top_k), i.e. indexer_top_k/kpool whole -// pools. The reference requires indexer_top_k % kpool == 0, so the division is exact. -// -// Pools are defined on positions, but everything the graph indexes is addressed by -// *cell*, and a cell index is whatever llama_kv_cache::find_slot handed out: cells of -// one pool are not adjacent, not ordered, and under a unified cache not even owned by -// the same sequence. The mapping therefore cannot be derived in the graph. It is built -// here, host side, from the cache's cells, and passed in as plain input tensors, as -// qwen4exp's QSA does for its compression blocks. -// -// A position only names a pool inside ONE sequence, so the map is per SEQUENCE. A -// non-unified cache gives every sequence its own stream and its own cells array, and -// the two are the same thing. A unified cache puts every sequence of the ubatch in one -// stream and one cells array, so the stream's pool table is PARTITIONED: each sequence -// gets a contiguous run of slots and rebases inside it. `pool_bias` is -INFINITY outside -// the query's own run, which is what stops a query from selecting a foreign pool. -// -// The runs are packed, not one full-width table per sequence. A full-width table per -// sequence would multiply the pool axis by the sequence count, and the indexer scores -// every pool against every query, so llama-embedding (which asks for n_seq_max 256 with -// a unified cache) reserved 286 GB of compute buffer that way. -// -// Nothing here may emit a negative index: ggml_set_rows asserts i1 >= 0 and -// ggml_get_rows has no sentinel, so unpopulated entries are clamped into range and -// neutralised by the additive masks instead. -// +// GLM-5-Next indexer pooling. the position -> cell map is built host side because +// find_slot's cell order is arbitrary. no input may hold a negative index: ggml_set_rows +// asserts i1 >= 0 and ggml_get_rows has no sentinel, so unusable entries are clamped into +// range and neutralised by the additive masks instead. -// number of pool slots the graph must allocate for `n_kv` cells shared by `n_seqs` -// sequences. -// -// pool ordinals are position-derived, then rebased on each sequence's lowest resident -// pool so sequences not starting at position 0 still land inside the array. rebasing -// can cost one slot at each end, hence 2 per sequence. -// -// n_kv/kpool is the budget the sequences share, which is exact while their cells are -// disjoint: a cell belongs to one pool of one sequence. llama_memory_seq_cp breaks that -// - a shared prefix cell is pooled by every sequence holding it - and then the runs are -// cut to the newest pools rather than the table being grown, see [TAG_KPOOL_PACK]. +// pool slots for `n_kv` cells shared by `n_seqs` sequences: n_kv/kpool, exact only while +// the sequences' cells are disjoint, plus 2 per sequence for rebasing. uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool, uint32_t n_seqs = 1); -// how many POOLS ggml_top_k selects. -// -// The reference (modular_glm5_next.py, Glm5NextTextIndexer.forward) scores the pool -// axis and takes select_k = min(index_topk // index_kpool, n_pools), then expands each -// selected pool to its members. This is that number, and top-k over pools rather than -// over cells is not an optimisation, it is the only correct spelling. -// -// The tempting alternative - broadcast a pool's score onto its kpool member cells and -// run one top-k of width indexer_top_k over cells - is WRONG, for a reason a -// set-similarity metric cannot see. Its argument is that a pool's members carry its -// score bit-exactly, so the cut must land on a pool boundary; that holds only if tie -// groups never span pools. They do: F.relu drives most pool scores to exactly 0.0 (on -// TinySparse, 92 cells tie at 0.0 for query 386 at layer 3), so the cut falls inside a -// tie group that straddles pools, and ggml_top_k - explicitly unordered among equals; -// the CPU op deliberately swaps the first two results, the CUDA op declares -// determinism::not_guaranteed - takes an arbitrary 1..kpool-1 members of the pool it -// cuts. Measured on TinySparse at 512 tokens the cell-level form leaves a partial pool -// on 7.51% of query rows at L3 and 5.93% at L7 (70 and 47 partial pools); the -// pool-level form leaves none. A pool-aligned top-k WIDTH does not help, because the -// ties are not aligned to anything. -// -// Selecting whole pools also removes the CPU/CUDA tie-break divergence structurally -// rather than making it unlikely: the backends may disagree about WHICH pools come out -// of a tie group, never about whether a pool is taken whole. -// -// NOT indexer_top_k + kpool - 1 either: that is the width of the reference's *output* -// buffer, tail included. The always-selected tail is forced in via `sel_mask`. +// select_k of modular_glm5_next.py, Glm5NextTextIndexer.forward. must run over POOLS, not +// cells: relu ties span pool boundaries, so a cell-level cut takes partial pools. uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t kpool); -// Fill the host-side inputs of the pooled indexer. -// -// `cell_pool` and `bias` are the per-CELL view of the same information and may be -// nullptr. They are kept because they are the independent spelling -// tests/test-glm5next-memory.cpp checks `pool_bias` and `cand_mask` against - but the -// graph selects at POOL granularity (see llama_kpool_select_k), so it passes nullptr -// for both. An input tensor with no consumer is never backed by the allocator, so -// building them anyway would write through a null buffer. -// -// cell_pool I32 [n_kv, n_stream] OPTIONAL -// cell -> its pool slot, or 0 when it has no usable pool. -// -// pool_cells I32 [kpool*n_pools, n_stream] -// pool slot, member -> the cell holding that position, or 0 when not resident. -// One stream's slots are cut into one contiguous run per sequence. -// Two consumers: the compressor gathers a pool's member keys and gates with it, -// and the top-k expands a selected POOL ordinal back into its kpool member CELLS -// with it - the reference's `pool_indices[batch_idx, selected]`. -// -// bias F32 [n_kv, n_tokens/n_stream, n_stream] OPTIONAL -// additive per-(cell, query) bias on a per-cell score: 0.0f in a complete pool -// whose last member the query can see, else -INFINITY, the trailing incomplete -// pool included. -// -// pool_bias F32 [n_pools, n_tokens/n_stream, n_stream] -// the same predicate per POOL, which is where the reference applies it: 0.0f when -// pool p is completely resident and its LAST member is visible to query q -// (`pool_valid & pool_visible`), else -INFINITY, the query's own trailing pool -// included so no budget is spent on it. Every slot outside the query's own -// sequence run is -INFINITY, so a query never selects another sequence's pool. -// -// Derived here rather than in the graph on purpose. Gathering `bias` at -// each pool's last member looks equivalent and is not: an incomplete or -// entirely absent pool has no resident last member, `pool_cells` points -// that slot at cell 0 to keep the gather in range, and the pool would then -// inherit cell 0's validity and compete for budget with a finite score. -// -// sel_mask F16/F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) -// what the top-k scatter starts from, replacing the ggml_fill(kq_mask, -INFINITY) -// opening the DSA mask build in llm_graph_context::build_attn: 0.0f for the -// query's own incomplete trailing pool, which GLM always attends to -// (index_kpool_always_select_tail), else -INFINITY, padding rows included. -// Forcing the tail in here rather than through the score keeps the budget a whole -// number of pools. -// -// cand_mask F16/F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) -// max(bias, sel_mask): the reference's candidate set, i.e. every cell it could -// return for this query. 0.0f in a complete visible pool OR in the query's own -// tail, else -INFINITY, padding rows included. -// -// ggml_top_k returns `select_k` pool ordinals even when fewer pools carry a -// finite score, and during prefill that is the NORMAL state, not a corner case: -// query q has only ~q/kpool complete visible pools against a budget of -// index_topk/kpool (on TinySparse, 20 of 20 query rows are under budget). The -// spilled ordinals are arbitrary among the -INFINITY ties and expanding them -// unmasks cells. Adding the causal KQ mask kills the spills that are empty, -// foreign or in the future. What it does not kill is a resident, causally visible -// cell in an INCOMPLETE pool below the tail - unreachable while positions are -// contiguous, reachable the moment a partial seq_rm leaves a hole. cand_mask -// kills exactly those, for one store per element in a loop that already computes -// both operands. -// -// (The other spill, an unfilled pool slot pointing at cell 0, is a provable -// no-op: the budget can only overflow once every finitely-scored pool is already -// selected, so cell 0 is either already selected through its own pool, or masked -// for the same reason it would have been anyway.) -// -// `kv` must be the ATTENTION (MLA) cache: its cells define the pools and are what the -// top-k indices are ultimately read against. llama_memory_hybrid gives the indexer -// cache the attention cache's slot layout, so the two agree cell for cell. +// `kv` must be the ATTENTION (MLA) cache; the indexer cache shares its slot layout. +// cell_pool I32 [n_kv, n_stream] per-cell view, optional, unused here +// pool_cells I32 [kpool*n_pools, n_stream] pool member -> cell, 0 if not resident +// bias F32 [n_kv, n_tps, n_stream] per-cell view, optional, unused here +// pool_bias F32 [n_pools, n_tps, n_stream] pool_valid & pool_visible, -INFINITY +// outside the query's own sequence run; computed, not gathered from `bias` at the +// last member, which an incomplete pool lacks and would inherit cell 0's validity +// sel_mask F16/F32 [n_kv, n_batch, 1, n_stream] 0.0f on the always-selected tail only +// cand_mask F16/F32 [n_kv, n_batch, 1, n_stream] max(bias, sel_mask); bounds the top-k +// spills that a partial seq_rm would otherwise let escape the candidate set void llama_kv_cache_set_input_kpool( const llama_kv_cache * kv, ggml_tensor * cell_pool, @@ -164,19 +44,9 @@ void llama_kv_cache_set_input_kpool( const llama_ubatch * ubatch, uint32_t kpool); -// One pooling map per ubatch, shared by every indexer layer: all four tensors, and -// k_idxs, depend only on the cells and the ubatch, never on the layer. Rebuilding them -// per layer costs O(n_kv * n_tokens) host writes each time (at 128 Ki cells, 512 tokens -// and 11 DSA layers, ~11 x 67M float stores per ubatch, which dominates prefill). -// -// Sharing `pool_bias`, `sel_mask` and `cand_mask` is correct only while every indexer -// layer sees the same candidate set. That holds for glm5next, whose indexer_types are -// all "full"; a model mixing in windowed indexers would need one map per window. -// -// `k_idxs` is present whenever the model has an indexer cache; the rest only when the -// graph actually scores. The indexer key and gate STORE is not gated on the sparse path -// - below n_select the selection is a no-op but the cells still have to be written, or -// the first ubatch to cross n_select would pool cells that were never filled. +// One pooling map per ubatch; rebuilding it per indexer layer costs O(n_kv * n_tokens) +// host writes and dominates prefill. sharing is valid only while every indexer layer sees +// the same candidate set - true for glm5next (indexer_types all "full"), not for windowed. class llm_graph_input_kpool : public llm_graph_input_i { public: llm_graph_input_kpool( @@ -192,10 +62,7 @@ class llm_graph_input_kpool : public llm_graph_input_i { ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] ggml_tensor * pool_bias = nullptr; // F32 [n_pools, n_tps, n_stream] - // pool_bias in the shape and type the fused lightning indexer wants for its mask. - // A ggml_cast node, not an input: built once per graph rather than once per DSA - // layer, since every indexer layer shares the same mask. The cast is exact - - // pool_bias only ever holds 0.0f or -INFINITY. nullptr when the fused path is off + // exact, since pool_bias only holds 0.0f or -INFINITY. nullptr if the fused path is off ggml_tensor * pool_bias_f16 = nullptr; // F16 [n_pools, n_tps, 1, n_stream] ggml_tensor * sel_mask = nullptr; // F16 [n_kv, n_batch, 1, n_stream] diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 983352d617c..ba5f2eca03f 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -366,14 +366,9 @@ class llama_kv_cache_context : public llama_memory_context_i { uint32_t get_n_kv() const; - // streams covered by the current slot info, matching the `ns` get_k/get_v use for - // their stream dimension. 1 for a unified cache. note: this is the stream RANGE - // s1 - s0 + 1, not n_seqs_unq; they differ when the active sequences are not a - // contiguous run of slots, i.e. exactly when a per-cell input sized from this would - // stop agreeing with a KQ mask + // the stream RANGE s1 - s0 + 1 that get_k/get_v use as `ns`, not n_seqs_unq uint32_t get_n_stream() const; - // the cache this context views, for host-side inputs that resolve cell -> position const llama_kv_cache * get_kv() const; ggml_type type_k() const; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 3d842f99f84..7050ecb81f3 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -69,15 +69,9 @@ llama_memory_hybrid::llama_memory_hybrid( : filter_recr )), mem_idx(filter_idx == nullptr ? nullptr : [&] { - // MQA with one key head of indexer_head_size, as llama_kv_cache_dsa shapes its - // lightning-indexer cache. n_embd_head_k_full is what n_embd_head_k(il) reads - // for a non-SWA layer, so an MLA model still has is_mla() true here, which - // suppresses the V allocation the indexer does not need. - // - // a *pooling* indexer (indexer_kpool > 0, glm5next only) needs a second head: - // its compressor gate is a projection of the same hidden state, so it must be - // cached alongside the key or the pool cannot be rebuilt once the member tokens - // leave the batch. every other arch leaves indexer_kpool 0 and is unchanged. + // a *pooling* indexer (indexer_kpool > 0, glm5next only) needs a second head for + // the compressor gate, or the pool cannot be rebuilt once its tokens leave the + // batch. every other arch leaves indexer_kpool 0 and is unchanged. const uint32_t n_head_idx = model.hparams.indexer_kpool > 0 ? 2 : 1; std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), n_head_idx); @@ -143,10 +137,9 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } - // the indexer is a side buffer addressed by the attention cache's cells, so it - // takes that slot layout rather than finding its own: allocating separately lets - // the two drift apart when the context is rewritten between turns, and the top-k - // indices, read against the attention mask, would then point at the wrong cells + // the indexer takes the attention cache's slot layout rather than finding its + // own: allocated separately the two drift apart when the context is rewritten + // between turns, and the top-k indices would then point at the wrong cells llama_kv_cache::slot_info_vec_t heads_idx; if (mem_idx) { heads_idx = heads_attn; @@ -241,8 +234,7 @@ std::map llama_memory_hybrid::memory_breakdo void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_write(io, seq_id, flags); - // indexer keys are not recomputable from the attention cache, so a restored - // session that skipped them would select the wrong cells + // indexer keys are not recomputable; skipping them here misselects on restore if (mem_idx) mem_idx->state_write(io, seq_id, flags); } mem_recr->state_write(io, seq_id, flags); @@ -283,11 +275,10 @@ llama_memory_hybrid_context::llama_memory_hybrid_context( bool optimize) : ctx_attn(mem->get_mem_attn()->init_update(lctx, optimize)), ctx_recr(mem->get_mem_recr()->init_update(lctx, optimize)), - // indexer keys carry no positional encoding, so a shift has nothing to correct in - // them, but the pending per-cell delta must still be cleared or the two caches - // disagree about whether a shift is outstanding. an indexer only exists for - // LLAMA_ROPE_TYPE_NONE archs, which is exactly when llama_kv_cache::update skips the - // K-shift graph and does only that + // indexer keys carry no positional encoding, but the pending per-cell delta must + // still be cleared or the two caches disagree about whether a shift is outstanding. + // safe because an indexer only exists for LLAMA_ROPE_TYPE_NONE archs, where + // llama_kv_cache::update skips the K-shift graph and does only that ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : mem->get_mem_idx()->init_update(lctx, optimize)), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } @@ -332,7 +323,7 @@ bool llama_memory_hybrid_context::apply() { res = res & ctx_idx->apply(); // a top-k over indexer cells is meaningful only if both caches cover the same - // window. only the batch context has slot infos to compare + // window if (!ubatches.empty()) { GGML_ASSERT(get_idx()->get_n_kv() == get_attn()->get_n_kv()); GGML_ASSERT(get_idx()->get_n_stream() == get_attn()->get_n_stream()); diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 92e576332e7..9cf79b8e871 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -40,8 +40,7 @@ class llama_memory_hybrid : public llama_memory_i { /* layer filters */ const layer_filter_cb & filter_attn = nullptr, const layer_filter_cb & filter_recr = nullptr, - /* optional per-token indexer key cache for sparse-attention - hybrids; absent unless filter_idx is given */ + /* optional indexer key cache; absent unless filter_idx */ const layer_filter_cb & filter_idx = nullptr, ggml_type type_idx = GGML_TYPE_F16); @@ -91,8 +90,6 @@ class llama_memory_hybrid : public llama_memory_i { private: const llama_hparams & hparams; - // indexer cache geometry: n_head_kv key heads of indexer_head_size, as - // llama_kv_cache_dsa builds its own llama_hparams hparams_idx; const std::unique_ptr mem_attn; @@ -121,8 +118,6 @@ class llama_memory_hybrid_context : public llama_memory_context_i { llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, std::vector ubatches, - // empty without an indexer. the indexer is addressed by the - // attention cache's cells, so it gets that cache's slots slot_info_vec_t sinfos_idx = {}); ~llama_memory_hybrid_context() = default; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0d1bfee0dca..a059fb5b966 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2434,8 +2434,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; - // null for every arch but the sparse-attention ones, which is what - // keeps the indexer cache from existing + // null except for sparse attention, which keeps the indexer cache from existing llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; ggml_type type_idx = GGML_TYPE_F16; if (arch == LLM_ARCH_FALCON_H1) { @@ -2457,17 +2456,14 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, }; if (arch == LLM_ARCH_GLM5NEXT && hparams.indexer_head_size > 0) { - // a unified cache is fine here: the pool map is per SEQUENCE, - // not per stream. see [TAG_KPOOL_SEQ_PARTITION] + // unified is fine, the pool map is per SEQUENCE. see [TAG_KPOOL_SEQ_PARTITION] // only the DSA layers carry an indexer key cache filter_idx = [&](uint32_t il) { return il < hparams.n_layer() && !hparams.is_recr(il); }; - // the gate cached next to the key feeds a softmax, so -ctk - // q8_0 would quantise something far more sensitive than a - // key. keep the indexer float + // the gate cached beside the key feeds a softmax, unlike -ctk q8_0's target type_idx = params.type_k; if (ggml_is_quantized(type_idx)) { LLAMA_LOG_WARN("%s: indexer key cache stays %s rather than %s: it also holds the compressor gates\n", @@ -2478,9 +2474,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { - // llama_memory_hybrid_iswa has no indexer cache; glm5next is - // swa_type NONE, but a sparse hybrid with SWA would silently - // lose its indexer + // llama_memory_hybrid_iswa has no indexer cache, so SWA would silently lose it GGML_ASSERT(filter_idx == nullptr && "hybrid-iswa cannot carry an indexer cache"); // Use hybrid-iswa for hybrid models with SWA diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 97b8f4d22ed..e1b0ecced40 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -350,27 +350,9 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param // do not quantize relative position bias (T5) quantize &= name.find("attn_rel_b.weight") == std::string::npos; - // glm5next: small, precision-sensitive tensors that must stay at source - // precision -- the mHC residual mixers, the lightning indexer (selection gate, - // learned k-pool position table, and the three indexer projections) and the KDA - // recurrence gates. ~1 GiB total on GLM-5.3-Flash, so the size cost is noise - // against a 100-240 GB quant, while quantizing them perturbs *which* pools the - // indexer selects and *how much* state each KDA step retains -- errors that - // compound over a sequence instead of averaging out. - // - // Note both spellings are required. The compressor tensors came in with the - // DeepSeek-V4 merge and are named with an UNDERSCORE (indexer_compressor_ape / - // _gate), while the projections use a DOT (indexer.proj / .attn_k / .attn_q_b). - // A single "indexer." prefix test silently misses the compressor pair. - // - // Deliberately NOT listed: attn_q_a / attn_kv_a_mqa / attn_k_b / attn_v_b. They - // are precision-sensitive too, but our release recipe pins them to q8_0 via - // --tensor-type, and that is the configuration the shipped quants were measured - // in (KLD 0.027 at Q5_K_XL). Forcing them full precision here would change quant - // sizes and invalidate those measurements. - // - // indexer.k_norm.weight needs no entry -- the generic "_norm.weight" rule above - // already excludes it. + // glm5next: quantizing these perturbs pool selection and KDA state retention, errors that + // compound over a sequence, for ~1 GiB. note the compressor tensors spell it with an + // UNDERSCORE and the projections with a DOT, so one "indexer." prefix test is not enough if (arch == LLM_ARCH_GLM5NEXT) { static const char * const glm5next_full_precision[] = { "hc_attn_fn", diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index d37f49a5398..23f62f4890d 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2257,12 +2257,9 @@ 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; - // glm4 tokenizer.json sets "ignore_merges": true. without it greedy BPE - // cannot reach some vocab entries: " 王" (Ġçİĭ, 102322) needs (Ġ,çİĭ)=242943 - // but (Ġ,ç)=27944 wins first, so it stops three tokens short. triggered by - // whitespace before CJK, inflating mixed Chinese-English ~13%. - // NOT applied to chatglm-bpe, which shares this pre_type but is a different - // tokenizer (ChatGLM3) not confirmed to declare the flag. + // glm4 tokenizer.json sets "ignore_merges": true; without it greedy BPE cannot + // reach some vocab entries, inflating mixed Chinese-English ~13%. chatglm-bpe + // shares this pre_type but is ChatGLM3 and not confirmed to declare the flag if (tokenizer_pre == "glm4") { ignore_merges = true; } diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index 4b6aec8046d..9284cd2a7d0 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -330,11 +330,8 @@ std::pair llm_build_delta_net_base::build_delta_ne cb(b, "b_in", il); cb(g, "g_in", il); - // the state is indexed [key, value]: ne0 is the key axis, as the k/q reductions - // below rely on, so KDA's per-key-channel decay must broadcast over ne1, not ne0. - // for GDA g->ne[0] is 1 and both spellings agree - // GDA: [1, 1, H_v, n_seqs] - // KDA: [S_k, 1, H_v, n_seqs] + // the state is indexed [key, value]: ne0 is the key axis, so KDA's per-key-channel + // decay must broadcast over ne1, not ne0 (for GDA g->ne[0] is 1 and both agree) g = ggml_reshape_4d(ctx0, g, g->ne[0], 1, H_v, n_seqs); b = ggml_reshape_4d(ctx0, b, 1, 1, H_v, n_seqs); diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 2cb8d4c20b3..0bcbdb95d18 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -4,22 +4,10 @@ #include "llama-memory-hybrid.h" #include "llama-kv-cache-kpool.h" -// -// GLM-5.3-Flash: hybrid KDA (linear) + DSA (nope-only MLA) attention, mHC -// hyper-connections, and a NextN block that is a full DSA decoder layer. -// -// ssm_a holds -exp(A_log), the kimi-k3 convention, so the decay gate reads -// exp(A_log) back as -ssm_a; bailingmoe3 stores +exp(A_log). indistinguishable at -// load time, so conversion/glm5next.py is the only place the sign is checked. -// - -// positions the indexer keeps: index_topk/index_kpool whole pools plus the -// always-selected tail pool minus one. below this many cached tokens every position is -// selected, so sparse selection is exactly the dense path built here. -// -// asserted, not measured: an off-by-one here is invisible to output comparison (the -// reference's own off-by-one on this width is bit-identical on both fixtures). the -// second assert is the parity harness's independent spelling, so the two must agree +// ssm_a holds -exp(A_log) (kimi-k3), not +exp(A_log) (bailingmoe3); converter checks + +// positions the indexer keeps; at or below this many the dense path IS the sparse one. +// asserted not measured (invisible to output); the second assert is an independent spelling static uint32_t glm5next_n_select(const llama_hparams & hparams) { GGML_ASSERT(hparams.indexer_kpool > 0); GGML_ASSERT(hparams.indexer_top_k >= hparams.indexer_kpool); @@ -37,13 +25,7 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); // indexer k_norm is a LayerNorm with bias; without this key it runs at eps 0 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); - // The reference HARDCODES nn.LayerNorm(index_head_dim, eps=1e-6); it does not - // read it from the config, and rms_norm_eps is 1e-5. Warned about rather than - // asserted, for two reasons. No output comparison can see it: seeding eps=1e-5 - // into the reference leaves the logits BIT-IDENTICAL at 512 tokens and moves - // the index jaccard by 1.7e-5 at 2048, against a bf16 floor of 0.63. And an - // assert would abort test-llama-archs, whose synthetic models have no reason - // to carry a converter contract + // warned not asserted: no output comparison sees it, and an assert breaks test-llama-archs if (hparams.f_norm_eps <= 0.0f || hparams.f_norm_eps > 2e-6f) { LLAMA_LOG_WARN("%s: indexer k_norm eps is %g, but the reference hardcodes 1e-6. " "this is invisible to every output comparison; check the converter\n", @@ -57,19 +39,15 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(hparams.n_lora_q > 0 && "glm5next requires a q LoRA"); GGML_ASSERT(hparams.n_rot() == 0 && "glm5next MLA is nope-only"); - // KDA. no GGUF key for linear_num_heads, so the KDA head count is - // attention.head_count, which also sizes the recurrent state via n_embd_r/s(). - // conversion/glm5next.py refuses a checkpoint where the two differ + // no linear_num_heads key: KDA head count is attention.head_count (converter enforces) 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); GGML_ASSERT(hparams.ssm_d_conv > 1); GGML_ASSERT(hparams.n_embd_head_kda > 0); - // required: absent, kimi-k3 selects the softplus branch, a different - // function, not a missing clamp + // required: absent, kimi-k3 selects the softplus branch, a different function ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f); - // DSA indexer 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); @@ -77,35 +55,20 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(hparams.indexer_kpool > 0); GGML_ASSERT(hparams.indexer_top_k % hparams.indexer_kpool == 0); - // below n_select resident positions the indexer selects every visible one, so - // the dense path is not an approximation of the sparse one, it IS it const uint32_t n_select = glm5next_n_select(hparams); LLAMA_LOG_INFO("%s: indexer selection width = %u cells (%u pools of %u, plus a %u-wide tail)\n", __func__, n_select, hparams.indexer_top_k/hparams.indexer_kpool, hparams.indexer_kpool, hparams.indexer_kpool - 1); - // 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 > 0); - // trunk residual is hc_mult streams wide (deepseek4); lm_head still sees - // n_embd, the streams are averaged first -- so n_embd_out stays n_embd. - // - // deepseek4 sets this to hc_mult*n_embd, and inheriting that here was a bug: - // our t_embd is build_norm(build_hc_mean(...)), which is [n_embd, n_tokens], - // but n_embd_out() would report 4*n_embd, so llama-context.cpp reads - // n_outputs*n_embd_out floats out of a tensor holding a quarter of that. - // The assert there sizes the DESTINATION, so nothing catches the short SOURCE. - // Only --embeddings / llama_get_embeddings* reach it, which is why plain - // generation never showed it. - // - // deepseek4 needs the wide value to size its MTP `h` input; when our NextN - // graph starts consuming it, give MTP its own width rather than widening this. + // n_embd_out stays n_embd: lm_head sees the stream mean. deepseek4's hc_mult*n_embd + // makes llama-context.cpp overread t_embd, and the assert there sizes the destination hparams.n_embd_out_impl = 0; - // MoE ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); @@ -123,8 +86,6 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all); - // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-k3 and bailingmoe3. - // a scalar head_count_kv would make every layer look like DSA, so require both uint32_t n_recr = 0; for (uint32_t il = 0; il < hparams.n_layer(); ++il) { hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0; @@ -132,8 +93,7 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { } GGML_ASSERT(n_recr > 0 && n_recr < hparams.n_layer() && "glm5next needs a per-layer attention.head_count_kv array"); - // every glm5next indexer is full; glm-dsa gates its indexer on this - // predicate and the generic loader only zero-fills the array + // every glm5next indexer is full; the generic loader only zero-fills the array for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { hparams.is_indexer_full_impl[il] = !hparams.is_recr_impl[il]; } @@ -162,7 +122,6 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { const int64_t hc_dim = (int64_t) hparams.dsv4_hc_mult * n_embd; const int64_t hc_mix_dim = (2 + (int64_t) hparams.dsv4_hc_mult) * hparams.dsv4_hc_mult; - // the trunk and the NextN block can be split across two GGUFs in either direction const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); 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); @@ -262,22 +221,14 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), {n_embd}, flags); layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), {n_embd}, flags); - // absent in the checkpoint: NextN shares the trunk's embeddings and - // lm_head. only accepted if an export adds them + // absent in the checkpoint: NextN shares the trunk's embeddings and lm_head layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); } } } -// -// KDA layer -// -// one depthwise conv over the concatenated q|k|v channels, as in the reference: it -// leaves the conv state one contiguous block, which is what build_conv_state needs to -// snapshot for recurrent-state rollback. three separate convs would be numerically -// identical but would restate the rollback write three times -// +// one conv over concatenated q|k|v: keeps the conv state one block for rollback ggml_tensor * llama_model_glm5next::graph::build_kda_layer( const llama_layer & layer, llm_graph_input_rs * inp_rs, @@ -305,7 +256,6 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); ggml_tensor * conv_in = build_conv_state(inp_rs, conv_states_all, qkv, d_conv, 3*d_inner, il); - // stored separately (kimi-linear, kimi-k3), stacked back into the single kernel ggml_tensor * conv_w = ggml_concat(ctx0, ggml_concat(ctx0, ggml_reshape_2d(ctx0, layer.ssm_q_conv, d_conv, d_inner), @@ -326,9 +276,7 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( Vcur = ggml_view_4d(ctx0, conv_out, head_dim, n_head, n_seq_tokens, n_seqs, nb_head, nb_qkv, nb_qkv*n_seq_tokens, ggml_row_size(conv_out->type, 2*d_inner)); - // 1e-6 is the reference's own constant, not the model's norm eps. ggml_l2_norm - // divides by max(sqrt(sum), eps) where the reference uses sqrt(sum + eps); at - // head_dim 128 the clamp never binds, so close but not bit-exact + // 1e-6 is the reference's own constant, not the model's norm eps Qcur = ggml_l2_norm(ctx0, Qcur, 1e-6f); Kcur = ggml_l2_norm(ctx0, Kcur, 1e-6f); cb(Qcur, "kda_q_norm", il); @@ -336,9 +284,7 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( // the 1/sqrt(head_dim) query scale is applied inside build_delta_net, after this norm - // forget gate. gate_lower_bound is a multiplicative scale, not a clamp: - // g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias)) - // ssm_a holds -exp(A_log), so exp(A_log) * y == -(y * ssm_a) + // g = lower_bound * sigmoid(exp(A_log)*(f_b(f_a(x)) + dt_bias)); it scales, not clamps ggml_tensor * g = ggml_mul_mat(ctx0, layer.ssm_f_b, ggml_mul_mat(ctx0, layer.ssm_f_a, inp)); g = ggml_add(ctx0, g, layer.ssm_dt_b); g = ggml_reshape_3d(ctx0, g, head_dim, n_head, n_tokens); @@ -358,17 +304,14 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( ggml_tensor * out = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g, beta, state, il); - // the fallbacks return a permuted view, the fused op a contiguous one; cont - // either way rather than depend on which ran + // the fallbacks return a permuted view, the fused op a contiguous one ggml_tensor * o = ggml_cont_3d(ctx0, out, head_dim, n_head, n_tokens); cb(o, "kda_scan_out", il); - // low-rank output gate (kimi-k3 has a single full-rank ssm_g instead) ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_g_b, ggml_mul_mat(ctx0, layer.ssm_g_a, inp)); gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens); - // RMS over head_dim only, one weight shared by every head, then a plain sigmoid - // gate: not the SiLU that FusedRMSNormGated defaults to + // plain sigmoid gate, not the SiLU that FusedRMSNormGated defaults to ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, gate)); cb(gated, "kda_normed", il); @@ -379,42 +322,12 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( return cur; } -// -// the lightning indexer, pooled -// -// Writes this layer's indexer key and compressor gate into the indexer cache, and -// then - when the cache is large enough for selection to bind - returns the I32 -// ATTENTION-cache cell indices this layer's queries select, ready for the scatter in -// build_attn_sparse. `qr` is the shared q LoRA residual, i.e. q_a_norm(q_a_proj(x)), -// which the indexer's own wq_b consumes; `cur` is the layer input. -// -// The store is NOT gated on the sparse path and the scoring is. Gating both the same -// way leaves every cell written below n_select - the first 2051 positions of every -// sequence on the real model - with no indexer state at all, and the first ubatch to -// cross n_select then pools cells that were never written. -// -// Three points where the reference implementations disagree, resolved 2-of-3 by -// reading transformers (modular_glm5_next.py Glm5NextTextIndexer), sglang -// (dsa_indexer_kpool.py) and vLLM (glm5next/nvidia/attention.py): -// -// * the weights_proj GEMM runs in fp32. sglang gives it params_dtype=fp32 and feeds -// x.float(); vLLM does the same and says why - bf16 head-gates move logits by -// ~1e-2, enough to flip near-tie pool rankings on long context. transformers is -// the outlier and rounds the activation to bf16. -// * k_norm is a LayerNorm WITH BIAS at eps 1e-6, hardcoded in transformers and in -// vLLM; sglang leaves torch's 1e-5 default. It is NOT f_norm_rms_eps (1e-5 here) -// and NOT ggml's 0 default. The value comes from the GGUF and load_arch_hparams -// warns when it is not 1e-6, because no output comparison can see it. -// * the ReLU between the QK dot and the head weighting is explicit in transformers -// and in sglang's non-pooled indexer; in the pooled path both engines hand it to -// the DeepGEMM MQA-logits kernel, so neither validates it. It is easy to drop and -// is written out here. -// -// No Hadamard rotation: sglang and vLLM rotate q and k by H128 before their fp8 -// kernel, but H is orthogonal so (Hq).(Hk) == q.k. It exists to spread magnitude ahead -// of fp8 quantisation; scoring in f32 here it would only cost accuracy. transformers, -// the semantic reference, has none. -// +// the store is NOT gated on the sparse path, the scoring is: gating both leaves cells +// below n_select with no indexer state, which the first ubatch to cross n_select pools. +// * weights_proj runs in fp32; bf16 head-gates flip near-tie pool rankings (vLLM, sglang) +// * k_norm is a LayerNorm WITH BIAS at eps 1e-6, not f_norm_rms_eps (transformers, vLLM) +// * the ReLU between the QK dot and the head weighting is real (modular_glm5_next.py) +// no Hadamard rotation: H is orthogonal so (Hq).(Hk) == q.k; it only helps fp8. ggml_tensor * llama_model_glm5next::graph::build_indexer( const llama_layer & layer, llm_graph_input_kpool * inp_kp, @@ -428,22 +341,17 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( const auto * mctx_idx = inp_kp->mctx_idx; - // a genuine LayerNorm: weight AND bias. glm-dsa runs this same norm at eps 0 today GGML_ASSERT(layer.indexer_k_norm_b != nullptr && "the indexer k_norm is a LayerNorm with bias"); ggml_tensor * ik = build_norm(ggml_mul_mat(ctx0, layer.indexer_attn_k, cur), layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il); cb(ik, "indexer_k", il); - // the pooling gate is a SECOND, INDEPENDENT projection of the hidden state, not a - // reuse of the indexer key. it has to be cached beside the key: the compressor - // mixes a pool's member keys with a softmax over these gates, and a pool is only - // rebuilt once its members have left the batch + // a SECOND, INDEPENDENT projection, not a reuse of the key, and cached beside it ggml_tensor * gate = ggml_mul_mat(ctx0, layer.indexer_comp_wgate, cur); cb(gate, "indexer_gate", il); - // {d_idx, 2, n_tokens}: head 0 is the key, head 1 the gate. llama_memory_hybrid - // allocates the second head exactly for this and only when indexer_kpool > 0 + // {d_idx, 2, n_tokens}: head 0 is the key, head 1 the gate ggml_tensor * packed = ggml_concat(ctx0, ggml_reshape_3d(ctx0, ik, d_idx, 1, n_tokens), ggml_reshape_3d(ctx0, gate, d_idx, 1, n_tokens), 1); @@ -453,7 +361,6 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( return nullptr; } - // {d_idx, 2, n_kv, n_stream} ggml_tensor * kbuf = mctx_idx->get_k(ctx0, il); const int64_t n_kv = kbuf->ne[2]; @@ -466,17 +373,10 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( GGML_ASSERT(kbuf->nb[1] == (size_t) d_idx*kbuf->nb[0] && "key and gate must be adjacent in a cell"); GGML_ASSERT(n_tokens == n_tps*n_stream); - // one cell's key and gate as a single row, so that the members of a pool are - // gathered once rather than twice: {2*d_idx, n_kv, n_stream} ggml_tensor * kg_rows = ggml_view_3d(ctx0, kbuf, 2*d_idx, n_kv, n_stream, kbuf->nb[2], kbuf->nb[3], 0); - // gather each pool's members. pool_cells names a cell per (pool, slot); slots that - // are not resident hold 0 rather than a negative sentinel, because ggml_get_rows - // has none. The resulting garbage pools are neutralised by pool_bias below, never - // by a NaN: unlike the reference, no -inf ever enters the compressor softmax. - // ggml_get_rows always yields F32, so the pooling runs in F32 even though the - // indexer cache is F16 + // non-resident slots hold 0, not a sentinel; garbage pools die to pool_bias, not NaN ggml_tensor * members = ggml_get_rows(ctx0, kg_rows, inp_kp->pool_cells); cb(members, "indexer_pool_members", il); @@ -487,10 +387,7 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_tensor * mem_g = ggml_view_4d(ctx0, members, d_idx, r, n_pools, n_stream, nb_mem, nb_mem*r, members->nb[2], d_idx*members->nb[0]); - // d_idx independent r-way softmaxes over the SLOT axis, so the slot axis has to be - // dim 0. ape is added PRE-softmax and is indexed by LOGICAL SLOT, so it broadcasts - // over pools and streams; pool_cells is built in position order, so slot m is - // position p % kpool and the two agree by construction + // r-way softmaxes over the SLOT axis, so it must be dim 0; ape is added PRE-softmax ggml_tensor * keys_t = ggml_cont(ctx0, ggml_permute(ctx0, mem_k, 1, 0, 2, 3)); ggml_tensor * gate_t = ggml_cont(ctx0, ggml_permute(ctx0, mem_g, 1, 0, 2, 3)); @@ -500,40 +397,27 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_tensor * probs = ggml_soft_max(ctx0, gate_t); cb(probs, "indexer_pool_probs", il); - // per-channel weighted average over the pool's members -> {d_idx, n_pools, 1, n_stream} ggml_tensor * pool_k = ggml_sum_rows(ctx0, ggml_mul(ctx0, keys_t, probs)); pool_k = ggml_reshape_4d(ctx0, pool_k, d_idx, n_pools, 1, n_stream); cb(pool_k, "indexer_pool_k", il); - // {d_idx, n_ihead, n_tps, n_stream}. no rope: n_rot() is 0 for the whole text tower + // no rope: n_rot() is 0 for the whole text tower ggml_tensor * iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr); iq = ggml_reshape_4d(ctx0, iq, d_idx, n_ihead, n_tps, n_stream); cb(iq, "indexer_q", il); - // sign-unconstrained head weights: no softmax, no abs, no relu. Both scale - // constants - the reference's softmax_scale = d_idx^-0.5 and its n_heads^-0.5 head - // factor - are folded in here, on an {n_ihead, n_tokens} tensor rather than on the - // {n_pools, n_tps, n_ihead} score tensor. relu is positively homogeneous and both - // constants are positive, so this is exactly the same function, and it is what the - // engines and the in-tree glm-dsa both do. - // - // GGML_PREC_F32 is not cosmetic: on vLLM a bf16 head gate moves a logit by ~1e-2, - // which is enough to swap two near-tied pools, and the top-k below is a hard cut + // sign-unconstrained head weights: no softmax, no abs, no relu; both scale constants + // fold in here. GGML_PREC_F32 is not cosmetic: bf16 can swap two near-tied pools ggml_tensor * w = ggml_mul_mat(ctx0, layer.indexer_proj, cur); ggml_mul_mat_set_prec(w, GGML_PREC_F32); w = ggml_reshape_4d(ctx0, w, n_ihead, n_tps, 1, n_stream); w = ggml_scale(ctx0, w, 1.0f/sqrtf(float(d_idx*n_ihead))); cb(w, "indexer_weights", il); - // both paths end with pool_bias added: -INFINITY on every pool the reference's - // `pool_valid & pool_visible` rejects, the query's own trailing pool included, so - // that no budget is spent on it ggml_tensor * pool_score = nullptr; if (cparams.fused_lid) { - // one node for the whole dot product -> relu -> head sum -> mask chain, and the - // mask add comes for free. pool_k stays f32, so the kernel takes its f32 vector - // path rather than the f16 wmma path, which would undo the GGML_PREC_F32 above + // pool_k stays f32 so the kernel takes its f32 path; f16 wmma would undo the prec ggml_tensor * pool_kf = ggml_reshape_4d(ctx0, pool_k, d_idx, 1, n_pools, n_stream); pool_score = ggml_lightning_indexer(ctx0, iq, pool_kf, w, inp_kp->pool_bias_f16); @@ -543,18 +427,13 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( pool_score = ggml_reshape_3d(ctx0, pool_score, n_pools, n_tps, n_stream); } else { - // {n_pools, n_tps, n_ihead, n_stream}: pool_k is MQA and broadcasts over the heads ggml_tensor * kq = ggml_mul_mat(ctx0, pool_k, ggml_permute(ctx0, iq, 0, 2, 1, 3)); - // {n_ihead, n_tps, n_pools, n_stream}, contiguous for the relu and the head sum. - // the ReLU sits BETWEEN the per-head dot product and the head weighting: moving it - // to either side is a different function, because the head weights are sign-free - // and the sum is not a convex combination + // the ReLU sits BETWEEN the per-head dot and the head weighting; either side differs kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); ggml_tensor * score = ggml_relu(ctx0, kq); cb(score, "indexer_score", il); - // {1, n_tps, n_pools, n_stream} -> {n_pools, n_tps, n_stream} pool_score = ggml_sum_rows(ctx0, ggml_mul(ctx0, score, w)); pool_score = ggml_cont(ctx0, ggml_permute(ctx0, pool_score, 2, 1, 0, 3)); pool_score = ggml_reshape_3d(ctx0, pool_score, n_pools, n_tps, n_stream); @@ -563,22 +442,8 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( cb(pool_score, "indexer_pool_score", il); } - // Top-k over POOLS at index_topk/index_kpool, then expand each selected pool to its - // members. This is the reference's own two-step (topk over the pool axis, then - // selected_indices = pool_indices[batch_idx, selected]) and it is NOT - // interchangeable with a single top-k of width index_topk over member cells. - // - // The cell-level form is the tempting one and the argument for it is false. It says - // a pool's members carry its score bit-exactly, so the cut must land on a pool - // boundary. But F.relu drives most pool scores to exactly 0.0, so tie groups SPAN - // pools, the cut falls inside one, and ggml_top_k - explicitly unordered among - // equals - takes an arbitrary 1..kpool-1 members of the pool it lands in. A - // pool-aligned top-k WIDTH does not save it: the ties are not aligned to anything. - // Measured on TinySparse at 512 tokens, the cell-level form leaves a partial pool - // on 7.51% of query rows at L3 and 5.93% at L7 (70 and 47 partial pools); this form - // leaves none. The index-set jaccard does not separate them - it scored the broken - // form at 0.9958/0.9764 against a bf16 noise floor of 0.9779/0.8385, i.e. ABOVE the - // floor - which is why scripts/glm5next_pool_integrity.py exists + // top-k over POOLS then expand, as in the reference: a cell-level top-k is wrong + // because relu ties span pool boundaries and ggml_top_k splits the pool it lands in const int64_t select_k = llama_kpool_select_k(n_pools, hparams.indexer_top_k, r); GGML_ASSERT(select_k > 0 && select_k <= n_pools); @@ -586,14 +451,10 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_tensor * sel = ggml_cont(ctx0, ggml_top_k(ctx0, pool_score, (int) select_k)); cb(sel, "indexer_top_k_pools", il); - // Expand pools to members: gather whole rows of `kpool` cells out of pool_cells. - // The query axis folds into the gather's row axis, which is what lets ONE - // ggml_get_rows serve every query, while the stream axis stays where get_rows wants - // it (src0 dim 2 is indexed by the index tensor's dim 1) + // the query axis folds into the gather's row axis, so ONE get_rows serves every query ggml_tensor * pc3 = ggml_reshape_3d(ctx0, inp_kp->pool_cells, r, n_pools, n_stream); ggml_tensor * sel_flat = ggml_reshape_2d(ctx0, sel, select_k*n_tps, n_stream); - // {r, select_k*n_tps, n_stream} -> {r*select_k, n_tps, n_stream} ggml_tensor * top_k = ggml_get_rows(ctx0, pc3, sel_flat); GGML_ASSERT(top_k->type == GGML_TYPE_I32 && "pool_cells is I32, so the gather stays I32"); top_k = ggml_reshape_3d(ctx0, top_k, r*select_k, n_tps, n_stream); @@ -602,16 +463,8 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( return top_k; } -// -// DSA layer. `scoring` false takes the dense limit: full MLA over the whole cache, no -// kpool, no top-k, which is what sparse selection collapses to below -// glm5next_n_select() resident positions -// -// absorbed form, as in deepseek2/deepseek32/glm-dsa: q_nope is pushed through wk_b so -// q.k is taken against the 512-wide latent the cache actually holds (is_mla() drops the -// V allocation and V becomes a view of K). the naive form would re-expand the latent to -// n_head 256-wide k/v every step and needs a V cache this layout does not have -// +// absorbed form (deepseek2/glm-dsa): q_nope goes through wk_b so q.k is taken against +// the latent the cache holds; the naive form needs a V cache this layout lacks ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( const llama_layer & layer, llm_graph_input_attn_k * inp_attn, @@ -622,21 +475,15 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); const int64_t kv_lora_rank = hparams.n_lora_kv; - // nope-only: the rope half is zero-width, so no split, no concat and no rope - // anywhere in the text tower GGML_ASSERT(hparams.n_rot() == 0); - // scale is over the MLA head size, as in the reference, not over the post-absorption - // width: 1/sqrt(kv_lora_rank) = 1/sqrt(512) would be a different model + // scale is over the MLA head size, as in the reference, not the absorbed width const float kq_scale = 1.0f/sqrtf(float(qk_head_dim)); 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, "dsa_q_a_norm", il); - // the indexer shares this q LoRA residual with the main MLA path and consumes it - // with its own wq_b, so it is built here rather than being handed the layer input - // twice. it also writes the indexer cache, which happens on the dense path too ggml_tensor * top_k = inp_kp ? build_indexer(layer, inp_kp, cur, qr, scoring, il) : nullptr; ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr); @@ -647,14 +494,10 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); cb(kv, "dsa_kv_a_norm", il); - // {qk_head_dim, n_tokens, n_head} q = ggml_permute(ctx0, q, 0, 2, 1, 3); - // {qk_head_dim, kv_lora_rank, n_head} x {qk_head_dim, n_tokens, n_head} q = ggml_mul_mat(ctx0, layer.wk_b, q); - // {kv_lora_rank, n_head, n_tokens}. deepseek2 gets this contiguous for free from the - // concat with the roped half, which does not exist here q = ggml_cont(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3)); cb(q, "dsa_q_absorbed", il); @@ -697,8 +540,7 @@ ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( int il) const { const auto & layer = model.layers[il]; - // the leading dense layers clamp the same way the experts do: the reference - // routes both through one Glm5NextTextMLP, so swiglu_limit is not MoE-only + // the leading dense layers clamp like the experts: one Glm5NextTextMLP serves both if (il < (int) hparams.n_layer_dense_lead) { return build_ffn(cur, layer.ffn_up, nullptr, nullptr, @@ -707,8 +549,7 @@ ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); } - // noaux_tc: exp_probs_b biases top-k selection only, the weights are the - // unbiased sigmoid scores. n_group is 1, so the group mask is a no-op + // noaux_tc: exp_probs_b biases top-k selection only; the weights stay unbiased ggml_tensor * moe_out = build_moe_ffn(cur, layer.ffn_gate_inp, layer.ffn_up_exps, @@ -728,8 +569,7 @@ ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); cb(shexp, "ffn_shexp", il); - // shared expert unscaled: routed_scaling_factor is applied inside build_moe_ffn, - // after norm_topk_prob, to the routed weights only + // shared expert unscaled: routed_scaling_factor applies to the routed weights only return ggml_add(ctx0, moe_out, shexp); } @@ -740,28 +580,10 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp = build_inp_embd(model.tok_embd); ggml_tensor * inp_out_ids = build_inp_out_ids(); - // MLA absorption leaves a K-only cache holding the latent, so the attention half of - // the hybrid memory is the _k variant, as in bailingmoe3 llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); - // One pooling map for the whole ubatch. pool_cells, pool_bias, sel_mask and - // cand_mask depend only on the cells and the ubatch, never on the layer, so - // rebuilding them per DSA layer would cost O(n_kv * n_tokens) host writes eleven - // times over - at 128 Ki cells and 512 tokens that dominates prefill on its own. - // - // The map is built whenever the model HAS an indexer cache, because the indexer - // key and gate store runs unconditionally. Only `scoring` is gated: below - // index_topk + index_kpool - 1 resident positions the reference selects every - // visible position, so the dense build_attn is not an approximation there, it is - // the same function. - // - // Gated on n_ctx and not on the ubatch's n_kv, even though n_kv is what actually - // decides whether the budget binds. n_kv grows as the cache fills, so gating on it - // would flip the graph's topology partway through a run. n_ctx is fixed for the - // lifetime of the context, so the topology is decided once. The cost is that a - // context configured larger than n_select runs the indexer even while the cache is - // still short, where it selects every visible pool: wasted work, never a wrong - // answer, and llama_kpool_select_k clamps the budget to the pools that exist. + // one map for the whole ubatch; nothing in it depends on the layer. gated on n_ctx, + // not n_kv, which grows and would flip the graph topology mid-run llm_graph_input_kpool * inp_kp = nullptr; bool indexer_scoring = false; { @@ -820,8 +642,7 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa &post, &comb, il); cb(cur, "hc_ffn_pre", il); - // expand before the sublayer so op offload does not pull the mHC state - // onto the expert weights' backend, as in deepseek4 + // expand before the sublayer so op offload does not pull mHC state to the experts ggml_build_forward_expand(gf, residual); ggml_build_forward_expand(gf, post); ggml_build_forward_expand(gf, comb); @@ -865,8 +686,7 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa } std::unique_ptr llama_model_glm5next::build_arch_graph(const llm_graph_params & params) const { - // llama_init_from_model accepts an MTP context whenever n_layer_nextn > 0, - // which every glm5next checkpoint has; without this it silently runs the trunk + // without this, an MTP context (accepted whenever n_layer_nextn > 0) runs the trunk GGML_ASSERT(params.gtype != LLM_GRAPH_TYPE_DECODER_MTP && "glm5next NextN graph not implemented yet"); return std::make_unique(*this, params); diff --git a/src/models/models.h b/src/models/models.h index 1182fd7207b..cfd725f5fb3 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1173,10 +1173,8 @@ 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; - // llm_build_delta_net_base is a method-only mixin over llm_graph_context (no data - // members, no new virtuals), so this graph's layout and behaviour are unchanged. - // deepseek4 has no recurrent layers; it is here so glm5next, which derives from - // this graph for the mHC residual, can reach build_delta_net for its KDA layers + // method-only mixin; here only so glm5next, deriving from this graph, reaches + // build_delta_net. deepseek4 itself has no recurrent layers struct graph : public llm_build_delta_net_base { graph(const llm_graph_params & params) : llm_build_delta_net_base(params) {} graph(const llama_model & model, const llm_graph_params & params); @@ -1299,7 +1297,6 @@ struct llama_model_deepseek4 : public llama_model_base { ggml_tensor * comb, int il) const; - // unweighted mean over the streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens] static ggml_tensor * build_hc_mean( ggml_context * ctx, ggml_tensor * x); @@ -1345,9 +1342,7 @@ struct llama_model_glm5next : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; - // glm5next's mHC is DeepSeek-V4's hyper-connection block (same wide residual, - // 24-row mixer split, activations, Sinkhorn); only the final collapse differs - // (unweighted mean, not a learned gated head), so derive rather than restate + // glm5next's mHC is DeepSeek-V4's block, collapsed by unweighted mean, not a gated head struct graph : public llama_model_deepseek4::graph { graph(const llama_model & model, const llm_graph_params & params); @@ -1366,11 +1361,7 @@ struct llama_model_glm5next : public llama_model_base { ggml_tensor * cur, int il); - // `scoring` false keeps the dense path: at or below - // index_topk + index_kpool - 1 resident positions the indexer selects - // every visible position, so dense is not an approximation there, it is - // the same function without the pooling work. The indexer STORE still - // runs; only the selection is skipped + // `scoring` false keeps the dense path, the same function below n_select ggml_tensor * build_dsa_layer( const llama_layer & layer, llm_graph_input_attn_k * inp_attn, @@ -1379,10 +1370,7 @@ struct llama_model_glm5next : public llama_model_base { ggml_tensor * cur, int il) const; - // writes this layer's indexer key and compressor gate into the indexer - // cache - always - and then, when `scoring`, returns the I32 attention - // cache CELL indices this layer's queries select, already expanded from - // whole pools: [kpool*select_k, n_tps, n_stream] + // always stores the key and gate; when `scoring`, returns the selected CELL indices ggml_tensor * build_indexer( const llama_layer & layer, llm_graph_input_kpool * inp_kp, diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 913f50e59b1..a5a87eb0f88 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -686,9 +686,7 @@ ggml_tensor * clip_graph::build_ffn( } break; case FFN_SILU_CLAMP: { - // the gate is bounded above only, the up projection on both sides, and both - // before the activation. ggml_swiglu_oai clamps the same way but then adds - // one to the up branch, which is a gpt-oss detail this model does not share + // not ggml_swiglu_oai: it clamps the same way but adds one to the up branch GGML_ASSERT(gate && "FFN_SILU_CLAMP is a gated activation"); const float limit = hparams.swiglu_limit; GGML_ASSERT(limit > 0.0f); @@ -1749,16 +1747,13 @@ struct clip_model_loader { { hparams.rope_theta = 10000.0f; hparams.n_merge = 2; - // the reference asks for BICUBIC (resample = PILImageResampling.BICUBIC); - // the bilinear here came from the GLM-4V case this was copied from. neither - // filter matches torchvision's exactly, so this is closer in kind, not exact + // the reference asks for PILImageResampling.BICUBIC, which this only approximates hparams.image_resize_algo = RESIZE_ALGO_BICUBIC; get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); - // drives the clamp in both the per-block MLP and the merger get_f32(KEY_VISION_SWIGLU_LIMIT, hparams.swiglu_limit); hparams.ffn_op = FFN_SILU_CLAMP; log_ffn_op = "silu_clamp"; - // min_pixels/max_pixels of the GLM-5.3-Flash preprocessor, in tokens + // the preprocessor's min_pixels/max_pixels, in tokens hparams.set_limit_image_tokens(16, 8000); hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup } break; diff --git a/tools/mtmd/models/glm4v.cpp b/tools/mtmd/models/glm4v.cpp index ad36e7a483d..78f6d7622e4 100644 --- a/tools/mtmd/models/glm4v.cpp +++ b/tools/mtmd/models/glm4v.cpp @@ -42,8 +42,7 @@ ggml_cgraph * clip_graph_glm4v::build() { cb(inp, "patch_bias", -1); // pos-conv norm - // Note: GLM-OCR does not have a post-conv norm, and build_norm still normalizes when the - // weight is null, so the whole call must be skipped rather than just the affine scale + // GLM-OCR has none, and build_norm still normalizes on a null weight, so skip the call if (model.norm_embd_w != nullptr) { inp = build_norm(inp, model.norm_embd_w, model.norm_embd_b, norm_t, eps, -1); } diff --git a/tools/mtmd/models/glm5next-vision.cpp b/tools/mtmd/models/glm5next-vision.cpp index e1130f1cc5d..5d16cb2c367 100644 --- a/tools/mtmd/models/glm5next-vision.cpp +++ b/tools/mtmd/models/glm5next-vision.cpp @@ -1,12 +1,8 @@ #include "models.h" -// GLM-5.3-Flash reuses the GLM-OCR ViT unchanged: no learned position embeddings, no -// post-conv norm, q/k norms and biases throughout, and an RMSNorm called post_layernorm. +// the GLM-OCR ViT plus a clamp on the SwiGLU gate and up projections, which the per-block MLP +// and the merger both pick up from hparams.ffn_op // ref: https://huggingface.co/zai-org/GLM-5.3-Flash/blob/main/modeling_glm5_next.py -// -// The one difference is the clamp on the SwiGLU gate and up projections. It applies to -// the per-block MLP and to the merger alike, and both read hparams.ffn_op, so selecting -// FFN_SILU_CLAMP for this projector type is enough to cover the pair. ggml_cgraph * clip_graph_glm5next::build() { GGML_ASSERT(hparams.ffn_op == FFN_SILU_CLAMP); GGML_ASSERT(model.norm_embd_w == nullptr); diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 5dc80faae4e..4627d4f68b2 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1588,8 +1588,7 @@ mtmd_image_preproc_out mtmd_image_preprocessor_muse_glimmer::preprocess(const cl // mtmd_image_preprocessor_glm5next // -// for a still image the reference's temporal_factor cancels on both sides of its budget comparison, -// leaving a plain pixel area against image_min_pixels / image_max_pixels (n_tokens * factor**2) +// for a still image the reference's temporal_factor cancels out, leaving pixel area vs min/max clip_image_size mtmd_image_preprocessor_glm5next::smart_resize(const clip_hparams & hparams, const clip_image_size & size) { const int factor = hparams.patch_size * hparams.n_merge; GGML_ASSERT(factor > 0); @@ -1598,7 +1597,7 @@ clip_image_size mtmd_image_preprocessor_glm5next::smart_resize(const clip_hparam const int height = size.height; const int width = size.width; if (height <= 0 || width <= 0) { - // an empty bitmap is reachable through the public API, same tolerance as calc_size_preserved_ratio + // an empty bitmap is reachable through the public API return { 0, 0 }; } @@ -1612,7 +1611,6 @@ clip_image_size mtmd_image_preprocessor_glm5next::smart_resize(const clip_hparam int aligned_height = align(height); int aligned_width = align(width); - // upscale an image that is too small to spend the minimum token budget if ((int64_t) aligned_height * aligned_width < min_pixels) { const double scale = std::sqrt((double) min_pixels / ((double) height * (double) width)); aligned_height = align(std::max(1, (int64_t) std::ceil(height * scale))); @@ -1620,8 +1618,7 @@ clip_image_size mtmd_image_preprocessor_glm5next::smart_resize(const clip_hparam } if ((int64_t) aligned_height * aligned_width > max_pixels) { - // binary search the tallest content height whose aligned canvas still fits the budget. the Qwen - // sqrt(area / max_pixels) scale leaves budget unspent: aligning both edges is not monotone in it + // aligning both edges is not monotone in the Qwen sqrt(area / max_pixels) scale, so search int low = 1; int high = height; aligned_height = factor; @@ -1657,7 +1654,7 @@ mtmd_image_preprocessor_glm5next::geometry mtmd_image_preprocessor_glm5next::get double scale = std::min((double) canvas.height / height, (double) canvas.width / width); if ((int64_t) height * (int64_t) width >= hparams.image_min_pixels) { - // an image already spending the minimum budget is only shrunk, never upscaled to fill the canvas + // already spending the minimum budget, so shrink only, never upscale to fill the canvas scale = std::min(1.0, scale); } diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index 4e0e707b2b8..41126a98c51 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -240,11 +240,8 @@ struct mtmd_image_preprocessor_muse_glimmer : mtmd_image_preprocessor { mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; }; -// GLM-5-Next / GLM-5.3-Flash dynamic resize -// ref: Glm5NextImageProcessor.{smart_resize,resize} of the 2026-08-26 adaptation -// -// unlike the Qwen-style smart_resize in mtmd_image_preprocessor_dyn_size, an over-budget image is fitted -// by binary search on the content height, then pasted top-left rather than centred and stretched to fill +// ref: Glm5NextImageProcessor.{smart_resize,resize}. unlike the Qwen-style smart_resize in +// mtmd_image_preprocessor_dyn_size, an over-budget image is pasted top-left, not centred struct mtmd_image_preprocessor_glm5next : mtmd_image_preprocessor { mtmd_image_preprocessor_glm5next(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; @@ -254,7 +251,7 @@ struct mtmd_image_preprocessor_glm5next : mtmd_image_preprocessor { clip_image_size content; // resized image, placed at the top-left of the canvas }; - // pure arithmetic, exposed so tests can reach it without a clip_ctx + // static so tests can reach it without a clip_ctx static clip_image_size smart_resize(const clip_hparams & hparams, const clip_image_size & size); static geometry get_geometry(const clip_hparams & hparams, const clip_image_size & size); }; From b5517b15dcea347ed88230d937dd5306209c345b Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 28 Aug 2026 06:57:19 +0000 Subject: [PATCH 31/36] tests : add the glm5next fixture --- tests/test-llama-archs.cpp | 36 ++++++++++++++++++++++++++++++++++-- 1 file changed, 34 insertions(+), 2 deletions(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index d58d90952eb..6e19954c683 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -118,7 +118,8 @@ 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_MISTRAL4) { + || arch == LLM_ARCH_MISTRAL4 + || arch == LLM_ARCH_GLM5NEXT) { n_embd = 128; n_head = 1; n_ff = 192; @@ -174,6 +175,17 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } 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 if (arch == LLM_ARCH_GLM5NEXT) { + // head_count doubles as the KDA head count, so it stays uniform; the kv array is what + // marks the recurrent layers, and the loader asserts it holds both a zero and a nonzero + GGML_ASSERT(n_layer >= 2); + std::vector n_head_kv_per_layer; + n_head_kv_per_layer.reserve(n_layer); + for (uint32_t il = 0; il < n_layer; il++) { + n_head_kv_per_layer.push_back(il == 1 ? 0 : n_head_kv); + } + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head); + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_kv_per_layer); } else { ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head); ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(1) : n_head_kv); @@ -213,12 +225,21 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types); } + } else if (arch == LLM_ARCH_GLM5NEXT) { + // nope-only MLA: the cache holds the bare latent, so no rope width is added on top of + // the kv LoRA rank and n_rot has to be an explicit 0, not the head size default + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(512)); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512)); + ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(0)); + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192)); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128)); } else if (arch == LLM_ARCH_MINIMAX_M3) { // partial rotary: n_rot must not exceed the indexer key length (64) ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); } ms.add_kv(LLM_KV_ATTENTION_CLAMP_KQV, 1.0f); - ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS, 1e-5f); + // glm5next warns on anything but the 1e-6 its indexer k_norm hardcodes + ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS, arch == LLM_ARCH_GLM5NEXT ? 1e-6f : 1e-5f); ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, 1e-5f); ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS, 1e-5f); ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS, uint32_t(8)); @@ -277,6 +298,16 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 10.0f); ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); + } else if (arch == LLM_ARCH_GLM5NEXT) { + ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); // build_hc_pre asserts exactly 4 streams + ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); + // the only arch that pools indexer keys; top_k must be a whole number of pools, and + // the resulting selection width has to stay under n_ctx or the sparse path goes unused + ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, uint32_t(4)); + // glm5next reads these unconditionally; the if (moe) block below never sets them + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } ms.add_kv(LLM_KV_TOKENIZER_MODEL, "no_vocab"); // ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT, n_embd); @@ -432,6 +463,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_EXAONE_MOE: case LLM_ARCH_BAILINGMOE: case LLM_ARCH_BAILINGMOE2: From f30bed88717059d8a4728864c88f8abad8d329a0 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 28 Aug 2026 06:58:46 +0000 Subject: [PATCH 32/36] kv-cache : group the context accessors after type_v --- src/llama-kv-cache.cpp | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ca2d12d5f20..8af06a570e6 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2779,14 +2779,6 @@ uint32_t llama_kv_cache_context::get_n_kv() const { return n_kv; } -uint32_t llama_kv_cache_context::get_n_stream() const { - return sinfos[i_cur].s1 - sinfos[i_cur].s0 + 1; -} - -const llama_kv_cache * llama_kv_cache_context::get_kv() const { - return kv; -} - ggml_type llama_kv_cache_context::type_k() const { return kv->type_k(); } @@ -2795,6 +2787,14 @@ ggml_type llama_kv_cache_context::type_v() const { return kv->type_v(); } +uint32_t llama_kv_cache_context::get_n_stream() const { + return sinfos[i_cur].s1 - sinfos[i_cur].s0 + 1; +} + +const llama_kv_cache * llama_kv_cache_context::get_kv() const { + return kv; +} + ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il) const { return kv->get_k(ctx, il, n_kv, sinfos[i_cur]); } From 00699716c275498ff84d71e329178fe21cba56a6 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 30 Aug 2026 03:28:18 +0000 Subject: [PATCH 33/36] Faster inference --- src/llama-graph.cpp | 26 ++++++ src/llama-kv-cache-kpool.cpp | 159 ++++++++++++++++++++++++++++++++++- src/llama-kv-cache-kpool.h | 33 ++++++++ src/llama-kv-cache.cpp | 49 +++++++++++ src/llama-kv-cache.h | 28 ++++++ src/llama-memory-hybrid.cpp | 17 +++- src/models/glm5next.cpp | 40 +++++++-- 7 files changed, 341 insertions(+), 11 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index e7ab7fc7bec..8677083f55e 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3614,6 +3614,32 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( inp->cand_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, n_kv, n_tps, 1, n_stream); ggml_set_input(inp->cand_mask); ggml_set_name(inp->cand_mask, "kpool_cand_mask"); + + // pooled-key cache. n_new_max is an exact bound on the pools one ubatch can + // complete: a sequence's tokens are a contiguous position run, so a run of L tokens + // closes at most L/kpool + 1 pools. it is FIXED for the whole decode phase so the + // graph shape does not depend on how many pools happened to close this step. + // after a position mutation every cached pooled key is stale, so this one graph has + // to be able to re-emit all of them. a shape change here forces a rebuild, which is + // the point: the wide shape is used for exactly one ubatch and then goes away. + const bool rebuild = mctx_attn->get_kv()->get_kpool_dirty(); + const int64_t n_new_max = rebuild ? n_pools : n_tps/kpool + n_ps; + + inp->n_new_max = (uint32_t) n_new_max; + inp->rebuild = rebuild; + + inp->pool_reps = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_pools, n_stream); + ggml_set_input(inp->pool_reps); + ggml_set_name(inp->pool_reps, "kpool_pool_reps"); + + inp->new_pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool*n_new_max, n_stream); + ggml_set_input(inp->new_pool_cells); + ggml_set_name(inp->new_pool_cells, "kpool_new_pool_cells"); + + // I64 because ggml_set_rows takes its row indices as I64 + inp->new_pool_reps = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new_max*n_stream); + ggml_set_input(inp->new_pool_reps); + ggml_set_name(inp->new_pool_reps, "kpool_new_pool_reps"); } return (llm_graph_input_kpool *) res->add_input(std::move(inp)); diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index 3fc1eaca616..485260586b8 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -71,6 +71,12 @@ void llama_kv_cache_set_input_kpool( ggml_tensor * pool_bias, ggml_tensor * sel_mask, ggml_tensor * cand_mask, + ggml_tensor * pool_reps, + ggml_tensor * new_pool_cells, + ggml_tensor * new_pool_reps, + const uint32_t * strm_of, + int64_t kv_size, + bool rebuild, const llama_ubatch * ubatch, uint32_t kpool) { GGML_ASSERT(kv != nullptr); @@ -136,6 +142,40 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(bias->ne[0] == n_kv && bias->ne[1] == n_tps && bias->ne[2] == n_ns); } + // pooled-key cache inputs travel together or not at all + const bool kcache = pool_reps != nullptr; + + GGML_ASSERT((new_pool_cells != nullptr) == kcache); + GGML_ASSERT((new_pool_reps != nullptr) == kcache); + + int64_t n_new_max = 0; + + if (kcache) { + GGML_ASSERT(ggml_backend_buffer_is_host(pool_reps ->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_cells->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_reps ->buffer)); + + GGML_ASSERT(pool_reps ->type == GGML_TYPE_I32); + GGML_ASSERT(new_pool_cells->type == GGML_TYPE_I32); + GGML_ASSERT(new_pool_reps ->type == GGML_TYPE_I64); + + GGML_ASSERT(ggml_is_contiguous(pool_reps)); + GGML_ASSERT(ggml_is_contiguous(new_pool_cells)); + GGML_ASSERT(ggml_is_contiguous(new_pool_reps)); + + GGML_ASSERT(pool_reps->ne[0] == n_pools && pool_reps->ne[1] == n_ns); + GGML_ASSERT(new_pool_cells->ne[0] % r == 0 && new_pool_cells->ne[1] == n_ns); + + n_new_max = new_pool_cells->ne[0]/r; + + GGML_ASSERT(new_pool_reps->ne[0] == n_new_max*n_ns); + GGML_ASSERT(strm_of != nullptr && kv_size > 0); + } + + int32_t * dst_pool_reps = kcache ? (int32_t *) pool_reps ->data : nullptr; + int32_t * dst_new_cells = kcache ? (int32_t *) new_pool_cells->data : nullptr; + int64_t * dst_new_reps = kcache ? (int64_t *) new_pool_reps ->data : nullptr; + int32_t * dst_cell_pool = cell_pool ? (int32_t *) cell_pool->data : nullptr; int32_t * dst_pool_cells = (int32_t *) pool_cells->data; float * dst_bias = bias ? (float *) bias->data : nullptr; @@ -167,6 +207,25 @@ void llama_kv_cache_set_input_kpool( std::fill(cur_pool_cells, cur_pool_cells + r*n_pools, 0); std::fill(cur_pool_bias, cur_pool_bias + n_tps*n_pools, -INFINITY); + int32_t * cur_pool_reps = kcache ? dst_pool_reps + s*n_pools : nullptr; + int32_t * cur_new_cells = kcache ? dst_new_cells + s*(r*n_new_max) : nullptr; + int64_t * cur_new_reps = kcache ? dst_new_reps + s*n_new_max : nullptr; + + // count of real entries emitted for this stream; the rest is padding + int64_t n_new = 0; + + // members of any complete pool in this stream, used to pad the fixed-size write. + // recomputing a complete pool is idempotent, so a repeat is always safe. + const int32_t * any_rep_src = nullptr; + + if (kcache) { + // a pool with no rep gathers row 0. that row's pooled third may hold another + // pool's key, but such a pool is always -INFINITY in pool_bias, so the value is + // discarded before it can score. + std::fill(cur_pool_reps, cur_pool_reps + n_pools, 0); + std::fill(cur_new_cells, cur_new_cells + r*n_new_max, 0); + } + // the token loop writes rows < n_tps in full; only the padding rows need clearing if (mask_f16) { kpool_mask_fill((ggml_fp16_t *) (cur_sel_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); @@ -293,6 +352,61 @@ void llama_kv_cache_set_input_kpool( } } + if (kcache) { + // the pooled key of a complete pool lives in the row of its LAST member, + // the cell holding pos % r == r-1. that slot is only meaningful once the + // pool is complete: for a partial pool it is still 0 from the fill above, + // and cell 0 is a real cell whose own pooled key we must not overwrite. + for (int64_t p = 0; p < n_run; ++p) { + if (filled[p] == (int32_t) r) { + cur_pool_reps[run_off[ps] + p] = part_pool_cells[p*r + (r - 1)]; + + if (any_rep_src == nullptr) { + any_rep_src = part_pool_cells + p*r; + } + } + } + + // recompute exactly the complete pools this ubatch wrote into. touched[] is + // over the run, so the cost is O(tokens), not O(n_kv). + std::vector touched(n_run, 0); + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + + if (ubatch->seq_id[i][0] != seq_of_pool) { + continue; + } + + const int64_t bo = ubatch->pos[i]/r - b_base; + + if (bo >= 0 && bo < n_run) { + touched[bo] = 1; + } + } + + for (int64_t p = 0; p < n_run; ++p) { + // in rebuild mode every cached pooled key is stale, so re-emit all of + // them, not just the pools this ubatch closed + if ((!touched[p] && !rebuild) || filled[p] != (int32_t) r) { + continue; + } + + // n_new_max = n_tps/kpool + n_ps bounds this while a sequence's tokens in + // one ubatch are a contiguous position run, which llama-batch.cpp + // enforces. dropping a completed pool would silently serve a stale key, + // so fail loudly instead of clamping. + GGML_ASSERT(n_new < n_new_max && "k-pool: more pools completed than the fixed bound"); + + std::copy(part_pool_cells + p*r, part_pool_cells + (p + 1)*r, + cur_new_cells + n_new*r); + + cur_new_reps[n_new] = (int64_t) strm_of[s]*kv_size + part_pool_cells[p*r + (r - 1)]; + + n_new++; + } + } + for (int64_t ii = 0; ii < n_tps; ++ii) { const int64_t i = s*n_tps + ii; @@ -350,6 +464,25 @@ void llama_kv_cache_set_input_kpool( // exactly one partition per row, or a query reads another sequence's pools GGML_ASSERT(n_done == n_tps && "every query must belong to a sequence of the ubatch"); + + if (kcache) { + // the write has a fixed row count, so the unused slots must name a destination + // that is safe to overwrite. two cases, both provably harmless: + // - some complete pool exists: repeat it. recomputing a complete pool yields + // the value already there, so the duplicate write is a no-op in effect. + // - none exists: no cell in this stream is the last member of a complete pool, + // so no pooled third is read (every pool is -INFINITY in pool_bias). cell 0 + // is then free. + for (int64_t p = n_new; p < n_new_max; ++p) { + if (any_rep_src) { + std::copy(any_rep_src, any_rep_src + r, cur_new_cells + p*r); + cur_new_reps[p] = (int64_t) strm_of[s]*kv_size + any_rep_src[r - 1]; + } else { + // cells already 0 from the fill; pool r copies of cell 0 into cell 0 + cur_new_reps[p] = (int64_t) strm_of[s]*kv_size; + } + } + } } } @@ -363,8 +496,32 @@ void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { return; } + // the pooled key is written into the INDEXER cache, whose slot layout the attention + // cache defines, so the stream map and cell count come from the indexer side + std::vector strm_of; + + if (pool_reps) { + strm_of.resize(mctx_idx->get_n_stream()); + + for (uint32_t s = 0; s < strm_of.size(); ++s) { + strm_of[s] = mctx_idx->get_strm(s); + } + } + llama_kv_cache_set_input_kpool( mctx_attn->get_kv(), /* cell_pool */ nullptr, pool_cells, /* bias */ nullptr, pool_bias, - sel_mask, cand_mask, ubatch, kpool); + sel_mask, cand_mask, + pool_reps, new_pool_cells, new_pool_reps, + strm_of.empty() ? nullptr : strm_of.data(), + pool_reps ? (int64_t) mctx_idx->get_kv()->get_size() : 0, + rebuild, + ubatch, kpool); + + // every pool has just been re-emitted, so the cache is consistent again. cleared here + // rather than in build_inp_kpool because a graph that is built but not evaluated must + // not clear it. + if (rebuild) { + mctx_attn->get_kv()->clear_kpool_dirty(); + } } diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index 7bd9d862a29..ff20a6e6667 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -33,6 +33,10 @@ uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t // sel_mask F16/F32 [n_kv, n_batch, 1, n_stream] 0.0f on the always-selected tail only // cand_mask F16/F32 [n_kv, n_batch, 1, n_stream] max(bias, sel_mask); bounds the top-k // spills that a partial seq_rm would otherwise let escape the candidate set +// pool_reps / new_pool_cells / new_pool_reps are optional (nullptr when the pooled-key +// cache is off). an entry is emitted only for a pool with filled == kpool: an incomplete +// pool's last-member slot is 0, and cell 0 is a legitimate cell, so writing it would +// clobber another pool's cached key. void llama_kv_cache_set_input_kpool( const llama_kv_cache * kv, ggml_tensor * cell_pool, @@ -41,6 +45,18 @@ void llama_kv_cache_set_input_kpool( ggml_tensor * pool_bias, ggml_tensor * sel_mask, ggml_tensor * cand_mask, + ggml_tensor * pool_reps, + ggml_tensor * new_pool_cells, + ggml_tensor * new_pool_reps, + // strm_of[s] is the physical stream behind view s and kv_size the per-stream cell + // count, so a global row is strm_of[s]*kv_size + cell. only read when pool_reps is + // set. the indexer cache shares the attention cache's slot layout, so one cell + // index addresses both. + const uint32_t * strm_of, + int64_t kv_size, + // re-emit every complete pool, not only the ones this ubatch closed. set after a + // position mutation; the graph is built with n_new_max == n_pools to hold them. + bool rebuild, const llama_ubatch * ubatch, uint32_t kpool); @@ -62,6 +78,23 @@ class llm_graph_input_kpool : public llm_graph_input_i { ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] ggml_tensor * pool_bias = nullptr; // F32 [n_pools, n_tps, n_stream] + // pooled-key cache. the pooled value of a pool lives in the third head of the row of + // the cell holding its LAST member (pos % kpool == kpool-1), so it is a pure function + // of cell content: independent of sequence and of pool ordinal, which is what lets a + // seq_cp share it and a rebase leave it alone. + ggml_tensor * pool_reps = nullptr; // I32 [n_pools, n_stream] stream-local rep cell + ggml_tensor * new_pool_cells = nullptr; // I32 [kpool*n_new_max, n_stream] members to (re)pool + ggml_tensor * new_pool_reps = nullptr; // I64 [n_new_max*n_stream] GLOBAL dest row + + // n_new_max is fixed for the whole decode phase on purpose. making the shape depend on + // how many pools happened to close this step would flip the graph topology every kpool + // tokens and defeat graph reuse / force CUDA-graph recapture. + uint32_t n_new_max = 0; + + // this graph was built to re-emit every pool after a position mutation, not just the + // ones this ubatch closed. decided at build time because it sets n_new_max above. + bool rebuild = false; + // exact, since pool_bias only holds 0.0f or -INFINITY. nullptr if the fused path is off ggml_tensor * pool_bias_f16 = nullptr; // F16 [n_pools, n_tps, 1, n_stream] diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 8af06a570e6..62f51287b44 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1217,6 +1217,18 @@ bool llama_kv_cache::get_has_shift() const { return result; } +void llama_kv_cache::set_kpool_dirty() { + kpool_dirty = true; +} + +bool llama_kv_cache::get_kpool_dirty() const { + return kpool_dirty; +} + +void llama_kv_cache::clear_kpool_dirty() const { + kpool_dirty = false; +} + ggml_type llama_kv_cache::type_k() const { return layers[0].k->type; } @@ -1351,6 +1363,30 @@ ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggm return ggml_set_rows(ctx, k, k_cur, k_idxs); } +ggml_tensor * llama_kv_cache::cpy_k_part(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, + int64_t n_embd, int64_t i_off) const { + const int32_t ikv = map_layer_ids.at(il); + + ggml_tensor * k = layers[ikv].k; + + const int64_t n_embd_gqa = k->ne[0]; + const int64_t kv_size = get_size(); + const int64_t n_stream = k->ne[2]; + + GGML_ASSERT(i_off >= 0 && i_off + n_embd <= n_embd_gqa); + GGML_ASSERT(k_cur->ne[0] == n_embd); + + // merge the streams: k_idxs are global, exactly as in cpy_k + ggml_tensor * k2 = ggml_reshape_2d(ctx, k, n_embd_gqa, kv_size*n_stream); + + // a row-slice view of every cell. ggml_set_rows needs contiguous rows in the DEST, + // which ggml_is_contiguous_rows() grants for a view whose ne[0] slice is contiguous. + ggml_tensor * dst = ggml_view_2d(ctx, k2, n_embd, kv_size*n_stream, + k2->nb[1], ggml_row_size(k2->type, i_off)); + + return ggml_set_rows(ctx, dst, k_cur, k_idxs); +} + ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const { GGML_UNUSED(sinfo); @@ -2791,6 +2827,14 @@ uint32_t llama_kv_cache_context::get_n_stream() const { return sinfos[i_cur].s1 - sinfos[i_cur].s0 + 1; } +uint32_t llama_kv_cache_context::get_strm(uint32_t s) const { + const auto & sinfo = sinfos[i_cur]; + + GGML_ASSERT(s < sinfo.strm.size()); + + return sinfo.strm[s]; +} + const llama_kv_cache * llama_kv_cache_context::get_kv() const { return kv; } @@ -2807,6 +2851,11 @@ ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_ return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]); } +ggml_tensor * llama_kv_cache_context::cpy_k_part(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, + int64_t n_embd, int64_t i_off) const { + return kv->cpy_k_part(ctx, k_cur, k_idxs, il, n_embd, i_off); +} + ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const { return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 0eed6a68f93..2ba032a34bf 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -160,6 +160,16 @@ class llama_kv_cache : public llama_memory_i { bool get_has_shift() const; + // GLM-5-Next pooled-key cache. a pool groups cells by absolute position, so anything + // that mutates positions in place (seq_add / seq_div) regroups the pools while every + // cached pooled key still looks complete. set here, consumed and cleared by + // llama_kv_cache_set_input_kpool, which then emits every pool instead of only the new + // ones. sticky by design: a flag that never clears degrades to the pre-cache cost, a + // flag that clears too early is silently wrong. + void set_kpool_dirty(); + bool get_kpool_dirty() const; + void clear_kpool_dirty() const; + ggml_type type_k() const; ggml_type type_v() const; @@ -193,6 +203,13 @@ class llama_kv_cache : public llama_memory_i { ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const; ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const; + // write n_embd elements at element offset i_off inside each row, leaving the rest of + // the row untouched. k_idxs are GLOBAL rows (stream-merged), as in cpy_k. + // returns the ggml_set_rows result so a later gather can be chained off it and become + // a real graph edge rather than relying on build order. + ggml_tensor * cpy_k_part(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, + int64_t n_embd, int64_t i_off) const; + // // preparation API // @@ -261,6 +278,10 @@ class llama_kv_cache : public llama_memory_i { bool v_trans = true; // the value tensor is transposed + // see set_kpool_dirty. mutable because the only consumer runs from set_input, which + // holds the cache by const pointer; nothing else observes it. + mutable bool kpool_dirty = false; + const uint32_t n_seq_max = 1; const uint32_t n_stream = 1; @@ -394,6 +415,11 @@ class llama_kv_cache_context : public llama_memory_context_i { // the stream RANGE s1 - s0 + 1 that get_k/get_v use as `ns`, not n_seqs_unq uint32_t get_n_stream() const; + // physical stream backing view index s, i.e. sinfo.strm[s]. a global cache row for a + // cell j of that view is get_strm(s)*kv->get_size() + j, the convention set_input_k_idxs + // uses. needed by the k-pool code, which writes cells that are not ubatch tokens. + uint32_t get_strm(uint32_t s) const; + const llama_kv_cache * get_kv() const; ggml_type type_k() const; @@ -410,6 +436,8 @@ class llama_kv_cache_context : public llama_memory_context_i { // - v_cur [n_embd_head_v, n_head_v, n_tokens] // - v_idxs [n_tokens] or [n_tokens*n_embd_v_gqa] depending if V cache is transposed ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const; + ggml_tensor * cpy_k_part(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, + int64_t n_embd, int64_t i_off) const; ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const; // create destination indices for each head of the current batch for where it would be written in the KV cache diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 7050ecb81f3..d0372becf68 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -72,7 +72,11 @@ llama_memory_hybrid::llama_memory_hybrid( // a *pooling* indexer (indexer_kpool > 0, glm5next only) needs a second head for // the compressor gate, or the pool cannot be rebuilt once its tokens leave the // batch. every other arch leaves indexer_kpool 0 and is unchanged. - const uint32_t n_head_idx = model.hparams.indexer_kpool > 0 ? 2 : 1; + // + // the third head is the pooled key of the pool this cell ENDS, i.e. it is written + // only for cells at pos % kpool == kpool-1. caching it turns the per-step pooling + // from O(n_kv) into O(pools completed by this ubatch): see build_indexer. + const uint32_t n_head_idx = model.hparams.indexer_kpool > 0 ? 3 : 1; std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), n_head_idx); hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; @@ -197,12 +201,23 @@ void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + // pools group cells by absolute position, so a shift regroups them; every pooled key + // cached in the indexer rows is stale even though each still looks complete. a shift + // that is a whole number of pools regroups nothing, which is the common context-shift + // case, so it is worth not paying for. + if (mem_idx && hparams.indexer_kpool > 0 && shift % (llama_pos) hparams.indexer_kpool != 0) { + mem_attn->set_kpool_dirty(); + } mem_attn->seq_add(seq_id, p0, p1, shift); if (mem_idx) mem_idx->seq_add(seq_id, p0, p1, shift); mem_recr->seq_add(seq_id, p0, p1, shift); } void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + // as seq_add, and a divide regroups for any d != 1 + if (mem_idx && hparams.indexer_kpool > 0 && d != 1) { + mem_attn->set_kpool_dirty(); + } mem_attn->seq_div(seq_id, p0, p1, d); if (mem_idx) mem_idx->seq_div(seq_id, p0, p1, d); mem_recr->seq_div(seq_id, p0, p1, d); diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 0bcbdb95d18..4605a17d76d 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -355,7 +355,11 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_tensor * packed = ggml_concat(ctx0, ggml_reshape_3d(ctx0, ik, d_idx, 1, n_tokens), ggml_reshape_3d(ctx0, gate, d_idx, 1, n_tokens), 1); - ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, packed, inp_kp->k_idxs, il)); + // key and gate are the first two heads of a three-head row; the third is the pooled key, + // written later from a different set of cells, so this store must leave it alone + ggml_build_forward_expand(gf, + mctx_idx->cpy_k_part(ctx0, ggml_reshape_2d(ctx0, packed, 2*d_idx, n_tokens), + inp_kp->k_idxs, il, 2*d_idx, 0)); if (!scoring) { return nullptr; @@ -368,23 +372,27 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( const int64_t n_tps = n_tokens/n_stream; const int64_t n_pools = inp_kp->pool_cells->ne[0]/r; - GGML_ASSERT(kbuf->ne[0] == d_idx && kbuf->ne[1] == 2 && - "the pooled indexer cache needs a key head and a gate head"); - GGML_ASSERT(kbuf->nb[1] == (size_t) d_idx*kbuf->nb[0] && "key and gate must be adjacent in a cell"); + GGML_ASSERT(kbuf->ne[0] == d_idx && kbuf->ne[1] == 3 && + "the pooled indexer cache needs a key head, a gate head and a pooled head"); + GGML_ASSERT(kbuf->nb[1] == (size_t) d_idx*kbuf->nb[0] && "key, gate and pooled must be adjacent in a cell"); GGML_ASSERT(n_tokens == n_tps*n_stream); ggml_tensor * kg_rows = ggml_view_3d(ctx0, kbuf, 2*d_idx, n_kv, n_stream, kbuf->nb[2], kbuf->nb[3], 0); - // non-resident slots hold 0, not a sentinel; garbage pools die to pool_bias, not NaN - ggml_tensor * members = ggml_get_rows(ctx0, kg_rows, inp_kp->pool_cells); + // pool only what this ubatch closed. the count is fixed (see build_inp_kpool), so the + // decode graph keeps one shape; unused slots repeat a complete pool, whose recompute is + // idempotent. non-resident slots hold 0, not a sentinel; garbage pools die to pool_bias. + const int64_t n_new_max = inp_kp->new_pool_cells->ne[0]/r; + + ggml_tensor * members = ggml_get_rows(ctx0, kg_rows, inp_kp->new_pool_cells); cb(members, "indexer_pool_members", il); const size_t nb_mem = members->nb[1]; - ggml_tensor * mem_k = ggml_view_4d(ctx0, members, d_idx, r, n_pools, n_stream, + ggml_tensor * mem_k = ggml_view_4d(ctx0, members, d_idx, r, n_new_max, n_stream, nb_mem, nb_mem*r, members->nb[2], 0); - ggml_tensor * mem_g = ggml_view_4d(ctx0, members, d_idx, r, n_pools, n_stream, + ggml_tensor * mem_g = ggml_view_4d(ctx0, members, d_idx, r, n_new_max, n_stream, nb_mem, nb_mem*r, members->nb[2], d_idx*members->nb[0]); // r-way softmaxes over the SLOT axis, so it must be dim 0; ape is added PRE-softmax @@ -397,7 +405,21 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_tensor * probs = ggml_soft_max(ctx0, gate_t); cb(probs, "indexer_pool_probs", il); - ggml_tensor * pool_k = ggml_sum_rows(ctx0, ggml_mul(ctx0, keys_t, probs)); + ggml_tensor * pool_new = ggml_sum_rows(ctx0, ggml_mul(ctx0, keys_t, probs)); + pool_new = ggml_reshape_2d(ctx0, pool_new, d_idx, n_new_max*n_stream); + cb(pool_new, "indexer_pool_new", il); + + // store into the third head of each representative cell's row. the write is indexed by + // GLOBAL row (stream-merged) while pool_reps below is stream-local, so the read cannot + // be chained off the write tensor; expand it into the graph first and let build order + // sequence them, exactly as the key/gate store above does. + ggml_build_forward_expand(gf, + mctx_idx->cpy_k_part(ctx0, pool_new, inp_kp->new_pool_reps, il, d_idx, 2*d_idx)); + + ggml_tensor * pooled_rd = ggml_view_3d(ctx0, kbuf, d_idx, n_kv, n_stream, + kbuf->nb[2], kbuf->nb[3], 2*d_idx*kbuf->nb[0]); + + ggml_tensor * pool_k = ggml_get_rows(ctx0, pooled_rd, inp_kp->pool_reps); pool_k = ggml_reshape_4d(ctx0, pool_k, d_idx, n_pools, 1, n_stream); cb(pool_k, "indexer_pool_k", il); From d07e71ede795b6ab60bb46d9212a6c584e4b2272 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 30 Aug 2026 06:24:53 +0000 Subject: [PATCH 34/36] Add MTP support --- src/llama-memory-recurrent.cpp | 12 ++- src/llama-model.cpp | 24 +++++- src/models/glm5next.cpp | 146 +++++++++++++++++++++++++++++++-- src/models/models.h | 8 ++ 4 files changed, 179 insertions(+), 11 deletions(-) diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 57919accf09..402384d0340 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -185,6 +185,12 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos // could be fatal return false; } + // a cache with no resident recurrent layers holds no state that could be partially + // erased, so the restriction does not apply to it. this is the glm5next MTP draft + // context, which runs only the NextN block and filters every KDA layer out. + const bool has_state = std::any_of(s_l.begin(), s_l.end(), + [](const ggml_tensor * t) { return t != nullptr; }); + if (0 <= seq_id) { int32_t & tail_id = cells[seq_id].tail; if (tail_id >= 0) { @@ -200,14 +206,16 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos cell.pos = p0 - 1; return true; } - return false; + if (has_state) { + return false; + } } // invalidate tails which will be cleared if (p0 <= cell.pos && cell.pos < p1) { tail_id = -1; } } - } else { + } else if (seq_id < 0) { // seq_id is negative, then the range should include everything or nothing if (p0 != p1 && (p0 != 0 || p1 != std::numeric_limits::max())) { //printf("[DEBUG] inside `llama_memory_recurrent::seq_rm`: `seq_id` is negative, so returning false\n"); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d37e8d619d4..d3dbc257c97 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2469,18 +2469,34 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_MINIMAX_01 || arch == LLM_ARCH_GLM5NEXT) { - filter_attn = [&](uint32_t il) { + // the MTP draft context runs the NextN block and nothing else, so it + // gets a cache for that one layer. the trunk never runs it and so keeps + // the layer range it always had. handing the draft the trunk's cache is + // not just waste: the KDA layers would take cells it can never roll + // back, since a draft context is built with n_rs_seq = 0, and then a + // rejected draft fails seq_rm outright. + const bool mtp_ctx = arch == LLM_ARCH_GLM5NEXT && + cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && + hparams.n_layer_all > hparams.n_layer(); + + filter_attn = [&, mtp_ctx](uint32_t il) { + if (mtp_ctx) { + return il >= hparams.n_layer() && il < hparams.n_layer_all; + } return il < hparams.n_layer() && !hparams.is_recr(il); }; - filter_recr = [&](uint32_t il) { - return il < hparams.n_layer() && hparams.is_recr(il); + filter_recr = [&, mtp_ctx](uint32_t il) { + return !mtp_ctx && il < hparams.n_layer() && hparams.is_recr(il); }; if (arch == LLM_ARCH_GLM5NEXT && hparams.indexer_head_size > 0) { // unified is fine, the pool map is per SEQUENCE. see [TAG_KPOOL_SEQ_PARTITION] // only the DSA layers carry an indexer key cache - filter_idx = [&](uint32_t il) { + filter_idx = [&, mtp_ctx](uint32_t il) { + if (mtp_ctx) { + return il >= hparams.n_layer() && il < hparams.n_layer_all; + } return il < hparams.n_layer() && !hparams.is_recr(il); }; diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 4605a17d76d..4c6db1eb32f 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -86,10 +86,12 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all); + // to n_layer_all, not n_layer(): the NextN block is an attention block, and the loop + // below plus the memory layer filters both read the entry for it uint32_t n_recr = 0; - for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0; - n_recr += hparams.is_recr_impl[il]; + n_recr += il < hparams.n_layer() ? hparams.is_recr_impl[il] : 0; } GGML_ASSERT(n_recr > 0 && n_recr < hparams.n_layer() && "glm5next needs a per-layer attention.head_count_kv array"); @@ -686,7 +688,11 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); } - if (inp_out_ids) { + // when unmasked nextn embeddings are requested, t_h_nextn must keep all rows, so the + // early output masking has to be skipped (it is applied after the final norm instead) + const bool mask_early = !cparams.embeddings_nextn || cparams.embeddings_nextn_masked; + + if (inp_out_ids && mask_early) { // flattened: get_rows needs one token's streams to be one contiguous row ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); inpL = ggml_reshape_3d(ctx0, ggml_get_rows(ctx0, flat, inp_out_ids), n_embd, hc, n_outputs); @@ -697,6 +703,15 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa cb(cur, "hc_mean", -1); cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + + // post-norm hidden state feeds the NextN/MTP draft head + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (inp_out_ids && !mask_early) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "result_norm", -1); res->t_embd = cur; @@ -707,9 +722,130 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa ggml_build_forward_expand(gf, cur); } +// NextN draft head. Unlike glm-dsa and deepseek32, whose MTP blocks skip the DSA indexer, +// blk.45 here ships a full indexer and is an ordinary DSA layer, so it runs the same +// build_layer_attn/build_layer_ffn the trunk does. What it does NOT run is the mHC mixer: +// the NextN block keeps a plain residual and has no hc_* tensors (see load_arch_tensors). +llama_model_glm5next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) + : graph(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "glm5next MTP supports a single NextN block"); + + const int il = hparams.n_layer() + cparams.nextn_layer_offset; + + GGML_ASSERT(cparams.nextn_layer_offset >= 0 && + cparams.nextn_layer_offset < (int) hparams.n_layer_nextn && + "nextn_layer_offset out of range [0, n_layer_nextn)"); + + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm && + "glm5next MTP block tensors missing; convert without --no-mtp"); + GGML_ASSERT(!layer.hc_attn_fn && "the NextN block has no mHC mixer"); + + 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(); + + llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); + + // the NextN block is DSA, so nothing here consumes the recurrent half. s_copy would then + // never enter the graph, never be allocated, and set_input would still read its buffer. + ggml_build_forward_expand(gf, inp_mem->get_recr()->s_copy); + + llm_graph_input_kpool * inp_kp = nullptr; + bool indexer_scoring = false; + { + const auto * mctx_hyb = static_cast(mctx); + + if (mctx_hyb->get_idx() != nullptr) { + indexer_scoring = cparams.n_ctx > glm5next_n_select(hparams); + + inp_kp = build_inp_kpool(mctx_hyb, + inp_mem->get_attn()->get_kq_mask(), indexer_scoring); + } + } + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, + ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0), layer.nextn.eh_proj_s); + cb(cur, "mtp_eh_proj", il); + + // plain pre-norm residual block, the trunk's helpers minus the mHC wrapper + ggml_tensor * residual = cur; + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + cur = build_layer_attn(model, inp_mem, inp_kp, indexer_scoring, cur, il); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, residual); + 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); + + cur = build_layer_ffn(model, cur, il); + cb(cur, "mtp_ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "mtp_post_ffn", il); + + ggml_tensor * head_norm_w = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm; + GGML_ASSERT(head_norm_w && "glm5next MTP: no nextn.shared_head_norm and no output_norm"); + cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1); + + // the post-norm hidden state seeds the next MTP step + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "mtp_shared_head_norm", -1); + + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; + GGML_ASSERT(head_w && "glm5next MTP: no nextn.shared_head_head and no output"); + cur = build_lora_mm(head_w, cur, head_s); + cb(cur, "result_output", -1); + + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} + std::unique_ptr llama_model_glm5next::build_arch_graph(const llm_graph_params & params) const { - // without this, an MTP context (accepted whenever n_layer_nextn > 0) runs the trunk - GGML_ASSERT(params.gtype != LLM_GRAPH_TYPE_DECODER_MTP && "glm5next NextN graph not implemented yet"); + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } diff --git a/src/models/models.h b/src/models/models.h index c47babd180b..699e58f246d 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1349,6 +1349,9 @@ struct llama_model_glm5next : public llama_model_base { struct graph : public llama_model_deepseek4::graph { graph(const llama_model & model, const llm_graph_params & params); + // builds nothing: lets graph_mtp reuse the block helpers below without the trunk + graph(const llm_graph_params & params) : llama_model_deepseek4::graph(params) {} + // not const: the delta-net helpers append to the graph through the base ggml_tensor * build_layer_attn( const llama_model & model, @@ -1388,6 +1391,11 @@ struct llama_model_glm5next : public llama_model_base { int il) const; }; + // NextN draft head. reuses the trunk's block helpers, so it stays a `graph` + 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 5796547f37f5943513dfa130065ec88f9e30e0f5 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Sun, 30 Aug 2026 22:33:57 +0000 Subject: [PATCH 35/36] glm5next: invalidate pooled keys when a shift splits a pool seq_add only skipped the pooled-key rebuild when the shift itself was a multiple of kpool. That is not sufficient: a pool straddling p0 or p1 keeps some members and moves the rest, so it is regrouped no matter how the shift is aligned, and its cached pooled key goes stale while still looking complete. Both callers pass an arbitrary bound. The server's context shift uses n_keep + n_discard and its prompt-cache reuse uses the match head, so this is reachable in normal use: with --keep 39 and n_ctx 8192, n_discard is 4076 and p0 is 4115, which is a multiple-of-4 shift starting mid-pool. Also require both bounds to be pool-aligned. A negative p0 or p1 means "from the start" / "to the end", which no pool can straddle. --- src/llama-memory-hybrid.cpp | 23 ++++++++++++++++++----- 1 file changed, 18 insertions(+), 5 deletions(-) diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index d0372becf68..e61d4c94bd4 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -202,11 +202,24 @@ void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { // pools group cells by absolute position, so a shift regroups them; every pooled key - // cached in the indexer rows is stale even though each still looks complete. a shift - // that is a whole number of pools regroups nothing, which is the common context-shift - // case, so it is worth not paying for. - if (mem_idx && hparams.indexer_kpool > 0 && shift % (llama_pos) hparams.indexer_kpool != 0) { - mem_attn->set_kpool_dirty(); + // cached in the indexer rows is then stale even though each pool still looks complete. + // it survives only if WHOLE pools move: the shift must be a multiple of kpool and the + // range must start and end on a pool boundary. a pool straddling p0 or p1 keeps some + // members and moves the rest, which regroups it however aligned the shift itself is, + // and both callers pass an arbitrary bound -- the server's context shift uses + // n_keep + n_discard, its prompt-cache reuse uses the match head. + if (mem_idx && hparams.indexer_kpool > 0) { + const llama_pos r = (llama_pos) hparams.indexer_kpool; + + // p0 < 0 means "from the start" and p1 < 0 means "to the end", so neither is a + // boundary a pool can straddle + const bool whole_pools = shift % r == 0 && + (p0 <= 0 || p0 % r == 0) && + (p1 < 0 || p1 % r == 0); + + if (!whole_pools) { + mem_attn->set_kpool_dirty(); + } } mem_attn->seq_add(seq_id, p0, p1, shift); if (mem_idx) mem_idx->seq_add(seq_id, p0, p1, shift); From 949f7efb097eb20ef36fecdb1afaebff9a4ae7ed Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 31 Aug 2026 10:39:19 +0000 Subject: [PATCH 36/36] glm5next: cut comments Second reduction pass over the arch's comments: 405 comment lines on the branch's own added lines down to 234, no code changes. Deletes rather than reshortens. What stayed is limited to things whose absence would let a reader make a specific mistake: reference constants and sign conventions, the ordering and precision constraints the graph relies on, and the shapes of ggml tensors, whose type carries none. --- conversion/glm5next.py | 48 ++++++------------------ gguf-py/gguf/constants.py | 6 --- src/llama-graph.cpp | 19 ++-------- src/llama-graph.h | 1 - src/llama-kv-cache-kpool.cpp | 72 ++++++++++-------------------------- src/llama-kv-cache-kpool.h | 56 +++++++++------------------- src/llama-kv-cache.h | 22 ++++------- src/llama-memory-hybrid.cpp | 36 +++++------------- src/llama-model.cpp | 14 ++----- src/models/glm5next.cpp | 71 +++++++++-------------------------- src/models/models.h | 3 +- tests/test-llama-archs.cpp | 10 ++--- tools/mtmd/mtmd-image.cpp | 3 -- 13 files changed, 95 insertions(+), 266 deletions(-) diff --git a/conversion/glm5next.py b/conversion/glm5next.py index f2db5867765..09ca8f73561 100644 --- a/conversion/glm5next.py +++ b/conversion/glm5next.py @@ -14,7 +14,6 @@ @ModelBase.register("Glm5NextForConditionalGeneration", "Glm5NextForCausalLM") -# [TAG_HF_EXAMPLE_MISSING] class Glm5NextModel(TextModel): """GLM-5.3-Flash text tower: hybrid KDA + DSA attention, nope-only MLA, mHC hyper-connections, and a NextN block with its own DSA attention and indexer. @@ -33,8 +32,7 @@ def __init__(self, *args, **kwargs): self.block_count = self.hparams["num_hidden_layers"] + nextn_layers self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) - # two independent spellings of the same partition; disagreement means an - # unexpected config + # two spellings of the same partition; disagreement means an odd config from_types = {il for il, t in enumerate(self.hparams["layer_types"]) if t == "deepseek_sparse_attention"} from_list = set(self.hparams["linear_attn_config"]["full_attn_layers"]) if from_types != from_list: @@ -47,7 +45,6 @@ def __init__(self, *args, **kwargs): raise ValueError("mlp_layer_types does not match first_k_dense_replace") def index_tensors(self, remote_hf_model_id: str | None = None): - # runs before TextModel.__init__ has hoisted text_config to the root hp = self.hparams.get("text_config", self.hparams) type(self)._main_layers = hp["num_hidden_layers"] return super().index_tensors(remote_hf_model_id=remote_hf_model_id) @@ -58,14 +55,12 @@ def set_vocab(self): def is_full_attention(self, bid: int) -> bool: return bid >= self.hparams["num_hidden_layers"] or bid in self._full_attn_layers - # -- metadata --------------------------------------------------------- def set_gguf_parameters(self): hp = self.hparams linear_cfg = hp["linear_attn_config"] - # checked here, not in the loader: head_count_kv is overwritten below with - # the per-layer 1/0 recurrence marker + # checked here, not in the loader: head_count_kv becomes the per-layer recurrence marker if hp["num_attention_heads"] != hp.get("num_key_value_heads"): raise ValueError("glm5next expects MHA-shaped head counts before MLA absorption") if hp["qk_rope_head_dim"] != 0 or not hp.get("mla_use_nope"): @@ -77,8 +72,7 @@ def set_gguf_parameters(self): if hp["index_topk"] % hp["index_kpool"] != 0: raise ValueError("glm5next index_topk must be a whole number of kpool pools") - # no GGUF key carries these and the graph cannot express them off, so refuse - # rather than write a silently wrong file + # no GGUF key carries these and the graph cannot express them off, so refuse to write if not hp.get("index_kpool_compress"): raise ValueError("glm5next without the indexer kpool compressor is not supported") if not hp.get("index_kpool_always_select_tail"): @@ -88,8 +82,6 @@ def set_gguf_parameters(self): if set(hp["indexer_types"]) != {"full"}: raise ValueError("glm5next expects every indexer to be full") - # drop both: head_dim is 0 in the config, head_count_kv is written as a - # per-layer array below hp.pop("head_dim", None) hp.pop("num_key_value_heads", None) @@ -101,7 +93,6 @@ def set_gguf_parameters(self): self.gguf_writer.add_head_count_kv( [1 if self.is_full_attention(il) else 0 for il in range(self.block_count)]) - # --- MLA --- kv_lora_rank = hp["kv_lora_rank"] qk_rope_head_dim = hp["qk_rope_head_dim"] self.gguf_writer.add_q_lora_rank(hp["q_lora_rank"]) @@ -112,29 +103,23 @@ def set_gguf_parameters(self): self.gguf_writer.add_key_length_mla(hp["qk_nope_head_dim"] + qk_rope_head_dim) self.gguf_writer.add_value_length_mla(hp["v_head_dim"]) - # indexer k_norm is a LayerNorm with bias at a fixed 1e-6, not the model's - # RMS eps. glm-dsa omits this key and runs that norm at eps 0 + # indexer k_norm is a LayerNorm with bias at a fixed 1e-6, not the model's RMS eps self.gguf_writer.add_layer_norm_eps(1e-6) - # --- KDA --- self.gguf_writer.add_ssm_conv_kernel(linear_cfg["short_conv_kernel_size"]) self.gguf_writer.add_kda_head_dim(linear_cfg["head_dim"]) - # not a clamp: scales the sigmoid decay gate. required, a missing key - # silently selects the softplus branch instead + # not a clamp: scales the sigmoid decay gate. required, else the softplus branch is chosen self.gguf_writer.add_kda_gate_lower_bound(linear_cfg["gate_lower_bound"]) - # --- DSA indexer --- 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"]) - # --- mHC --- 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 --- n_ff_exp = hp["moe_intermediate_size"] self.gguf_writer.add_expert_feed_forward_length(n_ff_exp) self.gguf_writer.add_expert_shared_feed_forward_length(n_ff_exp * hp["n_shared_experts"]) @@ -143,8 +128,7 @@ def set_gguf_parameters(self): self.gguf_writer.add_expert_weights_scale(hp["routed_scaling_factor"]) self.gguf_writer.add_expert_weights_norm(hp["norm_topk_prob"]) - # one limit for the whole model. no dense-FFN clamp key exists, so the - # expert arrays are sized for every layer to cover the leading dense ones + # no dense-FFN clamp key exists, so the expert arrays cover the leading dense layers too swiglu_limit = float(hp["swiglu_limit"]) self.gguf_writer.add_swiglu_clamp_exp([swiglu_limit] * self.block_count) self.gguf_writer.add_swiglu_clamp_shexp([swiglu_limit] * self.block_count) @@ -152,7 +136,6 @@ def set_gguf_parameters(self): if not self.no_mtp and (nextn_layers := hp.get("num_nextn_predict_layers", 0)): self.gguf_writer.add_nextn_predict_layers(nextn_layers) - # -- tensors ---------------------------------------------------------- @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: @@ -173,31 +156,28 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return name, gen def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: - # --- KDA conv1d: HF [d_inner, 1, d_conv] -> ggml ne [d_conv, 1, d_inner, 1] --- + # KDA conv1d: HF [d_inner, 1, d_conv] -> ggml ne [d_conv, 1, d_inner, 1] if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")): d_inner = data_torch.shape[0] d_conv = data_torch.shape[-1] data_torch = data_torch.reshape(1, d_inner, 1, d_conv) - # ssm_a holds -exp(A_log), the kimi-k3 convention (bailingmoe3 stores - # +exp(A_log)); the wrong sign turns decay into an unchecked growing state + # ssm_a holds -exp(A_log), the kimi-k3 convention (bailingmoe3 stores +exp(A_log)); + # the wrong sign turns decay into an unchecked growing state if name.endswith(".A_log"): - # eager: the sign is the point of the check, and A_log is one per head decay = LazyTorchTensor.to_eager(torch.exp(data_torch.float())) if not bool(torch.isfinite(decay).all() and (decay > 0).all()): raise ValueError(f"{name}: exp(A_log) must be finite and positive") data_torch = -decay - # dt_bias -> the name SSM_DT's mapping expects if name.endswith(".dt_bias"): name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" - # bare tensors in the checkpoint, but the GGUF names carry .weight + # bare in the checkpoint, but the GGUF names carry .weight if re.search(r"\.hc_(attn|ffn)_(fn|base|scale)$", name) or name.endswith( (".index_kpool_compress_gate", ".index_kpool_compress_ape")): name += ".weight" - # --- routed experts --- if ".mlp.experts." in name: n_experts = self.hparams["n_routed_experts"] assert bid is not None @@ -218,7 +198,6 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(torch.stack(tensors, dim=0), merged_name, bid) return - # --- MLA absorption --- if name.endswith(".kv_b_proj.weight"): assert bid is not None n_head = self.hparams["num_attention_heads"] @@ -234,8 +213,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(data_torch, name, bid) def tensor_force_quant(self, name, new_name, bid, n_dims): - # learned position table, one row per pooled key; pinned for the same - # reason POS_EMBD is in base.py + # pinned as POS_EMBD is in base.py if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_APE, bid): return gguf.GGMLQuantizationType.F32 return super().tensor_force_quant(name, new_name, bid, n_dims) @@ -265,9 +243,7 @@ def set_gguf_parameters(self): super().set_gguf_parameters() assert self.hparams_vision is not None - # Glm4VVisionModel bypasses Qwen3VLVisionModel entirely, which is also where - # the merge size is written, so no GLM4V-family mmproj carries this key and - # clip.cpp falls back to a hardcoded 2. Write it rather than rely on that + # no GLM4V-family mmproj carries this key and clip.cpp falls back to a hardcoded 2 self.gguf_writer.add_vision_spatial_merge_size(int(self.hparams_vision["spatial_merge_size"])) self.gguf_writer.add_vision_swiglu_limit(float(self.hparams_vision["swiglu_limit"])) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index a67d067043d..277fb13aebf 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -4001,14 +4001,12 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.OUTPUT, MODEL_TENSOR.ATTN_NORM, MODEL_TENSOR.FFN_NORM, - # mHC, layered on top of the per-layer norms above 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, @@ -4023,7 +4021,6 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.SSM_A, MODEL_TENSOR.SSM_DT, MODEL_TENSOR.SSM_NORM, - # DSA (MLA full-attention layers) MODEL_TENSOR.ATTN_Q_A, MODEL_TENSOR.ATTN_Q_B, MODEL_TENSOR.ATTN_Q_A_NORM, @@ -4032,14 +4029,12 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.ATTN_K_B, MODEL_TENSOR.ATTN_V_B, MODEL_TENSOR.ATTN_OUT, - # DSA indexer, with the kpool key compressor MODEL_TENSOR.INDEXER_K_NORM, MODEL_TENSOR.INDEXER_PROJ, MODEL_TENSOR.INDEXER_ATTN_K, MODEL_TENSOR.INDEXER_ATTN_Q_B, MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE, MODEL_TENSOR.INDEXER_COMPRESSOR_APE, - # FFN: dense on the leading blocks, MoE elsewhere MODEL_TENSOR.FFN_GATE, MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, @@ -4051,7 +4046,6 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_GATE_SHEXP, MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, - # NextN/MTP, a full DSA decoder layer with its own indexer MODEL_TENSOR.NEXTN_EH_PROJ, MODEL_TENSOR.NEXTN_EMBED_TOKENS, MODEL_TENSOR.NEXTN_ENORM, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index cc97afac4a5..35083a349c8 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3581,7 +3581,6 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[3] == n_stream); - // the selection terms below exist only for real queries GGML_ASSERT(kq_mask->ne[1] == n_tps && "the pooled indexer needs an unpadded KQ mask"); inp->pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool*n_pools, n_stream); @@ -3609,13 +3608,9 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( ggml_set_input(inp->cand_mask); ggml_set_name(inp->cand_mask, "kpool_cand_mask"); - // pooled-key cache. n_new_max is an exact bound on the pools one ubatch can - // complete: a sequence's tokens are a contiguous position run, so a run of L tokens - // closes at most L/kpool + 1 pools. it is FIXED for the whole decode phase so the - // graph shape does not depend on how many pools happened to close this step. - // after a position mutation every cached pooled key is stale, so this one graph has - // to be able to re-emit all of them. a shape change here forces a rebuild, which is - // the point: the wide shape is used for exactly one ubatch and then goes away. + // n_new_max is an exact bound (a contiguous run of L tokens closes at most L/kpool + 1 + // pools), FIXED for the decode phase so the graph shape does not track pools-closed-this- + // step; after a position mutation every cached key is stale, so it must re-emit all const bool rebuild = mctx_attn->get_kv()->get_kpool_dirty(); const int64_t n_new_max = rebuild ? n_pools : n_tps/kpool + n_ps; @@ -3630,7 +3625,6 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( ggml_set_input(inp->new_pool_cells); ggml_set_name(inp->new_pool_cells, "kpool_new_pool_cells"); - // I64 because ggml_set_rows takes its row indices as I64 inp->new_pool_reps = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new_max*n_stream); ggml_set_input(inp->new_pool_reps); ggml_set_name(inp->new_pool_reps, "kpool_new_pool_reps"); @@ -3661,7 +3655,6 @@ ggml_tensor * llm_graph_context::build_attn_sparse( const auto * mctx_cur = inp->mctx; - // store to KV cache { const auto & k_idxs = inp->get_k_idxs(); @@ -3679,26 +3672,22 @@ ggml_tensor * llm_graph_context::build_attn_sparse( // ggml_set_rows writes THROUGH, and sel_mask is shared per ubatch: scatter into a copy ggml_tensor * mask_all = ggml_dup(ctx0, sel_mask); - // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] mask_all = ggml_view_4d(ctx0, mask_all, 1, mask_all->ne[0], mask_all->ne[1], mask_all->ne[3], mask_all->nb[0], mask_all->nb[1], mask_all->nb[2], 0); - // [n_select, n_tps, n_stream] -> [n_select, n_tps, n_stream, 1] ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[2], 1, top_k->nb[1], top_k->nb[2], top_k->ne[2]*top_k->nb[2], 0); // a constant 0, never the cell's bias: scattering -inf would ERASE a zero granted to the - // tail (cand_mask rejects over-budget picks below). f32: CUDA only does SET_ROWS for f32 + // tail. f32 because CUDA only does SET_ROWS for f32 ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); zeros = ggml_fill(ctx0, zeros, 0.0f); ggml_tensor * mask_top_k = ggml_set_rows(ctx0, mask_all, zeros, top_k_3d); - // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] mask_top_k = ggml_view_4d(ctx0, mask_top_k, mask_top_k->ne[1], mask_top_k->ne[2], 1, mask_top_k->ne[3], mask_top_k->nb[2], mask_top_k->nb[3], mask_top_k->nb[3], 0); - // the reference's `selected_valid` gather, additively; cand_mask is candidates UNION tail mask_top_k = ggml_add(ctx0, mask_top_k, cand_mask); // ggml_flash_attn_ext asserts an f16 mask, and ggml_add would yield src0's f32 diff --git a/src/llama-graph.h b/src/llama-graph.h index 258a66dd53e..33beda21ff6 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -32,7 +32,6 @@ class llama_memory_recurrent_context; class llama_memory_hybrid_context; class llama_memory_hybrid_iswa_context; -// defined in llama-kv-cache-kpool.h, which includes this header, so forward declared only class llm_graph_input_kpool; // certain models (typically multi-modal) can produce different types of graphs diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index 485260586b8..eab9b9d4081 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -58,7 +58,6 @@ static void kpool_mask_row( const bool tail = pos_at[j] >= tail_start; cur_sel [j] = vis && tail ? v_sel : v_mask; - // the candidate set, which the top-k budget may overrun but must never escape cur_cand[j] = vis && (pooled || tail) ? v_sel : v_mask; } } @@ -103,8 +102,8 @@ void llama_kv_cache_set_input_kpool( const int64_t r = kpool; const int64_t n_tokens = ubatch->n_tokens; - // [TAG_KPOOL_SEQ_PARTITION] positions are unambiguous only within one sequence, so - // one pool map per SEQUENCE, not per stream + // [TAG_KPOOL_SEQ_PARTITION] positions are unambiguous only within one sequence, so one + // pool map per SEQUENCE, not per stream GGML_ASSERT(n_ns == 1 || (int64_t) ubatch->n_seqs_unq == n_ns); const int64_t n_ps = (int64_t) ubatch->n_seqs_unq/n_ns; @@ -142,7 +141,6 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(bias->ne[0] == n_kv && bias->ne[1] == n_tps && bias->ne[2] == n_ns); } - // pooled-key cache inputs travel together or not at all const bool kcache = pool_reps != nullptr; GGML_ASSERT((new_pool_cells != nullptr) == kcache); @@ -211,22 +209,17 @@ void llama_kv_cache_set_input_kpool( int32_t * cur_new_cells = kcache ? dst_new_cells + s*(r*n_new_max) : nullptr; int64_t * cur_new_reps = kcache ? dst_new_reps + s*n_new_max : nullptr; - // count of real entries emitted for this stream; the rest is padding int64_t n_new = 0; - // members of any complete pool in this stream, used to pad the fixed-size write. - // recomputing a complete pool is idempotent, so a repeat is always safe. + // pads the fixed-size write; recomputing a complete pool is idempotent, so a repeat is safe const int32_t * any_rep_src = nullptr; if (kcache) { - // a pool with no rep gathers row 0. that row's pooled third may hold another - // pool's key, but such a pool is always -INFINITY in pool_bias, so the value is - // discarded before it can score. + // a pool with no rep gathers row 0; such a pool is -INFINITY in pool_bias, so discarded std::fill(cur_pool_reps, cur_pool_reps + n_pools, 0); std::fill(cur_new_cells, cur_new_cells + r*n_new_max, 0); } - // the token loop writes rows < n_tps in full; only the padding rows need clearing if (mask_f16) { kpool_mask_fill((ggml_fp16_t *) (cur_sel_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); kpool_mask_fill((ggml_fp16_t *) (cur_cand_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); @@ -235,11 +228,8 @@ void llama_kv_cache_set_input_kpool( kpool_mask_fill((float *) (cur_cand_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); } - // [TAG_KPOOL_PACK] one packed run per sequence, sized on the pool range it holds. - // NOT one full-width table per sequence: the indexer scores every slot against - // every query, so that multiplies the score tensor by n_seq_max. - // llama_memory_seq_cp can ask for more slots than exist; then a sequence keeps its - // newest pools, the same cut a large hole already forces. + // [TAG_KPOOL_PACK] one packed run per sequence, NOT one full-width table: the indexer + // scores every slot against every query, which would multiply the score tensor by n_seq_max { int64_t n_want = 0; @@ -302,9 +292,8 @@ void llama_kv_cache_set_input_kpool( pos_at[j] = cells.is_empty(j) || !cells.seq_has(j, seq_of_pool) ? -1 : cells.pos_get(j); } - // anchoring at the absolute p/kpool follows vLLM and SGLang, not HF - // (valid_keys.argmax(-1)): it is the only anchor that keeps a pool's identity - // stable from the prefill that built it to the decodes that read it. + // anchor at the absolute p/kpool (vLLM, SGLang; not HF's valid_keys.argmax(-1)): the only + // anchor that keeps a pool's identity stable from prefill to the decodes that read it int64_t b_base = 0; { int64_t b_min = 0; @@ -353,10 +342,8 @@ void llama_kv_cache_set_input_kpool( } if (kcache) { - // the pooled key of a complete pool lives in the row of its LAST member, - // the cell holding pos % r == r-1. that slot is only meaningful once the - // pool is complete: for a partial pool it is still 0 from the fill above, - // and cell 0 is a real cell whose own pooled key we must not overwrite. + // a complete pool's key lives in the row of its LAST member; a partial pool leaves that + // slot 0, and cell 0 is a real cell for (int64_t p = 0; p < n_run; ++p) { if (filled[p] == (int32_t) r) { cur_pool_reps[run_off[ps] + p] = part_pool_cells[p*r + (r - 1)]; @@ -367,8 +354,7 @@ void llama_kv_cache_set_input_kpool( } } - // recompute exactly the complete pools this ubatch wrote into. touched[] is - // over the run, so the cost is O(tokens), not O(n_kv). + // touched[] is over the run, so the cost is O(tokens), not O(n_kv) std::vector touched(n_run, 0); for (int64_t ii = 0; ii < n_tps; ++ii) { @@ -386,16 +372,12 @@ void llama_kv_cache_set_input_kpool( } for (int64_t p = 0; p < n_run; ++p) { - // in rebuild mode every cached pooled key is stale, so re-emit all of - // them, not just the pools this ubatch closed if ((!touched[p] && !rebuild) || filled[p] != (int32_t) r) { continue; } - // n_new_max = n_tps/kpool + n_ps bounds this while a sequence's tokens in - // one ubatch are a contiguous position run, which llama-batch.cpp - // enforces. dropping a completed pool would silently serve a stale key, - // so fail loudly instead of clamping. + // bounded while a sequence's ubatch tokens are a contiguous run (llama-batch.cpp enforces); + // fail loudly, clamping would serve a stale key GGML_ASSERT(n_new < n_new_max && "k-pool: more pools completed than the fixed bound"); std::copy(part_pool_cells + p*r, part_pool_cells + (p + 1)*r, @@ -416,16 +398,13 @@ void llama_kv_cache_set_input_kpool( const llama_pos q = ubatch->pos[i]; - // q >= 0 is what makes the unsigned range test below a range test GGML_ASSERT(q >= 0); n_done++; - // index_kpool_always_select_tail, which lands selection on pool boundaries const llama_pos tail_start = (q + 1)/r*r; - // the reference tests visibility at a pool's LAST member, so a pool the - // query straddles is dropped whole + // the reference tests visibility at a pool's LAST member, so a straddled pool drops whole const int64_t bo_vis = std::max(0, tail_start/r - b_base); float * cur_bias = dst_bias ? dst_bias + i*n_kv : nullptr; @@ -449,8 +428,6 @@ void llama_kv_cache_set_input_kpool( } } - // the query's own sequence run only; every other slot keeps the -INFINITY - // of the fill above, which is what keeps a foreign pool out of the budget float * q_pool_bias = cur_pool_bias + ii*n_pools + run_off[ps]; for (int64_t p = 0; p < n_run; ++p) { @@ -466,19 +443,13 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(n_done == n_tps && "every query must belong to a sequence of the ubatch"); if (kcache) { - // the write has a fixed row count, so the unused slots must name a destination - // that is safe to overwrite. two cases, both provably harmless: - // - some complete pool exists: repeat it. recomputing a complete pool yields - // the value already there, so the duplicate write is a no-op in effect. - // - none exists: no cell in this stream is the last member of a complete pool, - // so no pooled third is read (every pool is -INFINITY in pool_bias). cell 0 - // is then free. + // the fixed row count means unused slots must name a safe destination: repeat a complete + // pool (recompute is a no-op), or cell 0 when none exists (nothing reads its pooled third) for (int64_t p = n_new; p < n_new_max; ++p) { if (any_rep_src) { std::copy(any_rep_src, any_rep_src + r, cur_new_cells + p*r); cur_new_reps[p] = (int64_t) strm_of[s]*kv_size + any_rep_src[r - 1]; } else { - // cells already 0 from the fill; pool r copies of cell 0 into cell 0 cur_new_reps[p] = (int64_t) strm_of[s]*kv_size; } } @@ -487,17 +458,14 @@ void llama_kv_cache_set_input_kpool( } void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { - // unconditional: the key and gate STORE runs on the dense path too. gating it the - // way the scoring is gated would leave every cell below n_select with no indexer - // state, and the first ubatch to cross n_select would pool cells never written + // unconditional: the key/gate STORE runs on the dense path too, or cells below n_select + // would have no indexer state when the first ubatch crosses it mctx_idx->set_input_k_idxs(k_idxs, ubatch); if (pool_cells == nullptr) { return; } - // the pooled key is written into the INDEXER cache, whose slot layout the attention - // cache defines, so the stream map and cell count come from the indexer side std::vector strm_of; if (pool_reps) { @@ -518,9 +486,7 @@ void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { rebuild, ubatch, kpool); - // every pool has just been re-emitted, so the cache is consistent again. cleared here - // rather than in build_inp_kpool because a graph that is built but not evaluated must - // not clear it. + // cleared here, not in build_inp_kpool: a graph built but not evaluated must not clear it if (rebuild) { mctx_attn->get_kv()->clear_kpool_dirty(); } diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index ff20a6e6667..b2f5cf3aac7 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -10,33 +10,21 @@ struct llama_ubatch; class llama_kv_cache; class llama_kv_cache_context; -// GLM-5-Next indexer pooling. the position -> cell map is built host side because -// find_slot's cell order is arbitrary. no input may hold a negative index: ggml_set_rows -// asserts i1 >= 0 and ggml_get_rows has no sentinel, so unusable entries are clamped into -// range and neutralised by the additive masks instead. +// GLM-5-Next indexer pooling. no input may hold a negative index (ggml_set_rows asserts +// i1 >= 0, ggml_get_rows has no sentinel), so unusable entries are clamped and masked. -// pool slots for `n_kv` cells shared by `n_seqs` sequences: n_kv/kpool, exact only while -// the sequences' cells are disjoint, plus 2 per sequence for rebasing. +// n_kv/kpool (exact only while the sequences' cells are disjoint) plus 2 per seq for rebasing uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool, uint32_t n_seqs = 1); -// select_k of modular_glm5_next.py, Glm5NextTextIndexer.forward. must run over POOLS, not -// cells: relu ties span pool boundaries, so a cell-level cut takes partial pools. +// select_k of Glm5NextTextIndexer.forward: must run over POOLS, a cell cut takes partial pools uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t kpool); // `kv` must be the ATTENTION (MLA) cache; the indexer cache shares its slot layout. -// cell_pool I32 [n_kv, n_stream] per-cell view, optional, unused here -// pool_cells I32 [kpool*n_pools, n_stream] pool member -> cell, 0 if not resident -// bias F32 [n_kv, n_tps, n_stream] per-cell view, optional, unused here -// pool_bias F32 [n_pools, n_tps, n_stream] pool_valid & pool_visible, -INFINITY -// outside the query's own sequence run; computed, not gathered from `bias` at the -// last member, which an incomplete pool lacks and would inherit cell 0's validity -// sel_mask F16/F32 [n_kv, n_batch, 1, n_stream] 0.0f on the always-selected tail only -// cand_mask F16/F32 [n_kv, n_batch, 1, n_stream] max(bias, sel_mask); bounds the top-k -// spills that a partial seq_rm would otherwise let escape the candidate set -// pool_reps / new_pool_cells / new_pool_reps are optional (nullptr when the pooled-key -// cache is off). an entry is emitted only for a pool with filled == kpool: an incomplete -// pool's last-member slot is 0, and cell 0 is a legitimate cell, so writing it would -// clobber another pool's cached key. +// pool_cells pool member -> cell, 0 if not resident +// pool_bias computed, NOT gathered at the last member, which an incomplete pool lacks +// cand_mask bounds top-k spills a partial seq_rm would let escape +// pool_reps / new_pool_cells / new_pool_reps are nullptr when the cache is off, and an entry +// is emitted only for filled == kpool: cell 0 is real, so writing its 0 slot would clobber void llama_kv_cache_set_input_kpool( const llama_kv_cache * kv, ggml_tensor * cell_pool, @@ -48,21 +36,15 @@ void llama_kv_cache_set_input_kpool( ggml_tensor * pool_reps, ggml_tensor * new_pool_cells, ggml_tensor * new_pool_reps, - // strm_of[s] is the physical stream behind view s and kv_size the per-stream cell - // count, so a global row is strm_of[s]*kv_size + cell. only read when pool_reps is - // set. the indexer cache shares the attention cache's slot layout, so one cell - // index addresses both. + // a global row is strm_of[s]*kv_size + cell; only read when pool_reps is set const uint32_t * strm_of, int64_t kv_size, - // re-emit every complete pool, not only the ones this ubatch closed. set after a - // position mutation; the graph is built with n_new_max == n_pools to hold them. + // re-emit every complete pool, not only those this ubatch closed; set after a position mutation bool rebuild, const llama_ubatch * ubatch, uint32_t kpool); -// One pooling map per ubatch; rebuilding it per indexer layer costs O(n_kv * n_tokens) -// host writes and dominates prefill. sharing is valid only while every indexer layer sees -// the same candidate set - true for glm5next (indexer_types all "full"), not for windowed. +// one map per ubatch; valid only while every indexer layer sees the same candidate set class llm_graph_input_kpool : public llm_graph_input_i { public: llm_graph_input_kpool( @@ -78,21 +60,17 @@ class llm_graph_input_kpool : public llm_graph_input_i { ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] ggml_tensor * pool_bias = nullptr; // F32 [n_pools, n_tps, n_stream] - // pooled-key cache. the pooled value of a pool lives in the third head of the row of - // the cell holding its LAST member (pos % kpool == kpool-1), so it is a pure function - // of cell content: independent of sequence and of pool ordinal, which is what lets a - // seq_cp share it and a rebase leave it alone. + // pooled-key cache: a pool's value lives in the row of its LAST member, a pure function of + // cell content (seq_cp shares it, rebase leaves it alone) ggml_tensor * pool_reps = nullptr; // I32 [n_pools, n_stream] stream-local rep cell ggml_tensor * new_pool_cells = nullptr; // I32 [kpool*n_new_max, n_stream] members to (re)pool ggml_tensor * new_pool_reps = nullptr; // I64 [n_new_max*n_stream] GLOBAL dest row - // n_new_max is fixed for the whole decode phase on purpose. making the shape depend on - // how many pools happened to close this step would flip the graph topology every kpool - // tokens and defeat graph reuse / force CUDA-graph recapture. + // fixed for the decode phase: a shape tracking pools-closed-this-step would flip the graph + // topology every kpool tokens and force CUDA-graph recapture uint32_t n_new_max = 0; - // this graph was built to re-emit every pool after a position mutation, not just the - // ones this ubatch closed. decided at build time because it sets n_new_max above. + // set at build time: re-emit every pool after a position mutation bool rebuild = false; // exact, since pool_bias only holds 0.0f or -INFINITY. nullptr if the fused path is off diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 2ba032a34bf..279ea79a3e5 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -160,12 +160,8 @@ class llama_kv_cache : public llama_memory_i { bool get_has_shift() const; - // GLM-5-Next pooled-key cache. a pool groups cells by absolute position, so anything - // that mutates positions in place (seq_add / seq_div) regroups the pools while every - // cached pooled key still looks complete. set here, consumed and cleared by - // llama_kv_cache_set_input_kpool, which then emits every pool instead of only the new - // ones. sticky by design: a flag that never clears degrades to the pre-cache cost, a - // flag that clears too early is silently wrong. + // GLM-5-Next pooled-key cache: seq_add/seq_div regroup pools while every cached key still + // looks complete. sticky by design, a flag that clears too early is silently wrong void set_kpool_dirty(); bool get_kpool_dirty() const; void clear_kpool_dirty() const; @@ -203,10 +199,8 @@ class llama_kv_cache : public llama_memory_i { ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const; ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const; - // write n_embd elements at element offset i_off inside each row, leaving the rest of - // the row untouched. k_idxs are GLOBAL rows (stream-merged), as in cpy_k. - // returns the ggml_set_rows result so a later gather can be chained off it and become - // a real graph edge rather than relying on build order. + // writes n_embd elements at element offset i_off, leaving the rest of the row untouched. + // returns the ggml_set_rows result so a later gather chains off it as a real graph edge ggml_tensor * cpy_k_part(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, int64_t n_embd, int64_t i_off) const; @@ -278,8 +272,7 @@ class llama_kv_cache : public llama_memory_i { bool v_trans = true; // the value tensor is transposed - // see set_kpool_dirty. mutable because the only consumer runs from set_input, which - // holds the cache by const pointer; nothing else observes it. + // see set_kpool_dirty. mutable: its only consumer runs from set_input, holding a const cache mutable bool kpool_dirty = false; const uint32_t n_seq_max = 1; @@ -415,9 +408,8 @@ class llama_kv_cache_context : public llama_memory_context_i { // the stream RANGE s1 - s0 + 1 that get_k/get_v use as `ns`, not n_seqs_unq uint32_t get_n_stream() const; - // physical stream backing view index s, i.e. sinfo.strm[s]. a global cache row for a - // cell j of that view is get_strm(s)*kv->get_size() + j, the convention set_input_k_idxs - // uses. needed by the k-pool code, which writes cells that are not ubatch tokens. + // physical stream behind view s; a global row for cell j is get_strm(s)*get_size() + j, + // the convention set_input_k_idxs uses. needed by k-pool, which writes non-ubatch cells uint32_t get_strm(uint32_t s) const; const llama_kv_cache * get_kv() const; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index e61d4c94bd4..2c7e9c8c162 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -69,13 +69,9 @@ llama_memory_hybrid::llama_memory_hybrid( : filter_recr )), mem_idx(filter_idx == nullptr ? nullptr : [&] { - // a *pooling* indexer (indexer_kpool > 0, glm5next only) needs a second head for - // the compressor gate, or the pool cannot be rebuilt once its tokens leave the - // batch. every other arch leaves indexer_kpool 0 and is unchanged. - // - // the third head is the pooled key of the pool this cell ENDS, i.e. it is written - // only for cells at pos % kpool == kpool-1. caching it turns the per-step pooling - // from O(n_kv) into O(pools completed by this ubatch): see build_indexer. + // a *pooling* indexer needs a second head for the compressor gate, or the pool cannot be + // rebuilt once its tokens leave the batch; the third head is the pooled key of the pool + // this cell ENDS (pos % kpool == kpool-1) const uint32_t n_head_idx = model.hparams.indexer_kpool > 0 ? 3 : 1; std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), n_head_idx); @@ -141,9 +137,8 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } - // the indexer takes the attention cache's slot layout rather than finding its - // own: allocated separately the two drift apart when the context is rewritten - // between turns, and the top-k indices would then point at the wrong cells + // the indexer takes the attention cache's slot layout: allocated separately the two drift + // apart when the context is rewritten, and top-k would point at wrong cells llama_kv_cache::slot_info_vec_t heads_idx; if (mem_idx) { heads_idx = heads_attn; @@ -201,18 +196,13 @@ void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { - // pools group cells by absolute position, so a shift regroups them; every pooled key - // cached in the indexer rows is then stale even though each pool still looks complete. - // it survives only if WHOLE pools move: the shift must be a multiple of kpool and the - // range must start and end on a pool boundary. a pool straddling p0 or p1 keeps some - // members and moves the rest, which regroups it however aligned the shift itself is, - // and both callers pass an arbitrary bound -- the server's context shift uses - // n_keep + n_discard, its prompt-cache reuse uses the match head. + // pools group cells by absolute position, so a cached pooled key survives a shift only if + // WHOLE pools move: the shift must be a multiple of kpool AND the range must start and + // end on a pool boundary. both callers pass an arbitrary bound. if (mem_idx && hparams.indexer_kpool > 0) { const llama_pos r = (llama_pos) hparams.indexer_kpool; - // p0 < 0 means "from the start" and p1 < 0 means "to the end", so neither is a - // boundary a pool can straddle + // p0 < 0 means "from the start", p1 < 0 "to the end": neither can be straddled const bool whole_pools = shift % r == 0 && (p0 <= 0 || p0 % r == 0) && (p1 < 0 || p1 % r == 0); @@ -227,7 +217,6 @@ void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p } void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { - // as seq_add, and a divide regroups for any d != 1 if (mem_idx && hparams.indexer_kpool > 0 && d != 1) { mem_attn->set_kpool_dirty(); } @@ -303,10 +292,7 @@ llama_memory_hybrid_context::llama_memory_hybrid_context( bool optimize) : ctx_attn(mem->get_mem_attn()->init_update(lctx, optimize)), ctx_recr(mem->get_mem_recr()->init_update(lctx, optimize)), - // indexer keys carry no positional encoding, but the pending per-cell delta must - // still be cleared or the two caches disagree about whether a shift is outstanding. - // safe because an indexer only exists for LLAMA_ROPE_TYPE_NONE archs, where - // llama_kv_cache::update skips the K-shift graph and does only that + // the pending per-cell delta must still be cleared or the caches disagree about a shift ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : mem->get_mem_idx()->init_update(lctx, optimize)), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } @@ -350,8 +336,6 @@ bool llama_memory_hybrid_context::apply() { if (ctx_idx) { res = res & ctx_idx->apply(); - // a top-k over indexer cells is meaningful only if both caches cover the same - // window if (!ubatches.empty()) { GGML_ASSERT(get_idx()->get_n_kv() == get_attn()->get_n_kv()); GGML_ASSERT(get_idx()->get_n_stream() == get_attn()->get_n_stream()); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 2fd38abd939..e8c509f77c3 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2457,11 +2457,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; - // null unless the arch has an indexer cache llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; ggml_type type_idx = GGML_TYPE_F16; - // qwen4exp uses the dedicated llama_memory_hybrid_idx; glm5next carries its - // indexer in llama_memory_hybrid via filter_idx/type_idx + // qwen4exp uses llama_memory_hybrid_idx; glm5next carries its indexer in llama_memory_hybrid const bool needs_mem_idx = (arch == LLM_ARCH_QWEN4EXP); if (arch == LLM_ARCH_FALCON_H1) { filter_attn = [&](uint32_t) { return true; }; @@ -2474,12 +2472,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_MINIMAX_01 || arch == LLM_ARCH_GLM5NEXT) { - // the MTP draft context runs the NextN block and nothing else, so it - // gets a cache for that one layer. the trunk never runs it and so keeps - // the layer range it always had. handing the draft the trunk's cache is - // not just waste: the KDA layers would take cells it can never roll - // back, since a draft context is built with n_rs_seq = 0, and then a - // rejected draft fails seq_rm outright. + // the draft runs only the NextN block, so it gets a cache for that one layer. the trunk's + // cache would let the KDA layers take cells it can never roll back (n_rs_seq = 0), and a + // rejected draft then fails seq_rm const bool mtp_ctx = arch == LLM_ARCH_GLM5NEXT && cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_all > hparams.n_layer(); @@ -2497,7 +2492,6 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, if (arch == LLM_ARCH_GLM5NEXT && hparams.indexer_head_size > 0) { // unified is fine, the pool map is per SEQUENCE. see [TAG_KPOOL_SEQ_PARTITION] - // only the DSA layers carry an indexer key cache filter_idx = [&, mtp_ctx](uint32_t il) { if (mtp_ctx) { return il >= hparams.n_layer() && il < hparams.n_layer_all; diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 4c6db1eb32f..0ec4cbb42eb 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -6,8 +6,6 @@ // ssm_a holds -exp(A_log) (kimi-k3), not +exp(A_log) (bailingmoe3); converter checks -// positions the indexer keeps; at or below this many the dense path IS the sparse one. -// asserted not measured (invisible to output); the second assert is an independent spelling static uint32_t glm5next_n_select(const llama_hparams & hparams) { GGML_ASSERT(hparams.indexer_kpool > 0); GGML_ASSERT(hparams.indexer_top_k >= hparams.indexer_kpool); @@ -25,7 +23,6 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); // indexer k_norm is a LayerNorm with bias; without this key it runs at eps 0 ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); - // warned not asserted: no output comparison sees it, and an assert breaks test-llama-archs if (hparams.f_norm_eps <= 0.0f || hparams.f_norm_eps > 2e-6f) { LLAMA_LOG_WARN("%s: indexer k_norm eps is %g, but the reference hardcodes 1e-6. " "this is invisible to every output comparison; check the converter\n", @@ -39,7 +36,6 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(hparams.n_lora_q > 0 && "glm5next requires a q LoRA"); GGML_ASSERT(hparams.n_rot() == 0 && "glm5next MLA is nope-only"); - // no linear_num_heads key: KDA head count is attention.head_count (converter enforces) 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); GGML_ASSERT(hparams.ssm_d_conv > 1); @@ -65,8 +61,7 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); GGML_ASSERT(hparams.dsv4_hc_mult > 0); - // n_embd_out stays n_embd: lm_head sees the stream mean. deepseek4's hc_mult*n_embd - // makes llama-context.cpp overread t_embd, and the assert there sizes the destination + // n_embd_out stays n_embd: deepseek4's hc_mult*n_embd overreads t_embd hparams.n_embd_out_impl = 0; ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); @@ -86,8 +81,7 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all); - // to n_layer_all, not n_layer(): the NextN block is an attention block, and the loop - // below plus the memory layer filters both read the entry for it + // n_layer_all, not n_layer(): the NextN block is an attention block the filters read uint32_t n_recr = 0; for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0; @@ -95,7 +89,6 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { } GGML_ASSERT(n_recr > 0 && n_recr < hparams.n_layer() && "glm5next needs a per-layer attention.head_count_kv array"); - // every glm5next indexer is full; the generic loader only zero-fills the array for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { hparams.is_indexer_full_impl[il] = !hparams.is_recr_impl[il]; } @@ -149,7 +142,6 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), {n_embd}, flags); layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), {n_embd}, flags); - // the NextN block keeps the plain residual, so it has no mHC mixer if (il < n_layer) { layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", il), {hc_dim, hc_mix_dim}, flags); layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", il), {hc_mix_dim}, flags); @@ -195,7 +187,6 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", il), {n_embd, n_embd_indexer}, flags); layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", il), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, flags); - // key pooling: DeepSeek-V4 doubles the compressor width, GLM-5.3 does not layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", il), {n_embd, n_embd_indexer}, flags); layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", il), {n_embd_indexer, kpool}, flags); } @@ -223,7 +214,6 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), {n_embd}, flags); layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), {n_embd}, flags); - // absent in the checkpoint: NextN shares the trunk's embeddings and lm_head layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); } @@ -284,7 +274,6 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( cb(Qcur, "kda_q_norm", il); cb(Kcur, "kda_k_norm", il); - // the 1/sqrt(head_dim) query scale is applied inside build_delta_net, after this norm // g = lower_bound * sigmoid(exp(A_log)*(f_b(f_a(x)) + dt_bias)); it scales, not clamps ggml_tensor * g = ggml_mul_mat(ctx0, layer.ssm_f_b, ggml_mul_mat(ctx0, layer.ssm_f_a, inp)); @@ -306,7 +295,6 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( ggml_tensor * out = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g, beta, state, il); - // the fallbacks return a permuted view, the fused op a contiguous one ggml_tensor * o = ggml_cont_3d(ctx0, out, head_dim, n_head, n_tokens); cb(o, "kda_scan_out", il); @@ -325,11 +313,7 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( } // the store is NOT gated on the sparse path, the scoring is: gating both leaves cells -// below n_select with no indexer state, which the first ubatch to cross n_select pools. -// * weights_proj runs in fp32; bf16 head-gates flip near-tie pool rankings (vLLM, sglang) -// * k_norm is a LayerNorm WITH BIAS at eps 1e-6, not f_norm_rms_eps (transformers, vLLM) -// * the ReLU between the QK dot and the head weighting is real (modular_glm5_next.py) -// no Hadamard rotation: H is orthogonal so (Hq).(Hk) == q.k; it only helps fp8. +// below n_select with no indexer state ggml_tensor * llama_model_glm5next::graph::build_indexer( const llama_layer & layer, llm_graph_input_kpool * inp_kp, @@ -353,12 +337,10 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_tensor * gate = ggml_mul_mat(ctx0, layer.indexer_comp_wgate, cur); cb(gate, "indexer_gate", il); - // {d_idx, 2, n_tokens}: head 0 is the key, head 1 the gate ggml_tensor * packed = ggml_concat(ctx0, ggml_reshape_3d(ctx0, ik, d_idx, 1, n_tokens), ggml_reshape_3d(ctx0, gate, d_idx, 1, n_tokens), 1); - // key and gate are the first two heads of a three-head row; the third is the pooled key, - // written later from a different set of cells, so this store must leave it alone + // the third head is the pooled key, written later from other cells; leave it alone ggml_build_forward_expand(gf, mctx_idx->cpy_k_part(ctx0, ggml_reshape_2d(ctx0, packed, 2*d_idx, n_tokens), inp_kp->k_idxs, il, 2*d_idx, 0)); @@ -382,9 +364,6 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( ggml_tensor * kg_rows = ggml_view_3d(ctx0, kbuf, 2*d_idx, n_kv, n_stream, kbuf->nb[2], kbuf->nb[3], 0); - // pool only what this ubatch closed. the count is fixed (see build_inp_kpool), so the - // decode graph keeps one shape; unused slots repeat a complete pool, whose recompute is - // idempotent. non-resident slots hold 0, not a sentinel; garbage pools die to pool_bias. const int64_t n_new_max = inp_kp->new_pool_cells->ne[0]/r; ggml_tensor * members = ggml_get_rows(ctx0, kg_rows, inp_kp->new_pool_cells); @@ -411,10 +390,8 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( pool_new = ggml_reshape_2d(ctx0, pool_new, d_idx, n_new_max*n_stream); cb(pool_new, "indexer_pool_new", il); - // store into the third head of each representative cell's row. the write is indexed by - // GLOBAL row (stream-merged) while pool_reps below is stream-local, so the read cannot - // be chained off the write tensor; expand it into the graph first and let build order - // sequence them, exactly as the key/gate store above does. + // the write is by GLOBAL row while pool_reps is stream-local, so the read must not chain + // off the write tensor; expand first and let build order sequence them ggml_build_forward_expand(gf, mctx_idx->cpy_k_part(ctx0, pool_new, inp_kp->new_pool_reps, il, d_idx, 2*d_idx)); @@ -430,8 +407,7 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( iq = ggml_reshape_4d(ctx0, iq, d_idx, n_ihead, n_tps, n_stream); cb(iq, "indexer_q", il); - // sign-unconstrained head weights: no softmax, no abs, no relu; both scale constants - // fold in here. GGML_PREC_F32 is not cosmetic: bf16 can swap two near-tied pools + // sign-unconstrained head weights; PREC_F32 is load-bearing, bf16 swaps near-tied pools ggml_tensor * w = ggml_mul_mat(ctx0, layer.indexer_proj, cur); ggml_mul_mat_set_prec(w, GGML_PREC_F32); w = ggml_reshape_4d(ctx0, w, n_ihead, n_tps, 1, n_stream); @@ -466,16 +442,13 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( cb(pool_score, "indexer_pool_score", il); } - // top-k over POOLS then expand, as in the reference: a cell-level top-k is wrong - // because relu ties span pool boundaries and ggml_top_k splits the pool it lands in + // top-k over POOLS then expand: a cell-level top-k is wrong, relu ties span pool bounds const int64_t select_k = llama_kpool_select_k(n_pools, hparams.indexer_top_k, r); GGML_ASSERT(select_k > 0 && select_k <= n_pools); - // {select_k, n_tps, n_stream} of POOL ordinals ggml_tensor * sel = ggml_cont(ctx0, ggml_top_k(ctx0, pool_score, (int) select_k)); cb(sel, "indexer_top_k_pools", il); - // the query axis folds into the gather's row axis, so ONE get_rows serves every query ggml_tensor * pc3 = ggml_reshape_3d(ctx0, inp_kp->pool_cells, r, n_pools, n_stream); ggml_tensor * sel_flat = ggml_reshape_2d(ctx0, sel, select_k*n_tps, n_stream); @@ -487,8 +460,7 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( return top_k; } -// absorbed form (deepseek2/glm-dsa): q_nope goes through wk_b so q.k is taken against -// the latent the cache holds; the naive form needs a V cache this layout lacks +// absorbed form (deepseek2/glm-dsa): q_nope goes through wk_b; the naive form needs a V cache ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( const llama_layer & layer, llm_graph_input_attn_k * inp_attn, @@ -564,7 +536,6 @@ ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( int il) const { const auto & layer = model.layers[il]; - // the leading dense layers clamp like the experts: one Glm5NextTextMLP serves both if (il < (int) hparams.n_layer_dense_lead) { return build_ffn(cur, layer.ffn_up, nullptr, nullptr, @@ -606,8 +577,7 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); - // one map for the whole ubatch; nothing in it depends on the layer. gated on n_ctx, - // not n_kv, which grows and would flip the graph topology mid-run + // gated on n_ctx, not n_kv, which grows and would flip the graph topology mid-run llm_graph_input_kpool * inp_kp = nullptr; bool indexer_scoring = false; { @@ -627,7 +597,7 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa const int64_t hc = hparams.dsv4_hc_mult; - // hc_mult exact copies of the embedding: no scaling, no one-hot into stream 0 + // exact copies: no scaling, no one-hot into stream 0 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); @@ -688,23 +658,21 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); } - // when unmasked nextn embeddings are requested, t_h_nextn must keep all rows, so the - // early output masking has to be skipped (it is applied after the final norm instead) + // unmasked nextn embeddings need all rows, so early output masking is skipped here const bool mask_early = !cparams.embeddings_nextn || cparams.embeddings_nextn_masked; if (inp_out_ids && mask_early) { - // flattened: get_rows needs one token's streams to be one contiguous row + // get_rows needs one token's streams contiguous ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); inpL = ggml_reshape_3d(ctx0, ggml_get_rows(ctx0, flat, inp_out_ids), n_embd, hc, n_outputs); } - // no hc_head tensor here: unweighted mean, not DeepSeek-V4's learned gated head + // unweighted mean, not DeepSeek-V4's learned gated head cur = build_hc_mean(ctx0, inpL); cb(cur, "hc_mean", -1); cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); - // post-norm hidden state feeds the NextN/MTP draft head cb(cur, "h_nextn", -1); res->t_h_nextn = cur; @@ -722,10 +690,7 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa ggml_build_forward_expand(gf, cur); } -// NextN draft head. Unlike glm-dsa and deepseek32, whose MTP blocks skip the DSA indexer, -// blk.45 here ships a full indexer and is an ordinary DSA layer, so it runs the same -// build_layer_attn/build_layer_ffn the trunk does. What it does NOT run is the mHC mixer: -// the NextN block keeps a plain residual and has no hc_* tensors (see load_arch_tensors). +// NextN draft head: an ordinary DSA layer, but a plain residual and no hc_* tensors llama_model_glm5next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) : graph(params) { GGML_ASSERT(hparams.n_layer_nextn == 1 && "glm5next MTP supports a single NextN block"); @@ -772,8 +737,8 @@ llama_model_glm5next::graph_mtp::graph_mtp(const llama_model & model, const llm_ llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); - // the NextN block is DSA, so nothing here consumes the recurrent half. s_copy would then - // never enter the graph, never be allocated, and set_input would still read its buffer. + // the NextN block is DSA: nothing consumes the recurrent half, so s_copy never enters the + // graph while set_input would still read its buffer ggml_build_forward_expand(gf, inp_mem->get_recr()->s_copy); llm_graph_input_kpool * inp_kp = nullptr; @@ -799,7 +764,6 @@ llama_model_glm5next::graph_mtp::graph_mtp(const llama_model & model, const llm_ ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0), layer.nextn.eh_proj_s); cb(cur, "mtp_eh_proj", il); - // plain pre-norm residual block, the trunk's helpers minus the mHC wrapper ggml_tensor * residual = cur; cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); @@ -823,7 +787,6 @@ llama_model_glm5next::graph_mtp::graph_mtp(const llama_model & model, const llm_ GGML_ASSERT(head_norm_w && "glm5next MTP: no nextn.shared_head_norm and no output_norm"); cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1); - // the post-norm hidden state seeds the next MTP step cb(cur, "h_nextn", -1); res->t_h_nextn = cur; diff --git a/src/models/models.h b/src/models/models.h index 82057594161..972280fcd3c 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1176,8 +1176,7 @@ 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; - // method-only mixin; here only so glm5next, deriving from this graph, reaches - // build_delta_net. deepseek4 itself has no recurrent layers + // method-only mixin, so glm5next reaches build_delta_net; deepseek4 has no recurrent layers struct graph : public llm_build_delta_net_base { graph(const llm_graph_params & params) : llm_build_delta_net_base(params) {} graph(const llama_model & model, const llm_graph_params & params); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index a70f39111d7..69f73d7c687 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -176,8 +176,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { 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 if (arch == LLM_ARCH_GLM5NEXT) { - // head_count doubles as the KDA head count, so it stays uniform; the kv array is what - // marks the recurrent layers, and the loader asserts it holds both a zero and a nonzero + // head_count doubles as the KDA head count; the kv array marks the recurrent layers GGML_ASSERT(n_layer >= 2); std::vector n_head_kv_per_layer; n_head_kv_per_layer.reserve(n_layer); @@ -226,8 +225,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types); } } else if (arch == LLM_ARCH_GLM5NEXT) { - // nope-only MLA: the cache holds the bare latent, so no rope width is added on top of - // the kv LoRA rank and n_rot has to be an explicit 0, not the head size default + // nope-only MLA: the cache holds the bare latent, so n_rot must be an explicit 0 ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(512)); ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512)); ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(0)); @@ -304,8 +302,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); // build_hc_pre asserts exactly 4 streams ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); - // the only arch that pools indexer keys; top_k must be a whole number of pools, and - // the resulting selection width has to stay under n_ctx or the sparse path goes unused + // the only arch that pools indexer keys; top_k must be a whole number of pools and the + // selection width must stay under n_ctx or the sparse path goes unused ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, uint32_t(4)); // glm5next reads these unconditionally; the if (moe) block below never sets them ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 4627d4f68b2..45db71fdcfe 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1584,9 +1584,6 @@ mtmd_image_preproc_out mtmd_image_preprocessor_muse_glimmer::preprocess(const cl return output; } -// -// mtmd_image_preprocessor_glm5next -// // for a still image the reference's temporal_factor cancels out, leaving pixel area vs min/max clip_image_size mtmd_image_preprocessor_glm5next::smart_resize(const clip_hparams & hparams, const clip_image_size & size) {