From a505a75878db9ec69e5cf48da4002529e911d840 Mon Sep 17 00:00:00 2001 From: danielhanchen Date: Wed, 26 Aug 2026 06:30:38 +0000 Subject: [PATCH 01/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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/62] 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 238f88023f1c9dbc1838c964d482bc65590e0c1e Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 31 Aug 2026 18:23:05 -0600 Subject: [PATCH 35/62] tests: add glm5next --depth-sweep mode (CPU-vs-device logit divergence over long context) Runs one tiny random-weight glm5next MoE model on the CPU backend and the first non-CPU device with an identical token stream, decoding in 2048-token chunks and comparing last-position logits (NMSE + greedy argmax) after each. --ctx and --ub override the context length baked into the fixture and the microbatch, because the production '@@@@' collapse depends jointly on (depth, n_ctx, n_ubatch). Toy-scale repro hunt for ggml-org/llama.cpp#27754. --- tests/test-llama-archs.cpp | 106 ++++++++++++++++++++++++++++++++++++- 1 file changed, 104 insertions(+), 2 deletions(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 6e19954c683..d68a00fa223 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -374,6 +374,10 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) return true; } +// depth-sweep overrides (0 = leave defaults). Set only by test_depth_sweep(). +static uint32_t g_depth_sweep_n_ctx = 0; +static uint32_t g_depth_sweep_n_ub = 0; + static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { @@ -386,11 +390,12 @@ static std::pair get_model_and_ctx( model_params.split_mode = split_mode; llama_context_params ctx_params = llama_context_default_params(); - ctx_params.n_ctx = 0; + ctx_params.n_ctx = g_depth_sweep_n_ctx; // 0 = model default (original behavior) ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; if (!encode) { - ctx_params.n_ubatch = 64; + ctx_params.n_ubatch = g_depth_sweep_n_ub ? g_depth_sweep_n_ub : 64; + ctx_params.n_batch = std::max(2048, ctx_params.n_ubatch); } size_t tmp = seed; @@ -783,6 +788,88 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg return all_ok ? 0 : 1; } +// Long-context CPU-vs-device divergence sweep for the glm5next "@@@@" collapse +// (ggml-org/llama.cpp#27754 / #27752). Runs ONE tiny random-weight glm5next +// model on the CPU backend and on the first non-CPU device with an identical +// token stream, decoding in chunks; after each chunk the last-position logits +// are compared (NMSE + greedy argmax). On the real model the collapse depth +// depends on (depth, n_ctx, n_ubatch); this asks whether that reproduces at +// toy scale, in minutes instead of hours. +static int test_depth_sweep(const size_t seed, const uint32_t max_depth, + const uint32_t n_ctx_arg, const uint32_t n_ub) { + g_depth_sweep_n_ctx = n_ctx_arg; + g_depth_sweep_n_ub = n_ub; + + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_GLM5NEXT, /*moe=*/true); + // the fixture bakes context_length=128; the KV/pool structures must be + // sized for the sweep target instead + gguf_set_val_u32(gguf_ctx.get(), "glm5next.context_length", n_ctx_arg); + + ggml_backend_dev_t dev_gpu = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_buffer_type(dev) != ggml_backend_cpu_buffer_type()) { + dev_gpu = dev; + break; + } + } + if (dev_gpu == nullptr) { + printf("depth-sweep: no non-CPU device found\n"); + return 1; + } + printf("depth-sweep: glm5next-moe, n_ctx=%u, n_ubatch=%u, max_depth=%u, device=%s\n", + n_ctx_arg, n_ub ? n_ub : 64, max_depth, ggml_backend_dev_description(dev_gpu)); + + auto mc_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}); + auto mc_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {dev_gpu}); + const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(mc_cpu.first.get())); + + const std::vector tokens = get_tokens(max_depth, n_vocab, seed); + const uint32_t chunk = 2048; + + llama_batch batch = llama_batch_init(chunk, 0, 1); + bool diverged = false; + for (uint32_t pos0 = 0; pos0 < max_depth; pos0 += chunk) { + const uint32_t n = std::min(chunk, max_depth - pos0); + common_batch_clear(batch); + for (uint32_t i = 0; i < n; i++) { + common_batch_add(batch, tokens[pos0 + i], pos0 + i, {0}, i == n - 1); + } + if (llama_decode(mc_cpu.second.get(), batch)) { + printf("depth-sweep: CPU decode failed at depth %u\n", pos0 + n); + llama_batch_free(batch); + return 1; + } + if (llama_decode(mc_dev.second.get(), batch)) { + printf("depth-sweep: device decode failed at depth %u\n", pos0 + n); + llama_batch_free(batch); + return 1; + } + const float * lc = llama_get_logits_ith(mc_cpu.second.get(), n - 1); + const float * ld = llama_get_logits_ith(mc_dev.second.get(), n - 1); + double se = 0.0, ref = 0.0; + uint32_t amax_c = 0, amax_d = 0; + for (uint32_t j = 0; j < n_vocab; j++) { + const double d = (double) lc[j] - (double) ld[j]; + se += d * d; + ref += (double) lc[j] * (double) lc[j]; + if (lc[j] > lc[amax_c]) amax_c = j; + if (ld[j] > ld[amax_d]) amax_d = j; + } + const double nmse_val = ref > 0.0 ? se / ref : se; + const bool bad = nmse_val > 1e-3 || amax_c != amax_d; + printf("depth=%7u nmse=%.3e argmax_cpu=%u argmax_dev=%u%s\n", + pos0 + n, nmse_val, amax_c, amax_d, bad ? " <-- DIVERGED" : ""); + fflush(stdout); + if (bad) { + diverged = true; + } + } + llama_batch_free(batch); + printf("depth-sweep: %s\n", diverged ? "DIVERGENCE FOUND" : "no divergence up to max depth"); + return diverged ? 2 : 0; +} + int main(int argc, char ** argv) { // FIXME these tests are disabled in the CI for macOS-latest-cmake-arm64 because they are segfaulting common_init(); @@ -790,6 +877,9 @@ int main(int argc, char ** argv) { llm_arch arch = LLM_ARCH_UNKNOWN; size_t seed = rd(); + uint32_t depth_sweep = 0; + uint32_t sweep_ctx = 131072; + uint32_t sweep_ub = 0; ggml_log_level log_level = GGML_LOG_LEVEL_ERROR; std::string out; @@ -807,6 +897,15 @@ int main(int argc, char ** argv) { return 1; } } + if (strcmp(argv[i], "--depth-sweep") == 0 && i + 1 < argc) { + depth_sweep = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--ctx") == 0 && i + 1 < argc) { + sweep_ctx = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--ub") == 0 && i + 1 < argc) { + sweep_ub = std::stoul(argv[++i]); + } if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) { if (i + 1 < argc) { seed = std::stoull(argv[++i]); @@ -831,6 +930,9 @@ int main(int argc, char ** argv) { printf("%s: using seed %zu\n", __func__, seed); try { + if (depth_sweep > 0) { + return test_depth_sweep(seed, depth_sweep, sweep_ctx, sweep_ub); + } if (!out.empty()) { return save_models(arch, seed, log_level, out); } From ed1313c0bdd9df5b350b3a136ff10d799218bdd0 Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 31 Aug 2026 18:26:04 -0600 Subject: [PATCH 36/62] tests: depth-sweep gains --topk override and same-backend ubatch A/B (--ub2/--b-cpu) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real GLM-5.3-Flash uses indexer.top_k=2048 (fixture: 8); --topk overrides it. --ub2 N with --b-cpu compares CPU@ub vs CPU@ub2 — the production collapse is microbatch-dependent, so a backend-independent defect in the per-ubatch indexer input construction shows up on this axis where CPU-vs-device cannot. --- tests/test-llama-archs.cpp | 37 +++++++++++++++++++++++++++++++------ 1 file changed, 31 insertions(+), 6 deletions(-) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index d68a00fa223..c1096bc402c 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -796,14 +796,19 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg // depends on (depth, n_ctx, n_ubatch); this asks whether that reproduces at // toy scale, in minutes instead of hours. static int test_depth_sweep(const size_t seed, const uint32_t max_depth, - const uint32_t n_ctx_arg, const uint32_t n_ub) { + const uint32_t n_ctx_arg, const uint32_t n_ub, + const uint32_t n_ub2, const bool b_on_cpu, + const uint32_t topk_override) { g_depth_sweep_n_ctx = n_ctx_arg; - g_depth_sweep_n_ub = n_ub; gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_GLM5NEXT, /*moe=*/true); // the fixture bakes context_length=128; the KV/pool structures must be // sized for the sweep target instead gguf_set_val_u32(gguf_ctx.get(), "glm5next.context_length", n_ctx_arg); + if (topk_override > 0) { + // real GLM-5.3-Flash: indexer.top_k = 2048, kpool = 4 (fixture: 8/4) + gguf_set_val_u32(gguf_ctx.get(), "glm5next.attention.indexer.top_k", topk_override); + } ggml_backend_dev_t dev_gpu = nullptr; for (size_t i = 0; i < ggml_backend_dev_count(); i++) { @@ -817,11 +822,18 @@ static int test_depth_sweep(const size_t seed, const uint32_t max_depth, printf("depth-sweep: no non-CPU device found\n"); return 1; } - printf("depth-sweep: glm5next-moe, n_ctx=%u, n_ubatch=%u, max_depth=%u, device=%s\n", - n_ctx_arg, n_ub ? n_ub : 64, max_depth, ggml_backend_dev_description(dev_gpu)); + // side A: CPU @ n_ub. side B: (cpu|device) @ (n_ub2 ? n_ub2 : n_ub). + const uint32_t ub_a = n_ub ? n_ub : 64; + const uint32_t ub_b = n_ub2 ? n_ub2 : ub_a; + printf("depth-sweep: glm5next-moe, n_ctx=%u, top_k=%s, max_depth=%u | A=cpu@ub%u B=%s@ub%u\n", + n_ctx_arg, topk_override ? std::to_string(topk_override).c_str() : "fixture(8)", + max_depth, ub_a, b_on_cpu ? "cpu" : ggml_backend_dev_description(dev_gpu), ub_b); + g_depth_sweep_n_ub = ub_a; auto mc_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}); - auto mc_dev = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {dev_gpu}); + g_depth_sweep_n_ub = ub_b; + auto mc_dev = b_on_cpu ? get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}) + : get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {dev_gpu}); const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(mc_cpu.first.get())); const std::vector tokens = get_tokens(max_depth, n_vocab, seed); @@ -880,6 +892,9 @@ int main(int argc, char ** argv) { uint32_t depth_sweep = 0; uint32_t sweep_ctx = 131072; uint32_t sweep_ub = 0; + uint32_t sweep_ub2 = 0; + uint32_t sweep_topk = 0; + bool sweep_b_cpu = false; ggml_log_level log_level = GGML_LOG_LEVEL_ERROR; std::string out; @@ -906,6 +921,15 @@ int main(int argc, char ** argv) { if (strcmp(argv[i], "--ub") == 0 && i + 1 < argc) { sweep_ub = std::stoul(argv[++i]); } + if (strcmp(argv[i], "--ub2") == 0 && i + 1 < argc) { + sweep_ub2 = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--b-cpu") == 0) { + sweep_b_cpu = true; + } + if (strcmp(argv[i], "--topk") == 0 && i + 1 < argc) { + sweep_topk = std::stoul(argv[++i]); + } if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) { if (i + 1 < argc) { seed = std::stoull(argv[++i]); @@ -931,7 +955,8 @@ int main(int argc, char ** argv) { try { if (depth_sweep > 0) { - return test_depth_sweep(seed, depth_sweep, sweep_ctx, sweep_ub); + return test_depth_sweep(seed, depth_sweep, sweep_ctx, sweep_ub, + sweep_ub2, sweep_b_cpu, sweep_topk); } if (!out.empty()) { return save_models(arch, seed, log_level, out); From 2841c7bcf73e2cef7d3d9607b20afa9e4e4169ef Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 31 Aug 2026 18:48:23 -0600 Subject: [PATCH 37/62] tests: depth-sweep --iheads/--iklen overrides (real glm5next: 32 heads, key 128) --- tests/test-llama-archs.cpp | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index c1096bc402c..c726dcbe49f 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -377,6 +377,8 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) // depth-sweep overrides (0 = leave defaults). Set only by test_depth_sweep(). static uint32_t g_depth_sweep_n_ctx = 0; static uint32_t g_depth_sweep_n_ub = 0; +static uint32_t g_sweep_iheads = 0; +static uint32_t g_sweep_iklen = 0; static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, @@ -809,6 +811,12 @@ static int test_depth_sweep(const size_t seed, const uint32_t max_depth, // real GLM-5.3-Flash: indexer.top_k = 2048, kpool = 4 (fixture: 8/4) gguf_set_val_u32(gguf_ctx.get(), "glm5next.attention.indexer.top_k", topk_override); } + if (g_sweep_iheads > 0) { // real: 32 (fixture: 1) + gguf_set_val_u32(gguf_ctx.get(), "glm5next.attention.indexer.head_count", g_sweep_iheads); + } + if (g_sweep_iklen > 0) { // real: 128 (fixture: 64) + gguf_set_val_u32(gguf_ctx.get(), "glm5next.attention.indexer.key_length", g_sweep_iklen); + } ggml_backend_dev_t dev_gpu = nullptr; for (size_t i = 0; i < ggml_backend_dev_count(); i++) { @@ -930,6 +938,12 @@ int main(int argc, char ** argv) { if (strcmp(argv[i], "--topk") == 0 && i + 1 < argc) { sweep_topk = std::stoul(argv[++i]); } + if (strcmp(argv[i], "--iheads") == 0 && i + 1 < argc) { + g_sweep_iheads = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--iklen") == 0 && i + 1 < argc) { + g_sweep_iklen = std::stoul(argv[++i]); + } if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) { if (i + 1 < argc) { seed = std::stoull(argv[++i]); From 40f7a342db66314d66081e7f237b20e47164c605 Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 31 Aug 2026 19:10:40 -0600 Subject: [PATCH 38/62] tests: depth-sweep --layers/--dlead (glm5next layer count + leading dense blocks) --- tests/test-llama-archs.cpp | 14 ++++++++++++++ 1 file changed, 14 insertions(+) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index c726dcbe49f..3ce8bdc1d03 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -123,6 +123,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { n_embd = 128; n_head = 1; n_ff = 192; + if (arch == LLM_ARCH_GLM5NEXT && g_sweep_layers > 0) { + n_layer = g_sweep_layers; + } } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { n_layer = 3; } else if (arch == LLM_ARCH_CHAMELEON) { @@ -379,6 +382,8 @@ static uint32_t g_depth_sweep_n_ctx = 0; static uint32_t g_depth_sweep_n_ub = 0; static uint32_t g_sweep_iheads = 0; static uint32_t g_sweep_iklen = 0; +static uint32_t g_sweep_layers = 0; // glm5next fixture layer-count override +static uint32_t g_sweep_dlead = 0; // leading dense layers (real model: 3 of 45) static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, @@ -817,6 +822,9 @@ static int test_depth_sweep(const size_t seed, const uint32_t max_depth, if (g_sweep_iklen > 0) { // real: 128 (fixture: 64) gguf_set_val_u32(gguf_ctx.get(), "glm5next.attention.indexer.key_length", g_sweep_iklen); } + if (g_sweep_dlead > 0) { // real: 3 + gguf_set_val_u32(gguf_ctx.get(), "glm5next.leading_dense_block_count", g_sweep_dlead); + } ggml_backend_dev_t dev_gpu = nullptr; for (size_t i = 0; i < ggml_backend_dev_count(); i++) { @@ -944,6 +952,12 @@ int main(int argc, char ** argv) { if (strcmp(argv[i], "--iklen") == 0 && i + 1 < argc) { g_sweep_iklen = std::stoul(argv[++i]); } + if (strcmp(argv[i], "--layers") == 0 && i + 1 < argc) { + g_sweep_layers = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--dlead") == 0 && i + 1 < argc) { + g_sweep_dlead = std::stoul(argv[++i]); + } if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) { if (i + 1 < argc) { seed = std::stoull(argv[++i]); From f5d5216bd276c82c44938c075813d482c61c3528 Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 31 Aug 2026 22:01:17 -0600 Subject: [PATCH 39/62] ggml-metal: 64-bit output offsets in mul_mm batched dst indexing MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit int32 im*ne1*ne0 (and im*N*M on the tensor path) wraps past 2^31 elements. For glm5next's dense-masked DSA attention the KQ tensor is [n_kv, n_ubatch, 64 heads] f32, so at ub=512 head h's base offset wraps once h*512*n_kv >= 2^31 — first at head 63 for n_kv ~66.6K (the observed 65K collapse-boundary flattening), reaching head 39 by n_kv ~108.8K. Wrapped heads read back zeros (standalone probe: FAIL corners exactly matching 2^31/(ub*n_kv) onset per head; all 256 corners OK after this change). Mechanism for the deep-context '@' collapse in PR #27754/#27752. --- ggml/src/ggml-metal/kernels/mul_mm.metal | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal index ee848eed6d6..df3be003890 100644 --- a/ggml/src/ggml-metal/kernels/mul_mm.metal +++ b/ggml/src/ggml-metal/kernels/mul_mm.metal @@ -134,7 +134,8 @@ kernel void kernel_mul_mm( // Store result tile to output matrix (with batch offset) // cT.store handles bounds checking via tD's extents (M, N) - device float * dstBatch = (device float *)dst + im * N * M; + // int32 im*N*M wraps for KQ [n_kv, ub, 64] past 2^31 elements (glm5next collapse) + device float * dstBatch = (device float *)dst + (uint64_t)im * (uint64_t)N * (uint64_t)M; auto tD = tensor(dstBatch, dextents(M, N), array({1, M})); cT.store(tD.slice(ra, rb)); @@ -318,7 +319,7 @@ kernel void kernel_mul_mm( // if no bounds checks on the output are needed, we can directly write to device memory device float * C = (device float *) dst + (r0 + 32*(sgitg & 1)) + \ - (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0; + (uint64_t)(r1 + 16*(sgitg >> 1)) * (uint64_t)args.ne0 + (uint64_t)im*(uint64_t)args.ne1*(uint64_t)args.ne0; for (short i = 0; i < 8; i++) { simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false); @@ -337,7 +338,7 @@ kernel void kernel_mul_mm( if (sgitg == 0) { for (int j = tiitg; j < nr1; j += NR1) { - device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0; + device float * D = (device float *) dst + r0 + (uint64_t)(r1 + j)*(uint64_t)args.ne0 + (uint64_t)im*(uint64_t)args.ne1*(uint64_t)args.ne0; device float4 * D4 = (device float4 *) D; threadgroup float * C = temp_str + (j*NR0); From c31aa2d8b81fdb312d569f6ff9d27c26693a75bb Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 31 Aug 2026 23:20:39 -0600 Subject: [PATCH 40/62] llama-context: env-gated alloc-map dump (LLAMA_ALLOCDUMP_*) for wrap-landing investigation --- src/llama-context.cpp | 59 +++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 59 insertions(+) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 9cbc8759d24..e98bee6b3f3 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -17,8 +17,10 @@ #include #include #include +#include #include #include +#include // // llama_context @@ -1379,6 +1381,63 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll ret = GGML_STATUS_ALLOC_FAILED; return nullptr; } + + // fabley: alloc-map dump for the wrap-landing investigation (~/conference). + // Env-gated; dumps every tensor's VA pointer + owning buffer for graphs whose + // batched kq ne0 crosses LLAMA_ALLOCDUMP_MIN_NE0, at most every _STRIDE tokens. + { + static const char * fbl_path = getenv("LLAMA_ALLOCDUMP_FILE"); + if (fbl_path) { + static long long fbl_gate = getenv("LLAMA_ALLOCDUMP_MIN_NE0") ? atoll(getenv("LLAMA_ALLOCDUMP_MIN_NE0")) : 0; + static long long fbl_stride = getenv("LLAMA_ALLOCDUMP_STRIDE") ? atoll(getenv("LLAMA_ALLOCDUMP_STRIDE")) : 2048; + static long long fbl_last = -1; + long long kq_ne0 = -1; + for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { + const ggml_tensor * n = ggml_graph_node(gf, i); + if (strncmp(n->name, "kq-", 3) == 0 && n->ne[2] >= 8 && n->ne[0] > kq_ne0) { + kq_ne0 = n->ne[0]; + } + } + if (kq_ne0 >= fbl_gate && (fbl_last < 0 || kq_ne0 - fbl_last >= fbl_stride)) { + fbl_last = kq_ne0; + if (FILE * f = fopen(fbl_path, "a")) { + fprintf(f, "GRAPH kq_ne0=%lld n_nodes=%d n_tokens=%d\n", + kq_ne0, ggml_graph_n_nodes(gf), (int) ubatch.n_tokens); + std::unordered_set fbl_seen; + auto fbl_dump = [&](const char * tag, int idx, const ggml_tensor * t) { + if (t == nullptr || fbl_seen.count(t)) { + return; + } + fbl_seen.insert(t); + ggml_backend_buffer_t buf = t->buffer; + fprintf(f, "%s|%d|%s|%s|%p|%zu|%lld,%lld,%lld,%lld|%p|%zu|%s|%p\n", + tag, idx, t->name, ggml_op_name(t->op), t->data, ggml_nbytes(t), + (long long) t->ne[0], (long long) t->ne[1], + (long long) t->ne[2], (long long) t->ne[3], + buf ? ggml_backend_buffer_get_base(buf) : nullptr, + buf ? ggml_backend_buffer_get_size(buf) : (size_t) 0, + buf ? ggml_backend_buffer_name(buf) : "-", + (const void *) t->view_src); + }; + for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { + ggml_tensor * n = ggml_graph_node(gf, i); + fbl_dump("N", i, n); + if (n->view_src) { + fbl_dump("V", i, n->view_src); + } + for (int j = 0; j < GGML_MAX_SRC; j++) { + fbl_dump("S", i, n->src[j]); + if (n->src[j] && n->src[j]->view_src) { + fbl_dump("V", i, n->src[j]->view_src); + } + } + } + fprintf(f, "ENDGRAPH\n"); + fclose(f); + } + } + } + } } // set the input data for the input tensors From b4394623bd7c69d8b23cf3af2772fc638e38b855 Mon Sep 17 00:00:00 2001 From: Your Name Date: Mon, 31 Aug 2026 23:57:32 -0600 Subject: [PATCH 41/62] metal: gpuAddress diagnostic (codex patch) + allocdump gpu_data/gpu_base columns --- ggml/include/ggml-metal.h | 5 +++++ ggml/src/ggml-metal/ggml-metal-device.h | 1 + ggml/src/ggml-metal/ggml-metal-device.m | 10 ++++++++++ ggml/src/ggml-metal/ggml-metal.cpp | 11 +++++++++++ src/llama-context.cpp | 10 ++++++++-- 5 files changed, 35 insertions(+), 2 deletions(-) diff --git a/ggml/include/ggml-metal.h b/ggml/include/ggml-metal.h index 433838f0d6d..eb7dec7dbea 100644 --- a/ggml/include/ggml-metal.h +++ b/ggml/include/ggml-metal.h @@ -54,6 +54,11 @@ GGML_BACKEND_API bool ggml_backend_metal_supports_family(ggml_backend_t backend, // capture all command buffers committed the next time `ggml_backend_graph_compute` is called GGML_BACKEND_API void ggml_backend_metal_capture_next_compute(ggml_backend_t backend); +// Returns tensor's Metal GPU virtual address, or 0 for a non-Metal buffer. +// This is a diagnostic interface; ordinary callers should not depend on the +// relative placement of separately allocated MTLBuffers. +GGML_BACKEND_API uint64_t ggml_backend_metal_buffer_get_gpu_address(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor); + GGML_BACKEND_API ggml_backend_reg_t ggml_backend_metal_reg(void); #ifdef __cplusplus diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 003b688dbac..cef4a083ce7 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -329,6 +329,7 @@ ggml_metal_buffer_t ggml_metal_buffer_map (ggml_metal_device_t dev, void * ptr, void ggml_metal_buffer_free (ggml_metal_buffer_t buf); void * ggml_metal_buffer_get_base (ggml_metal_buffer_t buf); +uint64_t ggml_metal_buffer_get_gpu_address(ggml_metal_buffer_t buf, const struct ggml_tensor * tensor); bool ggml_metal_buffer_is_shared(ggml_metal_buffer_t buf); void ggml_metal_buffer_memset_tensor(ggml_metal_buffer_t buf, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 41ce90dc8a9..3de01d6b5c6 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -2141,6 +2141,16 @@ void ggml_metal_buffer_free(ggml_metal_buffer_t buf) { return buf->all_data; } +uint64_t ggml_metal_buffer_get_gpu_address(ggml_metal_buffer_t buf, const struct ggml_tensor * tensor) { + struct ggml_metal_buffer_id bid = ggml_metal_buffer_get_id(buf, tensor); + if (bid.metal == nil) { + return 0; + } + + id metal = bid.metal; + return metal.gpuAddress + bid.offs; +} + bool ggml_metal_buffer_is_shared(ggml_metal_buffer_t buf) { return buf->is_shared; } diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 9756d47050c..414851b1335 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -182,6 +182,17 @@ static bool ggml_backend_buffer_is_metal(ggml_backend_buffer_t buffer) { buffer->iface.free_buffer == ggml_backend_metal_buffer_private_free_buffer; } +uint64_t ggml_backend_metal_buffer_get_gpu_address( + ggml_backend_buffer_t buffer, + const struct ggml_tensor * tensor) { + if (buffer == nullptr || tensor == nullptr || !ggml_backend_buffer_is_metal(buffer)) { + return 0; + } + + ggml_metal_buffer_t ctx = (ggml_metal_buffer_t) buffer->context; + return ggml_metal_buffer_get_gpu_address(ctx, tensor); +} + // // buffer types // diff --git a/src/llama-context.cpp b/src/llama-context.cpp index e98bee6b3f3..2b996d22bae 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1,6 +1,7 @@ #include "llama-context.h" #include "ggml.h" +#include "ggml-metal.h" #include "llama-arch.h" #include "llama-graph.h" #include "llama-impl.h" @@ -1410,14 +1411,19 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } fbl_seen.insert(t); ggml_backend_buffer_t buf = t->buffer; - fprintf(f, "%s|%d|%s|%s|%p|%zu|%lld,%lld,%lld,%lld|%p|%zu|%s|%p\n", + const uint64_t gpu_data = buf ? ggml_backend_metal_buffer_get_gpu_address(buf, t) : 0; + const uint64_t gpu_base = gpu_data + ? gpu_data - ((uintptr_t) t->data - (uintptr_t) ggml_backend_buffer_get_base(buf)) + : 0; + fprintf(f, "%s|%d|%s|%s|%p|%zu|%lld,%lld,%lld,%lld|%p|%zu|%s|%p|0x%llx|0x%llx\n", tag, idx, t->name, ggml_op_name(t->op), t->data, ggml_nbytes(t), (long long) t->ne[0], (long long) t->ne[1], (long long) t->ne[2], (long long) t->ne[3], buf ? ggml_backend_buffer_get_base(buf) : nullptr, buf ? ggml_backend_buffer_get_size(buf) : (size_t) 0, buf ? ggml_backend_buffer_name(buf) : "-", - (const void *) t->view_src); + (const void *) t->view_src, + (unsigned long long) gpu_data, (unsigned long long) gpu_base); }; for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { ggml_tensor * n = ggml_graph_node(gf, i); From 42b9e5e8b44f09749034e5d3873073d7e673218c Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 00:20:35 -0600 Subject: [PATCH 42/62] allocdump: LLAMA_ALLOCDUMP_REQ_TOKENS full-ubatch filter, bidirectional stride (grokk-reply) --- src/llama-context.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 2b996d22bae..5f87fb5b901 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1391,7 +1391,10 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll if (fbl_path) { static long long fbl_gate = getenv("LLAMA_ALLOCDUMP_MIN_NE0") ? atoll(getenv("LLAMA_ALLOCDUMP_MIN_NE0")) : 0; static long long fbl_stride = getenv("LLAMA_ALLOCDUMP_STRIDE") ? atoll(getenv("LLAMA_ALLOCDUMP_STRIDE")) : 2048; + static long long fbl_reqtok = getenv("LLAMA_ALLOCDUMP_REQ_TOKENS") ? atoll(getenv("LLAMA_ALLOCDUMP_REQ_TOKENS")) : 0; static long long fbl_last = -1; + // grokk-reply: reserve (n_tokens=1) and tail batches are the wrong geometry + const bool fbl_tok_ok = fbl_reqtok == 0 || (long long) ubatch.n_tokens == fbl_reqtok; long long kq_ne0 = -1; for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { const ggml_tensor * n = ggml_graph_node(gf, i); @@ -1399,7 +1402,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll kq_ne0 = n->ne[0]; } } - if (kq_ne0 >= fbl_gate && (fbl_last < 0 || kq_ne0 - fbl_last >= fbl_stride)) { + if (fbl_tok_ok && kq_ne0 >= fbl_gate && (fbl_last < 0 || llabs(kq_ne0 - fbl_last) >= fbl_stride)) { fbl_last = kq_ne0; if (FILE * f = fopen(fbl_path, "a")) { fprintf(f, "GRAPH kq_ne0=%lld n_nodes=%d n_tokens=%d\n", From b5d526418a0ebf15d2ad9f1864268a48faae5cab Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 09:21:44 -0600 Subject: [PATCH 43/62] metal: LLAMA_METAL_CAPTURE_ARM_FILE one-shot capture arming + tensor-named debug groups (speed conference 008/009) --- ggml/src/ggml-metal/ggml-metal-context.m | 17 +++++++++++++++++ ggml/src/ggml-metal/ggml-metal-ops.cpp | 5 ++++- 2 files changed, 21 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index 1227ed39a09..2eb98191322 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -467,6 +467,23 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * ctx->n_nodes_per_cb = (ctx->n_nodes_1 + ctx->n_cb - 1) / ctx->n_cb; + // fabley (speed conference 008/009): env-gated arm-file — when + // LLAMA_METAL_CAPTURE_ARM_FILE is set and the file exists, consume it + // and capture THIS compute. Production without the env pays nothing. + { + static const char * arm_path = NULL; + static bool arm_checked = false; + if (!arm_checked) { + arm_checked = true; + arm_path = getenv("LLAMA_METAL_CAPTURE_ARM_FILE"); + } + if (arm_path && access(arm_path, F_OK) == 0) { + unlink(arm_path); + ctx->capture_compute = 1; + GGML_LOG_WARN("%s: capture armed via %s\n", __func__, arm_path); + } + } + if (ctx->capture_compute >= 0) { ctx->capture_compute--; } diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 75de0f6dd08..d04bbe56fd2 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -521,7 +521,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { int ggml_metal_op_encode(ggml_metal_op_t ctx, int idx) { if (ctx->use_capture) { - ggml_metal_encoder_debug_group_push(ctx->enc, ggml_op_desc(ctx->node(idx))); + // fabley: include the tensor name so traces attribute per graph stage + char dbg[192]; + snprintf(dbg, sizeof(dbg), "%s|%s", ctx->node(idx)->name, ggml_op_desc(ctx->node(idx))); + ggml_metal_encoder_debug_group_push(ctx->enc, dbg); } int res = ggml_metal_op_encode_impl(ctx, idx); From b5fe9105197b5c8d2c0f967e5e79a364d52487b6 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 11:45:09 -0600 Subject: [PATCH 44/62] graph: env-gated packed-MQA decode transform (speed-conf 019 Patch A); tests: fix sweep-globals order, add --heads override --- src/llama-graph.cpp | 43 +++++++++++++++++++++++++++++++++++--- tests/test-llama-archs.cpp | 22 ++++++++++++------- 2 files changed, 55 insertions(+), 10 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8677083f55e..d799b00f27c 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2605,12 +2605,40 @@ ggml_tensor * llm_graph_context::build_attn_mha( cur = ggml_reshape_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]); } else { - ggml_tensor * kq = ggml_mul_mat(ctx0, k, q); + // fabley speed-conf 019 Patch A: packed-MQA decode transform. + // At decode (n_tokens==1) with an MQA cache (one KV head) and many + // query heads, mul_mat broadcasts over ne2 and Metal re-reads the + // whole K cache once per head (measured 6.03 ms/layer @108K). + // Packing heads into ne1 selects mul_mm (r2=1, K tiled once), then + // permutes back so softmax/mask see the original layout. + // Env-gated: LLAMA_PACKED_MQA_DECODE=1. Prefill is never affected. + static const bool fbl_packed_mqa = [] { + const char * e = getenv("LLAMA_PACKED_MQA_DECODE"); + return e && atoi(e) != 0; + }(); + const bool fbl_pack = fbl_packed_mqa && + q->ne[1] == 1 && k->ne[2] == 1 && q->ne[2] > 8 && n_stream == 1; + + ggml_tensor * kq = nullptr; + if (fbl_pack) { + static bool fbl_logged = false; + if (!fbl_logged) { fbl_logged = true; fprintf(stderr, "packed-MQA decode path ACTIVE\n"); } + ggml_tensor * qp = ggml_cont(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3)); // [d, n_head, 1] + cb(qp, "q_packed", il); + kq = ggml_mul_mat(ctx0, k, qp); // [n_kv, n_head, 1] + ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 0, 2, 1, 3)); // [n_kv, 1, n_head] + cb(kq, "kq_packed", il); + } else { + kq = ggml_mul_mat(ctx0, k, q); + } cb(kq, "kq", il); // note: this op tends to require high floating point range // while for some models F16 is enough, for others it is not, so we default to F32 here - ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + if (!fbl_pack) { + ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + } if (arch == LLM_ARCH_GROK) { // need to do the following: @@ -2649,7 +2677,16 @@ ggml_tensor * llm_graph_context::build_attn_mha( cb(v, "v_cont", il); } - ggml_tensor * kqv = ggml_mul_mat(ctx0, v, kq); + ggml_tensor * kqv = nullptr; + if (fbl_pack) { + ggml_tensor * kqp = ggml_cont(ctx0, ggml_permute(ctx0, kq, 0, 2, 1, 3)); // [n_kv, n_head, 1] + cb(kqp, "kq_soft_max_packed", il); + kqv = ggml_mul_mat(ctx0, v, kqp); // [d_v, n_head, 1] + kqv = ggml_cont(ctx0, ggml_permute(ctx0, kqv, 0, 2, 1, 3)); // [d_v, 1, n_head] + cb(kqv, "kqv_packed", il); + } else { + kqv = ggml_mul_mat(ctx0, v, kq); + } cb(kqv, "kqv", il); // for MLA with the absorption optimization, we need to "decompress" from MQA back to MHA diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 3ce8bdc1d03..64333f123dd 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -64,6 +64,15 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } } +// depth-sweep overrides (0 = leave defaults). Set only by test_depth_sweep(). +static uint32_t g_depth_sweep_n_ctx = 0; +static uint32_t g_depth_sweep_n_ub = 0; +static uint32_t g_sweep_iheads = 0; +static uint32_t g_sweep_iklen = 0; +static uint32_t g_sweep_layers = 0; // glm5next fixture layer-count override +static uint32_t g_sweep_dlead = 0; // leading dense layers (real model: 3 of 45) +static uint32_t g_sweep_heads = 0; // glm5next fixture attention-head override + static void usage(char ** argv) { printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-v/--verbose]\n", argv[0]); } @@ -126,6 +135,9 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { if (arch == LLM_ARCH_GLM5NEXT && g_sweep_layers > 0) { n_layer = g_sweep_layers; } + if (arch == LLM_ARCH_GLM5NEXT && g_sweep_heads > 0) { + n_head = g_sweep_heads; + } } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { n_layer = 3; } else if (arch == LLM_ARCH_CHAMELEON) { @@ -377,13 +389,6 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) return true; } -// depth-sweep overrides (0 = leave defaults). Set only by test_depth_sweep(). -static uint32_t g_depth_sweep_n_ctx = 0; -static uint32_t g_depth_sweep_n_ub = 0; -static uint32_t g_sweep_iheads = 0; -static uint32_t g_sweep_iklen = 0; -static uint32_t g_sweep_layers = 0; // glm5next fixture layer-count override -static uint32_t g_sweep_dlead = 0; // leading dense layers (real model: 3 of 45) static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, @@ -955,6 +960,9 @@ int main(int argc, char ** argv) { if (strcmp(argv[i], "--layers") == 0 && i + 1 < argc) { g_sweep_layers = std::stoul(argv[++i]); } + if (strcmp(argv[i], "--heads") == 0 && i + 1 < argc) { + g_sweep_heads = std::stoul(argv[++i]); + } if (strcmp(argv[i], "--dlead") == 0 && i + 1 < argc) { g_sweep_dlead = std::stoul(argv[++i]); } From d33b0ab1d5069e1aa767be2361666c1aed1990e7 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 12:33:23 -0600 Subject: [PATCH 45/62] graph: split packed-MQA modes (kq/kqv/both) + f32 prec on packed KQV, historical cb names (grokk 020); tests: DSA n_head_kv=1 under --heads, FA off in harness (production parity) --- src/llama-graph.cpp | 39 +++++++++++++++++++++++++++----------- tests/test-llama-archs.cpp | 8 +++++++- 2 files changed, 35 insertions(+), 12 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index d799b00f27c..23a4a07ac46 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2612,32 +2612,45 @@ ggml_tensor * llm_graph_context::build_attn_mha( // Packing heads into ne1 selects mul_mm (r2=1, K tiled once), then // permutes back so softmax/mask see the original layout. // Env-gated: LLAMA_PACKED_MQA_DECODE=1. Prefill is never affected. - static const bool fbl_packed_mqa = [] { + // mode: off (default) | kq | kqv | both (legacy "1" == both) + static const int fbl_mode = [] { const char * e = getenv("LLAMA_PACKED_MQA_DECODE"); - return e && atoi(e) != 0; + if (!e) return 0; + if (strcmp(e, "kq") == 0) return 1; + if (strcmp(e, "kqv") == 0) return 2; + if (strcmp(e, "both") == 0 || atoi(e) != 0) return 3; + return 0; }(); - const bool fbl_pack = fbl_packed_mqa && + const bool fbl_gate = fbl_mode != 0 && q->ne[1] == 1 && k->ne[2] == 1 && q->ne[2] > 8 && n_stream == 1; + const bool fbl_pack = fbl_gate && (fbl_mode & 1) != 0; // KQ side + const bool fbl_pack_o = fbl_gate && (fbl_mode & 2) != 0; // KQV side ggml_tensor * kq = nullptr; - if (fbl_pack) { + if (fbl_gate) { static bool fbl_logged = false; - if (!fbl_logged) { fbl_logged = true; fprintf(stderr, "packed-MQA decode path ACTIVE\n"); } + if (!fbl_logged) { + fbl_logged = true; + fprintf(stderr, "packed-MQA decode path ACTIVE (mode=%d)\n", fbl_mode); + } + } + if (fbl_pack) { ggml_tensor * qp = ggml_cont(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3)); // [d, n_head, 1] cb(qp, "q_packed", il); kq = ggml_mul_mat(ctx0, k, qp); // [n_kv, n_head, 1] ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + cb(kq, "kq", il); // grokk 020: historical name on the mul_mm kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 0, 2, 1, 3)); // [n_kv, 1, n_head] - cb(kq, "kq_packed", il); + cb(kq, "kq_unpacked", il); } else { kq = ggml_mul_mat(ctx0, k, q); + cb(kq, "kq", il); } - cb(kq, "kq", il); // note: this op tends to require high floating point range // while for some models F16 is enough, for others it is not, so we default to F32 here if (!fbl_pack) { - ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + ggml_mul_mat_set_prec(kq, GGML_PREC_F32); // packed branch set it on the mul_mm above } if (arch == LLM_ARCH_GROK) { @@ -2678,16 +2691,20 @@ ggml_tensor * llm_graph_context::build_attn_mha( } ggml_tensor * kqv = nullptr; - if (fbl_pack) { + if (fbl_pack_o) { ggml_tensor * kqp = ggml_cont(ctx0, ggml_permute(ctx0, kq, 0, 2, 1, 3)); // [n_kv, n_head, 1] cb(kqp, "kq_soft_max_packed", il); kqv = ggml_mul_mat(ctx0, v, kqp); // [d_v, n_head, 1] + // packing moves this product from mul_mv (f32 accumulation) to mul_mm + // (f16 tile accumulation by default) - request f32 like the KQ side + ggml_mul_mat_set_prec(kqv, GGML_PREC_F32); + cb(kqv, "kqv", il); // historical name on the mul_mm kqv = ggml_cont(ctx0, ggml_permute(ctx0, kqv, 0, 2, 1, 3)); // [d_v, 1, n_head] - cb(kqv, "kqv_packed", il); + cb(kqv, "kqv_unpacked", il); } else { kqv = ggml_mul_mat(ctx0, v, kq); + cb(kqv, "kqv", il); } - cb(kqv, "kqv", il); // for MLA with the absorption optimization, we need to "decompress" from MQA back to MHA if (v_mla) { diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 64333f123dd..cac8dc09559 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -193,11 +193,14 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } 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 + // Under --heads (query-head override), DSA stays MQA: n_head_kv is held at 1 + // (real model: 64 query heads over one latent KV head). KDA layers stay 0. GGML_ASSERT(n_layer >= 2); + const uint32_t glm_dsa_kv = (g_sweep_heads > 0) ? 1u : n_head_kv; 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); + n_head_kv_per_layer.push_back(il == 1 ? 0 : glm_dsa_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); @@ -402,6 +405,9 @@ static std::pair get_model_and_ctx( model_params.split_mode = split_mode; llama_context_params ctx_params = llama_context_default_params(); + // match production glm5next: FA off (AUTO would enable it on the toy and + // silently bypass the non-FA attention branch under test) + ctx_params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; ctx_params.n_ctx = g_depth_sweep_n_ctx; // 0 = model default (original behavior) ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; From 19446fe71262cba43cc1f2a8d80d0cc7a80c0466 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 12:36:44 -0600 Subject: [PATCH 46/62] tests: scope depth-sweep FA-off via arg (other arch tests keep AUTO); graph: NON_FA branch marker (codex 026 dual-marker doctrine) --- src/llama-graph.cpp | 8 ++++++++ tests/test-llama-archs.cpp | 13 ++++++------- 2 files changed, 14 insertions(+), 7 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 23a4a07ac46..839e1a4179b 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2605,6 +2605,14 @@ ggml_tensor * llm_graph_context::build_attn_mha( cur = ggml_reshape_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]); } else { + { // codex 026: complementary branch marker — OFF cells must prove they + // reached the non-FA branch, not merely lack the packed marker + static bool fbl_branch_logged = false; + if (!fbl_branch_logged) { + fbl_branch_logged = true; + fprintf(stderr, "GLM_ATTN_PATH=NON_FA\n"); + } + } // fabley speed-conf 019 Patch A: packed-MQA decode transform. // At decode (n_tokens==1) with an MQA cache (one KV head) and many // query heads, mul_mat broadcasts over ne2 and Metal re-reads the diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index cac8dc09559..8ecc933f7c2 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -395,7 +395,8 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, - const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { + const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, + llama_flash_attn_type fa_type = LLAMA_FLASH_ATTN_TYPE_AUTO) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -405,9 +406,7 @@ static std::pair get_model_and_ctx( model_params.split_mode = split_mode; llama_context_params ctx_params = llama_context_default_params(); - // match production glm5next: FA off (AUTO would enable it on the toy and - // silently bypass the non-FA attention branch under test) - ctx_params.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + ctx_params.flash_attn_type = fa_type; // depth-sweep passes DISABLED (production parity); others keep AUTO ctx_params.n_ctx = g_depth_sweep_n_ctx; // 0 = model default (original behavior) ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; @@ -857,10 +856,10 @@ static int test_depth_sweep(const size_t seed, const uint32_t max_depth, max_depth, ub_a, b_on_cpu ? "cpu" : ggml_backend_dev_description(dev_gpu), ub_b); g_depth_sweep_n_ub = ub_a; - auto mc_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}); + auto mc_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, false, LLAMA_FLASH_ATTN_TYPE_DISABLED); g_depth_sweep_n_ub = ub_b; - auto mc_dev = b_on_cpu ? get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}) - : get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {dev_gpu}); + auto mc_dev = b_on_cpu ? get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, false, LLAMA_FLASH_ATTN_TYPE_DISABLED) + : get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {dev_gpu}, LLAMA_SPLIT_MODE_LAYER, false, LLAMA_FLASH_ATTN_TYPE_DISABLED); const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(mc_cpu.first.get())); const std::vector tokens = get_tokens(max_depth, n_vocab, seed); From 5f3d764f6e58ee3202a883ffdd3214de31e8416e Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 13:01:14 -0600 Subject: [PATCH 47/62] graph: strict packed-MQA mode parse - numeric 2 no longer aliases to both (grokk 028), unrecognized values warn and disable --- src/llama-graph.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 839e1a4179b..2dab35c1da9 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2623,10 +2623,12 @@ ggml_tensor * llm_graph_context::build_attn_mha( // mode: off (default) | kq | kqv | both (legacy "1" == both) static const int fbl_mode = [] { const char * e = getenv("LLAMA_PACKED_MQA_DECODE"); - if (!e) return 0; + if (!e || strcmp(e, "0") == 0 || strcmp(e, "off") == 0) return 0; if (strcmp(e, "kq") == 0) return 1; if (strcmp(e, "kqv") == 0) return 2; - if (strcmp(e, "both") == 0 || atoi(e) != 0) return 3; + if (strcmp(e, "both") == 0 || strcmp(e, "1") == 0) return 3; // "1" = legacy alias + // grokk 028 §3: bare atoi made "2" mean both — fail loud instead + fprintf(stderr, "LLAMA_PACKED_MQA_DECODE=%s unrecognized (want off|kq|kqv|both) -> off\n", e); return 0; }(); const bool fbl_gate = fbl_mode != 0 && From 37a483c6b08666dcf8a94fe975563cec5567d704 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 13:27:38 -0600 Subject: [PATCH 48/62] =?UTF-8?q?kv-cache/graph:=20Patch=20B=20=E2=80=94?= =?UTF-8?q?=20derived=20transposed=20MLA-V=20mirror=20(LLAMA=5FMLA=5FV=5FM?= =?UTF-8?q?IRROR);=20contracts=204.1-4.6+7;=20hybrid-k=20wrapper=20fills?= =?UTF-8?q?=20v=5Fidxs=20(the=20sub-input=20set=5Finput=20is=20bypassed=20?= =?UTF-8?q?there);=20trunk-sparse=20consumer,=20MTP-off=20documented?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- src/llama-graph.cpp | 26 ++++++++- src/llama-graph.h | 2 + src/llama-kv-cache.cpp | 112 ++++++++++++++++++++++++++++++++---- src/llama-kv-cache.h | 14 ++++- src/llama-memory-hybrid.cpp | 7 ++- src/llama-memory-hybrid.h | 3 +- src/llama-model.cpp | 6 +- 7 files changed, 152 insertions(+), 18 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 2dab35c1da9..ff540c3e0ff 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -505,6 +505,10 @@ bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) { void llm_graph_input_attn_k::set_input(const llama_ubatch * ubatch) { mctx->set_input_k_idxs(self_k_idxs, ubatch); + if (self_v_idxs) { + mctx->set_input_v_idxs(self_v_idxs, ubatch); // Patch B mirror + } + mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); } @@ -519,6 +523,9 @@ bool llm_graph_input_attn_k::can_reuse_impl(const llm_graph_params & params) { res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; + // Patch B: expanded mirror idxs scale with n_tokens + res &= !self_v_idxs || (params.ubatch.n_tokens > 0 && self_v_idxs->ne[0] % params.ubatch.n_tokens == 0); + res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams); return res; @@ -1140,6 +1147,10 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) { void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) { mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); + if (inp_attn->self_v_idxs) { + mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); // Patch B mirror + } + mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); const int64_t n_rs = mctx->get_recr()->get_n_rs(); @@ -2937,6 +2948,10 @@ static std::unique_ptr build_attn_inp_k_impl( inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch); + if (mctx_cur->get_has_v_mirror()) { + inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch); // Patch B + } + inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); inp->self_kq_mask_cnv = inp->self_kq_mask; } @@ -3736,6 +3751,11 @@ ggml_tensor * llm_graph_context::build_attn_sparse( const auto & k_idxs = inp->get_k_idxs(); ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + + // Patch B: keep the transposed mirror in sync (same latent rows, O(1)/token) + if (mctx_cur->get_has_v_mirror()) { + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il)); + } } const auto & kq_mask = inp->get_kq_mask(); @@ -3782,7 +3802,11 @@ ggml_tensor * llm_graph_context::build_attn_sparse( 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); + // Patch B: with the mirror, V is the persistent transposed copy — build_attn_mha + // sees v_trans and never builds the per-token cont(transpose(v)) (the 4.83 ms v_cont) + ggml_tensor * v = mctx_cur->get_has_v_mirror() + ? mctx_cur->get_v(ctx0, il) + : 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); diff --git a/src/llama-graph.h b/src/llama-graph.h index 258a66dd53e..dcc2dcaee11 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -382,10 +382,12 @@ class llm_graph_input_attn_k : public llm_graph_input_i { bool can_reuse_impl(const llm_graph_params & params); ggml_tensor * get_k_idxs() const { return self_k_idxs; } + ggml_tensor * get_v_idxs() const { return self_v_idxs; } ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; } ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch] + ggml_tensor * self_v_idxs = nullptr; // I64 expanded [n_batch*w] — Patch B mirror only ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 62f51287b44..9fa879531ec 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -80,7 +80,8 @@ llama_kv_cache::llama_kv_cache( const layer_filter_cb & filter, const layer_reuse_cb & reuse, const layer_share_cb & share, - const char * name_tag) : + const char * name_tag, + bool v_mirror_opt) : model(model), hparams(hparams), v_trans(v_trans), n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type), other(static_cast(mem_other)), @@ -163,6 +164,8 @@ llama_kv_cache::llama_kv_cache( const bool is_mla = hparams.is_mla(); + v_mirror = v_mirror_opt && is_mla && v_trans; // Patch B opt-in, MLA transposed caches only + for (uint32_t il = 0; il < n_layer; il++) { if (!hparams.has_kv(il)) { LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il); @@ -229,10 +232,15 @@ llama_kv_cache::llama_kv_cache( } const bool has_k = true; - const bool has_v = !is_mla; + const bool has_v = !is_mla || v_mirror; // Patch B: mirror allocates V despite MLA + // grokk 024/020: the mirror is a layout of K -> width n_embd_k_gqa (512), type_k + const uint32_t n_embd_v_alloc = v_mirror ? n_embd_k_gqa : n_embd_v_gqa; + if (v_mirror && n_embd_k_gqa > mirror_w_max) { + mirror_w_max = n_embd_k_gqa; + } ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr; - ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr; + ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, v_mirror ? type_k : type_v, n_embd_v_alloc, kv_size, n_stream) : nullptr; has_k && ggml_format_name(k, "cache_%sk_l%d", name_tag, il); has_v && ggml_format_name(v, "cache_%sv_l%d", name_tag, il); @@ -242,7 +250,7 @@ llama_kv_cache::llama_kv_cache( for (uint32_t s = 0; s < n_stream; ++s) { k_stream.push_back(has_k ? ggml_view_2d(ctx, k, n_embd_k_gqa, kv_size, k->nb[1], s*k->nb[2]) : nullptr); - v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr); + v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_alloc, kv_size, v->nb[1], s*v->nb[2]) : nullptr); } map_layer_ids[il] = layers.size(); @@ -851,6 +859,10 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co if (layer.v_stream[ssrc]) { ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]); } + + if (v_mirror) { + mirror_dirty = true; + } } } } @@ -893,9 +905,66 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co } } + if (v_mirror && mirror_dirty) { + LLAMA_LOG_DEBUG("%s: rebuilding derived V mirror from K\n", __func__); + + ggml_backend_sched_reset(sched); + + auto * res = lctx->get_gf_res_reserve(); + + res->reset(); + + auto * gf = build_graph_mirror(res, lctx); + if (!ggml_backend_sched_alloc_graph(sched, gf)) { + LLAMA_LOG_ERROR("%s: failed to allocate compute graph for V-mirror rebuild\n", __func__); + return updated; + } + + res->set_inputs(nullptr); + + if (lctx->graph_compute(gf, false) != GGML_STATUS_SUCCESS) { + LLAMA_LOG_ERROR("%s: failed to compute V-mirror rebuild\n", __func__); + return updated; + } + + mirror_dirty = false; + fprintf(stderr, "MIRROR_REBUILD SUCCESS\n"); + updated = true; + } + return updated; } +ggml_cgraph * llama_kv_cache::build_graph_mirror(llm_graph_result * res, llama_context * lctx) const { + GGML_ASSERT(!other); + GGML_UNUSED(lctx); + + auto * ctx = res->get_ctx(); + auto * gf = res->get_gf(); + + const uint64_t kv_size = get_size(); + + for (const auto & layer : layers) { + if (!layer.v) { + continue; + } + + const uint64_t w = hparams.n_embd_k_gqa(layer.il); + + for (uint32_t st = 0; st < n_stream; ++st) { + ggml_tensor * ksrc = ggml_view_2d(ctx, layer.k, w, kv_size, layer.k->nb[1], st*layer.k->nb[2]); + // mirror layout: row e holds element e of every cell (transposed K) + ggml_tensor * vdst = ggml_view_2d(ctx, layer.v, kv_size, w, + ggml_row_size(layer.v->type, kv_size), + st*ggml_row_size(layer.v->type, kv_size*w)); + + ggml_build_forward_expand(gf, ggml_cpy(ctx, ggml_transpose(ctx, ksrc), vdst)); + } + } + + return gf; +} + llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, bool cont) const { if (debug > 0) { @@ -1319,6 +1388,17 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k ggml_row_size(v->type, n_embd_v_gqa*kv_size)*sinfo.s0); } + if (v_mirror) { + // contracts 4.4/7: mirror view is 512-wide (K slab), one KV head + const uint64_t w = hparams.n_embd_k_gqa(il); + return ggml_view_4d(ctx, v, + n_kv, 1, w, ns, + ggml_row_size(v->type, kv_size*w), // head stride > element stride => v_trans in mha + ggml_row_size(v->type, kv_size), + ggml_row_size(v->type, kv_size*w), + ggml_row_size(v->type, kv_size*w)*sinfo.s0); + } + // note: v->nb[1] > v->nb[2] return ggml_view_4d(ctx, v, n_kv, hparams.n_head_kv(il), hparams.n_embd_head_v(il), ns, @@ -1461,7 +1541,7 @@ ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama if (!v_trans) { v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens); } else { - v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens*hparams.n_embd_v_gqa_max()); + v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens*(v_mirror ? mirror_w_max : hparams.n_embd_v_gqa_max())); } ggml_set_input(v_idxs); @@ -1545,7 +1625,7 @@ void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ub // note: the V cache is transposed when not using flash attention const int64_t kv_size = get_size(); - const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa_max(); + const int64_t n_embd_v_gqa = v_mirror ? mirror_w_max : hparams.n_embd_v_gqa_max(); for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { const int64_t offs = sinfo.strm[s]*kv_size*n_embd_v_gqa; @@ -2351,7 +2431,7 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); auto * v = layer.v_stream[cr.strm]; - if (!v) { + if (!v || v_mirror) { // contract 4.5: derived mirror never serialized continue; } @@ -2380,7 +2460,7 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); auto * v = layer.v_stream[cr.strm]; - if (!v) { + if (!v || v_mirror) { // contract 4.5: derived mirror never serialized continue; } @@ -2590,6 +2670,10 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } + if (v_mirror) { + mirror_dirty = true; // contract 4.6: restored K needs a mirror rebuild before decode + } + // For each layer, read the keys for each cell, one row is one cell, read as one contiguous block for (const auto & layer : layers) { const uint32_t il = layer.il; @@ -2637,7 +2721,7 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); auto * v = layer.v_stream[strm]; - if (!v) { + if (!v || v_mirror) { // contract 4.5: mirror not on the wire continue; } @@ -2680,7 +2764,7 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); auto * v = layer.v_stream[strm]; - if (!v) { + if (!v || v_mirror) { // contract 4.5: mirror not on the wire continue; } @@ -2762,8 +2846,8 @@ llama_kv_cache_context::llama_kv_cache_context( llama_context * lctx, bool do_shift, stream_copy_info sc_info) : status(LLAMA_MEMORY_STATUS_SUCCESS), kv(kv), lctx(lctx), do_shift(do_shift), sc_info(std::move(sc_info)) { - if (!do_shift && this->sc_info.empty()) { - status = LLAMA_MEMORY_STATUS_NO_UPDATE; + if (!do_shift && this->sc_info.empty() && !kv->get_mirror_dirty()) { + status = LLAMA_MEMORY_STATUS_NO_UPDATE; // contract 4.6: a dirty mirror must vote for update } } @@ -2843,6 +2927,10 @@ ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il) cons return kv->get_k(ctx, il, n_kv, sinfos[i_cur]); } +bool llama_kv_cache_context::get_has_v_mirror() const { + return kv->get_has_v_mirror(); +} + ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) const { return kv->get_v(ctx, il, n_kv, sinfos[i_cur]); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 2ba032a34bf..9c39a75fcfd 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -114,7 +114,8 @@ class llama_kv_cache : public llama_memory_i { const layer_reuse_cb & reuse, const layer_share_cb & share, // a model can hold more than one cache, so the tensor names have to stay unique - const char * name_tag = ""); + const char * name_tag = "", + bool v_mirror_opt = false); ~llama_kv_cache() = default; @@ -199,6 +200,11 @@ class llama_kv_cache : public llama_memory_i { ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; + // speed-conf Patch B: derived transposed MLA-V mirror (022/024 contracts) + bool get_has_v_mirror() const { return v_mirror; } + bool get_mirror_dirty() const { return mirror_dirty; } + uint32_t mirror_width_max() const { return mirror_w_max; } + // store k_cur and v_cur in the cache based on the provided head location 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; @@ -277,6 +283,9 @@ class llama_kv_cache : public llama_memory_i { }; bool v_trans = true; // the value tensor is transposed + bool v_mirror = false; // derived transposed MLA-V mirror (never serialized; contract 4.5) + bool mirror_dirty = false; // set on state restore; rebuilt from K in update() (contract 4.6) + uint32_t mirror_w_max = 0; // max n_embd_k_gqa over mirrored layers (contract 7 idx sizing) // see set_kpool_dirty. mutable because the only consumer runs from set_input, which // holds the cache by const pointer; nothing else observes it. @@ -350,6 +359,8 @@ class llama_kv_cache : public llama_memory_i { float freq_scale, uint32_t il) const; + ggml_cgraph * build_graph_mirror(llm_graph_result * res, llama_context * lctx) const; + ggml_cgraph * build_graph_shift( llm_graph_result * res, llama_context * lctx) const; @@ -428,6 +439,7 @@ class llama_kv_cache_context : public llama_memory_context_i { // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il) const; + bool get_has_v_mirror() const; // store k_cur and v_cur in the cache based on the provided head location // note: the heads in k_cur and v_cur should be laid out contiguously in memory diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index d0372becf68..27de7f0d1c3 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -33,7 +33,8 @@ llama_memory_hybrid::llama_memory_hybrid( const layer_filter_cb & filter_attn, const layer_filter_cb & filter_recr, const layer_filter_cb & filter_idx, - ggml_type type_idx) : + ggml_type type_idx, + bool attn_v_mirror) : hparams(model.hparams), hparams_idx(model.hparams), mem_attn(new llama_kv_cache( @@ -54,7 +55,9 @@ llama_memory_hybrid::llama_memory_hybrid( [&](int32_t il) { return !hparams.is_recr(il); } : filter_attn, nullptr, - nullptr + nullptr, + "", + attn_v_mirror // Patch B: mem_attn only (contract 4.1); idx cache stays false )), mem_recr(new llama_memory_recurrent( model, diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 9cf79b8e871..c9036023259 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -42,7 +42,8 @@ class llama_memory_hybrid : public llama_memory_i { const layer_filter_cb & filter_recr = nullptr, /* optional indexer key cache; absent unless filter_idx */ const layer_filter_cb & filter_idx = nullptr, - ggml_type type_idx = GGML_TYPE_F16); + ggml_type type_idx = GGML_TYPE_F16, + bool attn_v_mirror = false); ~llama_memory_hybrid() = default; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d3dbc257c97..4c50d38c0ca 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2581,7 +2581,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* filter_attn */ std::move(filter_attn), /* filter_recr */ std::move(filter_recr), /* filter_idx */ std::move(filter_idx), - /* type_idx */ type_idx); + /* type_idx */ type_idx, + /* attn_v_mirror */ [] { + const char * e = getenv("LLAMA_MLA_V_MIRROR"); + return e && atoi(e) != 0; + }()); } } else { llama_kv_cache::layer_filter_cb filter = nullptr; From c4cc25e0bfd07403793268759431fe145cecbadf Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 13:36:55 -0600 Subject: [PATCH 49/62] graph/kv-cache: Patch B hardening (grokk 033 2.1/2.3/2.4 + codex 034 2-4): dense build_attn mirror wiring, wrapper can_reuse width check, strict V parser + MLA_V_MIRROR ACTIVE marker, construction guards, stream-copy keeps mirror clean --- src/llama-graph.cpp | 15 ++++++++++++++- src/llama-kv-cache.cpp | 24 ++++++++++++++++++++---- src/llama-kv-cache.h | 1 + src/llama-model.cpp | 5 ++++- 4 files changed, 39 insertions(+), 6 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index ff540c3e0ff..fc351d144f9 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1185,6 +1185,12 @@ bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) { res &= inp_rs->head == mctx->get_recr()->get_head(); res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z(); + if (inp_attn->self_v_idxs && + (int64_t) inp_attn->self_v_idxs->ne[0] != + (int64_t) params.ubatch.n_tokens * mctx->get_attn()->mirror_width_max()) { + return false; // Patch B (grokk 033 2.3) + } + return res; } @@ -2994,13 +3000,20 @@ ggml_tensor * llm_graph_context::build_attn( const auto & k_idxs = inp->get_k_idxs(); ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + + // Patch B (grokk 033 2.1): dense consumer maintains the mirror too + if (mctx_cur->get_has_v_mirror()) { + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il)); + } } const auto & kq_mask = inp->get_kq_mask(); 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 * v = mctx_cur->get_has_v_mirror() + ? mctx_cur->get_v(ctx0, il) + : 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, kq_mask, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il); diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 9fa879531ec..2c2bf908003 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -165,6 +165,12 @@ llama_kv_cache::llama_kv_cache( const bool is_mla = hparams.is_mla(); v_mirror = v_mirror_opt && is_mla && v_trans; // Patch B opt-in, MLA transposed caches only + if (v_mirror_opt && !v_mirror) { + fprintf(stderr, "MLA_V_MIRROR requested but DISABLED (is_mla=%d v_trans=%d)\n", (int) is_mla, (int) v_trans); + } + if (v_mirror) { + GGML_ASSERT((type_k == GGML_TYPE_F16 || type_k == GGML_TYPE_F32) && "V mirror v0 requires F16/F32 type_k"); + } for (uint32_t il = 0; il < n_layer; il++) { if (!hparams.has_kv(il)) { @@ -258,6 +264,15 @@ llama_kv_cache::llama_kv_cache( layers.push_back({ il, k, v, k_stream, v_stream, }); } + if (v_mirror) { + for (const auto & layer : layers) { + if (layer.v) { + GGML_ASSERT(hparams.n_embd_k_gqa(layer.il) == mirror_w_max && "V mirror v0 requires uniform mirrored K width"); + } + } + fprintf(stderr, "MLA_V_MIRROR ACTIVE width=%u type=%s\n", mirror_w_max, ggml_type_name(type_k)); + } + if (reuse) { LLAMA_LOG_DEBUG("%s: reusing layers:\n", __func__); @@ -857,12 +872,9 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co ggml_backend_tensor_copy(layer.k_stream[ssrc], layer.k_stream[sdst]); if (layer.v_stream[ssrc]) { + // mirror copied with the stream -> stays clean (codex 034 s4) ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]); } - - if (v_mirror) { - mirror_dirty = true; - } } } } @@ -2931,6 +2943,10 @@ bool llama_kv_cache_context::get_has_v_mirror() const { return kv->get_has_v_mirror(); } +uint32_t llama_kv_cache_context::mirror_width_max() const { + return kv->mirror_width_max(); +} + ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) const { return kv->get_v(ctx, il, n_kv, sinfos[i_cur]); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 9c39a75fcfd..e3826d7657f 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -440,6 +440,7 @@ class llama_kv_cache_context : public llama_memory_context_i { ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il) const; bool get_has_v_mirror() const; + uint32_t mirror_width_max() const; // store k_cur and v_cur in the cache based on the provided head location // note: the heads in k_cur and v_cur should be laid out contiguously in memory diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4c50d38c0ca..42a78b79baa 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2584,7 +2584,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* type_idx */ type_idx, /* attn_v_mirror */ [] { const char * e = getenv("LLAMA_MLA_V_MIRROR"); - return e && atoi(e) != 0; + if (!e || strcmp(e, "0") == 0 || strcmp(e, "off") == 0) return false; + if (strcmp(e, "1") == 0 || strcmp(e, "on") == 0) return true; + fprintf(stderr, "LLAMA_MLA_V_MIRROR=%s unrecognized (want off|0|on|1) -> off\n", e); + return false; }()); } } else { From f86a0da7c38cf2ec6da99f2c9a0b91f30647213e Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 13:45:01 -0600 Subject: [PATCH 50/62] graph: exact v_idxs width check in attn_k sub-input; name mirror update v_mirror_upd in both consumers (grokk 036) --- src/llama-graph.cpp | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index fc351d144f9..c2fdd7ace20 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -523,8 +523,10 @@ bool llm_graph_input_attn_k::can_reuse_impl(const llm_graph_params & params) { res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; - // Patch B: expanded mirror idxs scale with n_tokens - res &= !self_v_idxs || (params.ubatch.n_tokens > 0 && self_v_idxs->ne[0] % params.ubatch.n_tokens == 0); + // Patch B (grokk 036): exact width check, mirroring the hybrid-k wrapper + res &= !self_v_idxs || + (int64_t) self_v_idxs->ne[0] == + (int64_t) params.ubatch.n_tokens * mctx->mirror_width_max(); res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams); @@ -3003,7 +3005,9 @@ ggml_tensor * llm_graph_context::build_attn( // Patch B (grokk 033 2.1): dense consumer maintains the mirror too if (mctx_cur->get_has_v_mirror()) { - ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il)); + ggml_tensor * vupd = mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il); + cb(vupd, "v_mirror_upd", il); + ggml_build_forward_expand(gf, vupd); } } @@ -3767,7 +3771,9 @@ ggml_tensor * llm_graph_context::build_attn_sparse( // Patch B: keep the transposed mirror in sync (same latent rows, O(1)/token) if (mctx_cur->get_has_v_mirror()) { - ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il)); + ggml_tensor * vupd = mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il); + cb(vupd, "v_mirror_upd", il); + ggml_build_forward_expand(gf, vupd); } } From 0519ed89e9eb7dff20f787ba717cd6c1c421d851 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 15:36:15 -0600 Subject: [PATCH 51/62] kv-cache: scope V-mirror opt-in to nope-only GLM5Next (codex 051 s3) The env is parsed in the generic hybrid construction path; is_mla && v_trans alone does not prove K carries no rope-key half. Fail closed on any other arch and print arch/n_rot in the DISABLED diagnostic. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- src/llama-kv-cache.cpp | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 2c2bf908003..ef0c862ef6f 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -164,9 +164,16 @@ llama_kv_cache::llama_kv_cache( const bool is_mla = hparams.is_mla(); - v_mirror = v_mirror_opt && is_mla && v_trans; // Patch B opt-in, MLA transposed caches only + // Patch B opt-in. Scope guard (codex 051 s3): mirroring the FULL K slab as V + // is only proven for nope-only GLM5Next, where K carries no rope-key half. + // is_mla alone does not establish that contract; fail closed elsewhere so a + // globally inherited env cannot enable the mirror on an unproven MLA arch. + const bool mirror_supported = model.arch == LLM_ARCH_GLM5NEXT && hparams.n_rot(0) == 0; + + v_mirror = v_mirror_opt && mirror_supported && is_mla && v_trans; if (v_mirror_opt && !v_mirror) { - fprintf(stderr, "MLA_V_MIRROR requested but DISABLED (is_mla=%d v_trans=%d)\n", (int) is_mla, (int) v_trans); + fprintf(stderr, "MLA_V_MIRROR requested but DISABLED (arch=%d n_rot=%u is_mla=%d v_trans=%d)\n", + (int) model.arch, hparams.n_rot(0), (int) is_mla, (int) v_trans); } if (v_mirror) { GGML_ASSERT((type_k == GGML_TYPE_F16 || type_k == GGML_TYPE_F32) && "V mirror v0 requires F16/F32 type_k"); From dd80b059284f7682b80836314dd0a047d0cf391b Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 18:15:22 -0600 Subject: [PATCH 52/62] tests: single-token decode step in depth-sweep so packed decode branches are exercised (grokk 054 s3) The chunked sweep never sends n_tokens==1, so fbl_gate never opened and packed-path comparisons were vacuous. Decode one token after the sweep and compare full logits CPU-vs-device; the ACTIVE marker grep now has a real decode behind it. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- tests/test-llama-archs.cpp | 28 ++++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 8ecc933f7c2..70aade36bbf 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -903,6 +903,34 @@ static int test_depth_sweep(const size_t seed, const uint32_t max_depth, diverged = true; } } + // grokk 054 s3: the chunked sweep never takes the n_tokens==1 decode gate, + // so packed/view decode branches were invisible to it. One explicit + // single-token decode exercises them; grep the ACTIVE marker after this. + if (!diverged) { + common_batch_clear(batch); + common_batch_add(batch, tokens[0], max_depth, {0}, true); + if (llama_decode(mc_cpu.second.get(), batch) || llama_decode(mc_dev.second.get(), batch)) { + printf("depth-sweep-1tok: decode failed\n"); + llama_batch_free(batch); + return 1; + } + const float * lc = llama_get_logits_ith(mc_cpu.second.get(), 0); + const float * ld = llama_get_logits_ith(mc_dev.second.get(), 0); + double se = 0.0, ref = 0.0; + uint32_t amax_c = 0, amax_d = 0; + for (uint32_t j = 0; j < n_vocab; j++) { + const double d = (double) lc[j] - (double) ld[j]; + se += d * d; + ref += (double) lc[j] * (double) lc[j]; + if (lc[j] > lc[amax_c]) amax_c = j; + if (ld[j] > ld[amax_d]) amax_d = j; + } + const double nmse_val = ref > 0.0 ? se / ref : se; + const bool bad = nmse_val > 1e-3 || amax_c != amax_d; + printf("depth-sweep-1tok: depth=%u nmse=%.3e argmax_cpu=%u argmax_dev=%u%s\n", + max_depth + 1, nmse_val, amax_c, amax_d, bad ? " <-- DIVERGED" : ""); + if (bad) diverged = true; + } llama_batch_free(batch); printf("depth-sweep: %s\n", diverged ? "DIVERGENCE FOUND" : "no divergence up to max depth"); return diverged ? 2 : 0; From 92177ae919b9d6f7d238a30ab45c86384d60d1a6 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 18:53:24 -0600 Subject: [PATCH 53/62] tests: --dump-1tok writes the decode-1 device+cpu logits for cross-run oracles Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- tests/test-llama-archs.cpp | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 70aade36bbf..8d48bec27c8 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -67,6 +67,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { // depth-sweep overrides (0 = leave defaults). Set only by test_depth_sweep(). static uint32_t g_depth_sweep_n_ctx = 0; static uint32_t g_depth_sweep_n_ub = 0; +static std::string g_dump_1tok; static uint32_t g_sweep_iheads = 0; static uint32_t g_sweep_iklen = 0; static uint32_t g_sweep_layers = 0; // glm5next fixture layer-count override @@ -930,6 +931,16 @@ static int test_depth_sweep(const size_t seed, const uint32_t max_depth, printf("depth-sweep-1tok: depth=%u nmse=%.3e argmax_cpu=%u argmax_dev=%u%s\n", max_depth + 1, nmse_val, amax_c, amax_d, bad ? " <-- DIVERGED" : ""); if (bad) diverged = true; + if (!g_dump_1tok.empty()) { + FILE * fh = fopen(g_dump_1tok.c_str(), "wb"); + if (fh) { + fwrite(ld, sizeof(float), n_vocab, fh); // device logits + fwrite(lc, sizeof(float), n_vocab, fh); // cpu logits + fclose(fh); + printf("depth-sweep-1tok: dumped %u+%u logits to %s\n", + n_vocab, n_vocab, g_dump_1tok.c_str()); + } + } } llama_batch_free(batch); printf("depth-sweep: %s\n", diverged ? "DIVERGENCE FOUND" : "no divergence up to max depth"); @@ -944,6 +955,7 @@ int main(int argc, char ** argv) { llm_arch arch = LLM_ARCH_UNKNOWN; size_t seed = rd(); uint32_t depth_sweep = 0; + std::string dump_1tok; uint32_t sweep_ctx = 131072; uint32_t sweep_ub = 0; uint32_t sweep_ub2 = 0; @@ -969,6 +981,9 @@ int main(int argc, char ** argv) { if (strcmp(argv[i], "--depth-sweep") == 0 && i + 1 < argc) { depth_sweep = std::stoul(argv[++i]); } + if (strcmp(argv[i], "--dump-1tok") == 0 && i + 1 < argc) { + dump_1tok = argv[++i]; + } if (strcmp(argv[i], "--ctx") == 0 && i + 1 < argc) { sweep_ctx = std::stoul(argv[++i]); } @@ -1023,6 +1038,7 @@ int main(int argc, char ** argv) { printf("%s: using seed %zu\n", __func__, seed); try { + g_dump_1tok = dump_1tok; if (depth_sweep > 0) { return test_depth_sweep(seed, depth_sweep, sweep_ctx, sweep_ub, sweep_ub2, sweep_b_cpu, sweep_topk); From fbc35427795805ca3a8fe45bee69428e7b5bd6be Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 18:53:24 -0600 Subject: [PATCH 54/62] =?UTF-8?q?glm5next:=20D0=20gathered=20DSA=20decode?= =?UTF-8?q?=20(speed=20conference=20056/057)=20=E2=80=94=20LLAMA=5FGLM5=5F?= =?UTF-8?q?GATHERED=5FDSA?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Execute the selection instead of masking it: concat the indexer's physical top-k cells with the always-selected tail, gather the <=2051 latent K rows (dim-1 view, F32 return cast back, reshape to n_head_kv=1 for packed-KQ composition), gather the cand+causal cell mask and add a per-SLOT validity mask (expanded pool_bias of the selected pools; a gather, unlike the dense scatter, is not idempotent - padded picks name cell 0 and must die by slot). V is the same compact tensor; the V-mirror is still updated, not read. Decode-1, single-stream, strict env, GATHERED_DSA ACTIVE marker. Toy oracle: dense-vs-gathered Metal logits NMSE 5.9e-15, top-5 identical; duplicate-cell-0 adversarial (select_k > valid pools) matches dense. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- src/llama-graph.cpp | 127 ++++++++++++++++++++++++++++++++++- src/llama-graph.h | 29 +++++++- src/llama-kv-cache-kpool.cpp | 25 +++++++ src/llama-kv-cache-kpool.h | 9 +++ src/models/glm5next.cpp | 59 ++++++++++++++-- src/models/models.h | 7 +- 6 files changed, 247 insertions(+), 9 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index c2fdd7ace20..857910afd57 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3652,7 +3652,8 @@ llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k() const llm_graph_input_kpool * llm_graph_context::build_inp_kpool( const llama_memory_hybrid_context * mctx_cur, ggml_tensor * kq_mask, - bool scoring) const { + bool scoring, + bool gathered) const { const auto * mctx_attn = mctx_cur->get_attn(); const auto * mctx_idx = mctx_cur->get_idx(); @@ -3711,6 +3712,19 @@ 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"); + // D0 gathered decode inputs. ONLY when the gathered graph will consume them: + // an input tensor no node reads is never allocated, so the host fill would + // write through data == nullptr. + if (gathered) { + inp->tail_cells = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, kpool - 1, n_tps, n_stream); + ggml_set_input(inp->tail_cells); + ggml_set_name(inp->tail_cells, "kpool_tail_cells"); + + inp->tail_valid = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, kpool - 1, n_tps, n_stream); + ggml_set_input(inp->tail_valid); + ggml_set_name(inp->tail_valid, "kpool_tail_valid"); + } + // 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 @@ -3841,6 +3855,117 @@ ggml_tensor * llm_graph_context::build_attn_sparse( return cur; } + +ggml_tensor * llm_graph_context::build_attn_sparse_gathered( + 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 * slot_valid, + ggml_tensor * tail_cells, + ggml_tensor * tail_valid, + ggml_tensor * cand_mask, + float kq_scale, + int il) const { + // these nodes are added in the same order as build_attn / build_attn_sparse + 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, INCLUDING the V-mirror update: a gathered decode does not read + // the mirror, but the next dense graph (prefill, or env-off) will + { + const auto & k_idxs = inp->get_k_idxs(); + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + + if (mctx_cur->get_has_v_mirror()) { + ggml_tensor * vupd = mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il); + cb(vupd, "v_mirror_upd", il); + ggml_build_forward_expand(gf, vupd); + } + } + + const auto & kq_mask = inp->get_kq_mask(); + + // decode-1, single stream only (the gate in the model guarantees this) + GGML_ASSERT(top_k->ne[1] == 1 && top_k->ne[2] == 1); + GGML_ASSERT(kq_mask->ne[1] == 1 && kq_mask->ne[3] == 1); + GGML_ASSERT(slot_valid->ne[0] == top_k->ne[0]); + GGML_ASSERT(ggml_are_same_shape(tail_cells, tail_valid)); + + const int64_t n_kv = kq_mask->ne[0]; + const int64_t n_sel = top_k->ne[0] + tail_cells->ne[0]; + + static bool logged_active = false; + if (!logged_active) { + fprintf(stderr, "GATHERED_DSA ACTIVE n_sel_max=%lld n_kv=%lld\n", + (long long) n_sel, (long long) n_kv); + logged_active = true; + } + + // ids: selected cells then the always-selected tail. all entries are in [0, n_kv); + // invalid slots name cell 0 and die to the -inf slot mask below, BY SLOT. + ggml_tensor * ids = ggml_concat(ctx0, + ggml_reshape_1d(ctx0, top_k, top_k->ne[0]), + ggml_reshape_1d(ctx0, tail_cells, tail_cells->ne[0]), 0); + cb(ids, "gathered_ids", il); + + // per-CELL additive mask, gathered: candidate-set rejection + causal/empty/foreign. + // never gather the causal mask alone (grokk 057 s3): padded picks name cell 0 and + // cell 0 can be a real visible token; only cand_mask keeps those -inf. + ggml_tensor * cell_mask = ggml_add(ctx0, cand_mask, kq_mask); // F16 + ggml_tensor * mask_rows = ggml_reshape_3d(ctx0, cell_mask, 1, n_kv, 1); + ggml_tensor * mask_g = ggml_get_rows(ctx0, mask_rows, ids); // -> F32 [1, n_sel, 1] + mask_g = ggml_reshape_2d(ctx0, mask_g, n_sel, 1); + + // per-SLOT validity: expanded pool_bias of the selected pools, then the tail's + ggml_tensor * slot_mask = ggml_concat(ctx0, + ggml_reshape_1d(ctx0, slot_valid, slot_valid->ne[0]), + ggml_reshape_1d(ctx0, tail_valid, tail_valid->ne[0]), 0); + slot_mask = ggml_reshape_2d(ctx0, slot_mask, n_sel, 1); + + mask_g = ggml_add(ctx0, mask_g, slot_mask); + mask_g = ggml_reshape_4d(ctx0, mask_g, n_sel, 1, 1, 1); + cb(mask_g, "gathered_mask", il); + + // K rows: [512, 1, n_kv, 1] -> dim-1 view -> gather -> cast -> [512, 1, n_sel, 1]. + // the final reshape is load-bearing (grokk 057 s3): build_attn_mha permutes (0,2,1,3) + // and packed KQ keys on post-permute k->ne[2]==1 (n_head_kv). + ggml_tensor * kfull = mctx_cur->get_k(ctx0, il); + ggml_tensor * krows = ggml_view_3d(ctx0, kfull, + kfull->ne[0], kfull->ne[2], kfull->ne[3], + kfull->nb[2], kfull->nb[3], 0); + + ggml_tensor * kg = ggml_get_rows(ctx0, krows, ids); // F32 [512, n_sel, 1] + kg = ggml_cast(ctx0, kg, kfull->type); + kg = ggml_reshape_4d(ctx0, kg, kfull->ne[0], 1, n_sel, 1); + cb(kg, "gathered_k", il); + + // V is the same compact latent; v_trans false -> compact v_cont inside mha + ggml_tensor * cur = build_attn_mha(q_cur, kg, kg, kq_b, mask_g, 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 dcc2dcaee11..59c6d170a37 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1353,7 +1353,8 @@ struct llm_graph_context { llm_graph_input_kpool * build_inp_kpool( const llama_memory_hybrid_context * mctx_cur, ggml_tensor * kq_mask, - bool scoring) const; + bool scoring, + bool gathered = false) const; // build_attn, but masking with `top_k` over `sel_mask`; `cand_mask` drops over-budget picks ggml_tensor * build_attn_sparse( @@ -1373,6 +1374,32 @@ struct llm_graph_context { float kq_scale, int il) const; + // D0 gathered DSA (codex 056 / grokk 057): same contract as build_attn_sparse, but + // attention runs over the <= n_select+kpool-1 gathered rows instead of a dense masked + // n_kv. decode-1 / n_stream==1 only. slot_valid carries the expanded pool_bias of the + // selected pools (-INFINITY kills padded/duplicate picks BY SLOT; a gather, unlike the + // dense scatter, is not idempotent). K rows are gathered from the latent cache + // (F32 return, cast back), and V is the same compact tensor: v_trans is false, the + // per-token v_cont is ~2 MiB, and the V-mirror is not read (it IS still updated). + ggml_tensor * build_attn_sparse_gathered( + 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, 1, 1] physical cells + ggml_tensor * slot_valid, // F32 [n_select, 1, 1] 0/-inf per SLOT + ggml_tensor * tail_cells, // I32 [kpool-1, 1, 1] + ggml_tensor * tail_valid, // F32 [kpool-1, 1, 1] + ggml_tensor * cand_mask, // F16 [n_kv, 1, 1, 1] + float kq_scale, + int il) const; + // // pooling // diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index 485260586b8..c707ddd2b68 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -71,6 +71,8 @@ void llama_kv_cache_set_input_kpool( ggml_tensor * pool_bias, ggml_tensor * sel_mask, ggml_tensor * cand_mask, + ggml_tensor * tail_cells, + ggml_tensor * tail_valid, ggml_tensor * pool_reps, ggml_tensor * new_pool_cells, ggml_tensor * new_pool_reps, @@ -449,6 +451,28 @@ void llama_kv_cache_set_input_kpool( } } + // D0: emit the tail cells [tail_start, q] by slot; unused slots stay + // cell 0 / -INFINITY from the pre-fill below the loop + if (tail_cells) { + const int64_t r_t = r - 1; + int32_t * cur_tcell = (int32_t *) tail_cells->data + (s*n_tps + ii)*r_t; + float * cur_tvalid = (float *) tail_valid->data + (s*n_tps + ii)*r_t; + + for (int64_t t = 0; t < r_t; ++t) { + cur_tcell [t] = 0; + cur_tvalid[t] = -INFINITY; + } + for (int64_t j = 0; j < n_kv; ++j) { + const llama_pos p = pos_at[j]; + if (p >= tail_start && p <= q) { + const int64_t t = p - tail_start; + GGML_ASSERT(t < r_t); + cur_tcell [t] = (int32_t) j; + cur_tvalid[t] = 0.0f; + } + } + } + // 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]; @@ -512,6 +536,7 @@ void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { mctx_attn->get_kv(), /* cell_pool */ nullptr, pool_cells, /* bias */ nullptr, pool_bias, sel_mask, cand_mask, + tail_cells, tail_valid, 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, diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index ff20a6e6667..93e74bc4f47 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -45,6 +45,8 @@ void llama_kv_cache_set_input_kpool( ggml_tensor * pool_bias, ggml_tensor * sel_mask, ggml_tensor * cand_mask, + ggml_tensor * tail_cells, // optional, D0 gathered decode + ggml_tensor * tail_valid, // optional, with tail_cells ggml_tensor * pool_reps, ggml_tensor * new_pool_cells, ggml_tensor * new_pool_reps, @@ -101,6 +103,13 @@ class llm_graph_input_kpool : public llm_graph_input_i { 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] + // D0 gathered DSA (codex 056 / grokk 057): the always-selected incomplete-pool tail as + // explicit cells so a gathered graph can concat them to top_k. slot t holds the cell at + // pos tail_start+t; unused slots are cell 0 with tail_valid -INFINITY (get_rows has no + // sentinel). tail length is (q+1)%kpool in [0, kpool-1]. + ggml_tensor * tail_cells = nullptr; // I32 [kpool-1, n_tps, n_stream] + ggml_tensor * tail_valid = nullptr; // F32 [kpool-1, n_tps, 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 4c6db1eb32f..1ec0be1d683 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -330,13 +330,27 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( // * 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. +// D0 (speed conference 056-058): gathered sparse decode, strict opt-in +static bool glm5_gathered_dsa_enabled() { + static const bool on = []() { + const char * e = getenv("LLAMA_GLM5_GATHERED_DSA"); + if (e == nullptr) return false; + if (strcmp(e, "on") == 0 || strcmp(e, "1") == 0) return true; + if (strcmp(e, "off") == 0 || strcmp(e, "0") == 0) return false; + fprintf(stderr, "LLAMA_GLM5_GATHERED_DSA=%s unrecognized (want off|0|on|1) -> off\n", e); + return false; + }(); + return on; +} + 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 { + int il, + ggml_tensor ** out_slot_valid) 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; @@ -479,6 +493,22 @@ ggml_tensor * llama_model_glm5next::graph::build_indexer( 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); + // D0: per-slot validity of the selection. gather pool_bias at the selected pool + // ordinals and expand each pool to its kpool member slots. decode-1 only: pool_bias + // is per-token, and this path is gated to n_tps==1. + if (out_slot_valid != nullptr) { + GGML_ASSERT(n_tps == 1 && n_stream == 1 && "gathered DSA is decode-1/one-stream only"); + + ggml_tensor * pb = ggml_reshape_3d(ctx0, inp_kp->pool_bias, 1, n_pools, n_stream); + ggml_tensor * sb = ggml_get_rows(ctx0, pb, sel_flat); // F32 [1, select_k, 1] + + ggml_tensor * rep = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, r, select_k, n_stream); + sb = ggml_repeat(ctx0, sb, rep); // member-major, matches top_k + + *out_slot_valid = ggml_reshape_3d(ctx0, sb, r*select_k, n_tps, n_stream); + cb(*out_slot_valid, "indexer_slot_valid", il); + } + 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); @@ -508,7 +538,14 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); cb(qr, "dsa_q_a_norm", il); - ggml_tensor * top_k = inp_kp ? build_indexer(layer, inp_kp, cur, qr, scoring, il) : nullptr; + // D0 gate: strict env + decode-1 + one stream (mirrors fbl_gate's shape) + const bool gathered = glm5_gathered_dsa_enabled() && scoring && inp_kp != nullptr && + n_tokens == 1 && inp_kp->sel_mask != nullptr && inp_kp->sel_mask->ne[3] == 1 && + inp_kp->tail_cells != nullptr; + + ggml_tensor * slot_valid = nullptr; + ggml_tensor * top_k = inp_kp ? build_indexer(layer, inp_kp, cur, qr, scoring, il, + gathered ? &slot_valid : nullptr) : 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); @@ -529,7 +566,17 @@ 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); - if (top_k) { + if (top_k && gathered) { + GGML_ASSERT(slot_valid != nullptr && inp_kp->tail_cells != nullptr); + // sel_mask is not consumed by the gathered graph, but set_input still fills it + // (one shared host fill for all paths); expand the leaf so it stays allocated + ggml_build_forward_expand(gf, inp_kp->sel_mask); + cur = build_attn_sparse_gathered(inp_attn, + layer.wo, nullptr, nullptr, + q, k, k, nullptr, nullptr, layer.wv_b, + top_k, slot_valid, inp_kp->tail_cells, inp_kp->tail_valid, + inp_kp->cand_mask, kq_scale, il); + } else if (top_k) { cur = build_attn_sparse(inp_attn, layer.wo, nullptr, nullptr, q, k, k, nullptr, nullptr, layer.wv_b, @@ -617,7 +664,8 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa 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); + inp_mem->get_attn()->get_kq_mask(), indexer_scoring, + glm5_gathered_dsa_enabled() && n_tokens == 1); } } @@ -785,7 +833,8 @@ llama_model_glm5next::graph_mtp::graph_mtp(const llama_model & model, const llm_ 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); + inp_mem->get_attn()->get_kq_mask(), indexer_scoring, + glm5_gathered_dsa_enabled() && n_tokens == 1); } } diff --git a/src/models/models.h b/src/models/models.h index 699e58f246d..77e0f50ce1e 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1376,14 +1376,17 @@ struct llama_model_glm5next : public llama_model_base { ggml_tensor * cur, int il) const; - // always stores the key and gate; when `scoring`, returns the selected CELL indices + // always stores the key and gate; when `scoring`, returns the selected CELL indices. + // out_slot_valid (optional, D0): expanded pool_bias of the selected pools, + // F32 [kpool*select_k, n_tps, n_stream], 0/-inf per slot (decode-1 only) 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; + int il, + ggml_tensor ** out_slot_valid = nullptr) const; ggml_tensor * build_layer_ffn( const llama_model & model, From 4cbe31e30e8a7e70d616f2f833e95c6114cf8113 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 19:57:24 -0600 Subject: [PATCH 55/62] glm5next/tests: D0 P0 batch (codex 059 / grokk 060) - compact mask path: gather cand and kq at ids separately, add 2051-wide; the dense cand+kq add ran per DSA layer = eleven O(n_kv) ops per token - one gathered_decode boolean owns tail-input creation, slot_valid request and builder choice (the 058 crash class was producer/consumer gate disagreement); graph_mtp passes false until MTP has its own oracle - name the compute owners gathered_k_rows / gathered_k_cast / gathered_{cand,kq}_rows for the per-node cut; MQA-only assert on the view - fixture: four successive decode-1 steps (rolling tail 1->2->3->0, logits compared and dumped per step) and --skip-1tok Suite: dense-vs-gathered NMSE ~1e-15 all four steps, adversarial duplicate-cell-0 rolling green, h9/h64, env-off regression clean. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- src/llama-graph.cpp | 27 +++++++++---- src/models/glm5next.cpp | 19 ++++++--- tests/test-llama-archs.cpp | 83 ++++++++++++++++++++++---------------- 3 files changed, 81 insertions(+), 48 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 857910afd57..718128d35ee 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3920,13 +3920,19 @@ ggml_tensor * llm_graph_context::build_attn_sparse_gathered( ggml_reshape_1d(ctx0, tail_cells, tail_cells->ne[0]), 0); cb(ids, "gathered_ids", il); - // per-CELL additive mask, gathered: candidate-set rejection + causal/empty/foreign. + // per-CELL additive masks, gathered SEPARATELY then added compact (codex 059 P0: + // the dense cand+kq add ran once per DSA layer = eleven O(n_kv) ops per token). // never gather the causal mask alone (grokk 057 s3): padded picks name cell 0 and // cell 0 can be a real visible token; only cand_mask keeps those -inf. - ggml_tensor * cell_mask = ggml_add(ctx0, cand_mask, kq_mask); // F16 - ggml_tensor * mask_rows = ggml_reshape_3d(ctx0, cell_mask, 1, n_kv, 1); - ggml_tensor * mask_g = ggml_get_rows(ctx0, mask_rows, ids); // -> F32 [1, n_sel, 1] - mask_g = ggml_reshape_2d(ctx0, mask_g, n_sel, 1); + auto gather_mask = [&](ggml_tensor * m, const char * name) { + ggml_tensor * rows = ggml_reshape_3d(ctx0, m, 1, n_kv, 1); + ggml_tensor * out = ggml_get_rows(ctx0, rows, ids); // -> F32 [1, n_sel, 1] + cb(out, name, il); + return ggml_reshape_2d(ctx0, out, n_sel, 1); + }; + ggml_tensor * cand_g = gather_mask(cand_mask, "gathered_cand_rows"); + ggml_tensor * kq_g = gather_mask(kq_mask, "gathered_kq_rows"); + ggml_tensor * mask_g = ggml_add(ctx0, cand_g, kq_g); // per-SLOT validity: expanded pool_bias of the selected pools, then the tail's ggml_tensor * slot_mask = ggml_concat(ctx0, @@ -3942,13 +3948,18 @@ ggml_tensor * llm_graph_context::build_attn_sparse_gathered( // the final reshape is load-bearing (grokk 057 s3): build_attn_mha permutes (0,2,1,3) // and packed KQ keys on post-permute k->ne[2]==1 (n_head_kv). ggml_tensor * kfull = mctx_cur->get_k(ctx0, il); + GGML_ASSERT(kfull->ne[1] == 1 && "gathered DSA is MQA-only (n_head_kv == 1)"); ggml_tensor * krows = ggml_view_3d(ctx0, kfull, kfull->ne[0], kfull->ne[2], kfull->ne[3], kfull->nb[2], kfull->nb[3], 0); - ggml_tensor * kg = ggml_get_rows(ctx0, krows, ids); // F32 [512, n_sel, 1] - kg = ggml_cast(ctx0, kg, kfull->type); - kg = ggml_reshape_4d(ctx0, kg, kfull->ne[0], 1, n_sel, 1); + // name the actual compute owners (059 P0): the gather and the cast are what the + // per-node cut must price; the final reshape is a free view + ggml_tensor * kg32 = ggml_get_rows(ctx0, krows, ids); // F32 [512, n_sel, 1] + cb(kg32, "gathered_k_rows", il); + ggml_tensor * kg16 = ggml_cast(ctx0, kg32, kfull->type); + cb(kg16, "gathered_k_cast", il); + ggml_tensor * kg = ggml_reshape_4d(ctx0, kg16, kfull->ne[0], 1, n_sel, 1); cb(kg, "gathered_k", il); // V is the same compact latent; v_trans false -> compact v_cont inside mha diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 1ec0be1d683..b76bae6a72c 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -538,10 +538,9 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); cb(qr, "dsa_q_a_norm", il); - // D0 gate: strict env + decode-1 + one stream (mirrors fbl_gate's shape) - const bool gathered = glm5_gathered_dsa_enabled() && scoring && inp_kp != nullptr && - n_tokens == 1 && inp_kp->sel_mask != nullptr && inp_kp->sel_mask->ne[3] == 1 && - inp_kp->tail_cells != nullptr; + // D0 gate: the tail inputs exist iff build_inp_kpool's single gathered_decode + // boolean was true — resource presence IS the policy (codex 059 P0) + const bool gathered = inp_kp != nullptr && inp_kp->tail_cells != nullptr; ggml_tensor * slot_valid = nullptr; ggml_tensor * top_k = inp_kp ? build_indexer(layer, inp_kp, cur, qr, scoring, il, @@ -663,9 +662,16 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa if (mctx_hyb->get_idx() != nullptr) { indexer_scoring = cparams.n_ctx > glm5next_n_select(hparams); + // ONE owner for the D0 decision (codex 059 / grokk 060): the same boolean + // creates the tail inputs, requests slot_valid, and picks the builder. + // Producer/consumer gate disagreement was the 058 crash class. + const bool gathered_decode = glm5_gathered_dsa_enabled() && + indexer_scoring && n_tokens == 1 && + (cparams.kv_unified ? 1 : (int64_t) ubatch.n_seqs_unq) == 1; + inp_kp = build_inp_kpool(mctx_hyb, inp_mem->get_attn()->get_kq_mask(), indexer_scoring, - glm5_gathered_dsa_enabled() && n_tokens == 1); + gathered_decode); } } @@ -832,9 +838,10 @@ llama_model_glm5next::graph_mtp::graph_mtp(const llama_model & model, const llm_ if (mctx_hyb->get_idx() != nullptr) { indexer_scoring = cparams.n_ctx > glm5next_n_select(hparams); + // D0 stays OFF the NextN graph until MTP has its own oracle (grokk 060 s2) inp_kp = build_inp_kpool(mctx_hyb, inp_mem->get_attn()->get_kq_mask(), indexer_scoring, - glm5_gathered_dsa_enabled() && n_tokens == 1); + /* gathered */ false); } } diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 8d48bec27c8..c1c19ac3484 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -68,6 +68,7 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { static uint32_t g_depth_sweep_n_ctx = 0; static uint32_t g_depth_sweep_n_ub = 0; static std::string g_dump_1tok; +static bool g_skip_1tok = false; static uint32_t g_sweep_iheads = 0; static uint32_t g_sweep_iklen = 0; static uint32_t g_sweep_layers = 0; // glm5next fixture layer-count override @@ -904,43 +905,52 @@ static int test_depth_sweep(const size_t seed, const uint32_t max_depth, diverged = true; } } - // grokk 054 s3: the chunked sweep never takes the n_tokens==1 decode gate, - // so packed/view decode branches were invisible to it. One explicit - // single-token decode exercises them; grep the ACTIVE marker after this. - if (!diverged) { - common_batch_clear(batch); - common_batch_add(batch, tokens[0], max_depth, {0}, true); - if (llama_decode(mc_cpu.second.get(), batch) || llama_decode(mc_dev.second.get(), batch)) { - printf("depth-sweep-1tok: decode failed\n"); - llama_batch_free(batch); - return 1; - } - const float * lc = llama_get_logits_ith(mc_cpu.second.get(), 0); - const float * ld = llama_get_logits_ith(mc_dev.second.get(), 0); - double se = 0.0, ref = 0.0; - uint32_t amax_c = 0, amax_d = 0; - for (uint32_t j = 0; j < n_vocab; j++) { - const double d = (double) lc[j] - (double) ld[j]; - se += d * d; - ref += (double) lc[j] * (double) lc[j]; - if (lc[j] > lc[amax_c]) amax_c = j; - if (ld[j] > ld[amax_d]) amax_d = j; - } - const double nmse_val = ref > 0.0 ? se / ref : se; - const bool bad = nmse_val > 1e-3 || amax_c != amax_d; - printf("depth-sweep-1tok: depth=%u nmse=%.3e argmax_cpu=%u argmax_dev=%u%s\n", - max_depth + 1, nmse_val, amax_c, amax_d, bad ? " <-- DIVERGED" : ""); - if (bad) diverged = true; + // grokk 054 s3 / 060 s2: the chunked sweep never takes the n_tokens==1 gate. + // FOUR successive single-token decodes in one process exercise the packed/ + // gathered branches AND the stateful rolling tail (kpool tail length walks + // an ordered transition; starting right after a chunk boundary the expected + // tail lengths are (pos+1)%kpool = 1, 2, 3, 0 for kpool=4). + if (!diverged && !g_skip_1tok) { + FILE * dump = nullptr; if (!g_dump_1tok.empty()) { - FILE * fh = fopen(g_dump_1tok.c_str(), "wb"); - if (fh) { - fwrite(ld, sizeof(float), n_vocab, fh); // device logits - fwrite(lc, sizeof(float), n_vocab, fh); // cpu logits - fclose(fh); - printf("depth-sweep-1tok: dumped %u+%u logits to %s\n", - n_vocab, n_vocab, g_dump_1tok.c_str()); + dump = fopen(g_dump_1tok.c_str(), "wb"); + } + for (uint32_t step = 0; step < 4 && !diverged; step++) { + common_batch_clear(batch); + common_batch_add(batch, tokens[step], max_depth + step, {0}, true); + if (llama_decode(mc_cpu.second.get(), batch) || llama_decode(mc_dev.second.get(), batch)) { + printf("depth-sweep-1tok: decode failed at step %u\n", step); + if (dump) fclose(dump); + llama_batch_free(batch); + return 1; + } + const float * lc = llama_get_logits_ith(mc_cpu.second.get(), 0); + const float * ld = llama_get_logits_ith(mc_dev.second.get(), 0); + double se = 0.0, ref = 0.0; + uint32_t amax_c = 0, amax_d = 0; + for (uint32_t j = 0; j < n_vocab; j++) { + const double d = (double) lc[j] - (double) ld[j]; + se += d * d; + ref += (double) lc[j] * (double) lc[j]; + if (lc[j] > lc[amax_c]) amax_c = j; + if (ld[j] > ld[amax_d]) amax_d = j; + } + const double nmse_val = ref > 0.0 ? se / ref : se; + const bool bad = nmse_val > 1e-3 || amax_c != amax_d; + printf("depth-sweep-1tok: step=%u depth=%u tail_len=%u nmse=%.3e " + "argmax_cpu=%u argmax_dev=%u%s\n", + step, max_depth + step + 1, (max_depth + step + 1) % 4, + nmse_val, amax_c, amax_d, bad ? " <-- DIVERGED" : ""); + if (bad) diverged = true; + if (dump) { + fwrite(ld, sizeof(float), n_vocab, dump); + fwrite(lc, sizeof(float), n_vocab, dump); } } + if (dump) { + fclose(dump); + printf("depth-sweep-1tok: dumped 4 steps to %s\n", g_dump_1tok.c_str()); + } } llama_batch_free(batch); printf("depth-sweep: %s\n", diverged ? "DIVERGENCE FOUND" : "no divergence up to max depth"); @@ -956,6 +966,7 @@ int main(int argc, char ** argv) { size_t seed = rd(); uint32_t depth_sweep = 0; std::string dump_1tok; + bool skip_1tok = false; uint32_t sweep_ctx = 131072; uint32_t sweep_ub = 0; uint32_t sweep_ub2 = 0; @@ -984,6 +995,9 @@ int main(int argc, char ** argv) { if (strcmp(argv[i], "--dump-1tok") == 0 && i + 1 < argc) { dump_1tok = argv[++i]; } + if (strcmp(argv[i], "--skip-1tok") == 0) { + skip_1tok = true; + } if (strcmp(argv[i], "--ctx") == 0 && i + 1 < argc) { sweep_ctx = std::stoul(argv[++i]); } @@ -1039,6 +1053,7 @@ int main(int argc, char ** argv) { try { g_dump_1tok = dump_1tok; + g_skip_1tok = skip_1tok; if (depth_sweep > 0) { return test_depth_sweep(seed, depth_sweep, sweep_ctx, sweep_ub, sweep_ub2, sweep_b_cpu, sweep_topk); From 6056d0574958f60868460c957a1de9f94db52ac0 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 20:31:48 -0600 Subject: [PATCH 56/62] =?UTF-8?q?glm5next:=20PAD32=20=E2=80=94=20pad=20gat?= =?UTF-8?q?hered=20slots=20to=20the=20mul=5Fmm=20tile=20(codex=20063=20/?= =?UTF-8?q?=20grokk=20064)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit bc_inp = src0->ne[0] % 32 (ggml-metal-device.cpp:768): compact packed KQV's reduction dim is n_sel, and 2,051 is both %32!=0 and odd — the slow guarded path. Pad the tail-slot allocation to GGML_PAD(n_top + kpool-1, 32) with dead slots (cell 0 / -inf), the representation the duplicate-cell-0 adversarial proved exact. Marker renamed GATHERED_DSA GRAPH (it prints at construction, not execution) and reports n_sel_pad / n_finite_max. Toy: PAD(11,32)=32 with 21 dead slots live in every gather; rolling four tails dense-vs-gathered at e-15 (step 3 bit-identical), duplicate-cell-0 rolling MATCH x4, h9 and env-off clean. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- src/llama-graph.cpp | 28 +++++++++++++++++++++------- src/llama-graph.h | 5 +++-- src/llama-kv-cache-kpool.cpp | 15 +++++++++------ src/models/glm5next.cpp | 2 +- 4 files changed, 34 insertions(+), 16 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 718128d35ee..597957e3d78 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3715,12 +3715,21 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( // D0 gathered decode inputs. ONLY when the gathered graph will consume them: // an input tensor no node reads is never allocated, so the host fill would // write through data == nullptr. + // PAD32 (codex 063 / grokk 064): packed mul_mm's fast path needs the compact + // reduction dim % 32 == 0 (bc_inp law); pad the tail slots so + // n_top + n_tail_slots lands on the tile. Dead slots are cell 0 / -inf — + // the exact representation the duplicate-cell-0 adversarial proved. if (gathered) { - inp->tail_cells = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, kpool - 1, n_tps, n_stream); + const int64_t n_top = (int64_t) kpool * + llama_kpool_select_k((uint32_t) n_pools, hparams.indexer_top_k, kpool); + const int64_t n_sel_pad = GGML_PAD(n_top + kpool - 1, 32); + const int64_t n_tail = n_sel_pad - n_top; + + inp->tail_cells = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, n_tail, n_tps, n_stream); ggml_set_input(inp->tail_cells); ggml_set_name(inp->tail_cells, "kpool_tail_cells"); - inp->tail_valid = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, kpool - 1, n_tps, n_stream); + inp->tail_valid = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_tail, n_tps, n_stream); ggml_set_input(inp->tail_valid); ggml_set_name(inp->tail_valid, "kpool_tail_valid"); } @@ -3872,6 +3881,7 @@ ggml_tensor * llm_graph_context::build_attn_sparse_gathered( ggml_tensor * tail_cells, ggml_tensor * tail_valid, ggml_tensor * cand_mask, + uint32_t kpool, float kq_scale, int il) const { // these nodes are added in the same order as build_attn / build_attn_sparse @@ -3906,11 +3916,15 @@ ggml_tensor * llm_graph_context::build_attn_sparse_gathered( const int64_t n_kv = kq_mask->ne[0]; const int64_t n_sel = top_k->ne[0] + tail_cells->ne[0]; - static bool logged_active = false; - if (!logged_active) { - fprintf(stderr, "GATHERED_DSA ACTIVE n_sel_max=%lld n_kv=%lld\n", - (long long) n_sel, (long long) n_kv); - logged_active = true; + // printed at graph CONSTRUCTION (codex 063 s6) — reservation builds decode-1 + // graphs too, so this proves "gathered graph built", never "evaluated". + // Execution proof is a cb hit on gathered_k_rows/kqv_out plus the oracle. + static bool logged_graph = false; + if (!logged_graph) { + fprintf(stderr, "GATHERED_DSA GRAPH n_sel_pad=%lld n_finite_max=%lld n_kv=%lld\n", + (long long) n_sel, (long long) (top_k->ne[0] + kpool - 1), + (long long) n_kv); + logged_graph = true; } // ids: selected cells then the always-selected tail. all entries are in [0, n_kv); diff --git a/src/llama-graph.h b/src/llama-graph.h index 59c6d170a37..7951bc3f4e4 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1394,9 +1394,10 @@ struct llm_graph_context { ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] ggml_tensor * top_k, // I32 [n_select, 1, 1] physical cells ggml_tensor * slot_valid, // F32 [n_select, 1, 1] 0/-inf per SLOT - ggml_tensor * tail_cells, // I32 [kpool-1, 1, 1] - ggml_tensor * tail_valid, // F32 [kpool-1, 1, 1] + ggml_tensor * tail_cells, // I32 [n_tail_slots, 1, 1] (real tail + PAD32 dead slots) + ggml_tensor * tail_valid, // F32 [n_tail_slots, 1, 1] ggml_tensor * cand_mask, // F16 [n_kv, 1, 1, 1] + uint32_t kpool, float kq_scale, int il) const; diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index c707ddd2b68..6a972271162 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -454,11 +454,14 @@ void llama_kv_cache_set_input_kpool( // D0: emit the tail cells [tail_start, q] by slot; unused slots stay // cell 0 / -INFINITY from the pre-fill below the loop if (tail_cells) { - const int64_t r_t = r - 1; - int32_t * cur_tcell = (int32_t *) tail_cells->data + (s*n_tps + ii)*r_t; - float * cur_tvalid = (float *) tail_valid->data + (s*n_tps + ii)*r_t; - - for (int64_t t = 0; t < r_t; ++t) { + // width includes the PAD32 dead slots; the real tail occupies + // only [0, kpool-1). Invariant: every slot t >= kpool-1 stays + // cell 0 / -INFINITY from the pre-fill (grokk 064 s2). + const int64_t n_ts = tail_cells->ne[0]; + int32_t * cur_tcell = (int32_t *) tail_cells->data + (s*n_tps + ii)*n_ts; + float * cur_tvalid = (float *) tail_valid->data + (s*n_tps + ii)*n_ts; + + for (int64_t t = 0; t < n_ts; ++t) { cur_tcell [t] = 0; cur_tvalid[t] = -INFINITY; } @@ -466,7 +469,7 @@ void llama_kv_cache_set_input_kpool( const llama_pos p = pos_at[j]; if (p >= tail_start && p <= q) { const int64_t t = p - tail_start; - GGML_ASSERT(t < r_t); + GGML_ASSERT(t < r - 1); cur_tcell [t] = (int32_t) j; cur_tvalid[t] = 0.0f; } diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index b76bae6a72c..68b3373731f 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -574,7 +574,7 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( layer.wo, nullptr, nullptr, q, k, k, nullptr, nullptr, layer.wv_b, top_k, slot_valid, inp_kp->tail_cells, inp_kp->tail_valid, - inp_kp->cand_mask, kq_scale, il); + inp_kp->cand_mask, inp_kp->kpool, kq_scale, il); } else if (top_k) { cur = build_attn_sparse(inp_attn, layer.wo, nullptr, nullptr, From c4ec35f02e9c58ff2915432fb6fe4f1f78298bfa Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 18:10:27 -0600 Subject: [PATCH 57/62] =?UTF-8?q?graph:=20Patch=20C=20view-only=20?= =?UTF-8?q?=E2=80=94=20drop=20redundant=20conts=20on=20packed-path=20singl?= =?UTF-8?q?eton=20permutes=20(codex=20043/051,=20grokk=20045)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit All four packed-decode transforms swap a singleton axis with the head axis; the permuted view already satisfies ggml_is_contiguous (singleton strides are ignored), and ggml_cont always emitted a real Metal copy. Keep the logical permutes as zero-copy views and assert contiguity at graph build. cb names stay on the surviving tensors (051 s5). No !is_permuted assert: the kq unpack view is contiguous AND permuted by design (045 s1). Probe: kq packed_view 1.186 ms vs packed 1.203 vs packed_perm 2.167 at n_kv=108544 - the 0.96 ms copy is gone. Fixture: fixed-seed byte-identical outputs at heads 64/16/9/4 and default vs pre-C build. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk (cherry picked from commit d76dcb0b0b27ba5329b33fb81bccccbb0ae9417a) --- src/llama-graph.cpp | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 597957e3d78..3ccdcfcd371 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -2664,12 +2664,16 @@ ggml_tensor * llm_graph_context::build_attn_mha( } } if (fbl_pack) { - ggml_tensor * qp = ggml_cont(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3)); // [d, n_head, 1] + // Patch C (view-only): singleton-axis permutes are already contiguous; + // ggml_cont here was a real Metal copy. Keep the logical shape as a view. + ggml_tensor * qp = ggml_permute(ctx0, q, 0, 2, 1, 3); // [d, n_head, 1] view + GGML_ASSERT(ggml_is_contiguous(qp)); cb(qp, "q_packed", il); kq = ggml_mul_mat(ctx0, k, qp); // [n_kv, n_head, 1] ggml_mul_mat_set_prec(kq, GGML_PREC_F32); cb(kq, "kq", il); // grokk 020: historical name on the mul_mm - kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 0, 2, 1, 3)); // [n_kv, 1, n_head] + kq = ggml_permute(ctx0, kq, 0, 2, 1, 3); // [n_kv, 1, n_head] view (contiguous AND permuted) + GGML_ASSERT(ggml_is_contiguous(kq)); cb(kq, "kq_unpacked", il); } else { kq = ggml_mul_mat(ctx0, k, q); @@ -2721,14 +2725,16 @@ ggml_tensor * llm_graph_context::build_attn_mha( ggml_tensor * kqv = nullptr; if (fbl_pack_o) { - ggml_tensor * kqp = ggml_cont(ctx0, ggml_permute(ctx0, kq, 0, 2, 1, 3)); // [n_kv, n_head, 1] + ggml_tensor * kqp = ggml_permute(ctx0, kq, 0, 2, 1, 3); // [n_kv, n_head, 1] view + GGML_ASSERT(ggml_is_contiguous(kqp)); cb(kqp, "kq_soft_max_packed", il); kqv = ggml_mul_mat(ctx0, v, kqp); // [d_v, n_head, 1] // packing moves this product from mul_mv (f32 accumulation) to mul_mm // (f16 tile accumulation by default) - request f32 like the KQ side ggml_mul_mat_set_prec(kqv, GGML_PREC_F32); cb(kqv, "kqv", il); // historical name on the mul_mm - kqv = ggml_cont(ctx0, ggml_permute(ctx0, kqv, 0, 2, 1, 3)); // [d_v, 1, n_head] + kqv = ggml_permute(ctx0, kqv, 0, 2, 1, 3); // [d_v, 1, n_head] view + GGML_ASSERT(ggml_is_contiguous(kqv)); cb(kqv, "kqv_unpacked", il); } else { kqv = ggml_mul_mat(ctx0, v, kq); From 048799f098a19e52ec11a4dec93dbf9fead1233d Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 21:21:42 -0600 Subject: [PATCH 58/62] =?UTF-8?q?graph/context:=20R0=20=E2=80=94=20kpool?= =?UTF-8?q?=20can=5Freuse=20contract=20+=20split=20host=20timers=20(codex?= =?UTF-8?q?=20068/072,=20grokk=20070/073)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit llm_graph_input_kpool had no can_reuse override, so the inherited false vetoed reuse of every served GLM graph (every token re-ran build_graph + sched_alloc). The override rebinds both memory contexts and compares recomputed shapes and policy: k_idxs width, idx/attn n_kv (stored as build-time metadata so the contract survives E0c's optional masks), n_pools/mask/pool tensor shapes, rebuild==get_kpool_dirty(), n_new_max, and the recomputed PAD32 tail width (never a cached 2080). LLAMA_HOST_TIMERS=: decode-1 graph_rebuild_us and set_inputs_us as SEPARATE p50/p95 buckets (R0 shrinks only the first; E0c/E1 the second), printed every N samples with the reuse counter. Reuse debug log names the input class via typeid. Fixture: reused=48 with rebuild p50=0 in steady decode (G1 and dense); pad-256 boundary crossing = exactly one rebuild; logits bit-identical to LLAMA_GRAPH_REUSE_DISABLE=1 in all cases. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- src/llama-context.cpp | 34 +++++++++++++++++-- src/llama-context.h | 4 +++ src/llama-graph.cpp | 9 ++++- src/llama-kv-cache-kpool.cpp | 66 ++++++++++++++++++++++++++++++++++++ src/llama-kv-cache-kpool.h | 18 ++++++++-- 5 files changed, 126 insertions(+), 5 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 5f87fb5b901..ce44a552e17 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1,5 +1,7 @@ #include "llama-context.h" +#include + #include "ggml.h" #include "ggml-metal.h" #include "llama-arch.h" @@ -1348,6 +1350,12 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll // in order to correctly reuse a graph, it's full topology has to be uniquely determined by these parameters const auto gparams = graph_params(res, ubatch, mctx, gtype); + // R0 instrumentation (codex 072 s3): graph_rebuild_us and set_inputs_us are + // SEPARATE owners — R0 shrinks only the first, E0c/E1 only the second. + static const bool fbl_host_timers = getenv("LLAMA_HOST_TIMERS") != nullptr; + const int64_t fbl_t_rebuild0 = fbl_host_timers ? ggml_time_us() : 0; + bool fbl_reused = false; + if (!graph_reuse_disable && res->can_reuse(gparams)) { //LLAMA_LOG_DEBUG("%s: reusing previous graph\n", __func__); @@ -1359,6 +1367,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } n_reused++; + fbl_reused = true; } else { res->reset(); @@ -1451,12 +1460,33 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll // set the input data for the input tensors { - //const auto t_start_us = ggml_time_us(); + const int64_t fbl_t_set0 = fbl_host_timers ? ggml_time_us() : 0; // FIXME this call causes a crash if any model inputs were not used in the graph and were therefore not allocated res->set_inputs(&ubatch); - //LLAMA_LOG_INFO("graph set inputs time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0); + if (fbl_host_timers && ubatch.n_tokens == 1) { + // decode-1 only: prefill chunks would swamp the percentiles + const int64_t now = ggml_time_us(); + fbl_host_rebuild_us.push_back(fbl_reused ? 0 : (fbl_t_set0 - fbl_t_rebuild0)); + fbl_host_setinp_us.push_back(now - fbl_t_set0); + static const int64_t fbl_every = [](){ + const char * e = getenv("LLAMA_HOST_TIMERS"); + const long long v = e ? atoll(e) : 0; + return v > 0 ? v : 64; + }(); + if ((int64_t) fbl_host_rebuild_us.size() % fbl_every == 0) { + auto pct = [](std::vector v, double p) { + std::sort(v.begin(), v.end()); + return v.empty() ? (int64_t) 0 : v[(size_t) (p*(v.size() - 1))]; + }; + fprintf(stderr, "HOST_TIMERS n=%zu rebuild_us p50=%lld p95=%lld | set_inputs_us p50=%lld p95=%lld | reused=%d\n", + fbl_host_rebuild_us.size(), + (long long) pct(fbl_host_rebuild_us, 0.5), (long long) pct(fbl_host_rebuild_us, 0.95), + (long long) pct(fbl_host_setinp_us, 0.5), (long long) pct(fbl_host_setinp_us, 0.95), + n_reused); + } + } } const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b56..9b3f4a2073d 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -391,4 +391,8 @@ struct llama_context { mutable int32_t n_eval = 0; // number of eval calls mutable int32_t n_reused = 0; // number of times the previous graph was reused + + // R0 host-interval buckets (env LLAMA_HOST_TIMERS, decode-1 samples only) + std::vector fbl_host_rebuild_us; + std::vector fbl_host_setinp_us; }; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 3ccdcfcd371..7033ec60416 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1,5 +1,7 @@ #include "llama-graph.h" +#include + #include "llama-impl.h" #include "llama-model.h" #include "llama-batch.h" @@ -1439,7 +1441,7 @@ bool llm_graph_result::can_reuse(const llm_graph_params & params) { const bool cur = input->can_reuse(params); if (debug > 1) { - LLAMA_LOG_DEBUG("%s: can_reuse = %d\n", "placeholder", cur); + LLAMA_LOG_DEBUG("%s: can_reuse = %d\n", typeid(*input).name(), cur); } res = res && cur; @@ -3674,6 +3676,11 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( ggml_set_input(inp->k_idxs); ggml_set_name(inp->k_idxs, "kpool_k_idxs"); + // R0: build-time identity for can_reuse + inp->b_n_kv_idx = mctx_idx->get_n_kv(); + inp->b_n_kv_attn = mctx_attn->get_n_kv(); + inp->hparams_indexer_top_k = hparams.indexer_top_k; + if (scoring) { const int64_t n_kv = mctx_attn->get_n_kv(); diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index 6a972271162..e9a9ef3dd7b 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -2,6 +2,7 @@ #include "llama-batch.h" #include "llama-kv-cache.h" +#include "llama-memory-hybrid.h" #include "llama-kv-cells.h" #include @@ -513,6 +514,71 @@ void llama_kv_cache_set_input_kpool( } } + +bool llm_graph_input_kpool::can_reuse(const llm_graph_params & params) { + const auto * mctx_hyb = static_cast(params.mctx); + + const auto * attn = mctx_hyb->get_attn(); + const auto * idx = mctx_hyb->get_idx(); + + if (attn == nullptr || idx == nullptr) { + return false; + } + + // rebind so a reused graph's set_input reads the new contexts + mctx_attn = attn; + mctx_idx = idx; + + bool res = true; + + res &= k_idxs->ne[0] == (int64_t) params.ubatch.n_tokens; + res &= b_n_kv_idx == (int64_t) idx->get_n_kv(); + + if (pool_cells == nullptr) { + // non-scoring graph: only the key/gate store exists + return res; + } + + const int64_t n_kv = attn->get_n_kv(); + const int64_t n_stream = params.cparams.kv_unified ? 1 : params.ubatch.n_seqs_unq; + + if (n_stream <= 0 || params.ubatch.n_tokens % n_stream != 0 || + (int64_t) params.ubatch.n_seqs_unq % n_stream != 0) { + return false; + } + + const int64_t n_tps = params.ubatch.n_tokens/n_stream; + const int64_t n_ps = (int64_t) params.ubatch.n_seqs_unq/n_stream; + + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, kpool, (uint32_t) n_ps); + + res &= b_n_kv_attn == n_kv; + res &= sel_mask == nullptr || (sel_mask->ne[0] == n_kv && sel_mask->ne[1] == n_tps && + sel_mask->ne[3] == n_stream); + res &= pool_cells->ne[0] == (int64_t) kpool*n_pools && pool_cells->ne[1] == n_stream; + res &= pool_bias->ne[0] == n_pools && pool_bias->ne[1] == n_tps && pool_bias->ne[2] == n_stream; + + // policy: rebuild mode and the fixed pooled-key emission bound must match + // what a fresh build would choose (codex 068 s3) + const bool want_rebuild = mctx_attn->get_kv()->get_kpool_dirty(); + const int64_t want_new_max = want_rebuild ? n_pools : n_tps/(int64_t) kpool + n_ps; + + res &= rebuild == want_rebuild; + res &= (int64_t) n_new_max == want_new_max; + + if (tail_cells != nullptr) { + // gathered graph: recompute the PAD32 width — never cache 2080 (grokk 070 s3) + const int64_t n_top = (int64_t) kpool * + llama_kpool_select_k((uint32_t) n_pools, hparams_indexer_top_k, kpool); + const int64_t n_sel_pad = GGML_PAD(n_top + kpool - 1, 32); + + res &= tail_cells->ne[0] == n_sel_pad - n_top; + res &= n_tps == 1 && n_stream == 1; + } + + return res; +} + 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 diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index 93e74bc4f47..d2a1cd69370 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -76,6 +76,12 @@ class llm_graph_input_kpool : public llm_graph_input_i { void set_input(const llama_ubatch * ubatch) override; + // R0 (codex 068 / grokk 070): exact reuse contract. Rebinds the two memory + // contexts and compares every value that determines tensor shape or graph + // policy against the new params. Without this override the inherited + // can_reuse() == false vetoed reuse of EVERY served GLM graph. + bool can_reuse(const llm_graph_params & params) override; + 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] @@ -107,8 +113,16 @@ class llm_graph_input_kpool : public llm_graph_input_i { // explicit cells so a gathered graph can concat them to top_k. slot t holds the cell at // pos tail_start+t; unused slots are cell 0 with tail_valid -INFINITY (get_rows has no // sentinel). tail length is (q+1)%kpool in [0, kpool-1]. - ggml_tensor * tail_cells = nullptr; // I32 [kpool-1, n_tps, n_stream] - ggml_tensor * tail_valid = nullptr; // F32 [kpool-1, n_tps, n_stream] + // PAD32: physical width is n_sel_pad - n_top (>= kpool-1); only the finite + // tail prefix [0, kpool-1) is ever written with live cells. + ggml_tensor * tail_cells = nullptr; // I32 [n_tail_slots, n_tps, n_stream] + ggml_tensor * tail_valid = nullptr; // F32 [n_tail_slots, n_tps, n_stream] + + // build-time identity for can_reuse (explicit metadata, not inferred from + // mask tensors, so the contract survives E0c making the masks optional) + int64_t b_n_kv_attn = -1; + int64_t b_n_kv_idx = -1; + uint32_t hparams_indexer_top_k = 0; const llama_kv_cache_context * mctx_attn; const llama_kv_cache_context * mctx_idx; From 1f9064b251102c3d75aa7d6845f8103d0826882d Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 21:26:07 -0600 Subject: [PATCH 59/62] =?UTF-8?q?graph/context:=20E0a=20=E2=80=94=20live?= =?UTF-8?q?=20certificate-subsumption=20invariant,=20host-side=20class=20c?= =?UTF-8?q?ompare?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Under LLAMA_GLM5_E0A the gathered builder exposes the composed mask and the slot-certificate mask as named outputs; llama-context reads both back and compares 0-vs--inf class on the host. First version computed the delta with a graph-side exp/abs/sum chain, which mis-executed on one backend and produced impossible values (delta=-inf) — per the seen-to-fire doctrine the instrument was rebuilt backend-proof before any conclusion was drawn. Result: zero class mismatches on rolling tails and the duplicate-cell-0 adversarial (finite counts track the tail exactly); codex 066's algebra holds on all fixture-reachable cases. hole/seq_rm/seq_cp premises ride the env on real-server canaries, which is the point of a live invariant. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- src/llama-context.cpp | 44 +++++++++++++++++++++++++++++++++++++++++++ src/llama-graph.cpp | 17 +++++++++++++++++ 2 files changed, 61 insertions(+) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index ce44a552e17..28ad036d1b5 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1,6 +1,7 @@ #include "llama-context.h" #include +#include #include "ggml.h" #include "ggml-metal.h" @@ -1490,6 +1491,49 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); + + // E0a: report any layer where the gathered mask is not subsumed by the + // slot certificate (debug env; nodes exist only when the env is set) + static const bool fbl_e0a = getenv("LLAMA_GLM5_E0A") != nullptr; + if (fbl_e0a && status == GGML_STATUS_SUCCESS) { + ggml_cgraph * gfc = res->get_gf(); + std::map> layers; + for (int i = 0; i < ggml_graph_n_nodes(gfc); ++i) { + ggml_tensor * t = ggml_graph_node(gfc, i); + if (strncmp(t->name, "e0a_old-", 8) == 0) { + layers[atoi(t->name + 8)].first = t; + } else if (strncmp(t->name, "e0a_cert-", 9) == 0) { + layers[atoi(t->name + 9)].second = t; + } + } + for (auto & kv : layers) { + ggml_tensor * to = kv.second.first; + ggml_tensor * tc = kv.second.second; + if (!to || !tc) continue; + const int64_t n = to->ne[0]; + std::vector vo(n), vc(n); + ggml_backend_tensor_get(to, vo.data(), 0, n*sizeof(float)); + ggml_backend_tensor_get(tc, vc.data(), 0, n*sizeof(float)); + int bad = 0, finite = 0; + for (int64_t j = 0; j < n; ++j) { + const bool fo = vo[j] > -1e30f; + const bool fc = vc[j] > -1e30f; + finite += fc; + if (fo != fc) bad++; + } + if (bad) { + fprintf(stderr, "E0A VIOLATION layer=%d n=%lld class_mismatch=%d finite_cert=%d\n", + kv.first, (long long) n, bad, finite); + } else { + static bool logged = false; + if (!logged) { + fprintf(stderr, "E0A CHECK ACTIVE n=%lld finite=%d (layer %d clean)\n", + (long long) n, finite, kv.first); + logged = true; + } + } + } + } if (status != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); ret = status; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 7033ec60416..8b8c1061844 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3971,6 +3971,23 @@ ggml_tensor * llm_graph_context::build_attn_sparse_gathered( mask_g = ggml_reshape_4d(ctx0, mask_g, n_sel, 1, 1, 1); cb(mask_g, "gathered_mask", il); + // E0a (codex 066 s4 / grokk 067 s2): live certificate-subsumption invariant. + // Every mask value is exactly 0 or -inf, so exp() maps class to {1, 0} and + // sum(|exp(old) - exp(cert)|) is 0 iff cand_g+kq_g+slot == slot in class. + // Debug env only; read back and reported after compute in llama-context. + static const bool e0a_check = getenv("LLAMA_GLM5_E0A") != nullptr; + if (e0a_check) { + // backend-proof: expose BOTH raw masks as named graph outputs and do the + // 0-vs--inf class comparison on the host (a graph-side exp/abs/sum chain + // proved unreliable across backends for this shape) + ggml_tensor * cert = ggml_reshape_4d(ctx0, slot_mask, n_sel, 1, 1, 1); + ggml_format_name(cert, "e0a_cert-%d", il); + ggml_tensor * oldm = ggml_cont(ctx0, mask_g); + ggml_format_name(oldm, "e0a_old-%d", il); + ggml_build_forward_expand(gf, cert); + ggml_build_forward_expand(gf, oldm); + } + // K rows: [512, 1, n_kv, 1] -> dim-1 view -> gather -> cast -> [512, 1, n_sel, 1]. // the final reshape is load-bearing (grokk 057 s3): build_attn_mha permutes (0,2,1,3) // and packed KQ keys on post-permute k->ne[2]==1 (n_head_kv). From 37447da041600fedcb78a129f5641f3ff7a5d399 Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 22:10:31 -0600 Subject: [PATCH 60/62] context/kpool: timer + reuse-contract repairs (grokk 076/078, codex 077/079) - HOST_TIMERS: nonzero-only rebuild p50/p95/max + n_rebuild/n_reuse + burst-amortized + /256 serving estimate; p50 of a 255:1 zero series hid the one pad-256 rebuild entirely - perf_reset clears both host-timer vectors (instrument correctness must not depend on fresh processes) - kpool can_reuse: exhaustive optional-pair contract (sel<->cand paired + same shape; pool_reps/new_pool_cells/new_pool_reps shapes and pairedness; tail pair same shape) Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- src/llama-context.cpp | 26 +++++++++++++++++++++----- src/llama-kv-cache-kpool.cpp | 21 +++++++++++++++++++-- 2 files changed, 40 insertions(+), 7 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 28ad036d1b5..a6e89cefe97 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1481,11 +1481,25 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll std::sort(v.begin(), v.end()); return v.empty() ? (int64_t) 0 : v[(size_t) (p*(v.size() - 1))]; }; - fprintf(stderr, "HOST_TIMERS n=%zu rebuild_us p50=%lld p95=%lld | set_inputs_us p50=%lld p95=%lld | reused=%d\n", - fbl_host_rebuild_us.size(), - (long long) pct(fbl_host_rebuild_us, 0.5), (long long) pct(fbl_host_rebuild_us, 0.95), - (long long) pct(fbl_host_setinp_us, 0.5), (long long) pct(fbl_host_setinp_us, 0.95), - n_reused); + // rebuild stats over NONZERO samples only: p50 of a 255:1 zero-heavy + // series hides the one pad-256 rebuild entirely (grokk 076 s3). + std::vector nz; + int64_t sum_rebuild = 0; + for (int64_t v : fbl_host_rebuild_us) { + sum_rebuild += v; + if (v > 0) nz.push_back(v); + } + const size_t n_tot = fbl_host_rebuild_us.size(); + fprintf(stderr, "HOST_TIMERS n=%zu n_rebuild=%zu n_reuse=%zu | " + "rebuild_nz_us p50=%lld p95=%lld max=%lld | " + "burst_amort_us=%lld | serving_est_us=%lld | " + "set_inputs_us p50=%lld p95=%lld\n", + n_tot, nz.size(), n_tot - nz.size(), + (long long) pct(nz, 0.5), (long long) pct(nz, 0.95), + (long long) (nz.empty() ? 0 : *std::max_element(nz.begin(), nz.end())), + (long long) (n_tot ? sum_rebuild/(int64_t) n_tot : 0), + (long long) (pct(nz, 0.5)/256), + (long long) pct(fbl_host_setinp_us, 0.5), (long long) pct(fbl_host_setinp_us, 0.95)); } } } @@ -3409,6 +3423,8 @@ llama_perf_context_data llama_context::perf_get_data() const { } void llama_context::perf_reset() { + fbl_host_rebuild_us.clear(); + fbl_host_setinp_us.clear(); t_start_us = ggml_time_us(); t_eval_us = n_eval = 0; t_p_eval_us = n_p_eval = 0; diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index e9a9ef3dd7b..e67dd03253f 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -553,10 +553,26 @@ bool llm_graph_input_kpool::can_reuse(const llm_graph_params & params) { const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, kpool, (uint32_t) n_ps); res &= b_n_kv_attn == n_kv; - res &= sel_mask == nullptr || (sel_mask->ne[0] == n_kv && sel_mask->ne[1] == n_tps && - sel_mask->ne[3] == n_stream); + + // exhaustive optional-pair contract (codex 077 s4 / grokk 078 s1): + // presence, pairedness and exact shape for every optional tensor + res &= (sel_mask != nullptr) == (cand_mask != nullptr); + if (sel_mask != nullptr) { + res &= sel_mask->ne[0] == n_kv && sel_mask->ne[1] == n_tps && sel_mask->ne[3] == n_stream; + res &= ggml_are_same_shape(sel_mask, cand_mask); + } res &= pool_cells->ne[0] == (int64_t) kpool*n_pools && pool_cells->ne[1] == n_stream; res &= pool_bias->ne[0] == n_pools && pool_bias->ne[1] == n_tps && pool_bias->ne[2] == n_stream; + if (pool_reps != nullptr) { + res &= pool_reps->ne[0] == n_pools && pool_reps->ne[1] == n_stream; + res &= new_pool_cells != nullptr && new_pool_reps != nullptr; + res &= new_pool_cells->ne[0] == (int64_t) kpool*(int64_t) n_new_max && + new_pool_cells->ne[1] == n_stream; + res &= new_pool_reps->ne[0] == (int64_t) n_new_max*n_stream; + } else { + res &= new_pool_cells == nullptr && new_pool_reps == nullptr; + } + res &= (tail_cells != nullptr) == (tail_valid != nullptr); // policy: rebuild mode and the fixed pooled-key emission bound must match // what a fresh build would choose (codex 068 s3) @@ -567,6 +583,7 @@ bool llm_graph_input_kpool::can_reuse(const llm_graph_params & params) { res &= (int64_t) n_new_max == want_new_max; if (tail_cells != nullptr) { + res &= ggml_are_same_shape(tail_cells, tail_valid); // gathered graph: recompute the PAD32 width — never cache 2080 (grokk 070 s3) const int64_t n_top = (int64_t) kpool * llama_kpool_select_k((uint32_t) n_pools, hparams_indexer_top_k, kpool); From 03dc115c2b89b64ba9ce72bae127c5295026097a Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 22:25:10 -0600 Subject: [PATCH 61/62] context: rebuild event ledger replaces nonzero percentiles (grokk 081 s2, codex 082 s3) first_rebuild_us (cold graph after restore, never a pad cost) is separated from later events; rebuild_events=ordinal:us,... preserves chronology as the width-cadence oracle; percentiles over <=2 events removed. serving estimates move to the harness, which knows the horizon and the distance to the next pad boundary. Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- src/llama-context.cpp | 42 ++++++++++++++++++++++++++++++------------ 1 file changed, 30 insertions(+), 12 deletions(-) diff --git a/src/llama-context.cpp b/src/llama-context.cpp index a6e89cefe97..cbe23640388 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1481,24 +1481,42 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll std::sort(v.begin(), v.end()); return v.empty() ? (int64_t) 0 : v[(size_t) (p*(v.size() - 1))]; }; - // rebuild stats over NONZERO samples only: p50 of a 255:1 zero-heavy - // series hides the one pad-256 rebuild entirely (grokk 076 s3). - std::vector nz; + // rebuild EVENT LEDGER, chronological (codex 082 s3): the first + // event is the cold graph after restore, never a pad-256 cost + // (grokk 081 s2); later events are the width-cadence oracle. + // Percentiles over one or two events pretend to knowledge that + // does not exist, so raw ordinals+values are printed instead. + const size_t n_tot = fbl_host_rebuild_us.size(); int64_t sum_rebuild = 0; - for (int64_t v : fbl_host_rebuild_us) { + std::string events; + int64_t first_us = 0, later_max = 0; + size_t later_n = 0; + char buf[64]; + for (size_t j = 0; j < n_tot; ++j) { + const int64_t v = fbl_host_rebuild_us[j]; sum_rebuild += v; - if (v > 0) nz.push_back(v); + if (v > 0) { + if (first_us == 0) { + first_us = v; + } else { + later_n++; + later_max = std::max(later_max, v); + } + if (events.size() < 256) { + snprintf(buf, sizeof(buf), "%s%zu:%lld", + events.empty() ? "" : ",", j + 1, (long long) v); + events += buf; + } + } } - const size_t n_tot = fbl_host_rebuild_us.size(); fprintf(stderr, "HOST_TIMERS n=%zu n_rebuild=%zu n_reuse=%zu | " - "rebuild_nz_us p50=%lld p95=%lld max=%lld | " - "burst_amort_us=%lld | serving_est_us=%lld | " + "first_rebuild_us=%lld later_n=%zu later_max_us=%lld | " + "rebuild_events=%s | burst_amort_us=%lld | " "set_inputs_us p50=%lld p95=%lld\n", - n_tot, nz.size(), n_tot - nz.size(), - (long long) pct(nz, 0.5), (long long) pct(nz, 0.95), - (long long) (nz.empty() ? 0 : *std::max_element(nz.begin(), nz.end())), + n_tot, (size_t) (later_n + (first_us > 0)), n_tot - later_n - (first_us > 0), + (long long) first_us, later_n, (long long) later_max, + events.empty() ? "none" : events.c_str(), (long long) (n_tot ? sum_rebuild/(int64_t) n_tot : 0), - (long long) (pct(nz, 0.5)/256), (long long) pct(fbl_host_setinp_us, 0.5), (long long) pct(fbl_host_setinp_us, 0.95)); } } From 18dcfb368f6e85be0a0648a048d827f0782a0afb Mon Sep 17 00:00:00 2001 From: Your Name Date: Tue, 1 Sep 2026 23:09:56 -0600 Subject: [PATCH 62/62] glm5next: E0b+c MASKLESS gathered decode (codex 066, grokk 067) + diamond PASS MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The slot certificate subsumes cand+causal on the served lane (proven by E0a). The gathered graph now builds ONLY slot_valid||tail_valid as its mask; the dense sel/cand and the ordinary KQ mask are never created on that graph (build_inp_mem_hybrid_k(maskless) computed before the hybrid input, per 069/082), the kpool host fill is nullable-mask + explicit n_kv, and both hybrid-k and attn-k can_reuse fall back to a stored build-time n_kv witness when the mask is absent. Lane fail-closed: causal + no-ALiBi + no-SWA + scalar positions asserted in the builder. Oracle: maskless-vs-dense rolling tails NMSE <=7e-11 argmax-identical; duplicate-cell-0 adversarial and h9 clean; set_inputs_us p50 25->19 us at toy 8K with the two mask fills deleted. U/L/R/C diamond PASS at same- partition S: U==L==R==C, max_dlogprob 0.0 on every edge — the rewind, FULL serialization, and determinism paths are all exonerated, so the canary's live/cold divergence was purely the partition near-tie (084). Landmines fixed en route: b_n_kv member in the wrong base class; three unguarded sel_mask derefs in the host fill (n_padq the last, an anchor miss). Co-Authored-By: Claude Fable 5 Claude-Session: https://claude.ai/code/session_01YUN8B4NPpyNhziETmLiNpk --- src/llama-graph.cpp | 106 ++++++++++++++++++----------------- src/llama-graph.h | 6 +- src/llama-kv-cache-kpool.cpp | 54 ++++++++++++------ src/llama-kv-cache-kpool.h | 2 + src/models/glm5next.cpp | 21 +++---- 5 files changed, 105 insertions(+), 84 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8b8c1061844..f786ca9d5e3 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -511,7 +511,9 @@ void llm_graph_input_attn_k::set_input(const llama_ubatch * ubatch) { mctx->set_input_v_idxs(self_v_idxs, ubatch); // Patch B mirror } - mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + if (self_kq_mask) { + mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + } } bool llm_graph_input_attn_k::can_reuse(const llm_graph_params & params) { @@ -530,7 +532,12 @@ bool llm_graph_input_attn_k::can_reuse_impl(const llm_graph_params & params) { (int64_t) self_v_idxs->ne[0] == (int64_t) params.ubatch.n_tokens * mctx->mirror_width_max(); - res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams); + if (self_kq_mask) { + res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams); + } else { + // maskless graph: the width witness is the stored build-time n_kv + res &= b_n_kv == (int64_t) mctx->get_n_kv(); + } return res; } @@ -1099,7 +1106,9 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); - mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + if (inp_attn->self_kq_mask) { + mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + } if (inp_attn->self_k_rot) { mctx->get_attn()->set_input_k_rot(inp_attn->self_k_rot); @@ -1132,7 +1141,11 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) { res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens; //res &= inp_attn->self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there - res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams); + if (inp_attn->self_kq_mask) { + res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams); + } else { + res &= inp_attn->b_n_kv == (int64_t) mctx->get_attn()->get_n_kv(); + } res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs(); @@ -1155,7 +1168,9 @@ void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) { mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); // Patch B mirror } - mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + if (inp_attn->self_kq_mask) { + mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + } const int64_t n_rs = mctx->get_recr()->get_n_rs(); @@ -1179,7 +1194,11 @@ bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) { res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens; - res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams); + if (inp_attn->self_kq_mask) { + res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams); + } else { + res &= inp_attn->b_n_kv == (int64_t) mctx->get_attn()->get_n_kv(); + } res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs(); @@ -2955,7 +2974,8 @@ static std::unique_ptr build_attn_inp_k_impl( const llama_ubatch & ubatch, const llama_hparams & hparams, const llama_cparams & cparams, - const llama_kv_cache_context * mctx_cur) { + const llama_kv_cache_context * mctx_cur, + bool skip_kq_mask = false) { auto inp = std::make_unique(hparams, cparams, mctx_cur); @@ -2968,8 +2988,13 @@ static std::unique_ptr build_attn_inp_k_impl( inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch); // Patch B } - inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); - inp->self_kq_mask_cnv = inp->self_kq_mask; + // E0c: a maskless (gathered) graph never consumes the dense KQ mask; an + // unconsumed input is never allocated and the fill would fault (058 class) + if (!skip_kq_mask) { + inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); + inp->self_kq_mask_cnv = inp->self_kq_mask; + } + inp->b_n_kv = mctx_cur->get_n_kv(); } return inp; @@ -3646,11 +3671,11 @@ llm_graph_input_mem_hybrid * llm_graph_context::build_inp_mem_hybrid() const { return (llm_graph_input_mem_hybrid *) res->add_input(std::move(inp)); } -llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k() const { +llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k(bool maskless) const { const auto * mctx_cur = static_cast(mctx); auto inp_rs = build_rs_inp_impl (ctx0, ubatch, mctx_cur->get_recr()); - auto inp_attn = build_attn_inp_k_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_attn()); + auto inp_attn = build_attn_inp_k_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_attn(), maskless); auto inp = std::make_unique(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur); @@ -3695,10 +3720,11 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( 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); + GGML_ASSERT(kq_mask == nullptr || (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"); + GGML_ASSERT(kq_mask == nullptr || + (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); @@ -3716,7 +3742,10 @@ llm_graph_input_kpool * llm_graph_context::build_inp_kpool( ggml_set_name(inp->pool_bias_f16, "kpool_pool_bias_f16"); } - // lossless in f16 (only 0.0f and -INFINITY), and f16 + f32 -> f16 adds the KQ mask uncast + // lossless in f16 (only 0.0f and -INFINITY), and f16 + f32 -> f16 adds the KQ mask uncast. + // E0c: the maskless gathered graph consumes neither dense mask — creating + // them unconsumed would leave unallocated inputs for the host fill (058) + if (!gathered) { 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"); @@ -3724,6 +3753,7 @@ 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"); + } // D0 gathered decode inputs. ONLY when the gathered graph will consume them: // an input tensor no node reads is never allocated, so the host fill would @@ -3918,15 +3948,13 @@ ggml_tensor * llm_graph_context::build_attn_sparse_gathered( } } - const auto & kq_mask = inp->get_kq_mask(); - // decode-1, single stream only (the gate in the model guarantees this) GGML_ASSERT(top_k->ne[1] == 1 && top_k->ne[2] == 1); - GGML_ASSERT(kq_mask->ne[1] == 1 && kq_mask->ne[3] == 1); GGML_ASSERT(slot_valid->ne[0] == top_k->ne[0]); GGML_ASSERT(ggml_are_same_shape(tail_cells, tail_valid)); - const int64_t n_kv = kq_mask->ne[0]; + // explicit active width — never derived from a mask that no longer exists + const int64_t n_kv = mctx_cur->get_n_kv(); const int64_t n_sel = top_k->ne[0] + tail_cells->ne[0]; // printed at graph CONSTRUCTION (codex 063 s6) — reservation builds decode-1 @@ -3947,46 +3975,20 @@ ggml_tensor * llm_graph_context::build_attn_sparse_gathered( ggml_reshape_1d(ctx0, tail_cells, tail_cells->ne[0]), 0); cb(ids, "gathered_ids", il); - // per-CELL additive masks, gathered SEPARATELY then added compact (codex 059 P0: - // the dense cand+kq add ran once per DSA layer = eleven O(n_kv) ops per token). - // never gather the causal mask alone (grokk 057 s3): padded picks name cell 0 and - // cell 0 can be a real visible token; only cand_mask keeps those -inf. - auto gather_mask = [&](ggml_tensor * m, const char * name) { - ggml_tensor * rows = ggml_reshape_3d(ctx0, m, 1, n_kv, 1); - ggml_tensor * out = ggml_get_rows(ctx0, rows, ids); // -> F32 [1, n_sel, 1] - cb(out, name, il); - return ggml_reshape_2d(ctx0, out, n_sel, 1); - }; - ggml_tensor * cand_g = gather_mask(cand_mask, "gathered_cand_rows"); - ggml_tensor * kq_g = gather_mask(kq_mask, "gathered_kq_rows"); - ggml_tensor * mask_g = ggml_add(ctx0, cand_g, kq_g); + // E0b MASKLESS (codex 066, grokk 067, proven by E0a): the slot certificate + // subsumes cand+causal on this lane. Fail closed on the lane premises — + // outside them the dense masks can encode facts the certificate does not. + GGML_ASSERT(cparams.causal_attn && "maskless gathered decode requires causal attention"); + GGML_ASSERT(!hparams.use_alibi && "maskless gathered decode: no ALiBi"); + GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "maskless gathered decode: no SWA"); + GGML_ASSERT(!ubatch.is_pos_2d() && "maskless gathered decode: scalar positions only"); - // per-SLOT validity: expanded pool_bias of the selected pools, then the tail's - ggml_tensor * slot_mask = ggml_concat(ctx0, + ggml_tensor * mask_g = ggml_concat(ctx0, ggml_reshape_1d(ctx0, slot_valid, slot_valid->ne[0]), ggml_reshape_1d(ctx0, tail_valid, tail_valid->ne[0]), 0); - slot_mask = ggml_reshape_2d(ctx0, slot_mask, n_sel, 1); - - mask_g = ggml_add(ctx0, mask_g, slot_mask); mask_g = ggml_reshape_4d(ctx0, mask_g, n_sel, 1, 1, 1); cb(mask_g, "gathered_mask", il); - // E0a (codex 066 s4 / grokk 067 s2): live certificate-subsumption invariant. - // Every mask value is exactly 0 or -inf, so exp() maps class to {1, 0} and - // sum(|exp(old) - exp(cert)|) is 0 iff cand_g+kq_g+slot == slot in class. - // Debug env only; read back and reported after compute in llama-context. - static const bool e0a_check = getenv("LLAMA_GLM5_E0A") != nullptr; - if (e0a_check) { - // backend-proof: expose BOTH raw masks as named graph outputs and do the - // 0-vs--inf class comparison on the host (a graph-side exp/abs/sum chain - // proved unreliable across backends for this shape) - ggml_tensor * cert = ggml_reshape_4d(ctx0, slot_mask, n_sel, 1, 1, 1); - ggml_format_name(cert, "e0a_cert-%d", il); - ggml_tensor * oldm = ggml_cont(ctx0, mask_g); - ggml_format_name(oldm, "e0a_old-%d", il); - ggml_build_forward_expand(gf, cert); - ggml_build_forward_expand(gf, oldm); - } // K rows: [512, 1, n_kv, 1] -> dim-1 view -> gather -> cast -> [512, 1, n_sel, 1]. // the final reshape is load-bearing (grokk 057 s3): build_attn_mha permutes (0,2,1,3) diff --git a/src/llama-graph.h b/src/llama-graph.h index 7951bc3f4e4..e5afc17f945 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -344,8 +344,9 @@ class llm_graph_input_attn_kv : public llm_graph_input_i { ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch] ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa] - ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]; null on maskless gathered graphs (E0c) ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + int64_t b_n_kv = -1; // E0c: width witness when the mask is absent // note: assumes v_rot^2 == I ggml_tensor * self_k_rot = nullptr; @@ -390,6 +391,7 @@ class llm_graph_input_attn_k : public llm_graph_input_i { ggml_tensor * self_v_idxs = nullptr; // I64 expanded [n_batch*w] — Patch B mirror only ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + int64_t b_n_kv = -1; // E0c: width witness when the mask is absent ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] const llama_hparams hparams; @@ -1345,7 +1347,7 @@ struct llm_graph_context { // llm_graph_input_mem_hybrid * build_inp_mem_hybrid() const; - llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k() const; + llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k(bool maskless = false) const; llm_graph_input_mem_hybrid_iswa * build_inp_mem_hybrid_iswa() const; diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp index e67dd03253f..5d816955f21 100644 --- a/src/llama-kv-cache-kpool.cpp +++ b/src/llama-kv-cache-kpool.cpp @@ -74,6 +74,8 @@ void llama_kv_cache_set_input_kpool( ggml_tensor * cand_mask, ggml_tensor * tail_cells, ggml_tensor * tail_valid, + int64_t n_kv_arg, + int64_t n_ns_arg, ggml_tensor * pool_reps, ggml_tensor * new_pool_cells, ggml_tensor * new_pool_reps, @@ -87,22 +89,29 @@ void llama_kv_cache_set_input_kpool( GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(pool_bias ->buffer)); + GGML_ASSERT((sel_mask == nullptr) == (cand_mask == nullptr)); + if (sel_mask) { GGML_ASSERT(ggml_backend_buffer_is_host(sel_mask ->buffer)); GGML_ASSERT(ggml_backend_buffer_is_host(cand_mask ->buffer)); + } GGML_ASSERT(pool_cells->type == GGML_TYPE_I32); GGML_ASSERT(pool_bias ->type == GGML_TYPE_F32); + if (sel_mask) { 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"); - - 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)); + GGML_ASSERT(sel_mask->ne[0] == n_kv_arg && "explicit n_kv must match the masks"); + } + GGML_ASSERT(ggml_is_contiguous(pool_cells)); + GGML_ASSERT(ggml_is_contiguous(pool_bias)); - const int64_t n_kv = sel_mask->ne[0]; - const int64_t n_ns = sel_mask->ne[3]; + // E0c: n_kv is an explicit argument (grokk 067/070) — never inferred from + // a mask that may not exist + const int64_t n_kv = n_kv_arg; + const int64_t n_ns = sel_mask ? sel_mask->ne[3] : n_ns_arg; const int64_t r = kpool; const int64_t n_tokens = ubatch->n_tokens; @@ -117,13 +126,16 @@ void llama_kv_cache_set_input_kpool( 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)); + if (sel_mask) { + 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; - const int64_t n_padq = sel_mask->ne[1]; + // maskless gathered decode has no query-pad tensor; n_tps==1 on that lane + const int64_t n_padq = sel_mask ? sel_mask->ne[1] : n_tps; GGML_ASSERT(pool_bias->ne[1] == n_tps); GGML_ASSERT(n_padq >= n_tps); @@ -183,11 +195,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; - char * dst_sel_mask = (char *) sel_mask ->data; - char * dst_cand_mask = (char *) cand_mask ->data; + char * dst_sel_mask = sel_mask ? (char *) sel_mask ->data : nullptr; + char * dst_cand_mask = cand_mask ? (char *) cand_mask ->data : nullptr; - const bool mask_f16 = sel_mask->type == GGML_TYPE_F16; - const size_t mask_ts = ggml_type_size(sel_mask->type); + const bool mask_f16 = sel_mask ? (sel_mask->type == GGML_TYPE_F16) : true; + const size_t mask_ts = sel_mask ? ggml_type_size(sel_mask->type) : 0; // -1 marks a cell with no usable pool; host side only, never copied into cell_pool std::vector pool_of(n_kv); @@ -203,8 +215,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); - 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; + char * cur_sel_mask = dst_sel_mask ? dst_sel_mask + s*(n_padq*n_kv)*mask_ts : nullptr; + char * cur_cand_mask = dst_cand_mask ? dst_cand_mask + s*(n_padq*n_kv)*mask_ts : nullptr; float * cur_pool_bias = dst_pool_bias + s*(n_tps*n_pools); std::fill(cur_pool_cells, cur_pool_cells + r*n_pools, 0); @@ -230,7 +242,9 @@ void llama_kv_cache_set_input_kpool( } // the token loop writes rows < n_tps in full; only the padding rows need clearing - if (mask_f16) { + if (!cur_sel_mask) { + // maskless gathered graph: no dense mask to pad (E0c) + } else 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 { @@ -432,10 +446,12 @@ 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; - char * cur_sel = cur_sel_mask + ii*n_kv*mask_ts; - char * cur_cand = cur_cand_mask + ii*n_kv*mask_ts; + char * cur_sel = cur_sel_mask ? cur_sel_mask + ii*n_kv*mask_ts : nullptr; + char * cur_cand = cur_cand_mask ? cur_cand_mask + ii*n_kv*mask_ts : nullptr; - if (mask_f16) { + if (!sel_mask) { + // masks absent on the maskless gathered graph (E0c) + } else 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 { @@ -623,6 +639,8 @@ void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { /* cell_pool */ nullptr, pool_cells, /* bias */ nullptr, pool_bias, sel_mask, cand_mask, tail_cells, tail_valid, + (int64_t) mctx_attn->get_n_kv(), + (int64_t) (tail_cells ? tail_cells->ne[2] : 1), 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, diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h index d2a1cd69370..224f4881a72 100644 --- a/src/llama-kv-cache-kpool.h +++ b/src/llama-kv-cache-kpool.h @@ -47,6 +47,8 @@ void llama_kv_cache_set_input_kpool( ggml_tensor * cand_mask, ggml_tensor * tail_cells, // optional, D0 gathered decode ggml_tensor * tail_valid, // optional, with tail_cells + int64_t n_kv_arg, // explicit active width (E0c) + int64_t n_ns_arg, // explicit stream count when masks are null ggml_tensor * pool_reps, ggml_tensor * new_pool_cells, ggml_tensor * new_pool_reps, diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 68b3373731f..83543d726ea 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -567,9 +567,6 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( if (top_k && gathered) { GGML_ASSERT(slot_valid != nullptr && inp_kp->tail_cells != nullptr); - // sel_mask is not consumed by the gathered graph, but set_input still fills it - // (one shared host fill for all paths); expand the leaf so it stays allocated - ggml_build_forward_expand(gf, inp_kp->sel_mask); cur = build_attn_sparse_gathered(inp_attn, layer.wo, nullptr, nullptr, q, k, k, nullptr, nullptr, layer.wv_b, @@ -650,7 +647,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(); - llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); + // E0c: the gathered decision is computed BEFORE the hybrid input so a + // maskless graph never creates (or fills) the dense KQ mask (codex 066 s4) + const bool fbl_gathered_decode = glm5_gathered_dsa_enabled() && + cparams.n_ctx > glm5next_n_select(hparams) && n_tokens == 1 && + (cparams.kv_unified ? 1 : (int64_t) ubatch.n_seqs_unq) == 1; + + llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(fbl_gathered_decode); // 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 @@ -662,16 +665,10 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa if (mctx_hyb->get_idx() != nullptr) { indexer_scoring = cparams.n_ctx > glm5next_n_select(hparams); - // ONE owner for the D0 decision (codex 059 / grokk 060): the same boolean - // creates the tail inputs, requests slot_valid, and picks the builder. - // Producer/consumer gate disagreement was the 058 crash class. - const bool gathered_decode = glm5_gathered_dsa_enabled() && - indexer_scoring && n_tokens == 1 && - (cparams.kv_unified ? 1 : (int64_t) ubatch.n_seqs_unq) == 1; - + // ONE owner for the D0 decision; computed above, before the hybrid input inp_kp = build_inp_kpool(mctx_hyb, inp_mem->get_attn()->get_kq_mask(), indexer_scoring, - gathered_decode); + fbl_gathered_decode); } }