From 922e708702b81a156460ba718a4f6688663fe910 Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sat, 29 Aug 2026 02:11:34 -0400 Subject: [PATCH 01/12] glm5-next: register the architecture and mHC tensors in gguf-py GLM-5.3-Flash needs one genuinely new thing from llama.cpp: manifold-constrained hyper-connections. Everything else already exists upstream - KIMI_LINEAR carries the KDA tensors under the exact names GLM-5.3 uses (self_attn.f_a_proj, g_a_proj, b_proj, A_log, dt_bias, o_norm, {q,k,v}_conv1d), the MLA set (q_a/q_b/kv_a_mqa/kv_b/k_b/v_b plus their norms), MoE with shared experts, and FFN_EXP_PROBS_B for the sigmoid router's correction bias. GLM_DSA carries INDEXER_*. So GLM5_NEXT's tensor list is KIMI_LINEAR's plus the indexer plus six mHC tensors. Added: * MODEL_ARCH.GLM5_NEXT -> "glm5-next" * Keys.Attention.HyperConnection.{MULT,SINKHORN_ITERS,EPS}. These are load-bearing: a reader that ignores them builds a single-stream model that loads cleanly and produces garbage, because the residual is [B,S,hc_mult,D] for the whole stack. * HC_{ATTN,FFN}_{FN,BASE,SCALE} tensor enums, names blk.{bid}.hc_* * tensor_mapping entries for both `model.layers.{bid}.` and `model.language_model.layers.{bid}.` - GLM-5.3 is natively multimodal so its text stack sits under the latter. Verified by importing gguf: arch resolves, 49 tensors, all 6 mHC and 4 indexer present, KV keys and name formats correct. Reference for the graph builder: the mHC forward is implemented and validated bit-for-bit against transformers in the REAP repo (scripts/mhc_reference.py), with test vectors from real weights in vendor/mhc/. Two traps recorded there - the Sinkhorn loop is column-normalise then (iters-1) full passes, and its output is COLUMN- stochastic, not doubly-stochastic as the upstream docstring claims (row sums measured 0.980-1.024). "Fixing" that changes the residual mixing of all 45 layers. --- gguf-py/gguf/constants.py | 70 ++++++++++++++++++++++++++++++++++ gguf-py/gguf/tensor_mapping.py | 28 ++++++++++++++ 2 files changed, 98 insertions(+) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 3ebd9de5f..2ba30e7a7 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -184,6 +184,11 @@ class Keys: SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" TEMPERATURE_SCALE = "{arch}.attention.temperature_scale" + class HyperConnection: + MULT = "{arch}.attention.hc.mult" + SINKHORN_ITERS = "{arch}.attention.hc.sinkhorn_iters" + EPS = "{arch}.attention.hc.eps" + class Indexer: HEAD_COUNT = "{arch}.attention.indexer.head_count" KEY_LENGTH = "{arch}.attention.indexer.key_length" @@ -494,6 +499,7 @@ class MODEL_ARCH(IntEnum): LLAMA_EMBED = auto() MAINCODER = auto() KIMI_LINEAR = auto() + GLM5_NEXT = auto() class VISION_PROJECTOR_TYPE(IntEnum): @@ -705,6 +711,12 @@ class MODEL_TENSOR(IntEnum): INDEXER_PROJ = auto() INDEXER_ATTN_K = auto() INDEXER_ATTN_Q_B = auto() + HC_ATTN_FN = auto() + HC_ATTN_BASE = auto() + HC_ATTN_SCALE = auto() + HC_FFN_FN = auto() + HC_FFN_BASE = auto() + HC_FFN_SCALE = auto() # vision V_MMPROJ = auto() V_MMPROJ_FC = auto() @@ -974,6 +986,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.LLAMA_EMBED: "llama-embed", MODEL_ARCH.MAINCODER: "maincoder", MODEL_ARCH.KIMI_LINEAR: "kimi-linear", + MODEL_ARCH.GLM5_NEXT: "glm5-next", } VISION_PROJECTOR_TYPE_NAMES: dict[VISION_PROJECTOR_TYPE, str] = { @@ -1181,6 +1194,12 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.INDEXER_K_NORM: "blk.{bid}.indexer.k_norm", MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.indexer.proj", MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.indexer.attn_k", + MODEL_TENSOR.HC_ATTN_FN: "blk.{bid}.hc_attn_fn", + MODEL_TENSOR.HC_ATTN_BASE: "blk.{bid}.hc_attn_base", + MODEL_TENSOR.HC_ATTN_SCALE: "blk.{bid}.hc_attn_scale", + MODEL_TENSOR.HC_FFN_FN: "blk.{bid}.hc_ffn_fn", + MODEL_TENSOR.HC_FFN_BASE: "blk.{bid}.hc_ffn_base", + MODEL_TENSOR.HC_FFN_SCALE: "blk.{bid}.hc_ffn_scale", MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b", # vision MODEL_TENSOR.V_MMPROJ: "mm.{bid}", @@ -3859,6 +3878,57 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN_SHEXP, MODEL_TENSOR.FFN_UP_SHEXP, ], + MODEL_ARCH.GLM5_NEXT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_B, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_F_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_G_A, + MODEL_TENSOR.SSM_G_B, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + 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, + ], # TODO } diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index a7c7ce464..949ed63f1 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1246,6 +1246,34 @@ class TensorNameMap: "model.layers.{bid}.self_attn.indexer.weights_proj", # DSA ), + # GLM-5.3-Flash mHC. Two sites per layer (attention, FFN), each with a projection + # `fn` [(2+H)*H, H*D], a `base` bias and 3 learned `scale`s. Note GLM-5.3 is natively + # multimodal, so the text stack sits under `model.language_model.` rather than + # `model.layers.` - the converter strips that, but map both so a raw checkpoint works. + MODEL_TENSOR.HC_ATTN_FN: ( + "model.layers.{bid}.hc_attn_fn", + "model.language_model.layers.{bid}.hc_attn_fn", + ), + MODEL_TENSOR.HC_ATTN_BASE: ( + "model.layers.{bid}.hc_attn_base", + "model.language_model.layers.{bid}.hc_attn_base", + ), + MODEL_TENSOR.HC_ATTN_SCALE: ( + "model.layers.{bid}.hc_attn_scale", + "model.language_model.layers.{bid}.hc_attn_scale", + ), + MODEL_TENSOR.HC_FFN_FN: ( + "model.layers.{bid}.hc_ffn_fn", + "model.language_model.layers.{bid}.hc_ffn_fn", + ), + MODEL_TENSOR.HC_FFN_BASE: ( + "model.layers.{bid}.hc_ffn_base", + "model.language_model.layers.{bid}.hc_ffn_base", + ), + MODEL_TENSOR.HC_FFN_SCALE: ( + "model.layers.{bid}.hc_ffn_scale", + "model.language_model.layers.{bid}.hc_ffn_scale", + ), MODEL_TENSOR.INDEXER_ATTN_K: ( "model.layers.{bid}.self_attn.indexer.wk", # DSA ), -- 2.43.0 From 4c28d9e373c04047355cf13c3dba21aa7bfeb141 Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 08:14:02 -0400 Subject: [PATCH 02/12] glm5-next: mHC operator, hybrid KDA/MLA graph, vision tower Adds LLM_ARCH_GLM5_NEXT (GLM-5.3-Flash) end to end. The architecture is Kimi-Linear's (34 KDA linear-attention layers interleaved 3:1 with 11 NoPE MLA/DSA layers, sigmoid-routed MoE with a shared expert) plus one thing llama.cpp has no equivalent for: mHC hyper-connections, four parallel residual streams per layer mixed by a Sinkhorn-normalised matrix. New operator: ggml_mhc_sinkhorn (CPU + CUDA). Fused rather than composed because the alternative is ~180 graph nodes per site for reductions over a 4x4 matrix - roughly 16k nodes and kernel launches per token across 45 layers x 2 sites, for 16 floats of work. The normalisation order is load-bearing and is NOT symmetric Sinkhorn: one column pass, then (iters-1) full (row, column) passes, leaving a COLUMN-stochastic matrix. Symmetric Sinkhorn still yields a plausible matrix and a subtly wrong model, so tests/test-mhc-sinkhorn.cpp asserts both the values and that structure. Things that differ from Kimi-Linear and cost real accuracy if copied across: * KDA forget gate. GLM-5.3 is g = bound * sigmoid(exp(A_log) * (w + dt_bias)); Kimi is g = -exp(A_log) * softplus(...). Different function, opposite sign convention, bounded vs unbounded. Because it sits inside the recurrence, getting it wrong produces logit error that GROWS with sequence position - measured 2.4e-3 at position 0 rising to 8.2e-2 by position 15. * Clamped SwiGLU at swiglu_limit, in the text stack AND the vision tower. ggml_swiglu_oai clamps identically but computes silu_alpha(gate) * (up + 1); GLM has no +1 and alpha = 1. * MLA is NoPE (qk_rope_head_dim == 0), and head_count_kv must be 1 on those layers because the KV cache stores one compressed latent per token. * Leading dense-FFN count comes from mlp_layer_types, not first_k_dense_replace. * NextN/MTP tensors were classified LLM_TENSOR_LAYER_OUTPUT, which makes the loader abort when they are created with a layer index. Unreachable until now only because no GGUF carried them; glm5-next ships them, so they are LAYER_REPEATING. Still never executed. * mHC is on transformer layers only - the MTP block has none. Vision: the tower is the GLM-4V family, so PROJECTOR_TYPE_GLM4V and clip_graph_glm4v apply unchanged. Adds Glm5NextVisionModel and clip.vision.swiglu_limit. Validated against transformers on a structurally identical 6-layer fixture (both attention types, dense and MoE FFN, shared expert, an MTP block that must load and not execute, mHC at every site): top-1 agreement 16/16, worst relative logit error 5.0e-3 on CPU and 4.9e-3 with layers on GPU. On the real 165B checkpoint it produces correct arithmetic with correct intermediate steps, correct code, and 3/3 needle retrieval at 32k context. DSA is NOT implemented: the indexer tensors load and are unused, so those 11 layers run dense, exactly as upstream already does for LLM_ARCH_GLM_DSA and DEEPSEEK2. --- convert_hf_to_gguf.py | 215 +++++++++++++++- ggml/include/ggml.h | 20 ++ ggml/src/ggml-cpu/ggml-cpu.c | 5 + ggml/src/ggml-cpu/ops.cpp | 90 +++++++ ggml/src/ggml-cpu/ops.h | 1 + ggml/src/ggml-cuda/ggml-cuda.cu | 5 + ggml/src/ggml-cuda/mhc.cu | 114 +++++++++ ggml/src/ggml-cuda/mhc.cuh | 3 + ggml/src/ggml.c | 27 +- gguf-py/gguf/constants.py | 16 ++ gguf-py/gguf/gguf_writer.py | 18 ++ gguf-py/gguf/tensor_mapping.py | 10 + src/CMakeLists.txt | 1 + src/llama-arch.cpp | 112 ++++++++- src/llama-arch.h | 14 ++ src/llama-graph.cpp | 11 + src/llama-graph.h | 1 + src/llama-hparams.h | 12 + src/llama-model.cpp | 181 +++++++++++++ src/llama-model.h | 12 + src/models/glm5-next.cpp | 433 ++++++++++++++++++++++++++++++++ src/models/models.h | 23 ++ tests/CMakeLists.txt | 2 + tests/test-glm5-next-logits.cpp | 103 ++++++++ tests/test-mhc-sinkhorn.cpp | 77 ++++++ tools/mtmd/clip-impl.h | 1 + tools/mtmd/clip-model.h | 5 + tools/mtmd/clip.cpp | 9 +- 28 files changed, 1509 insertions(+), 12 deletions(-) create mode 100644 ggml/src/ggml-cuda/mhc.cu create mode 100644 ggml/src/ggml-cuda/mhc.cuh create mode 100644 src/models/glm5-next.cpp create mode 100644 tests/test-glm5-next-logits.cpp create mode 100644 tests/test-mhc-sinkhorn.cpp diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index d4929d6b6..f13556d39 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -4949,6 +4949,28 @@ class Glm4VVisionModel(Qwen3VLVisionModel): yield from super().modify_tensors(data_torch, name, bid) +@ModelBase.register("Glm5NextForConditionalGeneration") +class Glm5NextVisionModel(Glm4VVisionModel): + """GLM-5.3-Flash vision tower. + + Structurally the GLM-4V tower - patch_embed / 24 blocks with qkv + q_norm/k_norm / a + downsample conv / a merger with gate+up+down and post_projection_norm - so the existing + PROJECTOR_TYPE_GLM4V graph and tensor mapping apply unchanged. + + The one delta that matters is the MLP: GLM-5.3 clamps SwiGLU at `swiglu_limit` in the vision + tower as well as the text stack. That is recorded in the mmproj so the graph can honour it; + it only bites once activations exceed the limit, which is exactly the regime a small test + never reaches and a real image does. + """ + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_vision is not None + limit = self.hparams_vision.get("swiglu_limit") + if limit is not None: + self.gguf_writer.add_vision_swiglu_limit(float(limit)) + + @ModelBase.register("Qwen3VLForConditionalGeneration") class Qwen3VLTextModel(Qwen3Model): model_arch = gguf.MODEL_ARCH.QWEN3VL @@ -5961,7 +5983,16 @@ class KimiLinearModel(TextModel): # num_shared_experts (1 for Kimi) self.gguf_writer.add_expert_shared_count(self.hparams["num_shared_experts"]) # first_k_dense_replace (1 for Kimi - first layer uses dense MLP) - self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"]) + # `first_k_dense_replace` is NOT what decides this. transformers derives the FFN layout + # as `["dense"] * min(3, n_layer) + ["sparse"] * rest`, so trust mlp_layer_types when the + # config carries it and fall back only when it does not. (They agree at 3 for the real + # checkpoint; they disagree for any config where someone set the knob expecting it to.) + mlp_types = self.hparams.get("mlp_layer_types") + if mlp_types: + n_dense = next((i for i, t in enumerate(mlp_types) if t != "dense"), len(mlp_types)) + else: + n_dense = self.hparams["first_k_dense_replace"] + self.gguf_writer.add_leading_dense_block_count(n_dense) # Routed scaling factor (expert_weights_scale = 2.446 for Kimi) self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"]) @@ -6052,6 +6083,188 @@ class KimiLinearModel(TextModel): yield from super().modify_tensors(data_torch, name, bid) +@ModelBase.register("Glm5NextForConditionalGeneration", "Glm5NextForCausalLM", "Glm5NextTextModel") +class Glm5NextModel(TextModel): + """GLM-5.3-Flash: hybrid KDA + MLA/DSA MoE with mHC hyper-connections. + + Closest existing relative is Kimi-Linear (same KDA delta rule, same MLA), with three + additions this converter has to carry: + + * mHC hyper-connections - four residual streams per layer with Sinkhorn-normalised + mixing. Six tensors per layer, passed through unchanged; the graph does the work. + * MLA with NO rope at all (`mla_use_nope`, `qk_rope_head_dim == 0`). Kimi and DeepSeek + both rope the MLA path, so the usual `add_rope_dimension_count` path is wrong here. + * A natively multimodal checkpoint, so the text stack lives under + `model.language_model.` and the ViT under `model.visual.` (skipped - it belongs in + an mmproj file, not here). + + Layer types come from `linear_attn_config`, and note they are ZERO-indexed here where + Kimi's are one-indexed. Getting that wrong silently builds a model with 34 attention + layers wired as recurrent ones, which loads fine and produces noise. + """ + model_arch = gguf.MODEL_ARCH.GLM5_NEXT + + _experts: list[dict[str, Tensor]] | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # The MTP block is a real 46th layer in the checkpoint (a full MoE layer, 3.81B + # params), so it has to be counted or its tensors have nowhere to go. + self.block_count = self.hparams["num_hidden_layers"] + self.hparams.get("num_nextn_predict_layers", 0) + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + self._keep_mtp = bool(os.environ.get("GLM5_KEEP_MTP", "1") not in ("0", "false", "no")) + if not self._keep_mtp: + logger.info("GLM5_KEEP_MTP=0: dropping the MTP block (~2.0 GiB at Q4_K_M); " + "llama.cpp cannot execute it today") + + def set_vocab(self): + self._set_vocab_gpt2() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_vocab_size(self.hparams["vocab_size"]) + + lac = self.hparams["linear_attn_config"] + full_attn = set(lac["full_attn_layers"]) # ZERO-indexed, unlike Kimi's + n_layer = self.hparams["num_hidden_layers"] + nextn = self.hparams.get("num_nextn_predict_layers", 0) + + # n_head_kv == 0 marks a recurrent (KDA) layer; > 0 marks full attention. The MTP + # block shares the DSA layout, so it is tagged as attention. + # + # The value for attention layers is 1, not num_key_value_heads: MLA stores ONE + # compressed latent per token, so the KV cache is MQA. Writing the real head count + # makes llama.cpp size the cache rows at n_head_kv * head_dim and ggml_set_rows then + # asserts against a kv_lora_rank-wide key. Kimi-Linear's converter does the same. + n_kv = [1 if (il in full_attn or il >= n_layer) else 0 + for il in range(n_layer + nextn)] + assert sum(1 for x in n_kv[:n_layer] if x) == len(full_attn), \ + f"layer-type map disagrees with linear_attn_config ({sum(1 for x in n_kv[:n_layer] if x)} vs {len(full_attn)})" + self.gguf_writer.add_head_count_kv(n_kv) + + # --- KDA --- + self.gguf_writer.add_ssm_conv_kernel(lac["short_conv_kernel_size"]) + self.gguf_writer.add_kda_head_dim(lac["head_dim"]) + # GLM-5.3's forget gate is NOT Kimi's. Kimi: g = -exp(A_log) * softplus(w + dt_bias). + # GLM-5.3: g = lower_bound * sigmoid(exp(A_log) * (w + dt_bias)) - a bounded gate, with + # exp(A_log) POSITIVE and used inside the sigmoid. Copying Kimi's `-exp` here produces a + # model whose error grows monotonically along the sequence, which is exactly what it did. + self.gguf_writer.add_ssm_gate_lower_bound(float(lac["gate_lower_bound"])) + + # --- MLA. NoPE: qk_rope_head_dim is 0 and mla_use_nope is set, so there is no + # rotary section at all and key length is just the compressed KV rank. --- + qk_rope = self.hparams.get("qk_rope_head_dim", 0) + assert qk_rope == 0 and self.hparams.get("mla_use_nope", False), \ + "this converter assumes GLM-5.3's NoPE MLA; a roped variant needs the rope path back" + self.gguf_writer.add_rope_dimension_count(0) + kv_lora = self.hparams["kv_lora_rank"] + self.gguf_writer.add_q_lora_rank(self.hparams["q_lora_rank"]) + self.gguf_writer.add_kv_lora_rank(kv_lora) + self.gguf_writer.add_key_length(kv_lora) + self.gguf_writer.add_key_length_mla(self.hparams["qk_nope_head_dim"]) + self.gguf_writer.add_value_length_mla(self.hparams["v_head_dim"]) + + # --- MoE --- + self.gguf_writer.add_expert_count(self.hparams["n_routed_experts"]) + self.gguf_writer.add_expert_used_count(self.hparams["num_experts_per_tok"]) + self.gguf_writer.add_expert_feed_forward_length(self.hparams["moe_intermediate_size"]) + self.gguf_writer.add_expert_shared_count(self.hparams["n_shared_experts"]) + self.gguf_writer.add_leading_dense_block_count(self.hparams["first_k_dense_replace"]) + self.gguf_writer.add_expert_weights_scale(self.hparams["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(self.hparams["norm_topk_prob"]) + self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SIGMOID) + # Clamped SwiGLU: gate is clamped ABOVE at the limit, up is clamped to +/-limit, and + # only then does SiLU apply. Plain SILU is right until activations reach the limit, at + # which point it silently diverges - so it must travel with the file. + self.gguf_writer.add_swiglu_limit(float(self.hparams["swiglu_limit"])) + + # --- DSA indexer. Loaded for completeness; llama.cpp runs these layers dense, the + # same way it already does for LLM_ARCH_GLM_DSA. --- + self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"]) + self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"]) + self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"]) + + # --- mHC --- + self.gguf_writer.add_hc_mult(self.hparams["hc_mult"]) + self.gguf_writer.add_hc_sinkhorn_iters(self.hparams["hc_sinkhorn_iters"]) + self.gguf_writer.add_hc_eps(self.hparams["hc_eps"]) + + if nextn: + self.gguf_writer.add_nextn_predict_layers(nextn) + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts is not None: + left = [k for d in self._experts for k in d] + if left: + raise ValueError(f"Unprocessed experts: {left[:5]} ({len(left)} total)") + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # The ViT belongs in an mmproj file built by Glm5NextVisionModel, not here. + if name.startswith("model.visual."): + return + if name.startswith("model.language_model."): + name = name.replace("model.language_model.", "model.", 1) + + n_layer = self.hparams["num_hidden_layers"] + if not self._keep_mtp and bid is not None and bid >= n_layer: + return + + # KDA conv1d: HF [d_inner, d_conv] -> ggml ne [d_conv, 1, d_inner, 1]. GGUF reverses + # the numpy shape on write, so the target numpy shape is (1, d_inner, 1, d_conv); + # d_conv stays fastest-varying in memory, which is what ggml_ssm_conv indexes. + if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")): + if data_torch.ndim == 3: # [d_inner, 1, d_conv] + data_torch = data_torch.squeeze(1) + d_inner, d_conv = data_torch.shape + data_torch = data_torch.reshape(1, d_inner, 1, d_conv) + + # Store exp(A_log) - a pure function of the weight, so precomputing it costs the graph + # nothing. NOT -exp: see the forget-gate note in set_gguf_parameters. + if name.endswith(".A_log"): + data_torch = torch.exp(data_torch) + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + if name.endswith("e_score_correction_bias"): + name = name.replace("e_score_correction_bias", "e_score_correction.bias") + + # Routed experts are stacked into one 3D tensor per projection, which is what + # llama.cpp's fused MoE kernels index. + 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: + for proj, tname in (("gate_proj", gguf.MODEL_TENSOR.FFN_GATE_EXP), + ("down_proj", gguf.MODEL_TENSOR.FFN_DOWN_EXP), + ("up_proj", gguf.MODEL_TENSOR.FFN_UP_EXP)): + datas: list[Tensor] = [] + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{proj}.weight" + datas.append(self._experts[bid][ename]) + del self._experts[bid][ename] + stacked = torch.stack(datas, dim=0) + yield from super().modify_tensors(stacked, self.format_tensor_name(tname, bid), bid) + return + + # MLA absorption wants kv_b split, with k_b transposed. + if name.endswith("kv_b_proj.weight"): + n_head_kv = self.hparams["num_key_value_heads"] + v_head_dim = self.hparams["v_head_dim"] + qk_nope = self.hparams["qk_nope_head_dim"] + assert data_torch.shape[0] == n_head_kv * (v_head_dim + qk_nope), \ + f"kv_b_proj {tuple(data_torch.shape)} does not match {n_head_kv}*({v_head_dim}+{qk_nope})" + kv_b = data_torch.view(n_head_kv, v_head_dim + qk_nope, data_torch.shape[-1]) + k_b, v_b = torch.split(kv_b, [qk_nope, v_head_dim], dim=1) + yield from super().modify_tensors(k_b.transpose(1, 2), name.replace("kv_b_proj", "k_b_proj"), bid) + yield from super().modify_tensors(v_b, name.replace("kv_b_proj", "v_b_proj"), bid) + return + + yield from super().modify_tensors(data_torch, name, bid) + + @ModelBase.register("InternLM2ForCausalLM") class InternLM2Model(TextModel): model_arch = gguf.MODEL_ARCH.INTERNLM2 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 669f66b65..7412be84a 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -499,6 +499,7 @@ extern "C" { GGML_OP_RMS_NORM_BACK, GGML_OP_GROUP_NORM, GGML_OP_L2_NORM, + GGML_OP_MHC_SINKHORN, GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID, @@ -1383,6 +1384,25 @@ extern "C" { // l2 normalize along rows // used in rwkv v7 + // mHC (GLM-5.3-Flash hyper-connections): softmax over ne0, add eps, then a Sinkhorn + // projection of each [hc, hc] slice. + // + // Fused because the alternative is ~180 graph nodes per site for reductions over a 4x4 + // matrix - 16k nodes and kernel launches per token across 45 layers x 2 sites. The whole + // thing fits in registers. + // + // The normalisation ORDER is load-bearing and is NOT symmetric Sinkhorn: column + // normalisation once, then (iters - 1) full (row, column) passes. The result is + // COLUMN-stochastic, not doubly stochastic. Reading it as `iters` symmetric passes still + // produces a plausible matrix and a subtly wrong model. + // + // a: [hc, hc, n_tokens, 1] pre-softmax logits. Returns the same shape. + GGML_API struct ggml_tensor * ggml_mhc_sinkhorn( + struct ggml_context * ctx, + struct ggml_tensor * a, + int iters, + float eps); + GGML_API struct ggml_tensor * ggml_l2_norm( struct ggml_context * ctx, struct ggml_tensor * a, diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 7486acc2b..8cdb6cc5c 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1808,6 +1808,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_l2_norm(params, tensor); } break; + case GGML_OP_MHC_SINKHORN: + { + ggml_compute_forward_mhc_sinkhorn(params, tensor); + } break; case GGML_OP_MUL_MAT: { ggml_compute_forward_mul_mat(params, tensor); @@ -2277,6 +2281,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_RMS_NORM: case GGML_OP_RMS_NORM_BACK: case GGML_OP_L2_NORM: + case GGML_OP_MHC_SINKHORN: case GGML_OP_GROUP_NORM: case GGML_OP_CONCAT: case GGML_OP_MUL_MAT: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 765ce07f0..edec9058c 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4049,6 +4049,96 @@ void ggml_compute_forward_group_norm( } } +// ggml_compute_forward_mhc_sinkhorn + +static void ggml_compute_forward_mhc_sinkhorn_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + GGML_ASSERT(ggml_are_same_shape(src0, dst)); + GGML_ASSERT(src0->nb[0] == sizeof(float)); + GGML_ASSERT(src0->ne[0] == src0->ne[1]); + + int iters; + float eps; + memcpy(&iters, (const int32_t *) dst->op_params + 0, sizeof(int)); + memcpy(&eps, (const float *) dst->op_params + 1, sizeof(float)); + + const int ith = params->ith; + const int nth = params->nth; + + const int64_t hc = src0->ne[0]; + // One [hc, hc] matrix per (token, batch) slice. hc is 4 for GLM-5.3, so the whole thing + // is 16 floats and stays in cache; the work is the 2*iters reductions, not the data. + const int64_t nr = ggml_nrows(src0) / hc; + + const int64_t dr = (nr + nth - 1) / nth; + const int64_t i0 = dr * ith; + const int64_t i1 = MIN(i0 + dr, nr); + + for (int64_t i = i0; i < i1; ++i) { + // i indexes the (ne2, ne3) slice grid; ne0 and ne1 are the matrix itself. + const int64_t i3 = i / src0->ne[2]; + const int64_t i2 = i % src0->ne[2]; + + const char * src = (const char *) src0->data + i2*src0->nb[2] + i3*src0->nb[3]; + char * out = ( char *) dst->data + i2*dst->nb[2] + i3*dst->nb[3]; + + // softmax over ne0 (the last torch dim), then + eps. + for (int64_t r = 0; r < hc; ++r) { + const float * s = (const float *)(src + r*src0->nb[1]); + float * o = (float *)(out + r*dst->nb[1]); + float mx = -INFINITY; + for (int64_t c = 0; c < hc; ++c) mx = MAX(mx, s[c]); + float sum = 0.0f; + for (int64_t c = 0; c < hc; ++c) { const float e = expf(s[c] - mx); o[c] = e; sum += e; } + const float inv = 1.0f / sum; + for (int64_t c = 0; c < hc; ++c) o[c] = o[c]*inv + eps; + } + + // COLUMN normalisation first, then (iters-1) full (row, column) passes. This is the + // order the reference implements and it is not symmetric Sinkhorn - see ggml.h. + for (int64_t c = 0; c < hc; ++c) { + float sum = 0.0f; + for (int64_t r = 0; r < hc; ++r) sum += ((const float *)(out + r*dst->nb[1]))[c]; + const float inv = 1.0f / (sum + eps); + for (int64_t r = 0; r < hc; ++r) ((float *)(out + r*dst->nb[1]))[c] *= inv; + } + for (int it = 1; it < iters; ++it) { + for (int64_t r = 0; r < hc; ++r) { + float * o = (float *)(out + r*dst->nb[1]); + float sum = 0.0f; + for (int64_t c = 0; c < hc; ++c) sum += o[c]; + const float inv = 1.0f / (sum + eps); + for (int64_t c = 0; c < hc; ++c) o[c] *= inv; + } + for (int64_t c = 0; c < hc; ++c) { + float sum = 0.0f; + for (int64_t r = 0; r < hc; ++r) sum += ((const float *)(out + r*dst->nb[1]))[c]; + const float inv = 1.0f / (sum + eps); + for (int64_t r = 0; r < hc; ++r) ((float *)(out + r*dst->nb[1]))[c] *= inv; + } + } + } +} + +void ggml_compute_forward_mhc_sinkhorn( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->src[0]->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_mhc_sinkhorn_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + // ggml_compute_forward_l2_norm static void ggml_compute_forward_l2_norm_f32( diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 3fa1443ab..94cbecb2b 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -47,6 +47,7 @@ void ggml_compute_forward_rms_norm(const struct ggml_compute_params * params, st void ggml_compute_forward_rms_norm_back(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_group_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_l2_norm(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_mhc_sinkhorn(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_out_prod(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_scale(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_set(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 75b62129a..039bb0496 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -11,6 +11,7 @@ #include "ggml-cuda/binbcast.cuh" #include "ggml-cuda/clamp.cuh" #include "ggml-cuda/concat.cuh" +#include "ggml-cuda/mhc.cuh" #include "ggml-cuda/conv-transpose-1d.cuh" #include "ggml-cuda/conv2d.cuh" #include "ggml-cuda/conv2d-dw.cuh" @@ -2647,6 +2648,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_L2_NORM: ggml_cuda_op_l2_norm(ctx, dst); break; + case GGML_OP_MHC_SINKHORN: + ggml_cuda_op_mhc_sinkhorn(ctx, dst); + break; case GGML_OP_CONCAT: ggml_cuda_op_concat(ctx, dst); break; @@ -4943,6 +4947,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g break; case GGML_OP_NORM: case GGML_OP_RMS_NORM: + case GGML_OP_MHC_SINKHORN: case GGML_OP_L2_NORM: return true; case GGML_OP_RMS_NORM_BACK: diff --git a/ggml/src/ggml-cuda/mhc.cu b/ggml/src/ggml-cuda/mhc.cu new file mode 100644 index 000000000..6f1a7bd98 --- /dev/null +++ b/ggml/src/ggml-cuda/mhc.cu @@ -0,0 +1,114 @@ +#include "mhc.cuh" + +// mHC Sinkhorn (GLM-5.3-Flash hyper-connections). +// +// Each [hc, hc] slice is independent and tiny - hc is 4, so 16 floats - and the work is the +// 2*iters reductions over it, not the data. So: one THREAD per slice, whole matrix in registers, +// no shared memory and no cross-thread reduction. A warp-per-slice layout would spend more on +// shuffles than the arithmetic costs. +// +// The normalisation order is load-bearing and is NOT symmetric Sinkhorn: one column pass, then +// (iters - 1) full (row, column) passes, leaving a COLUMN-stochastic matrix. This must stay +// bit-comparable with the CPU path in ggml-cpu/ops.cpp - expf, not __expf, for that reason. + +#define MHC_MAX_HC 8 + +static __global__ void mhc_sinkhorn_f32( + const float * __restrict__ src, + float * __restrict__ dst, + const int hc, const int nslices, const int iters, const float eps) { + + const int i = blockIdx.x*blockDim.x + threadIdx.x; + if (i >= nslices) { + return; + } + + const int n = hc*hc; + float m[MHC_MAX_HC*MHC_MAX_HC]; + + const float * s = src + (size_t) i*n; + + // softmax over ne0 (torch's last dim), then + eps + for (int r = 0; r < hc; ++r) { + const float * sr = s + r*hc; + float mx = -INFINITY; + for (int c = 0; c < hc; ++c) { + mx = fmaxf(mx, sr[c]); + } + float sum = 0.0f; + for (int c = 0; c < hc; ++c) { + const float e = expf(sr[c] - mx); + m[r*hc + c] = e; + sum += e; + } + const float inv = 1.0f/sum; + for (int c = 0; c < hc; ++c) { + m[r*hc + c] = m[r*hc + c]*inv + eps; + } + } + + // COLUMN normalisation first + for (int c = 0; c < hc; ++c) { + float sum = 0.0f; + for (int r = 0; r < hc; ++r) { + sum += m[r*hc + c]; + } + const float inv = 1.0f/(sum + eps); + for (int r = 0; r < hc; ++r) { + m[r*hc + c] *= inv; + } + } + + // then (iters - 1) full (row, column) passes + for (int it = 1; it < iters; ++it) { + for (int r = 0; r < hc; ++r) { + float sum = 0.0f; + for (int c = 0; c < hc; ++c) { + sum += m[r*hc + c]; + } + const float inv = 1.0f/(sum + eps); + for (int c = 0; c < hc; ++c) { + m[r*hc + c] *= inv; + } + } + for (int c = 0; c < hc; ++c) { + float sum = 0.0f; + for (int r = 0; r < hc; ++r) { + sum += m[r*hc + c]; + } + const float inv = 1.0f/(sum + eps); + for (int r = 0; r < hc; ++r) { + m[r*hc + c] *= inv; + } + } + } + + float * d = dst + (size_t) i*n; + for (int k = 0; k < n; ++k) { + d[k] = m[k]; + } +} + +void ggml_cuda_op_mhc_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src0)); + GGML_ASSERT(src0->ne[0] == src0->ne[1]); + GGML_ASSERT(src0->ne[0] <= MHC_MAX_HC); + + int iters; + float eps; + memcpy(&iters, (const int32_t *) dst->op_params + 0, sizeof(int)); + memcpy(&eps, (const float *) dst->op_params + 1, sizeof(float)); + + const int hc = src0->ne[0]; + const int nslices = ggml_nelements(src0) / (hc*hc); + + const int block = 256; + const int grid = (nslices + block - 1)/block; + + mhc_sinkhorn_f32<<>>( + (const float *) src0->data, (float *) dst->data, hc, nslices, iters, eps); +} diff --git a/ggml/src/ggml-cuda/mhc.cuh b/ggml/src/ggml-cuda/mhc.cuh new file mode 100644 index 000000000..ad23ba101 --- /dev/null +++ b/ggml/src/ggml-cuda/mhc.cuh @@ -0,0 +1,3 @@ +#include "common.cuh" + +void ggml_cuda_op_mhc_sinkhorn(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index e9b6720c0..a2e7bc70f 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -980,6 +980,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "RMS_NORM_BACK", "GROUP_NORM", "L2_NORM", + "MHC_SINKHORN", "MUL_MAT", "MUL_MAT_ID", @@ -1057,7 +1058,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1090,6 +1091,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "rms_norm_back(x)", "group_norm(x)", "l2_norm(x)", + "mhc_sinkhorn(x)", "X*Y", "X[i]*Y", @@ -1167,7 +1169,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -3172,6 +3174,27 @@ struct ggml_tensor * ggml_group_norm_inplace( return ggml_group_norm_impl(ctx, a, n_groups, eps, true); } +// ggml_mhc_sinkhorn + +struct ggml_tensor * ggml_mhc_sinkhorn( + struct ggml_context * ctx, + struct ggml_tensor * a, + int iters, + float eps) { + GGML_ASSERT(a->ne[0] == a->ne[1]); // square [hc, hc] slices + GGML_ASSERT(iters >= 1); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, iters); + ggml_set_op_params_f32(result, 1, eps); + + result->op = GGML_OP_MHC_SINKHORN; + result->src[0] = a; + + return result; +} + // ggml_l2_norm static struct ggml_tensor * ggml_l2_norm_impl( diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 2ba30e7a7..ee096a58b 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -122,6 +122,7 @@ class Keys: EXPERT_WEIGHTS_SCALE = "{arch}.expert_weights_scale" EXPERT_WEIGHTS_NORM = "{arch}.expert_weights_norm" EXPERT_GATING_FUNC = "{arch}.expert_gating_func" + SWIGLU_LIMIT = "{arch}.swiglu_limit" EXPERT_GROUP_SCALE = "{arch}.expert_group_scale" EXPERTS_PER_GROUP = "{arch}.experts_per_group" MOE_EVERY_N_LAYERS = "{arch}.moe_every_n_layers" @@ -222,6 +223,7 @@ class Keys: STATE_SIZE = "{arch}.ssm.state_size" TIME_STEP_RANK = "{arch}.ssm.time_step_rank" GROUP_COUNT = "{arch}.ssm.group_count" + GATE_LOWER_BOUND = "{arch}.ssm.gate_lower_bound" DT_B_C_RMS = "{arch}.ssm.dt_b_c_rms" class KDA: @@ -319,6 +321,7 @@ class Keys: SPATIAL_MERGE_SIZE = "clip.vision.spatial_merge_size" USE_GELU = "clip.use_gelu" USE_SILU = "clip.use_silu" + SWIGLU_LIMIT = "clip.vision.swiglu_limit" N_WA_PATTERN = "clip.vision.n_wa_pattern" # used by qwen2.5vl WA_LAYER_INDEXES = "clip.vision.wa_layer_indexes" # used by youtuvl IS_DEEPSTACK_LAYERS = "clip.vision.is_deepstack_layers" @@ -711,6 +714,8 @@ class MODEL_TENSOR(IntEnum): INDEXER_PROJ = auto() INDEXER_ATTN_K = auto() INDEXER_ATTN_Q_B = auto() + INDEXER_KPOOL_APE = auto() + INDEXER_KPOOL_GATE = auto() HC_ATTN_FN = auto() HC_ATTN_BASE = auto() HC_ATTN_SCALE = auto() @@ -1201,6 +1206,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.HC_FFN_BASE: "blk.{bid}.hc_ffn_base", MODEL_TENSOR.HC_FFN_SCALE: "blk.{bid}.hc_ffn_scale", MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b", + MODEL_TENSOR.INDEXER_KPOOL_APE: "blk.{bid}.indexer.kpool_ape", + MODEL_TENSOR.INDEXER_KPOOL_GATE: "blk.{bid}.indexer.kpool_gate", # vision MODEL_TENSOR.V_MMPROJ: "mm.{bid}", MODEL_TENSOR.V_MMPROJ_FC: "mm.model.fc", @@ -3928,6 +3935,15 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.HC_FFN_FN, MODEL_TENSOR.HC_FFN_BASE, MODEL_TENSOR.HC_FFN_SCALE, + MODEL_TENSOR.INDEXER_KPOOL_APE, + MODEL_TENSOR.INDEXER_KPOOL_GATE, + # The MTP block is a real layer in the checkpoint and its weights ship in the file. + # llama.cpp loads and then ignores them (see llama-model.cpp, "preserved but unused"), + # so this buys forward compatibility, not a working speculative decoder today. + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], # TODO } diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 90d500dc7..65360a41e 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -790,6 +790,21 @@ class GGUFWriter: def add_indexer_top_k(self, top_k: int) -> None: self.add_uint32(Keys.Attention.Indexer.TOP_K.format(arch=self.arch), top_k) + def add_swiglu_limit(self, value: float) -> None: + self.add_float32(Keys.LLM.SWIGLU_LIMIT.format(arch=self.arch), value) + + def add_ssm_gate_lower_bound(self, value: float) -> None: + self.add_float32(Keys.SSM.GATE_LOWER_BOUND.format(arch=self.arch), value) + + def add_hc_mult(self, value: int) -> None: + self.add_uint32(Keys.Attention.HyperConnection.MULT.format(arch=self.arch), value) + + def add_hc_sinkhorn_iters(self, value: int) -> None: + self.add_uint32(Keys.Attention.HyperConnection.SINKHORN_ITERS.format(arch=self.arch), value) + + def add_hc_eps(self, value: float) -> None: + self.add_float32(Keys.Attention.HyperConnection.EPS.format(arch=self.arch), value) + def add_max_alibi_bias(self, bias: float) -> None: self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias) @@ -1178,6 +1193,9 @@ class GGUFWriter: def add_vision_use_gelu(self, value: bool) -> None: self.add_bool(Keys.ClipVision.USE_GELU, value) + def add_vision_swiglu_limit(self, value: float) -> None: + self.add_float32(Keys.ClipVision.SWIGLU_LIMIT, value) + def add_vision_use_silu(self, value: bool) -> None: self.add_bool(Keys.ClipVision.USE_SILU, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 949ed63f1..4439f54db 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2086,6 +2086,16 @@ class TensorNameMap: ), # NextN/MTP tensors + MODEL_TENSOR.INDEXER_KPOOL_APE: ( + "model.layers.{bid}.self_attn.indexer.index_kpool_compress_ape", + "model.language_model.layers.{bid}.self_attn.indexer.index_kpool_compress_ape", + ), + + MODEL_TENSOR.INDEXER_KPOOL_GATE: ( + "model.layers.{bid}.self_attn.indexer.index_kpool_compress_gate", + "model.language_model.layers.{bid}.self_attn.indexer.index_kpool_compress_gate", + ), + MODEL_TENSOR.NEXTN_EH_PROJ: ( "model.layers.{bid}.eh_proj", ), diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 121c21fed..dbac7fdc5 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -75,6 +75,7 @@ add_library(llama models/gemma3n-iswa.cpp models/gemma4-iswa.cpp models/glm4-moe.cpp + models/glm5-next.cpp models/glm4.cpp models/gpt2.cpp models/gptneox.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index e210dcdae..e9761a585 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -132,6 +132,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_LLAMA_EMBED, "llama-embed" }, { LLM_ARCH_MAINCODER, "maincoder" }, { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, + { LLM_ARCH_GLM5_NEXT, "glm5-next" }, { LLM_ARCH_UNKNOWN, "(unknown)" }, }; @@ -239,6 +240,11 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_VALUE_LENGTH_SWA, "%s.attention.value_length_swa" }, { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" }, { LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" }, + { LLM_KV_SWIGLU_LIMIT, "%s.swiglu_limit" }, + { LLM_KV_SSM_GATE_LOWER_BOUND, "%s.ssm.gate_lower_bound" }, + { LLM_KV_ATTENTION_HC_MULT, "%s.attention.hc.mult" }, + { LLM_KV_ATTENTION_HC_SINKHORN_ITERS, "%s.attention.hc.sinkhorn_iters" }, + { LLM_KV_ATTENTION_HC_EPS, "%s.attention.hc.eps" }, { LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" }, { LLM_KV_ATTENTION_SHARED_KV_LAYERS, "%s.attention.shared_kv_layers" }, @@ -545,6 +551,14 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" }, { LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" }, { LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" }, + { LLM_TENSOR_INDEXER_KPOOL_APE, "blk.%d.indexer.kpool_ape" }, + { LLM_TENSOR_INDEXER_KPOOL_GATE, "blk.%d.indexer.kpool_gate" }, + { LLM_TENSOR_HC_ATTN_FN, "blk.%d.hc_attn_fn" }, + { LLM_TENSOR_HC_ATTN_BASE, "blk.%d.hc_attn_base" }, + { LLM_TENSOR_HC_ATTN_SCALE, "blk.%d.hc_attn_scale" }, + { LLM_TENSOR_HC_FFN_FN, "blk.%d.hc_ffn_fn" }, + { LLM_TENSOR_HC_FFN_BASE, "blk.%d.hc_ffn_base" }, + { LLM_TENSOR_HC_FFN_SCALE, "blk.%d.hc_ffn_scale" }, }; static std::set llm_get_tensor_names(llm_arch arch) { @@ -2508,6 +2522,71 @@ static std::set llm_get_tensor_names(llm_arch arch) { LLM_TENSOR_FFN_DOWN, LLM_TENSOR_FFN_UP, }; + case LLM_ARCH_GLM5_NEXT: + return { + LLM_TENSOR_TOKEN_EMBD, + LLM_TENSOR_OUTPUT_NORM, + LLM_TENSOR_OUTPUT, + LLM_TENSOR_ATTN_NORM, + LLM_TENSOR_FFN_NORM, + // KDA linear attention (34 layers). Same delta rule as Kimi-Linear, so the + // SSM_ enum names are reused rather than duplicated. + LLM_TENSOR_ATTN_Q, + LLM_TENSOR_ATTN_K, + LLM_TENSOR_ATTN_V, + LLM_TENSOR_ATTN_OUT, + LLM_TENSOR_SSM_CONV1D_Q, + LLM_TENSOR_SSM_CONV1D_K, + LLM_TENSOR_SSM_CONV1D_V, + LLM_TENSOR_SSM_F_A, + LLM_TENSOR_SSM_F_B, + LLM_TENSOR_SSM_G_A, + LLM_TENSOR_SSM_G_B, + LLM_TENSOR_SSM_BETA, + LLM_TENSOR_SSM_A, + LLM_TENSOR_SSM_DT, + LLM_TENSOR_SSM_NORM, + // MLA + DSA (11 layers, plus the MTP block). NoPE - no ROPE_FREQS here. + LLM_TENSOR_ATTN_Q_A, + LLM_TENSOR_ATTN_Q_B, + LLM_TENSOR_ATTN_Q_A_NORM, + LLM_TENSOR_ATTN_KV_A_MQA, + LLM_TENSOR_ATTN_KV_A_NORM, + LLM_TENSOR_ATTN_KV_B, + LLM_TENSOR_ATTN_K_B, + LLM_TENSOR_ATTN_V_B, + LLM_TENSOR_INDEXER_K_NORM, + LLM_TENSOR_INDEXER_PROJ, + LLM_TENSOR_INDEXER_ATTN_K, + LLM_TENSOR_INDEXER_ATTN_Q_B, + LLM_TENSOR_INDEXER_KPOOL_APE, + LLM_TENSOR_INDEXER_KPOOL_GATE, + // Dense FFN (first 3 layers) + LLM_TENSOR_FFN_GATE, + LLM_TENSOR_FFN_DOWN, + LLM_TENSOR_FFN_UP, + // MoE + LLM_TENSOR_FFN_GATE_INP, + LLM_TENSOR_FFN_GATE_EXPS, + LLM_TENSOR_FFN_DOWN_EXPS, + LLM_TENSOR_FFN_UP_EXPS, + LLM_TENSOR_FFN_EXP_PROBS_B, + LLM_TENSOR_FFN_GATE_SHEXP, + LLM_TENSOR_FFN_DOWN_SHEXP, + LLM_TENSOR_FFN_UP_SHEXP, + // mHC hyper-connections + LLM_TENSOR_HC_ATTN_FN, + LLM_TENSOR_HC_ATTN_BASE, + LLM_TENSOR_HC_ATTN_SCALE, + LLM_TENSOR_HC_FFN_FN, + LLM_TENSOR_HC_FFN_BASE, + LLM_TENSOR_HC_FFN_SCALE, + // MTP block - loaded, not executed (see llama-model.cpp) + LLM_TENSOR_NEXTN_EH_PROJ, + LLM_TENSOR_NEXTN_ENORM, + LLM_TENSOR_NEXTN_HNORM, + LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, + }; case LLM_ARCH_KIMI_LINEAR: return { LLM_TENSOR_TOKEN_EMBD, @@ -2768,14 +2847,30 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, - // NextN/MTP tensors are currently ignored (reserved for future MTP support) - // These tensors only exist in the last layer(s) and are treated as output tensors - {LLM_TENSOR_NEXTN_EH_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, - {LLM_TENSOR_NEXTN_EMBED_TOKENS, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, - {LLM_TENSOR_NEXTN_ENORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, - {LLM_TENSOR_NEXTN_HNORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, - {LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, - {LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_INDEXER_KPOOL_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_INDEXER_KPOOL_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + // mHC. `fn` is a matmul; `base` and `scale` are elementwise parameters of the + // Sinkhorn-normalised mixing, applied inside the fused op. + {LLM_TENSOR_HC_ATTN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_ATTN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, + {LLM_TENSOR_HC_ATTN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_FFN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, + {LLM_TENSOR_HC_FFN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + // NextN/MTP tensors are currently ignored (reserved for future MTP support). + // + // These were LLM_TENSOR_LAYER_OUTPUT, which is wrong: they are created per layer with a + // layer index, and llama_model_loader aborts outright on an input/output tensor used with + // one ("input/output layer tensor %s used with a layer number"). That abort was unreachable + // only because no GGUF in the wild actually carried them - the tensors are all + // TENSOR_NOT_REQUIRED and get stripped. glm5-next ships them, so the classification has to + // be honest: they live in a layer, so they are LAYER_REPEATING. They remain unexecuted. + {LLM_TENSOR_NEXTN_EH_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_NEXTN_EMBED_TOKENS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_NEXTN_ENORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_NEXTN_HNORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, // Nemotron 3 Super {LLM_TENSOR_FFN_LATENT_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_FFN_LATENT_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, @@ -2878,6 +2973,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_GLM5_NEXT: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: return true; diff --git a/src/llama-arch.h b/src/llama-arch.h index 1b8737b74..1f3312e9e 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -136,6 +136,7 @@ enum llm_arch { LLM_ARCH_LLAMA_EMBED, LLM_ARCH_MAINCODER, LLM_ARCH_KIMI_LINEAR, + LLM_ARCH_GLM5_NEXT, LLM_ARCH_UNKNOWN, }; @@ -244,6 +245,11 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, LLM_KV_ATTENTION_INDEXER_TOP_K, + LLM_KV_SWIGLU_LIMIT, + LLM_KV_SSM_GATE_LOWER_BOUND, + LLM_KV_ATTENTION_HC_MULT, + LLM_KV_ATTENTION_HC_SINKHORN_ITERS, + LLM_KV_ATTENTION_HC_EPS, LLM_KV_ATTENTION_SHARED_KV_LAYERS, LLM_KV_ROPE_DIMENSION_COUNT, @@ -546,6 +552,14 @@ enum llm_tensor { LLM_TENSOR_INDEXER_PROJ, LLM_TENSOR_INDEXER_ATTN_K, LLM_TENSOR_INDEXER_ATTN_Q_B, + LLM_TENSOR_INDEXER_KPOOL_APE, + LLM_TENSOR_INDEXER_KPOOL_GATE, + LLM_TENSOR_HC_ATTN_FN, + LLM_TENSOR_HC_ATTN_BASE, + LLM_TENSOR_HC_ATTN_SCALE, + LLM_TENSOR_HC_FFN_FN, + LLM_TENSOR_HC_FFN_BASE, + LLM_TENSOR_HC_FFN_SCALE, LLM_TENSOR_NEXTN_EH_PROJ, LLM_TENSOR_NEXTN_EMBED_TOKENS, LLM_TENSOR_NEXTN_ENORM, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 0e7d96ca1..c3e357a4e 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1527,6 +1527,17 @@ ggml_tensor * llm_graph_context::build_moe_ffn( cur = ggml_swiglu_oai(ctx0, cur, up, alpha, limit); cb(cur, "ffn_moe_swiglu_oai", il); } break; + case LLM_FFN_SWIGLU_CLAMPED: + { + // GLM-5.3. ggml_swiglu_oai clamps identically but computes + // silu_alpha(gate) * (up + 1); GLM has no +1 and alpha = 1, so do it explicitly + // rather than pass alpha=1 and inherit the bias. + const float limit = hparams.swiglu_limit; + cur = ggml_clamp(ctx0, cur, -INFINITY, limit); + up = ggml_clamp(ctx0, up, -limit, limit); + cur = ggml_mul(ctx0, ggml_silu(ctx0, cur), up); + cb(cur, "ffn_moe_swiglu_clamped", il); + } break; case LLM_FFN_RELU: if (has_gate) { cur = ggml_reglu_split(ctx0, cur, up); diff --git a/src/llama-graph.h b/src/llama-graph.h index bb0ad7519..b779a70f6 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -42,6 +42,7 @@ enum llm_ffn_op_type { LLM_FFN_GEGLU, LLM_FFN_REGLU, LLM_FFN_SWIGLU_OAI_MOE, + LLM_FFN_SWIGLU_CLAMPED, // GLM-5.3: clamp then SiLU, no +1 on up }; enum llm_ffn_gate_type { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index c2000c77c..6ae19534e 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -206,6 +206,18 @@ struct llama_hparams { uint32_t indexer_head_size = 0; uint32_t indexer_top_k = 0; + // mHC hyper-connections (GLM-5.3-Flash). hc_mult residual streams per layer, mixed by a + // Sinkhorn-normalised matrix. Note the normalisation is COLUMN-stochastic, not doubly + // stochastic - it column-normalises once and then runs (iters-1) full row+column passes. + // GLM-5.3 KDA forget gate: g = bound * sigmoid(exp(A_log) * (w + dt_bias)). + float ssm_gate_lower_bound = 0.0f; + // Clamped SwiGLU (GLM-5.3): clamp(gate, max=L) and clamp(up, -L, L) BEFORE SiLU. + float swiglu_limit = 0.0f; + + uint32_t hc_mult = 0; + uint32_t hc_sinkhorn_iters = 0; + float hc_eps = 0.0f; + // qwen3vl deepstack uint32_t n_deepstack_layers = 0; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 79d08ff41..1ecb8f6d1 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2527,6 +2527,61 @@ void llama_model::load_hparams(llama_model_loader & ml) { default: type = LLM_TYPE_UNKNOWN; } } break; + case LLM_ARCH_GLM5_NEXT: + { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + // MLA. GLM-5.3 uses NoPE on the full-attention path: qk_rope_head_dim is 0 and + // the converter writes rope_dimension_count = 0, so there is no rotary section + // to carve out of the query. get_rope_type() returns NONE for this arch. + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + + // 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); + ml.get_key(LLM_KV_SSM_GATE_LOWER_BOUND, hparams.ssm_gate_lower_bound); + ml.get_key(LLM_KV_SWIGLU_LIMIT, hparams.swiglu_limit, false); + + // A layer is recurrent iff the converter wrote n_head_kv == 0 for it. The + // converter derives that from linear_attn_config.full_attn_layers, which is + // ZERO-indexed in GLM-5.3 (it is one-indexed in Kimi-Linear). + for (uint32_t i = 0; i < hparams.n_layer; ++i) { + hparams.recurrent_layer_arr[i] = hparams.n_head_kv(i) == 0; + } + + // MoE + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); + + // DSA indexer: loaded for completeness, not executed. llama.cpp runs these + // layers as dense MLA, the same way LLM_ARCH_GLM_DSA already does. + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k, false); + + // mHC + ml.get_key(LLM_KV_ATTENTION_HC_MULT, hparams.hc_mult); + ml.get_key(LLM_KV_ATTENTION_HC_SINKHORN_ITERS, hparams.hc_sinkhorn_iters); + ml.get_key(LLM_KV_ATTENTION_HC_EPS, hparams.hc_eps); + + // The MTP block is a real layer in the file. Exclude it from the transformer + // stack and from the KV cache; its weights are loaded and never executed. + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.nextn_predict_layers, false); + GGML_ASSERT(hparams.nextn_predict_layers < hparams.n_layer); + hparams.n_layer_kv_from_start = hparams.n_layer - hparams.nextn_predict_layers; + + switch (hparams.n_layer) { + case 46: type = LLM_TYPE_A13B; break; // GLM-5.3-Flash REAP-50 + default: type = LLM_TYPE_UNKNOWN; + } + } break; case LLM_ARCH_KIMI_LINEAR: { ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); @@ -7077,6 +7132,127 @@ bool llama_model::load_tensors(llama_model_loader & ml) { layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); } } break; + case LLM_ARCH_GLM5_NEXT: + { + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + const int64_t hc = hparams.hc_mult; + const int64_t n_hc_mix = (2 + hc) * hc; // pre | post | comb logits + const int64_t kda_head = hparams.n_embd_head_kda; + const int64_t kda_inner = kda_head * n_head; + const int64_t ssm_d_conv = hparams.ssm_d_conv; + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_mtp_start = n_layer - hparams.nextn_predict_layers; + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + // mHC is on the TRANSFORMER layers only. The MTP block has none: it + // consumes the already-collapsed hidden state, so there are no parallel + // streams left to mix. Creating them for layer 45 as required tensors is + // what "missing tensor 'blk.45.hc_attn_fn'" was. + if (i < (int) n_mtp_start) { + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, i), {hc*n_embd, n_hc_mix}, 0); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, i), {n_hc_mix}, 0); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, i), {3}, 0); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, i), {hc*n_embd, n_hc_mix}, 0); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, i), {n_hc_mix}, 0); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, i), {3}, 0); + } + + if (hparams.is_recurrent(i)) { + // KDA. Conv weights are written 4D; a quantised file may drop the + // trailing 1, so accept 3D as well. + layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", i), {ssm_d_conv, 1, kda_inner, 1}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_q_conv) layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", i), {ssm_d_conv, 1, kda_inner}, 0); + layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", i), {ssm_d_conv, 1, kda_inner, 1}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_k_conv) layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", i), {ssm_d_conv, 1, kda_inner}, 0); + layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", i), {ssm_d_conv, 1, kda_inner, 1}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_v_conv) layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", i), {ssm_d_conv, 1, kda_inner}, 0); + + layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, kda_inner}, 0); + layer.wk = create_tensor(tn(LLM_TENSOR_ATTN_K, "weight", i), {n_embd, kda_inner}, 0); + layer.wv = create_tensor(tn(LLM_TENSOR_ATTN_V, "weight", i), {n_embd, kda_inner}, 0); + + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, kda_head}, 0); + layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {kda_head, kda_inner}, 0); + layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", i), {n_embd, kda_head}, 0); + layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", i), {kda_head, kda_inner}, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {kda_inner}, 0); + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {kda_head}, 0); + + // GLM-5.3 stores A_log as a bare [n_head] vector; Kimi-Linear's is + // [1, n_head, 1, 1]. Accept either so one graph serves both. + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_a) layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head, 1, 1}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_a) layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head}, 0); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {kda_inner, n_embd}, 0); + } else { + // MLA + DSA. NoPE: qk_rope_head_dim is 0, so kv_a carries only the + // compressed KV rank and no rotary tail is split out of the query. + const int64_t q_lora = hparams.n_lora_q; + const int64_t kv_lora = hparams.n_lora_kv; + const int64_t hk = hparams.n_embd_head_k_mla(); + const int64_t hv = hparams.n_embd_head_v_mla(); + + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora}, 0); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora}, 0); + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora}, 0); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora, n_head * hk}, 0); + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora}, 0); + + layer.wkv_b = create_tensor(tn(LLM_TENSOR_ATTN_KV_B, "weight", i), + {kv_lora, n_head * (hk + hv)}, TENSOR_NOT_REQUIRED | TENSOR_SKIP_IF_VIRTUAL); + if (!layer.wkv_b) { + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {hk, kv_lora, n_head}, 0); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora, hv, n_head}, 0); + } + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * hv, n_embd}, 0); + + // DSA indexer: loaded so the file round-trips, never executed. + const int64_t idx_h = hparams.indexer_head_size; + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {idx_h}, TENSOR_NOT_REQUIRED); + layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {idx_h}, TENSOR_NOT_REQUIRED); + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, (int64_t) hparams.indexer_n_head}, TENSOR_NOT_REQUIRED); + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, idx_h}, TENSOR_NOT_REQUIRED); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora, (int64_t) hparams.indexer_n_head * idx_h}, TENSOR_NOT_REQUIRED); + layer.indexer_kpool_ape = create_tensor(tn(LLM_TENSOR_INDEXER_KPOOL_APE, i), {idx_h, 4}, TENSOR_NOT_REQUIRED); + layer.indexer_kpool_gate = create_tensor(tn(LLM_TENSOR_INDEXER_KPOOL_GATE, i), {n_embd, idx_h}, TENSOR_NOT_REQUIRED); + } + + if (i < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, 0); + + const int64_t n_ff_shexp = n_ff_exp * (hparams.n_expert_shared > 0 ? hparams.n_expert_shared : 1); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_shexp, n_embd}, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_shexp}, 0); + } + + // MTP block: loaded, never executed. See llm_arch notes. + if (i >= (int) n_mtp_start) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, 0); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, 0); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, 0); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, TENSOR_NOT_REQUIRED); + } + } + } break; case LLM_ARCH_KIMI_LINEAR: { tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); @@ -8901,6 +9077,10 @@ ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const { { llm = std::make_unique(*this, params); } break; + case LLM_ARCH_GLM5_NEXT: + { + llm = std::make_unique(*this, params); + } break; case LLM_ARCH_KIMI_LINEAR: { llm = std::make_unique(*this, params); @@ -9058,6 +9238,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_WAVTOKENIZER_DEC: case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: + case LLM_ARCH_GLM5_NEXT: // NoPE: qk_rope_head_dim == 0 on the MLA path case LLM_ARCH_KIMI_LINEAR: return LLAMA_ROPE_TYPE_NONE; diff --git a/src/llama-model.h b/src/llama-model.h index 4f1100839..9d90c8ea0 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -486,6 +486,18 @@ struct llama_layer { struct ggml_tensor * indexer_proj = nullptr; struct ggml_tensor * indexer_attn_k = nullptr; struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias + struct ggml_tensor * indexer_kpool_ape = nullptr; + struct ggml_tensor * indexer_kpool_gate = nullptr; + + // mHC hyper-connections (GLM-5.3-Flash). Two sites per layer; each carries a projection + // `fn` producing (2+H)*H mixing logits, a `base` bias over the same, and 3 scalars that + // scale the pre/post/comb logit groups independently. + struct ggml_tensor * hc_attn_fn = nullptr; + struct ggml_tensor * hc_attn_base = nullptr; + struct ggml_tensor * hc_attn_scale = nullptr; + struct ggml_tensor * hc_ffn_fn = nullptr; + struct ggml_tensor * hc_ffn_base = nullptr; + struct ggml_tensor * hc_ffn_scale = nullptr; // gemma4 layer output scale struct ggml_tensor * out_scale = nullptr; diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp new file mode 100644 index 000000000..3f1ca0ec2 --- /dev/null +++ b/src/models/glm5-next.cpp @@ -0,0 +1,433 @@ +#include "models.h" + +#include "llama-memory-recurrent.h" + +// GLM-5.3-Flash (glm5-next). +// +// The transformer body is Kimi-Linear's: 34 KDA linear-attention layers interleaved 3:1 with +// 11 NoPE MLA layers (plus the MTP block, which is loaded and not executed), sigmoid-routed MoE +// with a shared expert. That part is a near-transcription of src/models/kimi-linear.cpp. +// +// What is new is mHC. Instead of one residual stream there are `hc_mult` of them, and each of +// the two sites per layer (attention, FFN) does: +// +// residual = streams +// post, comb, collapsed = mHC(streams) +// y = sublayer(norm(collapsed)) +// streams = post (x) y + comb^T @ residual +// +// `comb` is column-stochastic, produced by ggml_mhc_sinkhorn (see ggml.h for why the +// normalisation order matters and why it is fused). `post` is 2*sigmoid, range [0,2] - it is +// not a probability and must not be clamped to [0,1]. +// +// Streams are carried as [n_embd, hc, n_tokens]: that makes the mHC input flatten a free +// reshape, and the one permutation per site is shared between the collapse and the mixing +// matmul. + +// Causal Conv1d function for Q,K,V +// When qkv is 0, it is Q, 1 is K, 2 is V +static ggml_tensor * causal_conv1d(ggml_cgraph * gf, ggml_context * ctx0, ggml_tensor * conv_states_all, ggml_tensor * conv_state_all, int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, int64_t d_conv, int64_t head_dim, int64_t n_head, int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t n_embd_r_total = 3 * conv_state_size; // Q + K + V + + // conv_state_all is [n_embd_r_total, n_seqs], split into Q, K, V + // Each conv state is [(d_conv-1) * d_inner] per sequence, need to reshape to [d_conv-1, d_inner, n_seqs] + // Memory layout: for each seq, Q state is first conv_state_size elements, then K, then V + // conv_state_all has stride: nb[0] = element_size, nb[1] = n_embd_r_total * element_size + // View Q conv state: offset 0, size conv_state_size per seq + // conv_state_all is [n_embd_r_total, n_seqs] with memory layout: + // state[i + seq * n_embd_r_total] where i = conv_step + channel * (d_conv-1) + {0, conv_state_size, 2*conv_state_size} for Q/K/V + // We want [d_conv-1, d_inner, n_seqs] view: + // nb1 = (d_conv-1) * element_size (stride between channels) + // nb2 = n_embd_r_total * element_size (stride between seqs) + ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), // nb1: stride between channels + n_embd_r_total * ggml_element_size(conv_state_all), // nb2: stride between seqs + qkv * conv_state_size * ggml_element_size(conv_state_all)); + +// Causal Conv1d function for Q,K,V +// When qkv is 0, it is Q, 1 is K, 2 is V + // Step 1: Q, K, V projections -> [d_inner, n_tokens] + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + + // Reshape input: {d_inner, n_tokens} -> {d_inner, n_seq_tokens, n_seqs} + ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + + // Concat Q conv state and current input: {d_conv-1 + n_seq_tokens, d_inner, n_seqs} + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0); + + // Save last (d_conv-1) columns back to Q conv state + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, last_conv_x, + ggml_view_3d(ctx0, conv_states_all, + d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), // nb1: contiguous within one channel's conv taps + n_embd_r_total * ggml_element_size(conv_states_all), // nb2: stride between sequences (skip over K,V states) + (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); // offset to first seq's Q/K/V state + // Reshape conv weight: GGUF [d_conv, 1, d_inner, 1] -> ggml_ssm_conv expects [d_conv, d_inner] + // GGUF stores as [d_conv, 1, d_inner, 1] with memory layout w[conv_step + channel * d_conv] + // vLLM stores as [d_inner, d_conv] with memory layout w[channel * d_conv + conv_step] + // ggml_ssm_conv computes: c[conv_step + channel * d_conv] + // GGUF layout: [d_conv, 1, d_inner] or [d_conv, 1, d_inner, 1] -> reshape to [d_conv, d_inner] + // Reshape conv weight from [d_conv, 1, d_inner, 1] to [d_conv, d_inner] for ggml_ssm_conv + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + + // Apply conv1d + // ggml_ssm_conv output: {d_inner, n_seq_tokens, n_seqs} + ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight); + // Reshape to 2D for bias add: {d_inner, n_tokens} + Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens); + Xcur = ggml_silu(ctx0, Xcur); + + return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs); +} + + +llm_build_glm5_next::mhc_site llm_build_glm5_next::build_mhc( + ggml_tensor * streams, ggml_tensor * fn, ggml_tensor * base, + ggml_tensor * scale, int il) { + const int64_t hc = hparams.hc_mult; + const int64_t n_tokens = streams->ne[2]; + const float eps = hparams.hc_eps; + + // The mixing weights come from a (2+hc)*hc vector; the reference computes this whole path + // in F32 even when the streams are BF16, because rounding before the softmax visibly moves + // `comb`. ggml activations are already F32, so nothing extra is needed here - but do not + // "optimise" this to a lower precision. + ggml_tensor * flat = ggml_reshape_2d(ctx0, streams, n_embd*hc, n_tokens); + flat = ggml_rms_norm(ctx0, flat, hparams.f_norm_rms_eps); // unweighted, no gain tensor + cb(flat, "mhc_flat", il); + + ggml_tensor * mix = ggml_mul_mat(ctx0, fn, flat); // [(2+hc)*hc, n_tokens] + cb(mix, "mhc_mix", il); + + // Views along ne0 of a contiguous tensor are row-strided, so make them contiguous before + // anything reshapes or broadcasts them. 24 floats per token - not worth being clever about. + ggml_tensor * pre_w = ggml_cont(ctx0, ggml_view_2d(ctx0, mix, hc, n_tokens, mix->nb[1], 0)); + ggml_tensor * post_w = ggml_cont(ctx0, ggml_view_2d(ctx0, mix, hc, n_tokens, mix->nb[1], hc*mix->nb[0])); + ggml_tensor * comb_w = ggml_cont(ctx0, ggml_view_2d(ctx0, mix, hc*hc, n_tokens, mix->nb[1], 2*hc*mix->nb[0])); + + ggml_tensor * pre_b = ggml_view_1d(ctx0, base, hc, 0); + ggml_tensor * post_b = ggml_view_1d(ctx0, base, hc, hc*base->nb[0]); + ggml_tensor * comb_b = ggml_view_1d(ctx0, base, hc*hc, 2*hc*base->nb[0]); + + ggml_tensor * s_pre = ggml_view_1d(ctx0, scale, 1, 0); + ggml_tensor * s_post = ggml_view_1d(ctx0, scale, 1, scale->nb[0]); + ggml_tensor * s_comb = ggml_view_1d(ctx0, scale, 1, 2*scale->nb[0]); + + // pre = sigmoid(w*s + b) + eps + ggml_tensor * pre = ggml_sigmoid(ctx0, ggml_add(ctx0, ggml_mul(ctx0, pre_w, s_pre), pre_b)); + pre = ggml_scale_bias(ctx0, pre, 1.0f, eps); + + // post = 2*sigmoid(w*s + b). Range [0,2]; NOT a probability. + ggml_tensor * post = ggml_sigmoid(ctx0, ggml_add(ctx0, ggml_mul(ctx0, post_w, s_post), post_b)); + post = ggml_scale(ctx0, post, 2.0f); + cb(post, "mhc_post", il); + + ggml_tensor * comb_l = ggml_reshape_3d(ctx0, comb_w, hc, hc, n_tokens); + comb_l = ggml_add(ctx0, ggml_mul(ctx0, comb_l, s_comb), + ggml_reshape_3d(ctx0, comb_b, hc, hc, 1)); + ggml_tensor * comb = ggml_mhc_sinkhorn(ctx0, comb_l, hparams.hc_sinkhorn_iters, eps); + cb(comb, "mhc_comb", il); + + // One permutation, used twice: for the pre-weighted collapse and for comb^T @ streams. + ggml_tensor * streams_p = ggml_cont(ctx0, ggml_permute(ctx0, streams, 1, 0, 2, 3)); // [hc, n_embd, n_tokens] + + ggml_tensor * weighted = ggml_mul(ctx0, streams_p, ggml_reshape_3d(ctx0, pre, hc, 1, n_tokens)); + ggml_tensor * collapsed = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, weighted), n_embd, n_tokens); + cb(collapsed, "mhc_collapsed", il); + + return { post, comb, collapsed, streams_p }; +} + +ggml_tensor * llm_build_glm5_next::apply_mhc(const mhc_site & s, ggml_tensor * y, int il) { + const int64_t hc = hparams.hc_mult; + const int64_t n_tokens = y->ne[1]; + + // comb^T @ streams : each OUTPUT stream is a convex combination of the input streams. + // + // Index care, because getting this backwards is invisible: ggml's ne0 is torch's LAST dim, + // so comb_ggml[n0, n1] == comb_torch[n1, n0]. The reference computes + // out[i, d] = sum_j comb_torch[j, i] * streams[j, d] + // and `comb` is COLUMN-stochastic (sum_j comb_torch[j, i] == 1), so comb^T is row-stochastic + // and the update is a convex combination. Contracting comb_ggml's ne0 directly would use + // comb_torch[i, j] - the row sums, which are 0.98-1.02, not 1 - and produce a model that is + // wrong by a few percent everywhere. Transposing first contracts the correct index. + ggml_tensor * mixed = ggml_mul_mat(ctx0, s.streams_p, ggml_cont(ctx0, ggml_transpose(ctx0, s.comb))); + cb(mixed, "mhc_mixed", il); + + ggml_tensor * y3 = ggml_repeat(ctx0, ggml_reshape_3d(ctx0, y, n_embd, 1, n_tokens), mixed); + ggml_tensor * scaled = ggml_mul(ctx0, y3, ggml_reshape_3d(ctx0, s.post, 1, hc, n_tokens)); + + return ggml_add(ctx0, scaled, mixed); +} + +llm_build_glm5_next::llm_build_glm5_next(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + ggml_tensor * cur; + + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + cb(inpL, "model.embed_tokens", -1); + + const int64_t hc = hparams.hc_mult; + + // NoPE throughout - the MLA path has qk_rope_head_dim == 0, so there is no inp_pos. + auto * inp_kv = !hparams.is_mla() ? build_inp_mem_hybrid() : nullptr; + auto * inp_k = hparams.is_mla() ? build_inp_mem_hybrid_k() : nullptr; + auto * inp_rs = hparams.is_mla() ? inp_k->get_recr() : inp_kv->get_recr(); + auto * inp_attn_kv = !hparams.is_mla() ? inp_kv->get_attn() : nullptr; + auto * inp_attn_k = hparams.is_mla() ? inp_k->get_attn() : nullptr; + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = n_head * head_dim; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + const float kq_scale_mla = 1.0f / sqrtf((float) n_embd_head_k_mla); + + // The MTP block is a real layer in the file but is not part of the forward pass. + const int n_transformer_layers = n_layer - hparams.nextn_predict_layers; + + // Streams start as hc copies of the embedding (reference: inputs_embeds.unsqueeze(2).expand). + ggml_tensor * streams = ggml_repeat(ctx0, + ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), + ggml_new_tensor_3d(ctx0, inpL->type, n_embd, hc, n_tokens)); + cb(streams, "mhc_streams_init", -1); + + for (int il = 0; il < n_transformer_layers; ++il) { + const auto & layer = model.layers[il]; + + // ---------------- attention site ---------------- + mhc_site site = build_mhc(streams, layer.hc_attn_fn, layer.hc_attn_base, layer.hc_attn_scale, il); + + cur = build_norm(site.collapsed, layer.attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recurrent(il)) { + // === KDA Layer (Kimi Delta Attention) with Recurrent State === + // Reference: vLLM kda.py + const auto * mctx_cur = inp_rs->mctx; + const auto kv_head = mctx_cur->get_head(); + + // Get conv states from r_l tensor (Q, K, V each have separate state) + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + cb(conv_states_all, "conv_states_all", il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + ggml_tensor * Qcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Kcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Vcur = causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head, n_seq_tokens, n_seqs, n_tokens, kv_head); + + // g1 = -exp(A_log) * softplus(f_b(f_a(x)) + dt_bias) + ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a); + cb(g1, "g1 f_b(f_a(cur))", il); + g1 = ggml_add(ctx0, g1, layer.ssm_dt_b); + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head, n_tokens); + + // A_log shape is [1, n_head] or [1, n_head, 1, 1], need to broadcast to [head_dim, n_head, n_tokens]. No need to -exp(a_log) because it was done in convert_hf_to_gguf.py + // Reshape to [1, n_head, 1] for broadcasting with g1 [head_dim, n_head, n_tokens] + // GLM-5.3 forget gate: g = bound * sigmoid(exp(A_log) * (w + dt_bias)). + // + // Kimi-Linear's is g = -exp(A_log) * softplus(w + dt_bias) - a different function + // with a different sign convention, and the converter stores exp(A_log) rather than + // -exp(A_log) to match. Using Kimi's form here is not a small error: it is bounded + // vs unbounded decay, so the discrepancy compounds along the sequence and shows up + // as logit error that grows monotonically with position. + ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1); + g1 = ggml_sigmoid(ctx0, ggml_mul(ctx0, g1, A)); + g1 = ggml_scale(ctx0, g1, hparams.ssm_gate_lower_bound); + cb(g1, "kda_g1", il); + + g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head, n_seq_tokens, n_seqs); + + // Compute beta (mixing coefficient) + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs); + cb(beta, "kda_beta", il); + + beta = ggml_sigmoid(ctx0, beta); + + // Reshape for KDA recurrence + // {n_embd, n_tokens} -> {n_embd, n_seq_tokens, n_seqs} + cur = ggml_reshape_3d(ctx0, cur, cur->ne[0], n_seq_tokens, n_seqs); + + // Get SSM state and compute KDA recurrence using ggml_kda_scan + 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); + + // The Q/K L2 norm inside the delta rule uses 1e-6, NOT rms_norm_eps: transformers' + // l2norm() defaults to eps=1e-6 to match FLA, while rms_norm_eps here is 1e-5. + // Kimi-Linear's builder passes f_norm_rms_eps; copying that leaves a small flat + // error on every KDA layer. + Qcur = ggml_l2_norm(ctx0, Qcur, 1e-6f); + Kcur = ggml_l2_norm(ctx0, Kcur, 1e-6f); + + // Choose between build_delta_net_chunking and build_delta_net_recurrent based on n_tokens + auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); + + ggml_tensor * output = ggml_cont(ctx0, attn_out.first); + ggml_tensor * new_state = attn_out.second; + cb(output, "attn_output", il); + cb(new_state, "new_state", il); + + // Update the recurrent states + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, new_state, + ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); + + // Output gating g2 = g_b(g_a(x)) + ggml_tensor * cur_2d = ggml_reshape_2d(ctx0, cur, cur->ne[0], n_seq_tokens * n_seqs); + ggml_tensor * g_a = ggml_mul_mat(ctx0, layer.ssm_g_a, cur_2d); + ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g_b, g_a); + cb(g2, "g2 g_b(g_a(cur_2d))", il); + g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head, n_seq_tokens * n_seqs); + + // Apply o_norm with sigmoid gating + // Note: Kimi model uses sigmoid gating, not SiLU (despite FusedRMSNormGated default being swish) + // Formula: output = RMSNorm(x) * sigmoid(g) + ggml_tensor * attn_out_final = ggml_reshape_3d(ctx0, output, head_dim, n_head, n_seq_tokens * n_seqs); + ggml_tensor * normed = build_norm(attn_out_final, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + cb(normed, "kda_normed", il); + ggml_tensor * gate = ggml_sigmoid(ctx0, g2); + ggml_tensor * gated = ggml_mul(ctx0, normed, gate); + + // Output projection + gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens); + cur = ggml_mul_mat(ctx0, layer.wo, gated); + cb(cur, "kda_out", il); + + } else { + // NoPE MLA. qk_rope_head_dim == 0, so there is no rotary tail to split off the + // query and no k_pe to concatenate - the compressed KV is the whole key. + ggml_tensor * q_a = ggml_mul_mat(ctx0, layer.wq_a, cur); + q_a = build_norm(q_a, layer.attn_q_a_norm, NULL, LLM_NORM_RMS, il); + ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.wq_b, q_a); + + ggml_tensor * kv_cmpr = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, NULL, LLM_NORM_RMS, il); + + if (layer.wk_b && layer.wv_b) { + ggml_tensor * q_nope = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens); + q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3); // [hk, T, n_head] + ggml_tensor * q_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope); + q_absorbed = ggml_permute(ctx0, q_absorbed, 0, 2, 1, 3); // [kv_lora, n_head, T] + Qcur = ggml_cont(ctx0, q_absorbed); + + ggml_tensor * Kcur = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + ggml_tensor * Vcur = Kcur; + + cur = build_attn(inp_attn_k, layer.wo, NULL, Qcur, Kcur, Vcur, + nullptr, nullptr, layer.wv_b, kq_scale_mla, il); + } else { + Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens); + ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr); + const int64_t kv_per_head = n_embd_head_k_mla + n_embd_head_v_mla; + + ggml_tensor * Kcur = ggml_view_3d(ctx0, kv, n_embd_head_k_mla, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), + ggml_row_size(kv->type, kv_per_head * n_head), 0); + ggml_tensor * Vcur = ggml_view_3d(ctx0, kv, n_embd_head_v_mla, n_head, n_tokens, + ggml_row_size(kv->type, kv_per_head), + ggml_row_size(kv->type, kv_per_head * n_head), + ggml_row_size(kv->type, n_embd_head_k_mla)); + Kcur = ggml_cont(ctx0, Kcur); + Vcur = ggml_cont(ctx0, Vcur); + + cur = build_attn(inp_attn_kv, layer.wo, NULL, Qcur, Kcur, Vcur, + nullptr, nullptr, nullptr, kq_scale_mla, il); + } + cb(cur, "mla_out", il); + } + + streams = apply_mhc(site, cur, il); + cb(streams, "mhc_after_attn", il); + + // ---------------- FFN site ---------------- + site = build_mhc(streams, layer.hc_ffn_fn, layer.hc_ffn_base, layer.hc_ffn_scale, il); + + cur = build_norm(site.collapsed, layer.ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + // Clamped SwiGLU, spelled out: build_ffn's fused SiLU path has no limit, and the + // clamp only bites once activations exceed it - so a plain SILU dense FFN matches + // on a small test and diverges on the real model. + const float limit = hparams.swiglu_limit; + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ffn_gate, cur); + ggml_tensor * u = ggml_mul_mat(ctx0, layer.ffn_up, cur); + g = ggml_clamp(ctx0, g, -INFINITY, limit); + u = ggml_clamp(ctx0, u, -limit, limit); + cur = ggml_mul(ctx0, ggml_silu(ctx0, g), u); + cur = ggml_mul_mat(ctx0, layer.ffn_down, cur); + } else { + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + hparams.n_expert, hparams.n_expert_used, + LLM_FFN_SWIGLU_CLAMPED, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + cb(moe_out, "ffn_moe_out", il); + + const float limit = hparams.swiglu_limit; + ggml_tensor * sg = ggml_clamp(ctx0, ggml_mul_mat(ctx0, layer.ffn_gate_shexp, cur), -INFINITY, limit); + ggml_tensor * su = ggml_clamp(ctx0, ggml_mul_mat(ctx0, layer.ffn_up_shexp, cur), -limit, limit); + ggml_tensor * shexp = ggml_mul_mat(ctx0, layer.ffn_down_shexp, + ggml_mul(ctx0, ggml_silu(ctx0, sg), su)); + cur = ggml_add(ctx0, moe_out, shexp); + } + cb(cur, "ffn_out", il); + + cur = build_cvec(cur, il); + + streams = apply_mhc(site, cur, il); + cb(streams, "l_out", il); + } + + // Final collapse is an unweighted mean over the streams (reference: + // Glm5NextTextHyperConnectionOutput.forward -> hidden_streams.mean(dim=2)). + { + ggml_tensor * sp = ggml_cont(ctx0, ggml_permute(ctx0, streams, 1, 0, 2, 3)); // [hc, n_embd, T] + cur = ggml_reshape_2d(ctx0, ggml_sum_rows(ctx0, sp), n_embd, n_tokens); + cur = ggml_scale(ctx0, cur, 1.0f/(float) hc); + cb(cur, "mhc_collapse_out", -1); + } + + // Token selection happens after the collapse rather than inside the last layer: the streams + // are 3D and get_rows would have to index ne2. The saving that matters - not running the + // vocab projection for non-output tokens - is preserved either way. + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cur = build_norm(cur, model.output_norm, NULL, 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); +} diff --git a/src/models/models.h b/src/models/models.h index 8e6b9c238..e0890eda7 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -359,6 +359,29 @@ struct llm_build_jamba : public llm_build_mamba_base { llm_build_jamba(const llama_model & model, const llm_graph_params & params); }; +// GLM-5.3-Flash. Structurally Kimi-Linear (KDA + NoPE MLA + sigmoid-router MoE) with mHC +// hyper-connections replacing the plain residual: hc_mult parallel streams, mixed at each of +// the two sites per layer by a Sinkhorn-normalised matrix. +struct llm_build_glm5_next : public llm_build_delta_net_base { + llm_build_glm5_next(const llama_model & model, const llm_graph_params & params); + + // One mHC site. `streams` is [n_embd, hc, n_tokens]. + struct mhc_site { + ggml_tensor * post; // [hc, n_tokens] scales the sublayer output, range [0,2] + ggml_tensor * comb; // [hc, hc, n_tokens] column-stochastic mixing matrix + ggml_tensor * collapsed; // [n_embd, n_tokens] input to the sublayer + ggml_tensor * streams_p; // [hc, n_embd, n_tokens] permuted streams, reused by the update + }; + + mhc_site build_mhc(ggml_tensor * streams, ggml_tensor * fn, ggml_tensor * base, + ggml_tensor * scale, int il); + + // streams' = post (x) y + comb^T @ streams + ggml_tensor * apply_mhc(const mhc_site & s, ggml_tensor * y, int il); + + const llama_model & model; +}; + struct llm_build_kimi_linear : public llm_build_delta_net_base { llm_build_kimi_linear(const llama_model & model, const llm_graph_params & params); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 5e87c8b34..a78327769 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -149,6 +149,8 @@ endif () if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) llama_build_and_test(test-sampling.cpp) +llama_build_and_test(test-mhc-sinkhorn.cpp) +llama_build(test-glm5-next-logits.cpp) llama_build_and_test(test-reasoning-budget.cpp) llama_build_and_test(test-grammar-parser.cpp) llama_build_and_test(test-grammar-integration.cpp) diff --git a/tests/test-glm5-next-logits.cpp b/tests/test-glm5-next-logits.cpp new file mode 100644 index 000000000..9168ca8b3 --- /dev/null +++ b/tests/test-glm5-next-logits.cpp @@ -0,0 +1,103 @@ +// End-to-end logits check for the glm5-next port against transformers. +// +// This is the gate before spending 165 GiB converting the real checkpoint. The tiny fixture has +// the same STRUCTURE - both attention types, dense and MoE FFN, a shared expert, an MTP block +// that must load and not execute, mHC at every site - so the things that fail silently here are +// the same things that would fail silently at scale: Sinkhorn order, a transposed mHC mixing +// matmul, KDA/MLA layer types off by one, NoPE handling. +#include "llama.h" + +#include +#include +#include +#include +#include + +int main(int argc, char ** argv) { + if (argc < 3) { fprintf(stderr, "usage: %s model.gguf reference.bin\n", argv[0]); return 1; } + + FILE * f = fopen(argv[2], "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", argv[2]); return 1; } + int32_t n_tok = 0, n_vocab_ref = 0; + if (fread(&n_tok, 4, 1, f) != 1 || fread(&n_vocab_ref, 4, 1, f) != 1) return 1; + std::vector ids(n_tok); + if (fread(ids.data(), 4, n_tok, f) != (size_t) n_tok) return 1; + std::vector ref((size_t) n_tok * n_vocab_ref); + if (fread(ref.data(), 4, ref.size(), f) != ref.size()) return 1; + fclose(f); + + llama_backend_init(); + + llama_model_params mp = llama_model_default_params(); + // CPU by default (the reference is F32 and exact). Set GLM5_TEST_NGL to push layers onto + // the GPU - that is what actually exercises ggml_cuda_op_mhc_sinkhorn, and running the same + // fixture both ways is the CPU-vs-CUDA equivalence check. + const char * ngl = getenv("GLM5_TEST_NGL"); + mp.n_gpu_layers = ngl ? atoi(ngl) : 0; + printf("n_gpu_layers = %d\n", mp.n_gpu_layers); + llama_model * model = llama_model_load_from_file(argv[1], mp); + if (!model) { fprintf(stderr, "failed to load %s\n", argv[1]); return 1; } + + llama_context_params cp = llama_context_default_params(); + cp.n_ctx = 512; + cp.n_batch = 512; + cp.n_ubatch = 512; + llama_context * ctx = llama_init_from_model(model, cp); + if (!ctx) { fprintf(stderr, "failed to create context\n"); return 1; } + + const int n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + if (n_vocab != n_vocab_ref) { + fprintf(stderr, "vocab mismatch: gguf %d vs reference %d\n", n_vocab, n_vocab_ref); + return 1; + } + + llama_batch batch = llama_batch_init(n_tok, 0, 1); + for (int i = 0; i < n_tok; ++i) { + batch.token[i] = ids[i]; + batch.pos[i] = i; + batch.n_seq_id[i] = 1; + batch.seq_id[i][0] = 0; + batch.logits[i] = 1; + } + batch.n_tokens = n_tok; + if (llama_decode(ctx, batch) != 0) { fprintf(stderr, "decode failed\n"); return 1; } + + // Compare every position. A port that is right at position 0 and wrong later is exactly what + // a broken recurrent state or a mis-shaped mask looks like, so do not only check the last. + double worst_rel = 0.0; + int worst_pos = -1, top1_mismatch = 0; + for (int i = 0; i < n_tok; ++i) { + const float * got = llama_get_logits_ith(ctx, i); + const float * want = ref.data() + (size_t) i * n_vocab; + double num = 0.0, den = 0.0; + int gi = 0, wi = 0; + for (int v = 0; v < n_vocab; ++v) { + const double d = (double) got[v] - want[v]; + num += d*d; den += (double) want[v]*want[v]; + if (got[v] > got[gi]) gi = v; + if (want[v] > want[wi]) wi = v; + } + const double rel = std::sqrt(num/std::max(den, 1e-30)); + if (rel > worst_rel) { worst_rel = rel; worst_pos = i; } + if (gi != wi) ++top1_mismatch; + printf(" pos %2d rel %.4e top1 got %6d want %6d%s\n", i, rel, gi, wi, gi==wi?"":" <-- MISMATCH"); + } + printf("positions : %d\n", n_tok); + printf("worst rel error : %.4e (position %d)\n", worst_rel, worst_pos); + printf("top-1 mismatches: %d/%d\n", top1_mismatch, n_tok); + + llama_batch_free(batch); + llama_free(ctx); + llama_model_free(model); + llama_backend_free(); + + // TOP-1 IS THE HARD GATE. Every structural bug this test found - the KDA forget gate, + // the MQA head count, the leading-dense count, the mHC mixing transpose - showed up first as + // top-1 mismatches or as error that GREW with position. What remains is ~5e-3, flat across + // positions: the graph reassociates sums differently from torch, and a randomly-initialised + // 6-layer network has poorly-conditioned logits, so a few e-3 is what this fixture is worth. + // The real check on the real weights is whether the model generates coherent text. + const bool ok = worst_rel < 1e-2 && top1_mismatch == 0; + printf("%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} diff --git a/tests/test-mhc-sinkhorn.cpp b/tests/test-mhc-sinkhorn.cpp new file mode 100644 index 000000000..bcbf1d52a --- /dev/null +++ b/tests/test-mhc-sinkhorn.cpp @@ -0,0 +1,77 @@ +// Checks ggml_mhc_sinkhorn against the torch reference that was itself validated bit-for-bit +// against transformers' Glm5NextTextHyperConnection on real checkpoint weights. +// +// The failure this exists to catch is not "the numbers are a bit off" - it is the normalisation +// ORDER. Symmetric Sinkhorn, or an extra/missing column pass, still yields a plausible matrix +// with sensible sums, and the model that results is wrong in a way no later test localises. +// The reference output here is COLUMN-stochastic (column sums 1.0, row sums 0.98-1.02); a +// doubly-stochastic result means the order is wrong. +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include + +int main(int argc, char ** argv) { + const char * path = argc > 1 ? argv[1] + : "/home/patrickd/glm-5.3-reap/vendor/mhc/sinkhorn_case.bin"; + FILE * f = fopen(path, "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", path); return 1; } + + int32_t hc, n, iters; float eps; + if (fread(&hc, 4, 1, f) != 1 || fread(&n, 4, 1, f) != 1 || + fread(&iters, 4, 1, f) != 1 || fread(&eps, 4, 1, f) != 1) return 1; + + std::vector logits((size_t) hc*hc*n), want((size_t) hc*hc*n); + if (fread(logits.data(), 4, logits.size(), f) != logits.size()) return 1; + if (fread(want.data(), 4, want.size(), f) != want.size()) return 1; + fclose(f); + printf("case: hc=%d slices=%d iters=%d eps=%g\n", hc, n, iters, (double) eps); + + ggml_init_params ip = { (size_t) 64*1024*1024, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + ggml_tensor * a = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, hc, hc, n); + memcpy(a->data, logits.data(), ggml_nbytes(a)); + + ggml_tensor * out = ggml_mhc_sinkhorn(ctx, a, iters, eps); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, out); + ggml_graph_compute_with_ctx(ctx, gf, 1); + + const float * got = (const float *) out->data; + double num = 0.0, den = 0.0, maxabs = 0.0; + for (size_t i = 0; i < want.size(); ++i) { + const double d = (double) got[i] - want[i]; + num += d*d; den += (double) want[i]*want[i]; + maxabs = std::max(maxabs, std::fabs(d)); + } + const double rel = std::sqrt(num/den); + + // Structural check: the reference is column-stochastic, not doubly stochastic. + double col_err = 0.0, row_spread = 0.0; + for (int s = 0; s < n; ++s) { + for (int c = 0; c < hc; ++c) { + double cs = 0.0; + for (int r = 0; r < hc; ++r) cs += got[(size_t) s*hc*hc + r*hc + c]; + col_err = std::max(col_err, std::fabs(cs - 1.0)); + } + for (int r = 0; r < hc; ++r) { + double rs = 0.0; + for (int c = 0; c < hc; ++c) rs += got[(size_t) s*hc*hc + r*hc + c]; + row_spread = std::max(row_spread, std::fabs(rs - 1.0)); + } + } + printf("rel error vs reference : %.3e max abs %.3e\n", rel, maxabs); + printf("max |col sum - 1| : %.3e (must be ~0: column-stochastic)\n", col_err); + printf("max |row sum - 1| : %.3e (must be NON-zero: not doubly stochastic)\n", row_spread); + + ggml_free(ctx); + const bool ok = rel < 1e-5 && col_err < 1e-4 && row_spread > 1e-3; + printf("%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 5fa487367..ce414f5d3 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -25,6 +25,7 @@ #define KEY_HAS_VISION_ENC "clip.has_vision_encoder" #define KEY_USE_GELU "clip.use_gelu" #define KEY_USE_SILU "clip.use_silu" +#define KEY_VISION_SWIGLU_LIMIT "clip.vision.swiglu_limit" #define KEY_N_EMBD "clip.%s.embedding_length" #define KEY_N_FF "clip.%s.feed_forward_length" diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 70270d6e7..62a971fc2 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -75,6 +75,11 @@ struct clip_hparams { ffn_op_type ffn_op = FFN_GELU; + // GLM-5.3 clamps SwiGLU in the VISION tower too: clamp(gate, max=L), clamp(up, -L, L), + // then SiLU. 0 means unclamped, which is every other model. The clamp only bites once + // activations exceed L - so it passes a small test and diverges on a real image. + 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 12517123e..5d495a3c7 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -580,7 +580,13 @@ ggml_tensor * clip_graph::build_ffn( // we only support parallel ffn for now switch (type_op) { case FFN_SILU: - if (gate) { + if (gate && hparams.swiglu_limit > 0.0f) { + const float lim = hparams.swiglu_limit; + cur = ggml_clamp(ctx0, cur, -INFINITY, lim); + tmp = ggml_clamp(ctx0, tmp, -lim, lim); + cur = ggml_mul(ctx0, ggml_silu(ctx0, cur), tmp); + cb(cur, "ffn_swiglu_clamped", il); + } else if (gate) { cur = ggml_swiglu_split(ctx0, cur, tmp); cb(cur, "ffn_swiglu", il); } else { @@ -1127,6 +1133,7 @@ struct clip_model_loader { log_ffn_op = "gelu"; } else if (use_silu) { hparams.ffn_op = FFN_SILU; + get_f32(KEY_VISION_SWIGLU_LIMIT, hparams.swiglu_limit, false); log_ffn_op = "silu"; } else { hparams.ffn_op = FFN_GELU_QUICK; -- 2.43.0 From 4463f62a1c305e1ffe97ac8180cd386e6a40b7be Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 08:50:53 -0400 Subject: [PATCH 03/12] glm5-next vision: declarative clip.use_mrope, and the position bug it fixes mtmd_decode_use_mrope inferred M-RoPE purely from the projector type, so GLM-5.3 inherited GLM-4V's behaviour: an image advanced position by max(nx, ny) while submitting n_tokens of them, and the KV cache rejected the batch with find_slot: non-consecutive token position 5 after 4 for sequence 0 with 256 new tokens The inference is wrong in general - the projector family and the text model's rope scheme are independent. GLM-5.3 uses the GLM-4V tower with a NoPE text stack (qk_rope_head_dim == 0, LLAMA_ROPE_TYPE_NONE), so it needs GLM4V's graph and Qwen-style position counting OFF. clip.use_mrope now carries this explicitly, with the existing projector-type table as the default when the key is absent, so no other model changes behaviour. Presence is detected with gguf_find_key rather than get_bool's return, because get_bool returns void and leaves the target untouched on a missing key - 'absent' and 'explicitly false' were otherwise indistinguishable. Verified end to end: on a synthetic image the pruned 4-bit model correctly reports a red square top-left, a blue circle bottom-right, and a black '42'. --- convert_hf_to_gguf.py | 7 +++++++ gguf-py/gguf/constants.py | 1 + gguf-py/gguf/gguf_writer.py | 3 +++ tools/mtmd/clip-impl.h | 1 + tools/mtmd/clip-model.h | 7 +++++++ tools/mtmd/clip.cpp | 12 ++++++++++++ tools/mtmd/clip.h | 2 ++ tools/mtmd/mtmd.cpp | 8 ++++++++ 8 files changed, 41 insertions(+) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index f13556d39..46e68d46b 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -4970,6 +4970,13 @@ class Glm5NextVisionModel(Glm4VVisionModel): if limit is not None: self.gguf_writer.add_vision_swiglu_limit(float(limit)) + # GLM-4V uses M-RoPE, so mtmd advances position by max(nx, ny) for an image. GLM-5.3 is + # NoPE end to end (qk_rope_head_dim == 0, LLAMA_ROPE_TYPE_NONE), so an image must advance + # position by its full token count. Inheriting GLM-4V's behaviour submits 256 tokens while + # claiming 16 positions, and the KV cache rejects it: + # "find_slot: non-consecutive token position 5 after 4 ... with 256 new tokens" + self.gguf_writer.add_vision_use_mrope(False) + @ModelBase.register("Qwen3VLForConditionalGeneration") class Qwen3VLTextModel(Qwen3Model): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index ee096a58b..b2c062d31 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -322,6 +322,7 @@ class Keys: USE_GELU = "clip.use_gelu" USE_SILU = "clip.use_silu" SWIGLU_LIMIT = "clip.vision.swiglu_limit" + USE_MROPE = "clip.use_mrope" N_WA_PATTERN = "clip.vision.n_wa_pattern" # used by qwen2.5vl WA_LAYER_INDEXES = "clip.vision.wa_layer_indexes" # used by youtuvl IS_DEEPSTACK_LAYERS = "clip.vision.is_deepstack_layers" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 65360a41e..14977136c 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -1193,6 +1193,9 @@ class GGUFWriter: def add_vision_use_gelu(self, value: bool) -> None: self.add_bool(Keys.ClipVision.USE_GELU, value) + def add_vision_use_mrope(self, value: bool) -> None: + self.add_bool(Keys.ClipVision.USE_MROPE, value) + def add_vision_swiglu_limit(self, value: float) -> None: self.add_float32(Keys.ClipVision.SWIGLU_LIMIT, value) diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index ce414f5d3..d8f6d4fca 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -26,6 +26,7 @@ #define KEY_USE_GELU "clip.use_gelu" #define KEY_USE_SILU "clip.use_silu" #define KEY_VISION_SWIGLU_LIMIT "clip.vision.swiglu_limit" +#define KEY_USE_MROPE "clip.use_mrope" #define KEY_N_EMBD "clip.%s.embedding_length" #define KEY_N_FF "clip.%s.feed_forward_length" diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 62a971fc2..2909b66e1 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -80,6 +80,13 @@ struct clip_hparams { // activations exceed L - so it passes a small test and diverges on a real image. float swiglu_limit = 0.0f; + // Whether the TEXT model consumes image positions as M-RoPE. Defaults per projector type; + // a model may override it, because the projector family and the rope scheme are independent. + // GLM-5.3 uses the GLM-4V tower with a NoPE text stack, so it needs GLM4V's graph and + // Qwen-style position counting turned OFF. + bool use_mrope = false; + bool use_mrope_set = false; + 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 5d495a3c7..7a65dd865 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1134,6 +1134,10 @@ struct clip_model_loader { } else if (use_silu) { hparams.ffn_op = FFN_SILU; get_f32(KEY_VISION_SWIGLU_LIMIT, hparams.swiglu_limit, false); + // get_bool returns void and leaves the target untouched when the key is + // absent, so detect presence explicitly rather than inferring it. + hparams.use_mrope_set = gguf_find_key(ctx_gguf.get(), KEY_USE_MROPE) >= 0; + get_bool(KEY_USE_MROPE, hparams.use_mrope, false); log_ffn_op = "silu"; } else { hparams.ffn_op = FFN_GELU_QUICK; @@ -2786,6 +2790,14 @@ int clip_n_output_tokens(const struct clip_ctx * ctx, struct clip_image_f32 * im return n_patches; } +bool clip_use_mrope(const struct clip_ctx * ctx) { + return ctx->model.hparams.use_mrope; +} + +bool clip_use_mrope_is_set(const struct clip_ctx * ctx) { + return ctx->model.hparams.use_mrope_set; +} + bool clip_image_encode(struct clip_ctx * ctx, const int n_threads, clip_image_f32 * img, float * vec) { clip_image_f32_batch imgs; clip_image_f32_ptr img_copy(clip_image_f32_init()); diff --git a/tools/mtmd/clip.h b/tools/mtmd/clip.h index a859b3865..5cf7aa129 100644 --- a/tools/mtmd/clip.h +++ b/tools/mtmd/clip.h @@ -99,6 +99,8 @@ void clip_build_img_from_pixels(const unsigned char * rgb_pixels, int nx, int ny struct ggml_tensor * clip_get_newline_tensor(const struct clip_ctx * ctx); +bool clip_use_mrope(const struct clip_ctx * ctx); +bool clip_use_mrope_is_set(const struct clip_ctx * ctx); bool clip_image_encode (struct clip_ctx * ctx, int n_threads, struct clip_image_f32 * img, float * vec); bool clip_image_batch_encode(struct clip_ctx * ctx, int n_threads, const struct clip_image_f32_batch * imgs, float * vec); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 35b4396fd..6a6141494 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -989,6 +989,14 @@ bool mtmd_decode_use_non_causal(mtmd_context * ctx) { } bool mtmd_decode_use_mrope(mtmd_context * ctx) { + // An explicit clip.use_mrope in the mmproj wins. The projector family and the text model's + // rope scheme are independent: GLM-5.3 uses the GLM-4V tower with a NoPE text stack, so it + // needs GLM4V's graph but Qwen-style position counting OFF. Without this it advances an + // image by max(nx, ny) positions while submitting n_tokens of them, and the KV cache + // rejects the batch outright. + if (clip_use_mrope_is_set(ctx->ctx_v)) { + return clip_use_mrope(ctx->ctx_v); + } switch (ctx->proj_type_v()) { case PROJECTOR_TYPE_QWEN2VL: case PROJECTOR_TYPE_QWEN25VL: -- 2.43.0 From e84b35888cc6e08416c69f7d233da4842c2567a2 Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 11:02:37 -0400 Subject: [PATCH 04/12] glm5-next: DSA scaffolding - indexer kpool hparam and opt-in flag Phase 2 groundwork per research/DSA_LLAMACPP.md. No behaviour change: dsa_enabled defaults false, so the 11 DSA layers keep running dense exactly as the shipped GGUFs were validated with, and as upstream already does for LLM_ARCH_GLM_DSA and DEEPSEEK2. Sparse selection will widen the KV row on those layers (kv_lora + indexer key + gate + valid), which is why it is gated rather than switched on: n_embd_k_gqa already supports per-layer variation, but a bug there breaks a path that currently works. GLM5_DSA=1 opts in. Phase 1 (the indexer forward itself) is validated in tests/test-dsa-indexer.cpp against the transformers oracle at 2.18e-06. --- convert_hf_to_gguf.py | 1 + gguf-py/gguf/gguf_writer.py | 3 + src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/llama-hparams.h | 10 +++ src/llama-model.cpp | 7 ++ tests/CMakeLists.txt | 1 + tests/test-dsa-indexer.cpp | 146 ++++++++++++++++++++++++++++++++++++ 8 files changed, 170 insertions(+) create mode 100644 tests/test-dsa-indexer.cpp diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 46e68d46b..2f37829cf 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -6190,6 +6190,7 @@ class Glm5NextModel(TextModel): self.gguf_writer.add_indexer_head_count(self.hparams["index_n_heads"]) self.gguf_writer.add_indexer_key_length(self.hparams["index_head_dim"]) self.gguf_writer.add_indexer_top_k(self.hparams["index_topk"]) + self.gguf_writer.add_indexer_kpool(self.hparams["index_kpool"]) # --- mHC --- self.gguf_writer.add_hc_mult(self.hparams["hc_mult"]) diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 14977136c..120ef1d81 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -787,6 +787,9 @@ class GGUFWriter: def add_indexer_key_length(self, length: int) -> None: self.add_uint32(Keys.Attention.Indexer.KEY_LENGTH.format(arch=self.arch), length) + def add_indexer_kpool(self, value: int) -> None: + self.add_uint32(Keys.Attention.Indexer.KPOOL.format(arch=self.arch), value) + def add_indexer_top_k(self, top_k: int) -> None: self.add_uint32(Keys.Attention.Indexer.TOP_K.format(arch=self.arch), top_k) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index e9761a585..15783b5e5 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -245,6 +245,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_HC_MULT, "%s.attention.hc.mult" }, { LLM_KV_ATTENTION_HC_SINKHORN_ITERS, "%s.attention.hc.sinkhorn_iters" }, { LLM_KV_ATTENTION_HC_EPS, "%s.attention.hc.eps" }, + { LLM_KV_ATTENTION_INDEXER_KPOOL, "%s.attention.indexer.kpool" }, { LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" }, { LLM_KV_ATTENTION_SHARED_KV_LAYERS, "%s.attention.shared_kv_layers" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 1f3312e9e..5cd242dcb 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -245,6 +245,7 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, LLM_KV_ATTENTION_INDEXER_TOP_K, + LLM_KV_ATTENTION_INDEXER_KPOOL, LLM_KV_SWIGLU_LIMIT, LLM_KV_SSM_GATE_LOWER_BOUND, LLM_KV_ATTENTION_HC_MULT, diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 6ae19534e..2e0a1b457 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -205,6 +205,16 @@ struct llama_hparams { uint32_t indexer_n_head = 0; uint32_t indexer_head_size = 0; uint32_t indexer_top_k = 0; + uint32_t indexer_kpool = 0; + + // DSA is OPT-IN and defaults OFF. + // + // The shipped GGUFs were validated with these layers running dense, which is also what + // upstream does for LLM_ARCH_GLM_DSA and DEEPSEEK2. Enabling sparse selection changes the + // KV cache row width on those layers (kv_lora + indexer key + gate + valid), so a bug here + // breaks a path that currently works. Default-off keeps the published artifacts' behaviour + // bit-identical while the sparse path is built and gated. + bool dsa_enabled = false; // mHC hyper-connections (GLM-5.3-Flash). hc_mult residual streams per layer, mixed by a // Sinkhorn-normalised matrix. Note the normalisation is COLUMN-stochastic, not doubly diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 1ecb8f6d1..4ad87752d 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2565,6 +2565,13 @@ void llama_model::load_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head, false); ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size, false); ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KPOOL, hparams.indexer_kpool, false); + // Opt-in via GLM5_DSA=1 until the sparse path is gated at ctx <= index_topk + // (where selection is a no-op and must match dense exactly) and by NIAH above it. + { + const char * e = getenv("GLM5_DSA"); + hparams.dsa_enabled = e && *e && *e != '0'; + } // mHC ml.get_key(LLM_KV_ATTENTION_HC_MULT, hparams.hc_mult); diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index a78327769..ff961d14d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -150,6 +150,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # these tests are disabled on Windows because they use internal functions not exported with LLAMA_API (when building with shared libraries) llama_build_and_test(test-sampling.cpp) llama_build_and_test(test-mhc-sinkhorn.cpp) +llama_build(test-dsa-indexer.cpp) llama_build(test-glm5-next-logits.cpp) llama_build_and_test(test-reasoning-budget.cpp) llama_build_and_test(test-grammar-parser.cpp) diff --git a/tests/test-dsa-indexer.cpp b/tests/test-dsa-indexer.cpp new file mode 100644 index 000000000..4fdebee30 --- /dev/null +++ b/tests/test-dsa-indexer.cpp @@ -0,0 +1,146 @@ +// DSA indexer forward, built from ggml ops and checked against the transformers oracle. +// +// Phase 1 of research/DSA_LLAMACPP.md. This validates the GRAPH FORMULATION in isolation, before +// any of it touches the model's attention path or the KV cache - so a wrong pooling offset or a +// dropped relu fails here, at 24 tokens, instead of as a quietly worse model at 128k. +// +// The four hazards this is written to catch (see scripts/dsa_reference.py): +// 1. Pooling starts at the FIRST REAL TOKEN, not slot 0. +// 2. relu sits between the per-head scores and the head-weighted sum. +// 3. The pool key is a PER-CHANNEL softmax over the kpool tokens of (gate + ape). +// 4. A pool is visible to a query only if its LAST token is <= the query position. +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include + +static std::vector rd(FILE * f, size_t n) { + std::vector v(n); + if (fread(v.data(), 4, n, f) != n) { fprintf(stderr, "short read\n"); exit(1); } + return v; +} + +int main(int argc, char ** argv) { + const char * path = argc > 1 ? argv[1] + : "/home/patrickd/glm-5.3-reap/vendor/dsa/dsa_case.bin"; + FILE * f = fopen(path, "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", path); return 1; } + + int32_t S, D, hd, nh, kp, ql, P; + int32_t hdr[7]; + if (fread(hdr, 4, 7, f) != 7) return 1; + S = hdr[0]; D = hdr[1]; hd = hdr[2]; nh = hdr[3]; kp = hdr[4]; ql = hdr[5]; P = hdr[6]; + printf("case: S=%d D=%d hd=%d nh=%d kpool=%d q_lora=%d pools=%d\n", S, D, hd, nh, kp, ql, P); + + auto x = rd(f, (size_t) S*D); + auto qres = rd(f, (size_t) S*ql); + auto wq_b = rd(f, (size_t) nh*hd*ql); + auto wk = rd(f, (size_t) hd*D); + auto knw = rd(f, (size_t) hd); + auto knb = rd(f, (size_t) hd); + auto wproj = rd(f, (size_t) nh*D); + auto kgate = rd(f, (size_t) hd*D); + auto kape = rd(f, (size_t) kp*hd); + auto want = rd(f, (size_t) S*P); + fclose(f); + + ggml_init_params ip = { (size_t) 512*1024*1024, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + auto mk2 = [&](int ne0, int ne1, std::vector & src) { + ggml_tensor * t = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, ne0, ne1); + memcpy(t->data, src.data(), ggml_nbytes(t)); + return t; + }; + auto mk1 = [&](int ne0, std::vector & src) { + ggml_tensor * t = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ne0); + memcpy(t->data, src.data(), ggml_nbytes(t)); + return t; + }; + + ggml_tensor * X = mk2(D, S, x); // [D, S] + ggml_tensor * QR = mk2(ql, S, qres); // [q_lora, S] + ggml_tensor * WQ = mk2(ql, nh*hd, wq_b); // [q_lora, nh*hd] + ggml_tensor * WK = mk2(D, hd, wk); // [D, hd] + ggml_tensor * KNW= mk1(hd, knw); + ggml_tensor * KNB= mk1(hd, knb); + ggml_tensor * WP = mk2(D, nh, wproj); // [D, nh] + ggml_tensor * KG = mk2(D, hd, kgate); // [D, hd] + ggml_tensor * AP = mk2(hd, kp, kape); // [hd, kpool] + + // q = wq_b(q_resid) -> [hd, nh, S] + ggml_tensor * q = ggml_mul_mat(ctx, WQ, QR); + q = ggml_reshape_3d(ctx, q, hd, nh, S); + + // k = LayerNorm(wk(x)) -> [hd, S]. LayerNorm, not RMS, and it has a bias. + ggml_tensor * k = ggml_mul_mat(ctx, WK, X); + k = ggml_norm(ctx, k, 1e-6f); + k = ggml_add(ctx, ggml_mul(ctx, k, KNW), KNB); + + // gate = kpool_gate(x) -> [hd, S] + ggml_tensor * gate = ggml_mul_mat(ctx, KG, X); + + // Pool: groups of kp consecutive tokens. All tokens valid here, so pools start at 0 and the + // count is ceil(S/kp); the first-real-token offset only matters with left padding. + const int n_pools = (S + kp - 1) / kp; + if (n_pools != P) { printf("pool count %d != expected %d\n", n_pools, P); return 1; } + + // [hd, kp, n_pools] views of k and gate + ggml_tensor * k3 = ggml_reshape_3d(ctx, k, hd, kp, n_pools); + ggml_tensor * g3 = ggml_reshape_3d(ctx, gate, hd, kp, n_pools); + + // logits = gate + ape (ape broadcast over pools); softmax over the kp axis, PER CHANNEL. + ggml_tensor * ape3 = ggml_reshape_3d(ctx, AP, hd, kp, 1); + ggml_tensor * lg = ggml_add(ctx, g3, ape3); + // soft_max reduces ne0, so bring kp to ne0, normalise, and put it back. + lg = ggml_cont(ctx, ggml_permute(ctx, lg, 1, 0, 2, 3)); // [kp, hd, n_pools] + lg = ggml_soft_max(ctx, lg); + lg = ggml_cont(ctx, ggml_permute(ctx, lg, 1, 0, 2, 3)); // [hd, kp, n_pools] + + ggml_tensor * pk = ggml_mul(ctx, lg, k3); + // sum over kp: bring it to ne0 and sum_rows + pk = ggml_cont(ctx, ggml_permute(ctx, pk, 1, 0, 2, 3)); // [kp, hd, n_pools] + pk = ggml_sum_rows(ctx, pk); // [1, hd, n_pools] + pk = ggml_reshape_2d(ctx, pk, hd, n_pools); // [hd, n_pools] + + // scores = relu( (q . pool_keys) * hd^-0.5 ) -> [n_pools, nh, S] + ggml_tensor * sc = ggml_mul_mat(ctx, pk, q); + sc = ggml_scale(ctx, sc, 1.0f/sqrtf((float) hd)); + sc = ggml_relu(ctx, sc); + + // weights = weights_proj(x) * nh^-0.5 -> [nh, S]; index_scores = weights . scores + ggml_tensor * wgt = ggml_mul_mat(ctx, WP, X); + wgt = ggml_scale(ctx, wgt, 1.0f/sqrtf((float) nh)); + ggml_tensor * wgt3 = ggml_reshape_3d(ctx, wgt, nh, 1, S); + ggml_tensor * scp = ggml_cont(ctx, ggml_permute(ctx, sc, 1, 0, 2, 3)); // [nh, n_pools, S] + ggml_tensor * idx = ggml_mul_mat(ctx, scp, wgt3); // [n_pools, 1, S] + idx = ggml_reshape_2d(ctx, idx, n_pools, S); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, idx); + ggml_graph_compute_with_ctx(ctx, gf, 4); + + const float * got = (const float *) idx->data; + const float LO = -1e30f; + double num = 0, den = 0; int cmp = 0, bad = 0; + for (int s = 0; s < S; ++s) { + for (int p = 0; p < n_pools; ++p) { + const float w = want[(size_t) s*n_pools + p]; + if (w <= LO) continue; // masked in the reference (causality) + const double d = (double) got[(size_t) s*n_pools + p] - w; + num += d*d; den += (double) w*w; ++cmp; + if (std::fabs(d) > 1e-3) ++bad; + } + } + const double rel = std::sqrt(num/std::max(den, 1e-30)); + printf("compared %d unmasked entries | rel error %.4e | %d over 1e-3\n", cmp, rel, bad); + + ggml_free(ctx); + const bool ok = rel < 1e-4; + printf("%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} -- 2.43.0 From bb28e3fa87da3bbdeb00a8079caffcc5001bcdf0 Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 11:02:46 -0400 Subject: [PATCH 05/12] fix: Indexer.KPOOL constant placement (padding-aligned block, not the outer class) --- gguf-py/gguf/constants.py | 1 + 1 file changed, 1 insertion(+) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index b2c062d31..b39b5c4c4 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -194,6 +194,7 @@ class Keys: HEAD_COUNT = "{arch}.attention.indexer.head_count" KEY_LENGTH = "{arch}.attention.indexer.key_length" TOP_K = "{arch}.attention.indexer.top_k" + KPOOL = "{arch}.attention.indexer.kpool" class Rope: DIMENSION_COUNT = "{arch}.rope.dimension_count" -- 2.43.0 From b8f614894266aa0db517d66e8eb1f578daf9ad5f Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 11:36:25 -0400 Subject: [PATCH 06/12] glm5-next: DSA indexer in the model graph (prefill), gated off Moves the validated indexer formulation out of the test and into llm_build_glm5_next, so the graph-side code lives with the model rather than only in tests/test-dsa-indexer.cpp (which still checks it against the transformers oracle at 2.18e-06). Computes index_scores for the CURRENT batch only - the prefill case. Decode needs the indexer key and gate of every cached token, which means widening the KV row on DSA layers to carry kv_lora + indexer key + gate + valid. That is Phase 2b; until then the scores are computed and discarded, and the whole path is behind hparams.dsa_enabled (GLM5_DSA=1), default off. Returns nullptr on a ragged final pool rather than mis-pooling: the reference pads the last pool and masks the missing slots, and silently grouping the wrong tokens is worse than not running. All three tests still pass: mhc-sinkhorn, dsa-indexer 2.18e-06, glm5-next-logits 16/16 top-1. --- src/models/glm5-next.cpp | 78 ++++++++++++++++++++++++++++++++++++++++ src/models/models.h | 4 +++ 2 files changed, 82 insertions(+) diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 3f1ca0ec2..363a2f12f 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -86,6 +86,69 @@ static ggml_tensor * causal_conv1d(ggml_cgraph * gf, ggml_context * ctx0, ggml_t } +// DSA index scores for the current batch. +// +// Validated formulation - tests/test-dsa-indexer.cpp checks exactly this against the +// transformers oracle at 2.18e-06. Four things here are easy to get wrong and invisible if you +// do (see scripts/dsa_reference.py): +// * k_norm is a LAYERNORM with a bias, not RMS. +// * The pool key is a PER-CHANNEL softmax over the kpool tokens of (gate + ape) - not a mean. +// * relu sits between the per-head scores and the head-weighted sum. +// * Pools start at the first real token; with no left padding that is slot 0. +// +// LIMITATION: this scores only the CURRENT batch, which is the prefill case. Decode needs the +// indexer key and gate of every cached token, which means widening the KV row on this layer - +// Phase 2b. Until then the scores are computed and discarded, so nothing consumes them. +ggml_tensor * llm_build_glm5_next::build_dsa_index_scores( + ggml_tensor * x, ggml_tensor * q_a, const llama_layer & layer, int il) { + const int64_t hd = hparams.indexer_head_size; + const int64_t nh = hparams.indexer_n_head; + const int64_t kp = hparams.indexer_kpool ? hparams.indexer_kpool : 4; + const int64_t S = x->ne[1]; + + ggml_tensor * q = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, q_a); + q = ggml_reshape_3d(ctx0, q, hd, nh, S); + + ggml_tensor * k = ggml_mul_mat(ctx0, layer.indexer_attn_k, x); + k = ggml_norm(ctx0, k, 1e-6f); // LayerNorm, eps 1e-6 + k = ggml_add(ctx0, ggml_mul(ctx0, k, layer.indexer_k_norm), layer.indexer_k_norm_b); + + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.indexer_kpool_gate, x); + + const int64_t n_pools = (S + kp - 1)/kp; + if (n_pools*kp != S) { + // Ragged tail: the reference pads the final pool and masks the missing slots. Not yet + // handled here, and silently mis-pooling would be worse than not running. + return nullptr; + } + + ggml_tensor * k3 = ggml_reshape_3d(ctx0, k, hd, kp, n_pools); + ggml_tensor * g3 = ggml_reshape_3d(ctx0, gate, hd, kp, n_pools); + ggml_tensor * ape3 = ggml_reshape_3d(ctx0, layer.indexer_kpool_ape, hd, kp, 1); + + // soft_max reduces ne0, so bring kp there and put it back. + ggml_tensor * lg = ggml_add(ctx0, g3, ape3); + lg = ggml_cont(ctx0, ggml_permute(ctx0, lg, 1, 0, 2, 3)); + lg = ggml_soft_max(ctx0, lg); + lg = ggml_cont(ctx0, ggml_permute(ctx0, lg, 1, 0, 2, 3)); + + ggml_tensor * pk = ggml_mul(ctx0, lg, k3); + pk = ggml_cont(ctx0, ggml_permute(ctx0, pk, 1, 0, 2, 3)); + pk = ggml_sum_rows(ctx0, pk); + pk = ggml_reshape_2d(ctx0, pk, hd, n_pools); + + ggml_tensor * sc = ggml_mul_mat(ctx0, pk, q); + sc = ggml_scale(ctx0, sc, 1.0f/sqrtf((float) hd)); + sc = ggml_relu(ctx0, sc); // load-bearing + + ggml_tensor * wgt = ggml_mul_mat(ctx0, layer.indexer_proj, x); + wgt = ggml_scale(ctx0, wgt, 1.0f/sqrtf((float) nh)); + ggml_tensor * wgt3 = ggml_reshape_3d(ctx0, wgt, nh, 1, S); + ggml_tensor * scp = ggml_cont(ctx0, ggml_permute(ctx0, sc, 1, 0, 2, 3)); + ggml_tensor * idx = ggml_mul_mat(ctx0, scp, wgt3); + return ggml_reshape_2d(ctx0, idx, n_pools, S); +} + llm_build_glm5_next::mhc_site llm_build_glm5_next::build_mhc( ggml_tensor * streams, ggml_tensor * fn, ggml_tensor * base, ggml_tensor * scale, int il) { @@ -320,6 +383,21 @@ llm_build_glm5_next::llm_build_glm5_next(const llama_model & model, const llm_gr q_a = build_norm(q_a, layer.attn_q_a_norm, NULL, LLM_NORM_RMS, il); ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.wq_b, q_a); + // DSA indexer. Formulation validated against transformers in + // tests/test-dsa-indexer.cpp (rel error 2.18e-06); see research/DSA_LLAMACPP.md. + // + // Gated OFF by default: consuming these scores means widening the KV row on this + // layer to carry the indexer key, gate and valid flag, and the shipped GGUFs were + // validated with these layers dense. Built here so the graph-side formulation lives + // with the model rather than only in a test. + if (hparams.dsa_enabled && layer.indexer_attn_k && layer.indexer_kpool_gate) { + ggml_tensor * iscores = build_dsa_index_scores(cur, q_a, layer, il); + if (iscores) { + cb(iscores, "dsa_index_scores", il); + ggml_build_forward_expand(gf, iscores); + } + } + ggml_tensor * kv_cmpr = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, NULL, LLM_NORM_RMS, il); diff --git a/src/models/models.h b/src/models/models.h index e0890eda7..fbd40b430 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -379,6 +379,10 @@ struct llm_build_glm5_next : public llm_build_delta_net_base { // streams' = post (x) y + comb^T @ streams ggml_tensor * apply_mhc(const mhc_site & s, ggml_tensor * y, int il); + // DSA index scores for the current batch (prefill). Validated in test-dsa-indexer.cpp. + ggml_tensor * build_dsa_index_scores(ggml_tensor * x, ggml_tensor * q_a, + const llama_layer & layer, int il); + const llama_model & model; }; -- 2.43.0 From a0571d11270e14daf0ee4ff4b89268c1362883ea Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 13:45:20 -0400 Subject: [PATCH 07/12] glm5-next: DSA Phase 2b - indexer state in the K row, gated off Decode needs the indexer key and gate of every cached token, and neither survives the token: both are projections of the hidden state, and the cached MLA latent is a different projection that cannot be inverted. So the state must be cached. Widens the K row on attention layers to kv_lora + key + gate when dsa_enabled, and zero-pads Q to match so the extra dimensions contribute nothing to attention scores. Costs ~50% on the QK matmul for 11 layers, which is noise against Phase 3's ~64x reduction at 128k and a straight loss if Phase 3 never lands - stated rather than hidden. The alternative, a separate cache stream, was priced and is WORSE: n_embd_r() is global rather than per-layer, so DSA layers would have to be marked recurrent (layer typing is consulted throughout) or a new cache type added. More shared-code surface, not less. Full analysis in research/DSA_LLAMACPP.md, including a correction to an earlier note that wrongly claimed the widening was contained in the graph builder. V is no longer aliased to K. It was the same pointer; widening K would have fed the indexer state into wv_b as if it were value content. The reference's 'valid' flag is dropped - llama.cpp already tracks populated cache slots, and two notions of validity can only disagree. Gate: at ctx <= index_topk selection is a no-op, so sparse must equal dense. DSA off worst rel 5.0232e-03, top-1 16/16 DSA on worst rel 5.0219e-03, top-1 16/16 The 1.3e-06 delta is float reassociation from the wider matmul. --- src/llama-hparams.cpp | 8 ++++++++ src/llama-hparams.h | 4 ++++ src/llama-model.cpp | 9 +++++++++ src/models/glm5-next.cpp | 26 ++++++++++++++++++++++++++ 4 files changed, 47 insertions(+) diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 002d15d41..a22b83c5c 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -86,6 +86,14 @@ uint32_t llama_hparams::n_embd_out() const { uint32_t llama_hparams::n_embd_head_k(uint32_t il) const { if (il < n_layer) { + // DSA (opt-in): the attention layers carry the indexer's key and gate alongside the + // compressed MLA latent, so their K row is wider. Q is zero-padded to match, which + // leaves attention scores unchanged - the extra dimensions contribute nothing. Costs + // ~50% on the QK matmul for those layers; see research/DSA_LLAMACPP.md for why the + // alternative (a separate cache stream) is more shared-code surface, not less. + if (dsa_enabled && n_embd_head_k_dsa != 0 && !is_recurrent(il)) { + return n_embd_head_k_dsa; + } return is_swa(il) ? n_embd_head_k_swa : n_embd_head_k_full; } diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 2e0a1b457..9c6011ee1 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -216,6 +216,10 @@ struct llama_hparams { // bit-identical while the sparse path is built and gated. bool dsa_enabled = false; + // Widened K-row width for DSA layers when dsa_enabled: kv_lora + indexer key + gate. + // Zero when DSA is off, which is the shipped configuration. + uint32_t n_embd_head_k_dsa = 0; + // mHC hyper-connections (GLM-5.3-Flash). hc_mult residual streams per layer, mixed by a // Sinkhorn-normalised matrix. Note the normalisation is COLUMN-stochastic, not doubly // stochastic - it column-normalises once and then runs (iters-1) full row+column passes. diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4ad87752d..0a4bc7bf1 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2571,6 +2571,15 @@ void llama_model::load_hparams(llama_model_loader & ml) { { const char * e = getenv("GLM5_DSA"); hparams.dsa_enabled = e && *e && *e != '0'; + if (hparams.dsa_enabled) { + // kv_lora + indexer key + gate. The +1 "valid" flag the reference packs + // alongside is redundant here: llama.cpp already tracks which cache + // slots are populated, so carrying it would waste a row element and + // invite the two notions of validity to disagree. + hparams.n_embd_head_k_dsa = hparams.n_lora_kv + 2*hparams.indexer_head_size; + LLAMA_LOG_INFO("%s: DSA enabled - K row %u -> %u on attention layers\n", + __func__, hparams.n_lora_kv, hparams.n_embd_head_k_dsa); + } } // mHC diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 363a2f12f..a7c94262b 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -411,6 +411,32 @@ llm_build_glm5_next::llm_build_glm5_next(const llama_model & model, const llm_gr ggml_tensor * Kcur = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); ggml_tensor * Vcur = Kcur; + // DSA (opt-in): carry the indexer key and gate in the K row so decode can score + // cached tokens. Q is zero-padded to match, so the extra dimensions contribute + // exactly nothing to the attention scores. + // + // V is deliberately NOT widened and no longer aliases K: wv_b expands from the + // compressed latent, so a wider V would feed it the indexer state as if it were + // value content. + if (hparams.dsa_enabled && layer.indexer_attn_k && layer.indexer_kpool_gate) { + const int64_t ihd = hparams.indexer_head_size; + + ggml_tensor * ik = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur); + ik = ggml_norm(ctx0, ik, 1e-6f); + ik = ggml_add(ctx0, ggml_mul(ctx0, ik, layer.indexer_k_norm), + layer.indexer_k_norm_b); + ggml_tensor * ig = ggml_mul_mat(ctx0, layer.indexer_kpool_gate, cur); + + ggml_tensor * extra = ggml_concat(ctx0, ik, ig, 0); // [2*ihd, T] + extra = ggml_reshape_3d(ctx0, extra, 2*ihd, 1, n_tokens); + Kcur = ggml_concat(ctx0, Kcur, extra, 0); // [kv_lora+2*ihd, 1, T] + + ggml_tensor * pad = ggml_new_tensor_3d(ctx0, Qcur->type, 2*ihd, n_head, n_tokens); + pad = ggml_scale(ctx0, pad, 0.0f); + Qcur = ggml_concat(ctx0, Qcur, pad, 0); + cb(Kcur, "dsa_k_widened", il); + } + cur = build_attn(inp_attn_k, layer.wo, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, kq_scale_mla, il); } else { -- 2.43.0 From 60523a29e2eec1660c79f8f83618b6ed626d8d93 Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 14:20:59 -0400 Subject: [PATCH 08/12] glm5-next: DSA Phase 3 index expansion, validated Expands selected POOL indices into TOKEN indices in ggml - the step that looked like a blocker. ggml_top_k returns I32 and ggml has no integer add or scale, so pool p -> [p*kpool, p*kpool+kpool) appeared inexpressible without a custom op. ggml_cpy converts I32 <-> F32, so the arithmetic routes through float and back: top_k -> cpy(F32) -> scale(kpool) -> broadcast-add arange(kpool) -> cpy(I32) -> get_rows Tested against a scalar reference rather than eyeballed: an off-by-one here selects NEIGHBOURING tokens, which still yields fluent output and would stay invisible until someone measured long-context retrieval carefully. Also records why Phase 2c (the additive mask) is skipped rather than deferred: it needs a per-token scatter along ne0, and ggml has none - top_k gives indices without values, argsort cannot extract a per-row k-th value, and set_rows scatters along ne1 with one shared index vector. It would have required the custom operator this port was glad not to need, to buy correctness parity that dense already delivers (3/3 needle retrieval at 32k, measured). Phase 3 now has both halves proven independently: indexer scoring at 2.18e-06 against the transformers oracle, and this expansion exactly. --- tests/CMakeLists.txt | 1 + tests/test-dsa-gather.cpp | 86 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 87 insertions(+) create mode 100644 tests/test-dsa-gather.cpp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index ff961d14d..2dbe79ba9 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -151,6 +151,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-sampling.cpp) llama_build_and_test(test-mhc-sinkhorn.cpp) llama_build(test-dsa-indexer.cpp) +llama_build(test-dsa-gather.cpp) llama_build(test-glm5-next-logits.cpp) llama_build_and_test(test-reasoning-budget.cpp) llama_build_and_test(test-grammar-parser.cpp) diff --git a/tests/test-dsa-gather.cpp b/tests/test-dsa-gather.cpp new file mode 100644 index 000000000..1228343f6 --- /dev/null +++ b/tests/test-dsa-gather.cpp @@ -0,0 +1,86 @@ +// Phase 3 building block: expand selected POOL indices into TOKEN indices, in ggml. +// +// This is the step that looked like a blocker. ggml_top_k returns I32 and ggml has no integer +// add or scale, so pool p -> tokens [p*kpool, p*kpool + kpool) seemed inexpressible. It is not: +// ggml_cpy converts I32 <-> F32, so the arithmetic goes through float and back. +// +// Checked here against a scalar reference because an off-by-one in this expansion selects +// neighbouring tokens - which still produces fluent output, and would be invisible until someone +// measured long-context retrieval carefully. +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include + +int main() { + const int n_pools = 6; + const int kpool = 4; + const int select_k = 3; + + // Scores over pools; the top-3 are pools 4, 1, 5 (values 9, 7, 6). + const std::vector scores = { 2.0f, 7.0f, 1.0f, 0.5f, 9.0f, 6.0f }; + + ggml_init_params ip = { (size_t) 32*1024*1024, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + ggml_tensor * S = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_pools); + memcpy(S->data, scores.data(), ggml_nbytes(S)); + + // 1. top-k -> I32 pool indices, "in no particular order" + ggml_tensor * sel = ggml_top_k(ctx, S, select_k); + + // 2. I32 -> F32 so arithmetic is available at all + ggml_tensor * self = ggml_cpy(ctx, sel, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, select_k)); + + // 3. pool p -> first token p*kpool, then broadcast-add arange(0..kpool-1) + ggml_tensor * first = ggml_scale(ctx, self, (float) kpool); + first = ggml_reshape_2d(ctx, first, 1, select_k); + ggml_tensor * ar = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kpool, 1); + for (int i = 0; i < kpool; ++i) ((float *) ar->data)[i] = (float) i; + + ggml_tensor * base = ggml_repeat(ctx, first, + ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kpool, select_k)); + ggml_tensor * tok = ggml_add(ctx, base, ar); // [kpool, select_k] + + // 4. back to I32, flattened - this is what ggml_get_rows consumes + ggml_tensor * flat = ggml_reshape_1d(ctx, tok, kpool*select_k); + ggml_tensor * toki = ggml_cpy(ctx, flat, + ggml_new_tensor_1d(ctx, GGML_TYPE_I32, kpool*select_k)); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, toki); + ggml_build_forward_expand(gf, sel); + ggml_graph_compute_with_ctx(ctx, gf, 1); + + const int32_t * pools = (const int32_t *) sel->data; + const int32_t * got = (const int32_t *) toki->data; + + std::vector want; + printf("selected pools:"); + for (int i = 0; i < select_k; ++i) { + printf(" %d", pools[i]); + for (int j = 0; j < kpool; ++j) want.push_back(pools[i]*kpool + j); + } + printf(" (expected top-3 of {2,7,1,0.5,9,6} = pools 4,1,5 in some order)\n"); + + bool ok = true; + printf("expanded tokens:"); + for (size_t i = 0; i < want.size(); ++i) { + printf(" %d", got[i]); + if (got[i] != want[i]) ok = false; + } + printf("\n"); + + // The selection itself must be the right SET, order aside. + std::vector sp(pools, pools + select_k); + std::sort(sp.begin(), sp.end()); + const std::vector expect_pools = {1, 4, 5}; + if (sp != expect_pools) { printf("wrong pools selected\n"); ok = false; } + + ggml_free(ctx); + printf("%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} -- 2.43.0 From 19a4efa2203cd7f4d4439e4b229a4842c285ade8 Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 18:18:47 -0400 Subject: [PATCH 09/12] glm5-next: DSA Phase 3 gathered attention, and the mutation test that gave it teeth Attention over gathered cache rows must equal dense attention over the full cache masked to those same rows: both reduce to exp(s_i)/sum_{j in S} exp(s_j). Agreement to 1.19e-07, which is summation order only -- ggml_top_k returns pools unordered, so the softmax denominator and the V-weighted sum accumulate in a different order than the dense pass. The first version of this test was worthless and I only found that by mutating it. With T=1 a [n_kv,1] mask transposes to [1,n_kv], so gathering along either axis coincides and a mask gather with NO transposes passed cleanly. T is now 3, and visibility varies per token: MTP/DFlash spec decode submits several tokens per step, so T>1 is the production path here, not an edge case. Mutants now caught, having been run to confirm each fails: - mask gathered without the outbound transpose -> visibility lands on the wrong token - V viewed at the wrong offset in the widened K row -> 1.267e+00, indexer state read as value - off-by-one in the pool->token expansion -> selects the neighbouring token everywhere Gathering V separately with the same indices does NOT fail, and should not -- it is equivalent. The hazard is a different index order or a different offset, not a separate gather. The cache is deliberately a strided view (kv_size 40 > n_kv 24) so get_rows is exercised through nb1 rather than a packed buffer, which is how llama_kv_cache::get_k actually hands it over. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016B6ZbadovyJKFiN1CwDBLA --- tests/CMakeLists.txt | 1 + tests/test-dsa-attn.cpp | 195 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 196 insertions(+) create mode 100644 tests/test-dsa-attn.cpp diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 2dbe79ba9..11ac1645f 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -152,6 +152,7 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) llama_build_and_test(test-mhc-sinkhorn.cpp) llama_build(test-dsa-indexer.cpp) llama_build(test-dsa-gather.cpp) +llama_build(test-dsa-attn.cpp) llama_build(test-glm5-next-logits.cpp) llama_build_and_test(test-reasoning-budget.cpp) llama_build_and_test(test-grammar-parser.cpp) diff --git a/tests/test-dsa-attn.cpp b/tests/test-dsa-attn.cpp new file mode 100644 index 000000000..aef497cda --- /dev/null +++ b/tests/test-dsa-attn.cpp @@ -0,0 +1,195 @@ +// DSA Phase 3: attention over GATHERED cache rows must equal dense attention over the full cache +// masked to those same rows. +// +// This is the invariant that makes sparse attention safe to ship. softmax over a gathered subset +// and softmax over the full row with -inf everywhere outside the subset are the same function: +// both reduce to exp(s_i) / sum_{j in S} exp(s_j). If the gather, the mask gather, or the V +// alignment is wrong, the two disagree. Nothing else in the pipeline would catch it -- a +// mis-gathered attention still produces fluent text. +// +// Three specific hazards this pins down: +// 1. The cache is a STRIDED VIEW (kv_size > n_kv), so get_rows must read through nb1, not +// assume a packed buffer. A packed-buffer assumption reads the right count of wrong rows. +// 2. The mask is [n_kv, T] and must be gathered along ne0, which get_rows cannot do -- it +// selects along ne1. It has to go through a transpose, and a missing transpose silently +// gathers along the token axis instead. +// 3. V is a VIEW of K at width kv_lora_rank. Gathered V rows must stay paired with the same +// gathered K rows; any independent gather of V permutes value content against its weights. +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include +#include + +int main() { + const int kv_lora = 4; + const int ihd = 2; + const int kw = kv_lora + 2*ihd; // widened K row: latent + indexer key + gate + const int n_head = 2; + const int T = 3; // decode; >1 because MTP/DFlash spec decode submits several + // tokens per step, and T==1 makes the mask transpose a no-op + const int n_kv = 24; + const int kv_size = 40; // cache is BIGGER than n_kv -> strided view + const int kpool = 4; + const int n_pools = n_kv / kpool; // 6 + const int select_k = 3; // -> 12 of 24 tokens survive + const int sel = select_k * kpool; + const float scale = 1.0f / std::sqrt((float) kv_lora); + + std::mt19937 rng(1234); + std::normal_distribution nd(0.0f, 1.0f); + + ggml_init_params ip = { (size_t) 256*1024*1024, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + // ---- the cache, as llama_kv_cache actually lays it out: [kw, kv_size] ---- + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kw, kv_size); + for (int i = 0; i < kw*kv_size; ++i) ((float *) cache->data)[i] = nd(rng); + + // Query, already absorbed: [kv_lora, n_head, T], zero-padded out to kw so the indexer + // dimensions contribute exactly nothing to the scores. + ggml_tensor * q = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, kw, n_head, T); + memset(q->data, 0, ggml_nbytes(q)); + for (int h = 0; h < n_head; ++h) + for (int t = 0; t < T; ++t) + for (int d = 0; d < kv_lora; ++d) + ((float *) q->data)[(t*n_head + h)*kw + d] = nd(rng); + + // ---- pool scores -> selected pools -> selected token indices ---- + std::vector pool_scores(n_pools); + for (auto & s : pool_scores) s = nd(rng); + + ggml_tensor * S = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_pools); + memcpy(S->data, pool_scores.data(), ggml_nbytes(S)); + + ggml_tensor * selp = ggml_top_k(ctx, S, select_k); // I32 + ggml_tensor * self = ggml_cpy(ctx, selp, ggml_new_tensor_1d(ctx, GGML_TYPE_F32, select_k)); + ggml_tensor * first = ggml_reshape_2d(ctx, ggml_scale(ctx, self, (float) kpool), 1, select_k); + ggml_tensor * off = ggml_reshape_2d(ctx, ggml_arange(ctx, 0.0f, (float) kpool, 1.0f), kpool, 1); + ggml_tensor * toks = ggml_add(ctx, ggml_repeat(ctx, first, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kpool, select_k)), + ggml_repeat(ctx, off, ggml_new_tensor_2d(ctx, GGML_TYPE_F32, kpool, select_k))); + ggml_tensor * idx = ggml_cpy(ctx, ggml_reshape_1d(ctx, toks, sel), + ggml_new_tensor_1d(ctx, GGML_TYPE_I32, sel)); + + // ================= GATHERED PATH ================= + // The cache as attention sees it, then gathered to the selected rows. + ggml_tensor * k2d = ggml_view_2d(ctx, cache, kw, n_kv, cache->nb[1], 0); + ggml_tensor * k_sel = ggml_get_rows(ctx, k2d, idx); // [kw, sel] + ggml_tensor * kg = ggml_reshape_3d(ctx, k_sel, kw, 1, sel); + ggml_tensor * vg = ggml_view_3d(ctx, kg, kv_lora, 1, sel, kg->nb[1], kg->nb[2], 0); + + ggml_tensor * qp = ggml_permute(ctx, q, 0, 2, 1, 3); // [kw, T, n_head] + ggml_tensor * kgp = ggml_permute(ctx, kg, 0, 2, 1, 3); // [kw, sel, 1] + ggml_tensor * vgp = ggml_permute(ctx, vg, 0, 2, 1, 3); // [kv_lora, sel, 1] + + // The real decode mask is not all-visible: unoccupied cache cells and causal structure put + // -inf inside pools the selector still picks. Gather it for real, through the transpose, + // rather than assuming every selected row is visible. + // Visibility VARIES BY TOKEN -- causal structure plus a few dead cache cells. If it did not + // vary, every column of the mask would be identical and a gather along the wrong axis would + // return the right numbers by accident. + ggml_tensor * mask_base = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_kv, T); + std::vector> vis(T, std::vector(n_kv, 1)); + for (int t = 0; t < T; ++t) { + for (int i = 0; i < n_kv; ++i) if (i > n_kv - T + t) vis[t][i] = 0; // causal tail + for (int i : {2, 3, 9, 17}) vis[t][i] = 0; // dead cells + vis[t][(5 + t) % n_kv] = 0; // token-specific + } + for (int t = 0; t < T; ++t) + for (int i = 0; i < n_kv; ++i) + ((float *) mask_base->data)[t*n_kv + i] = vis[t][i] ? 0.0f : -INFINITY; + + // get_rows selects along ne1, but the mask needs gathering along ne0 -- hence the transpose + // either side. Dropping either one gathers along the token axis and silently returns garbage + // that is still finite and still the right shape. + ggml_tensor * mask_g = ggml_cont(ctx, ggml_transpose(ctx, + ggml_get_rows(ctx, ggml_cont(ctx, ggml_transpose(ctx, mask_base)), idx))); + + ggml_tensor * kq_g = ggml_mul_mat(ctx, kgp, qp); // [sel, T, n_head] + kq_g = ggml_soft_max_ext(ctx, kq_g, mask_g, scale, 0.0f); + ggml_tensor * out_g = ggml_mul_mat(ctx, ggml_cont(ctx, ggml_transpose(ctx, vgp)), kq_g); + + // ================= DENSE PATH ================= + ggml_tensor * kd = ggml_reshape_3d(ctx, ggml_cont(ctx, k2d), kw, 1, n_kv); + ggml_tensor * vd = ggml_view_3d(ctx, kd, kv_lora, 1, n_kv, kd->nb[1], kd->nb[2], 0); + ggml_tensor * kdp = ggml_permute(ctx, kd, 0, 2, 1, 3); + ggml_tensor * vdp = ggml_permute(ctx, vd, 0, 2, 1, 3); + + // -inf everywhere the selection did not land. + std::vector sel_pools(select_k); + { + std::vector ord(n_pools); for (int i=0;i pool_scores[b]; }); + std::copy(ord.begin(), ord.begin()+select_k, sel_pools.begin()); + } + std::vector keep(n_kv, 0); + for (int p : sel_pools) for (int j = 0; j < kpool; ++j) keep[p*kpool + j] = 1; + + ggml_tensor * mask_d = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_kv, T); + for (int t = 0; t < T; ++t) + for (int i = 0; i < n_kv; ++i) + ((float *) mask_d->data)[t*n_kv + i] = (keep[i] && vis[t][i]) ? 0.0f : -INFINITY; + + ggml_tensor * kq_d = ggml_mul_mat(ctx, kdp, qp); + kq_d = ggml_soft_max_ext(ctx, kq_d, mask_d, scale, 0.0f); + ggml_tensor * out_d = ggml_mul_mat(ctx, ggml_cont(ctx, ggml_transpose(ctx, vdp)), kq_d); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, out_g); + ggml_build_forward_expand(gf, out_d); + ggml_build_forward_expand(gf, idx); + ggml_build_forward_expand(gf, mask_g); + ggml_graph_compute_with_ctx(ctx, gf, 4); + + // ---- 1. the gather picked the pools the reference picked ---- + std::vector want; + for (int p : sel_pools) for (int j = 0; j < kpool; ++j) want.push_back(p*kpool + j); + std::vector got((int32_t *) idx->data, (int32_t *) idx->data + sel); + std::vector want_s = want, got_s = got; + std::sort(want_s.begin(), want_s.end()); std::sort(got_s.begin(), got_s.end()); + if (want_s != got_s) { + printf("FAIL: token selection mismatch\n want:"); + for (int v : want_s) printf(" %d", v); + printf("\n got: "); + for (int v : got_s) printf(" %d", v); + printf("\n"); + return 1; + } + + // ---- 2. the gathered mask carries the right visibility, in the right order ---- + for (int t = 0; t < T; ++t) { + for (int i = 0; i < sel; ++i) { + const float want_m = vis[t][got[i]] ? 0.0f : -INFINITY; + const float got_m = ((float *) mask_g->data)[t*sel + i]; + if (!((std::isinf(want_m) && std::isinf(got_m)) || want_m == got_m)) { + printf("FAIL: gathered mask[t=%d][%d] (token %d) = %f, want %f\n", + t, i, got[i], got_m, want_m); + return 1; + } + } + } + + // ---- 3. gathered attention == dense attention masked to the same rows ---- + const int n = kv_lora * T * n_head; + double maxerr = 0.0; + for (int i = 0; i < n; ++i) { + maxerr = std::max(maxerr, (double) std::fabs(((float *) out_g->data)[i] - ((float *) out_d->data)[i])); + } + + printf("selected %d of %d tokens (%d of %d pools)\n", sel, n_kv, select_k, n_pools); + printf("max |gathered - dense_masked| = %.3e\n", maxerr); + + // Not bit-exact: top_k returns the pools in no particular order, so the softmax denominator + // and the V-weighted sum accumulate in a different order than the dense pass. The set is + // identical, so the difference is pure summation order. + if (!(maxerr < 1e-6)) { printf("FAIL: gathered attention diverges from dense\n"); return 1; } + + printf("OK\n"); + ggml_free(ctx); + return 0; +} -- 2.43.0 From 9dc34911f655766f92c298e07d97bca9be821c19 Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 21:50:21 -0400 Subject: [PATCH 10/12] glm5-next: DSA Phase 3 - sparse attention over the KV cache (opt-in) build_attn_dsa runs gathered attention against top-k selected pools, gated behind hparams.dsa_enabled with a dense build_attn fallback for every case it does not handle. Gathering happens at POOL granularity, not token granularity. Consecutive cache cells are contiguous, so the [D, n_kv] cache view re-views as [D*kpool, n_pools] and one get_rows with pool ids gathers whole pools. That removes the index arithmetic entirely - ggml has no ops for scaling I32 pool ids into token ids - and pool granularity is DSA's own, since select_k = indexer_top_k / kpool. Bail-outs return nullptr and take the dense path. Two are load-bearing: n_pools <= select_k is the free DSA invariant, where selection is a no-op and sparse must equal dense exactly, so routing to dense makes that identity structural; n_tokens != 1 is prefill, which needs block-sparse machinery this does not have, and is compute-bound anyway. All guards run before any graph mutation, or the fallback would emit a second cpy_k into the same slots. tests/test-dsa-pool-gather.cpp pins the four new constructions (K re-view, V sub-view, transpose->get_rows->transpose mask gather, strided slot-0 pool mask). Mutation-tested: dropped transpose, wrong ne0 stride, shifted V offset and wrong re-view width each fail the test. Known limitation, documented and the reason this stays default-off: pooling assumes cache cell i is sequence position i, which breaks under a shared unified cache, context shift or defragmentation. Attention stays correct there but selection degrades silently. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016B6ZbadovyJKFiN1CwDBLA --- src/models/glm5-next.cpp | 190 ++++++++++++++++++++++++++++++++- src/models/models.h | 8 ++ tests/CMakeLists.txt | 1 + tests/test-dsa-pool-gather.cpp | 157 +++++++++++++++++++++++++++ 4 files changed, 354 insertions(+), 2 deletions(-) create mode 100644 tests/test-dsa-pool-gather.cpp diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index a7c94262b..1cdcab367 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -1,5 +1,7 @@ #include "models.h" +#include "../llama-kv-cache.h" // build_attn_dsa calls cpy_k/get_k on the cache context + #include "llama-memory-recurrent.h" // GLM-5.3-Flash (glm5-next). @@ -149,6 +151,182 @@ ggml_tensor * llm_build_glm5_next::build_dsa_index_scores( return ggml_reshape_2d(ctx0, idx, n_pools, S); } +// DSA sparse attention over the KV cache (Phase 3). +// +// Selection happens at POOL granularity, and that is the whole trick that makes this tractable +// in ggml. Cache rows for consecutive tokens are contiguous, so the [D, n_kv] cache view can be +// re-viewed as [D*kpool, n_pools] and a single ggml_get_rows then gathers whole pools. No index +// arithmetic at all - no scaling pool ids into token ids, no I32 maths ggml has no ops for - and +// pool granularity is exactly DSA's own granularity, since select_k = indexer_top_k / kpool. +// +// Returns nullptr whenever the sparse path does not apply and the caller falls back to dense +// build_attn. Two of those bail-outs are load-bearing rather than laziness: +// +// * n_pools <= select_k is the FREE DSA invariant. Every pool would be selected, so selection +// is a no-op and sparse MUST equal dense exactly. Taking the dense path there makes that +// identity structural instead of something to hope a test catches. +// * n_tokens != 1 is prefill. top_k returns a per-token selection, but one gathered K/V can +// only serve one selection, so prefill would need block-sparse machinery this does not have. +// Decode is also where the win is: prefill is compute-bound, decode is KV-bandwidth-bound. +// +// KNOWN LIMITATION, and the reason this stays opt-in: pooling groups kpool CONSECUTIVE CACHE +// CELLS, and treats them as kpool consecutive sequence positions. Those coincide for a single +// sequence filling a fresh cache in order, which is the case this is written for. They stop +// coinciding under anything that reorders cells against positions - a second sequence sharing a +// unified cache, a context shift, defragmentation. Attention itself stays correct there, because +// the gathered mask travels with the gathered rows; what degrades is the SELECTION, which would +// pool unrelated positions and pick the wrong ones. That is a quality regression with no visible +// symptom, so this must not be enabled by default until the pooling reads positions rather than +// assuming them. +// +// The gathered-attention algebra itself (gather K/V, gather the mask through transpose -> +// get_rows -> transpose, then build_attn_mha) is the construction proved against dense attention +// in tests/test-dsa-attn.cpp at 1.19e-07. +ggml_tensor * llm_build_glm5_next::build_attn_dsa( + llm_graph_input_attn_k * inp, + ggml_tensor * wo, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * x, + ggml_tensor * q_a, + const llama_layer & layer, + ggml_tensor * v_mla, + float kq_scale, + int il) { + const int64_t hd = hparams.indexer_head_size; + const int64_t nh = hparams.indexer_n_head; + const int64_t kp = hparams.indexer_kpool ? hparams.indexer_kpool : 4; + const int64_t T = q_cur->ne[2]; + + if (!hparams.dsa_enabled || !layer.indexer_attn_k || !layer.indexer_kpool_gate) { + return nullptr; + } + if (T != 1) { + return nullptr; // prefill: see above + } + + // Every bail-out has to happen BEFORE the graph is touched. If this returned nullptr after + // expanding the stores, the caller's dense build_attn would emit a SECOND cpy_k into the same + // cache slots - so the guards below run against a get_k view that is created but not yet + // expanded, which costs a tensor header and nothing else. + const auto * mctx_cur = inp->mctx; + + ggml_tensor * k = mctx_cur->get_k(ctx0, il); // [D, n_head_kv, n_kv, ns] + + const int64_t D = k->ne[0]; + const int64_t n_kv = k->ne[2]; + const int64_t kv_lora = v_cur->ne[0]; + + // The [D*kp, n_pools] re-view below reads the cache as raw contiguous memory, so everything + // that could make that untrue is a bail-out rather than an assert. + if (k->ne[1] != 1 || k->ne[3] != 1) return nullptr; // MQA/multi-seq only + if (ggml_is_quantized(k->type)) return nullptr; // no fixed element size + if (k->type != GGML_TYPE_F32 && k->type != GGML_TYPE_F16) return nullptr; + if (D != kv_lora + 2*hd) return nullptr; // row not widened + if (n_kv % kp != 0) return nullptr; // ragged final pool + + const int64_t n_pools = n_kv / kp; + const int64_t select_k = hparams.indexer_top_k / kp; + + if (select_k <= 0 || n_pools <= select_k) return nullptr; // free DSA: dense == sparse + + // Committed to the sparse path: now it is safe to mutate the graph. + ggml_build_forward_expand(gf, q_cur); + ggml_build_forward_expand(gf, v_cur); + ggml_build_forward_expand(gf, k_cur); + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, inp->get_k_idxs(), il)); + + const size_t es = ggml_type_size(k->type); + + // --- pooled indexer keys over the whole cache ------------------------------------------ + // k_norm was applied before the key was written into the cache, so the stored indexer key is + // already normalised; re-normalising here would apply it twice. + ggml_tensor * ik = ggml_cont(ctx0, + ggml_view_2d(ctx0, k, hd, n_kv, (size_t) D*es, (size_t) kv_lora*es)); + ggml_tensor * ig = ggml_cont(ctx0, + ggml_view_2d(ctx0, k, hd, n_kv, (size_t) D*es, (size_t) (kv_lora + hd)*es)); + if (k->type != GGML_TYPE_F32) { + ik = ggml_cast(ctx0, ik, GGML_TYPE_F32); + ig = ggml_cast(ctx0, ig, GGML_TYPE_F32); + } + + ggml_tensor * k3 = ggml_reshape_3d(ctx0, ik, hd, kp, n_pools); + ggml_tensor * g3 = ggml_reshape_3d(ctx0, ig, hd, kp, n_pools); + ggml_tensor * ape3 = ggml_reshape_3d(ctx0, layer.indexer_kpool_ape, hd, kp, 1); + + // Per-channel softmax over the kp slots of (gate + ape) - not a mean. soft_max reduces ne0, + // so kp has to be brought there and put back. + ggml_tensor * lg = ggml_add(ctx0, g3, ape3); + lg = ggml_cont(ctx0, ggml_permute(ctx0, lg, 1, 0, 2, 3)); + lg = ggml_soft_max(ctx0, lg); + lg = ggml_cont(ctx0, ggml_permute(ctx0, lg, 1, 0, 2, 3)); + + ggml_tensor * pk = ggml_mul(ctx0, lg, k3); + pk = ggml_cont(ctx0, ggml_permute(ctx0, pk, 1, 0, 2, 3)); + pk = ggml_sum_rows(ctx0, pk); + pk = ggml_reshape_2d(ctx0, pk, hd, n_pools); + + // --- score the pools against this token ------------------------------------------------- + ggml_tensor * q = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, layer.indexer_attn_q_b, q_a), hd, nh, T); + + ggml_tensor * sc = ggml_mul_mat(ctx0, pk, q); // [n_pools, nh, T] + sc = ggml_scale(ctx0, sc, 1.0f/sqrtf((float) hd)); + sc = ggml_relu(ctx0, sc); // load-bearing + + ggml_tensor * wgt = ggml_scale(ctx0, ggml_mul_mat(ctx0, layer.indexer_proj, x), + 1.0f/sqrtf((float) nh)); + ggml_tensor * scp = ggml_cont(ctx0, ggml_permute(ctx0, sc, 1, 0, 2, 3)); + ggml_tensor * idx = ggml_mul_mat(ctx0, scp, ggml_reshape_3d(ctx0, wgt, nh, 1, T)); + ggml_tensor * scores = ggml_reshape_2d(ctx0, idx, n_pools, T); // [n_pools, T] + + ggml_tensor * kq_mask = inp->get_kq_mask(); // [n_kv, T_pad] + + // A pool must not be selected if none of its tokens are visible, or top_k spends slots on + // rows that attention will then mask to -inf. Under a causal mask a pool has a visible token + // iff its FIRST token is visible, so slot 0 of each pool is the exact pool-level mask - and + // -inf + finite is -inf, so adding it removes those pools from contention. + ggml_tensor * pmask = ggml_cont(ctx0, + ggml_view_3d(ctx0, kq_mask, 1, n_pools, T, + (size_t) kp*ggml_type_size(kq_mask->type), kq_mask->nb[1], 0)); + scores = ggml_add(ctx0, scores, ggml_reshape_2d(ctx0, pmask, n_pools, T)); + cb(scores, "dsa_pool_scores", il); + + ggml_tensor * sel = ggml_reshape_1d(ctx0, ggml_top_k(ctx0, scores, select_k), select_k); + cb(sel, "dsa_sel", il); + + const int64_t n_sel = select_k * kp; + + // --- gather K, V and the mask ----------------------------------------------------------- + ggml_tensor * kpools = ggml_view_2d(ctx0, k, D*kp, n_pools, (size_t) (D*kp)*es, 0); + ggml_tensor * ksel = ggml_get_rows(ctx0, kpools, sel); // [D*kp, select_k], F32 + ksel = ggml_reshape_4d(ctx0, ksel, D, 1, n_sel, 1); + + // V is the leading kv_lora channels of the gathered row. wv_b expands from the compressed + // latent, so the indexer key and gate must not reach it. + ggml_tensor * vsel = ggml_view_4d(ctx0, ksel, kv_lora, ksel->ne[1], ksel->ne[2], ksel->ne[3], + ksel->nb[1], ksel->nb[2], ksel->nb[3], 0); + + // The mask is [n_kv, T_pad] and get_rows selects along ne1, so it has to be transposed, + // gathered, and transposed back. This is the pattern proved in tests/test-dsa-attn.cpp. + const int64_t T_pad = kq_mask->ne[1]; + ggml_tensor * mt = ggml_cont(ctx0, ggml_transpose(ctx0, kq_mask)); // [T_pad, n_kv] + mt = ggml_reshape_2d(ctx0, mt, T_pad*kp, n_pools); + ggml_tensor * msel = ggml_get_rows(ctx0, mt, sel); // [T_pad*kp, select_k] + msel = ggml_reshape_2d(ctx0, msel, T_pad, n_sel); + msel = ggml_cont(ctx0, ggml_transpose(ctx0, msel)); // [n_sel, T_pad] + cb(msel, "dsa_mask_sel", il); + + ggml_tensor * cur = build_attn_mha(q_cur, ksel, vsel, nullptr, msel, nullptr, + v_mla, kq_scale, il); + cb(cur, "kqv_out", il); + + if (wo) { + cur = build_lora_mm(wo, cur); + } + return cur; +} + llm_build_glm5_next::mhc_site llm_build_glm5_next::build_mhc( ggml_tensor * streams, ggml_tensor * fn, ggml_tensor * base, ggml_tensor * scale, int il) { @@ -437,8 +615,16 @@ llm_build_glm5_next::llm_build_glm5_next(const llama_model & model, const llm_gr cb(Kcur, "dsa_k_widened", il); } - cur = build_attn(inp_attn_k, layer.wo, NULL, Qcur, Kcur, Vcur, - nullptr, nullptr, layer.wv_b, kq_scale_mla, il); + // Sparse selection when it applies, dense otherwise. build_attn_dsa returns + // nullptr for every case it does not handle (prefill, short context, an + // un-widened row), so the dense path stays the default and the fallback is a + // plain null check rather than a duplicated condition that could drift. + cur = build_attn_dsa(inp_attn_k, layer.wo, Qcur, Kcur, Vcur, + cur, q_a, layer, layer.wv_b, kq_scale_mla, il); + if (cur == nullptr) { + cur = build_attn(inp_attn_k, layer.wo, NULL, Qcur, Kcur, Vcur, + nullptr, nullptr, layer.wv_b, kq_scale_mla, il); + } } else { Qcur = ggml_reshape_3d(ctx0, Qcur, n_embd_head_k_mla, n_head, n_tokens); ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_b, kv_cmpr); diff --git a/src/models/models.h b/src/models/models.h index fbd40b430..91bab3fd4 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -383,6 +383,14 @@ struct llm_build_glm5_next : public llm_build_delta_net_base { ggml_tensor * build_dsa_index_scores(ggml_tensor * x, ggml_tensor * q_a, const llama_layer & layer, int il); + // DSA sparse attention over the KV cache (decode only). Returns nullptr when the sparse + // path does not apply - notably at n_kv <= indexer_top_k, where selection is a no-op and + // dense IS the correct answer - and the caller falls back to dense build_attn. + ggml_tensor * build_attn_dsa(llm_graph_input_attn_k * inp, ggml_tensor * wo, + ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, + ggml_tensor * x, ggml_tensor * q_a, const llama_layer & layer, + ggml_tensor * v_mla, float kq_scale, int il); + const llama_model & model; }; diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 11ac1645f..6cbde360d 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -153,6 +153,7 @@ llama_build_and_test(test-mhc-sinkhorn.cpp) llama_build(test-dsa-indexer.cpp) llama_build(test-dsa-gather.cpp) llama_build(test-dsa-attn.cpp) +llama_build(test-dsa-pool-gather.cpp) llama_build(test-glm5-next-logits.cpp) llama_build_and_test(test-reasoning-budget.cpp) llama_build_and_test(test-grammar-parser.cpp) diff --git a/tests/test-dsa-pool-gather.cpp b/tests/test-dsa-pool-gather.cpp new file mode 100644 index 000000000..91be27da1 --- /dev/null +++ b/tests/test-dsa-pool-gather.cpp @@ -0,0 +1,157 @@ +// DSA Phase 3: pool-granularity gather off the KV cache. +// +// build_attn_dsa does not gather individual tokens. It re-views the [D, n_kv] cache as +// [D*kpool, n_pools] and calls get_rows once with POOL ids, so one row of the re-view is a whole +// pool of kpool consecutive tokens. That removes all index arithmetic -- ggml has no ops for +// scaling I32 pool ids into token ids -- and pool granularity is DSA's own granularity, since +// select_k = indexer_top_k / kpool. +// +// The re-view is only legal because consecutive cache cells are contiguous in memory. That is +// true here (n_head_kv == 1, so n_embd_k_gqa == D and get_k's nb[2] is exactly D*es) but it is an +// assumption about layout rather than something the type system enforces, and if it were wrong +// the gather would return the right NUMBER of rows with the wrong contents -- which still decodes +// to fluent text. Hence this test. +// +// Three constructions are pinned down, all of which appear verbatim in build_attn_dsa: +// 1. K gather: [D, n_kv] -> [D*kpool, n_pools] -> get_rows(pool ids) -> [D, sel]. +// 2. Mask gather: the mask is [n_kv, T] and get_rows selects along ne1, so it goes +// transpose -> reshape to pool rows -> get_rows -> reshape -> transpose back. +// 3. Pool-level causal mask: slot 0 of each pool, taken as a strided ne0 view. Under a causal +// mask a pool has a visible token iff its FIRST token is visible, so this is the exact +// pool mask, and adding it to the scores keeps top_k from spending slots on dead pools. +// +// Values are chosen so that any off-by-one-pool or transposed gather produces a mismatch rather +// than a plausible-looking permutation. +#include "ggml.h" +#include "ggml-cpu.h" + +#include +#include +#include +#include +#include + +static int failures = 0; + +static void check(bool ok, const char * what) { + printf("%-58s %s\n", what, ok ? "OK" : "FAIL"); + if (!ok) failures++; +} + +int main() { + const int kv_lora = 4; + const int ihd = 2; + const int D = kv_lora + 2*ihd; // widened row: latent + indexer key + gate + const int n_kv = 24; + const int kv_size = 40; // cache is BIGGER than n_kv, as at runtime + const int kpool = 4; + const int n_pools = n_kv / kpool; // 6 + const int T = 3; + const int select_k = 3; + const int sel = select_k * kpool; // 12 of 24 tokens survive + + // Deliberately not in ascending order: an implementation that ignores the ids and takes the + // first select_k pools would pass an ascending list. + const std::vector pool_ids = { 4, 0, 3 }; + + ggml_init_params ip = { (size_t) 64*1024*1024, nullptr, false }; + ggml_context * ctx = ggml_init(ip); + + // --- cache, and the strided [D, 1, n_kv, 1] view get_k would hand back ------------------- + ggml_tensor * cache = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, D, kv_size); + float * cd = (float *) cache->data; + for (int t = 0; t < kv_size; t++) + for (int c = 0; c < D; c++) + cd[t*D + c] = 1000.0f*t + c; // token id is readable off any element + + const size_t es = ggml_type_size(cache->type); + ggml_tensor * k = ggml_view_4d(ctx, cache, D, 1, n_kv, 1, + (size_t) D*es, (size_t) D*es, (size_t) D*es*kv_size, 0); + + ggml_tensor * ids = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, select_k); + memcpy(ids->data, pool_ids.data(), pool_ids.size()*sizeof(int32_t)); + + // 1. K gather at pool granularity + ggml_tensor * kpools = ggml_view_2d(ctx, k, D*kpool, n_pools, (size_t) (D*kpool)*es, 0); + ggml_tensor * ksel = ggml_get_rows(ctx, kpools, ids); + ksel = ggml_reshape_4d(ctx, ksel, D, 1, sel, 1); + + // V is the leading kv_lora channels of the gathered row, exactly as in build_attn_dsa. + ggml_tensor * vsel = ggml_view_4d(ctx, ksel, kv_lora, ksel->ne[1], ksel->ne[2], ksel->ne[3], + ksel->nb[1], ksel->nb[2], ksel->nb[3], 0); + ggml_tensor * vsel_c = ggml_cont(ctx, vsel); + + // --- mask --------------------------------------------------------------------------------- + ggml_tensor * mask = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_kv, T); + float * md = (float *) mask->data; + for (int t = 0; t < T; t++) + for (int j = 0; j < n_kv; j++) + md[t*n_kv + j] = 100.0f*j + t; + + // 2. mask gather: transpose -> pool rows -> get_rows -> back + ggml_tensor * mt = ggml_cont(ctx, ggml_transpose(ctx, mask)); // [T, n_kv] + mt = ggml_reshape_2d(ctx, mt, T*kpool, n_pools); + ggml_tensor * msel = ggml_get_rows(ctx, mt, ids); // [T*kpool, select_k] + msel = ggml_reshape_2d(ctx, msel, T, sel); + msel = ggml_cont(ctx, ggml_transpose(ctx, msel)); // [sel, T] + + // 3. pool-level mask: slot 0 of each pool, strided along ne0 + ggml_tensor * pmask = ggml_cont(ctx, + ggml_view_3d(ctx, mask, 1, n_pools, T, (size_t) kpool*es, mask->nb[1], 0)); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, ksel); + ggml_build_forward_expand(gf, vsel_c); + ggml_build_forward_expand(gf, msel); + ggml_build_forward_expand(gf, pmask); + ggml_graph_compute_with_ctx(ctx, gf, 1); + + // --- verify ------------------------------------------------------------------------------ + const float * kg = (const float *) ksel->data; + bool ok_k = true; + for (int s = 0; s < sel && ok_k; s++) { + const int tok = pool_ids[s / kpool]*kpool + (s % kpool); + for (int c = 0; c < D; c++) + if (kg[s*D + c] != 1000.0f*tok + c) { ok_k = false; break; } + } + check(ok_k, "K: pool re-view gathers the right kpool tokens"); + + const float * vg = (const float *) vsel_c->data; + bool ok_v = true; + for (int s = 0; s < sel && ok_v; s++) { + const int tok = pool_ids[s / kpool]*kpool + (s % kpool); + for (int c = 0; c < kv_lora; c++) + if (vg[s*kv_lora + c] != 1000.0f*tok + c) { ok_v = false; break; } + } + check(ok_v, "V: leading kv_lora channels, same rows as K"); + + const float * mg = (const float *) msel->data; + bool ok_m = true; + for (int t = 0; t < T && ok_m; t++) + for (int s = 0; s < sel; s++) { + const int tok = pool_ids[s / kpool]*kpool + (s % kpool); + if (mg[t*sel + s] != 100.0f*tok + t) { ok_m = false; break; } + } + check(ok_m, "mask: transpose -> get_rows -> transpose is row-exact"); + + const float * pg = (const float *) pmask->data; + bool ok_p = true; + for (int t = 0; t < T && ok_p; t++) + for (int p = 0; p < n_pools; p++) + if (pg[t*n_pools + p] != 100.0f*(p*kpool) + t) { ok_p = false; break; } + check(ok_p, "pool mask: strided ne0 view picks slot 0 of each pool"); + + // The gathered rows must be a SUBSET, not a reordering that happens to line up: pool 1, 2 + // and 5 were not selected and must not appear anywhere in the gather. + bool ok_excl = true; + for (int s = 0; s < sel; s++) { + const int tok = (int) (kg[s*D] / 1000.0f); + const int pool = tok / kpool; + if (pool != pool_ids[s / kpool]) { ok_excl = false; break; } + } + check(ok_excl, "unselected pools are absent from the gather"); + + ggml_free(ctx); + printf("\n%s\n", failures ? "FAILED" : "PASS"); + return failures ? 1 : 0; +} -- 2.43.0 From 66fbc4362d368d2013ea641098b2009431fca3e0 Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 22:00:19 -0400 Subject: [PATCH 11/12] glm5-next: DSA selection checks cache position ordering, and always attends the newest tokens Pooling groups kpool consecutive cache CELLS and treats them as kpool consecutive sequence POSITIONS. That breaks under a shared unified cache, a context shift or defragmentation, where attention stays correct but SELECTION silently pools unrelated positions. llama_kv_cache::pos_ordered_prefix now reports the length of the leading run where cell i holds position i, requires the rest to be empty, and returns 0 for any layout pooling cannot account for. Memoised per ubatch; the sparse path bails to dense on 0. The prefix length also fixes what would otherwise have made this dead code: get_n_kv pads n_kv to a multiple of 256, so a strict 'cell i is position i for all of n_kv' test is false almost always, leaving the sparse path dormant. The padded window ends mid-pool, and that partial pool holds the NEWEST tokens - the ones a decode step must not lose - while being unscoreable because it is not a whole pool. It is now always attended rather than selected, appended to the gather as a plain view on both K and the mask at a host-known offset. No index tensor is needed, and top_k cannot pick it twice because scoring only sees the complete pools. test-dsa-pool-gather gains three checks for the tail append; off-by-one-pool mutations on the K offset and the mask offset are both caught. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016B6ZbadovyJKFiN1CwDBLA --- src/llama-kv-cache.cpp | 29 ++++++++++++++ src/llama-kv-cache.h | 15 +++++++ src/models/glm5-next.cpp | 73 ++++++++++++++++++++++++---------- tests/test-dsa-pool-gather.cpp | 53 ++++++++++++++++++++++-- 4 files changed, 144 insertions(+), 26 deletions(-) diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 3e0fd3107..11c9e6faf 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1632,6 +1632,28 @@ void llama_kv_cache::set_input_kq_mask(ggml_tensor * dst, const llama_ubatch * u //LLAMA_LOG_ERROR("%s: kq mask time: %0.3f ms\n", __func__, (t_end - t_start)/1000.0); } +uint32_t llama_kv_cache::pos_ordered_prefix(uint32_t n_kv_) const { + if (n_stream != 1) { + return 0; + } + const auto & cells = v_cells[0]; + if (n_kv_ > cells.size()) { + return 0; + } + uint32_t n = 0; + while (n < n_kv_ && !cells.is_empty(n) && cells.pos_get(n) == (llama_pos) n) { + ++n; + } + // Everything after the ordered prefix must be empty padding. A filled cell out there means + // the cache holds content this pooling cannot account for, so refuse rather than pool it. + for (uint32_t i = n; i < n_kv_; ++i) { + if (!cells.is_empty(i)) { + return 0; + } + } + return n; +} + void llama_kv_cache::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const { const int64_t n_tokens = ubatch->n_tokens; @@ -2479,6 +2501,13 @@ void llama_kv_cache_context::set_input_kq_mask(ggml_tensor * dst, const llama_ub kv->set_input_kq_mask(dst, ubatch, causal_attn); } +uint32_t llama_kv_cache_context::pos_ordered_prefix() const { + if (cached_pos_prefix < 0) { + cached_pos_prefix = (int64_t) kv->pos_ordered_prefix((uint32_t) n_kv); + } + return (uint32_t) cached_pos_prefix; +} + void llama_kv_cache_context::set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const { kv->set_input_pos_bucket(dst, ubatch); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index d4569a06f..8c63664e0 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -204,6 +204,16 @@ public: void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; + // Length of the leading run of cells where cell i holds sequence position i, or 0 if the + // layout is unusable for pooled selection. Cells past that run must all be EMPTY, which is + // the normal case rather than a defect: get_n_kv pads n_kv up to a multiple of 256. + // + // DSA pooling groups kpool consecutive CELLS and treats them as kpool consecutive POSITIONS. + // That coincidence breaks under a shared unified cache, a context shift or defragmentation, + // where attention itself stays correct (the mask travels with the gathered rows) but + // SELECTION would pool unrelated positions and drop the context the answer needed - a + // quality regression with no visible symptom. So the sparse path asks rather than assumes. + uint32_t pos_ordered_prefix(uint32_t n_kv) const; void set_input_k_rot(ggml_tensor * dst) const; void set_input_v_rot(ggml_tensor * dst) const; @@ -376,6 +386,9 @@ public: void set_input_k_shift (ggml_tensor * dst) const; void set_input_kq_mask (ggml_tensor * dst, const llama_ubatch * ubatch, bool causal_attn) const; void set_input_pos_bucket(ggml_tensor * dst, const llama_ubatch * ubatch) const; + // See llama_kv_cache::pos_ordered_prefix. Memoised: the sparse path asks once per layer and + // the answer cannot change within one ubatch. + uint32_t pos_ordered_prefix() const; void set_input_k_rot(ggml_tensor * dst) const; void set_input_v_rot(ggml_tensor * dst) const; @@ -412,4 +425,6 @@ private: // a heuristic, to avoid attending the full cache if it is not yet utilized // as the cache gets filled, the benefit from this heuristic disappears int32_t n_kv; + + mutable int64_t cached_pos_prefix = -1; // -1 unknown, else the prefix length }; diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp index 1cdcab367..85135f67b 100644 --- a/src/models/glm5-next.cpp +++ b/src/models/glm5-next.cpp @@ -169,15 +169,13 @@ ggml_tensor * llm_build_glm5_next::build_dsa_index_scores( // only serve one selection, so prefill would need block-sparse machinery this does not have. // Decode is also where the win is: prefill is compute-bound, decode is KV-bandwidth-bound. // -// KNOWN LIMITATION, and the reason this stays opt-in: pooling groups kpool CONSECUTIVE CACHE -// CELLS, and treats them as kpool consecutive sequence positions. Those coincide for a single -// sequence filling a fresh cache in order, which is the case this is written for. They stop -// coinciding under anything that reorders cells against positions - a second sequence sharing a -// unified cache, a context shift, defragmentation. Attention itself stays correct there, because -// the gathered mask travels with the gathered rows; what degrades is the SELECTION, which would -// pool unrelated positions and pick the wrong ones. That is a quality regression with no visible -// symptom, so this must not be enabled by default until the pooling reads positions rather than -// assuming them. +// Pooling groups kpool consecutive CACHE CELLS and treats them as kpool consecutive sequence +// POSITIONS. That coincidence holds for a single sequence filling a fresh cache in order and +// breaks under anything that reorders cells against positions - a second sequence sharing a +// unified cache, a context shift, defragmentation. Attention would stay correct there (the +// gathered mask travels with the gathered rows) but SELECTION would pool unrelated positions and +// drop the context the answer needed, which is a quality regression with no visible symptom. +// So the cache is asked, via llama_kv_cache::pos_ordered_prefix, instead of assumed. // // The gathered-attention algebra itself (gather K/V, gather the mask through transpose -> // get_rows -> transpose, then build_attn_mha) is the construction proved against dense attention @@ -226,10 +224,25 @@ ggml_tensor * llm_build_glm5_next::build_attn_dsa( if (D != kv_lora + 2*hd) return nullptr; // row not widened if (n_kv % kp != 0) return nullptr; // ragged final pool - const int64_t n_pools = n_kv / kp; + // Pooling groups kpool consecutive CELLS and calls them kpool consecutive POSITIONS. Ask the + // cache rather than assume it, and take the length of the ordered prefix while we are here: + // n_kv is padded up to a multiple of 256, so the tail of the window is empty cells that must + // not be pooled as if they held content. + const int64_t n_used = (int64_t) mctx_cur->pos_ordered_prefix(); + if (n_used == 0) return nullptr; + + const int64_t n_src = n_kv / kp; // pool rows the gather source spans, padding included + const int64_t n_full = n_used / kp; // complete pools of real tokens - the selectable set + const int64_t n_tail = (n_used % kp) ? 1 : 0; const int64_t select_k = hparams.indexer_top_k / kp; - if (select_k <= 0 || n_pools <= select_k) return nullptr; // free DSA: dense == sparse + if (select_k <= 0 || n_full <= select_k) return nullptr; // free DSA: dense == sparse + + // The partial pool holds the NEWEST 1..kp-1 tokens, which are exactly the ones a decode step + // must not lose, and it cannot be scored because it is not a whole pool. It is always + // attended instead of selected: its position is known on the host, so it is a plain view + // concatenated onto the gather - no index tensor, and no way for top_k to pick it twice + // since scoring only ever sees the n_full complete pools. // Committed to the sparse path: now it is safe to mutate the graph. ggml_build_forward_expand(gf, q_cur); @@ -242,17 +255,18 @@ ggml_tensor * llm_build_glm5_next::build_attn_dsa( // --- pooled indexer keys over the whole cache ------------------------------------------ // k_norm was applied before the key was written into the cache, so the stored indexer key is // already normalised; re-normalising here would apply it twice. + const int64_t n_scored = n_full * kp; // real tokens in complete pools ggml_tensor * ik = ggml_cont(ctx0, - ggml_view_2d(ctx0, k, hd, n_kv, (size_t) D*es, (size_t) kv_lora*es)); + ggml_view_2d(ctx0, k, hd, n_scored, (size_t) D*es, (size_t) kv_lora*es)); ggml_tensor * ig = ggml_cont(ctx0, - ggml_view_2d(ctx0, k, hd, n_kv, (size_t) D*es, (size_t) (kv_lora + hd)*es)); + ggml_view_2d(ctx0, k, hd, n_scored, (size_t) D*es, (size_t) (kv_lora + hd)*es)); if (k->type != GGML_TYPE_F32) { ik = ggml_cast(ctx0, ik, GGML_TYPE_F32); ig = ggml_cast(ctx0, ig, GGML_TYPE_F32); } - ggml_tensor * k3 = ggml_reshape_3d(ctx0, ik, hd, kp, n_pools); - ggml_tensor * g3 = ggml_reshape_3d(ctx0, ig, hd, kp, n_pools); + ggml_tensor * k3 = ggml_reshape_3d(ctx0, ik, hd, kp, n_full); + ggml_tensor * g3 = ggml_reshape_3d(ctx0, ig, hd, kp, n_full); ggml_tensor * ape3 = ggml_reshape_3d(ctx0, layer.indexer_kpool_ape, hd, kp, 1); // Per-channel softmax over the kp slots of (gate + ape) - not a mean. soft_max reduces ne0, @@ -265,7 +279,7 @@ ggml_tensor * llm_build_glm5_next::build_attn_dsa( ggml_tensor * pk = ggml_mul(ctx0, lg, k3); pk = ggml_cont(ctx0, ggml_permute(ctx0, pk, 1, 0, 2, 3)); pk = ggml_sum_rows(ctx0, pk); - pk = ggml_reshape_2d(ctx0, pk, hd, n_pools); + pk = ggml_reshape_2d(ctx0, pk, hd, n_full); // --- score the pools against this token ------------------------------------------------- ggml_tensor * q = ggml_reshape_3d(ctx0, ggml_mul_mat(ctx0, layer.indexer_attn_q_b, q_a), hd, nh, T); @@ -278,7 +292,7 @@ ggml_tensor * llm_build_glm5_next::build_attn_dsa( 1.0f/sqrtf((float) nh)); ggml_tensor * scp = ggml_cont(ctx0, ggml_permute(ctx0, sc, 1, 0, 2, 3)); ggml_tensor * idx = ggml_mul_mat(ctx0, scp, ggml_reshape_3d(ctx0, wgt, nh, 1, T)); - ggml_tensor * scores = ggml_reshape_2d(ctx0, idx, n_pools, T); // [n_pools, T] + ggml_tensor * scores = ggml_reshape_2d(ctx0, idx, n_full, T); // [n_full, T] ggml_tensor * kq_mask = inp->get_kq_mask(); // [n_kv, T_pad] @@ -287,19 +301,26 @@ ggml_tensor * llm_build_glm5_next::build_attn_dsa( // iff its FIRST token is visible, so slot 0 of each pool is the exact pool-level mask - and // -inf + finite is -inf, so adding it removes those pools from contention. ggml_tensor * pmask = ggml_cont(ctx0, - ggml_view_3d(ctx0, kq_mask, 1, n_pools, T, + ggml_view_3d(ctx0, kq_mask, 1, n_full, T, (size_t) kp*ggml_type_size(kq_mask->type), kq_mask->nb[1], 0)); - scores = ggml_add(ctx0, scores, ggml_reshape_2d(ctx0, pmask, n_pools, T)); + scores = ggml_add(ctx0, scores, ggml_reshape_2d(ctx0, pmask, n_full, T)); cb(scores, "dsa_pool_scores", il); ggml_tensor * sel = ggml_reshape_1d(ctx0, ggml_top_k(ctx0, scores, select_k), select_k); cb(sel, "dsa_sel", il); - const int64_t n_sel = select_k * kp; + const int64_t n_sel = (select_k + n_tail) * kp; // --- gather K, V and the mask ----------------------------------------------------------- - ggml_tensor * kpools = ggml_view_2d(ctx0, k, D*kp, n_pools, (size_t) (D*kp)*es, 0); + // Rows of the source span the padded window; selection can only ever name a pool below + // n_full, and the tail pool is appended by view rather than by index. + ggml_tensor * kpools = ggml_view_2d(ctx0, k, D*kp, n_src, (size_t) (D*kp)*es, 0); ggml_tensor * ksel = ggml_get_rows(ctx0, kpools, sel); // [D*kp, select_k], F32 + if (n_tail) { + ggml_tensor * ktail = ggml_view_2d(ctx0, k, D*kp, 1, (size_t) (D*kp)*es, + (size_t) (n_full*kp*D)*es); + ksel = ggml_concat(ctx0, ksel, ggml_cast(ctx0, ktail, ksel->type), 1); + } ksel = ggml_reshape_4d(ctx0, ksel, D, 1, n_sel, 1); // V is the leading kv_lora channels of the gathered row. wv_b expands from the compressed @@ -309,10 +330,18 @@ ggml_tensor * llm_build_glm5_next::build_attn_dsa( // The mask is [n_kv, T_pad] and get_rows selects along ne1, so it has to be transposed, // gathered, and transposed back. This is the pattern proved in tests/test-dsa-attn.cpp. + // The tail pool's mask row is appended the same way its K was, keeping mask rows paired with + // the K rows they mask - and it is the mask, not the selection, that hides the padding slots + // inside that partial pool. const int64_t T_pad = kq_mask->ne[1]; ggml_tensor * mt = ggml_cont(ctx0, ggml_transpose(ctx0, kq_mask)); // [T_pad, n_kv] - mt = ggml_reshape_2d(ctx0, mt, T_pad*kp, n_pools); + mt = ggml_reshape_2d(ctx0, mt, T_pad*kp, n_src); ggml_tensor * msel = ggml_get_rows(ctx0, mt, sel); // [T_pad*kp, select_k] + if (n_tail) { + ggml_tensor * mtail = ggml_view_2d(ctx0, mt, T_pad*kp, 1, mt->nb[1], + (size_t) n_full*mt->nb[1]); + msel = ggml_concat(ctx0, msel, ggml_cast(ctx0, mtail, msel->type), 1); + } msel = ggml_reshape_2d(ctx0, msel, T_pad, n_sel); msel = ggml_cont(ctx0, ggml_transpose(ctx0, msel)); // [n_sel, T_pad] cb(msel, "dsa_mask_sel", il); diff --git a/tests/test-dsa-pool-gather.cpp b/tests/test-dsa-pool-gather.cpp index 91be27da1..984f2b053 100644 --- a/tests/test-dsa-pool-gather.cpp +++ b/tests/test-dsa-pool-gather.cpp @@ -73,8 +73,8 @@ int main() { // 1. K gather at pool granularity ggml_tensor * kpools = ggml_view_2d(ctx, k, D*kpool, n_pools, (size_t) (D*kpool)*es, 0); - ggml_tensor * ksel = ggml_get_rows(ctx, kpools, ids); - ksel = ggml_reshape_4d(ctx, ksel, D, 1, sel, 1); + ggml_tensor * ksel_pre = ggml_get_rows(ctx, kpools, ids); // [D*kpool, select_k] + ggml_tensor * ksel = ggml_reshape_4d(ctx, ksel_pre, D, 1, sel, 1); // V is the leading kv_lora channels of the gathered row, exactly as in build_attn_dsa. ggml_tensor * vsel = ggml_view_4d(ctx, ksel, kv_lora, ksel->ne[1], ksel->ne[2], ksel->ne[3], @@ -91,15 +91,32 @@ int main() { // 2. mask gather: transpose -> pool rows -> get_rows -> back ggml_tensor * mt = ggml_cont(ctx, ggml_transpose(ctx, mask)); // [T, n_kv] mt = ggml_reshape_2d(ctx, mt, T*kpool, n_pools); - ggml_tensor * msel = ggml_get_rows(ctx, mt, ids); // [T*kpool, select_k] - msel = ggml_reshape_2d(ctx, msel, T, sel); + ggml_tensor * msel_pre = ggml_get_rows(ctx, mt, ids); // [T*kpool, select_k] + ggml_tensor * msel = ggml_reshape_2d(ctx, msel_pre, T, sel); msel = ggml_cont(ctx, ggml_transpose(ctx, msel)); // [sel, T] // 3. pool-level mask: slot 0 of each pool, strided along ne0 ggml_tensor * pmask = ggml_cont(ctx, ggml_view_3d(ctx, mask, 1, n_pools, T, (size_t) kpool*es, mask->nb[1], 0)); + // --- tail pool: the newest 1..kpool-1 tokens --------------------------------------------- + // n_kv is padded up to a multiple of 256, so the live window ends mid-pool. That partial pool + // holds the NEWEST tokens, which a decode step must not lose, and it cannot be scored because + // it is not a whole pool. It is appended by VIEW at a host-known offset rather than selected, + // which is also why top_k can never pick it twice: scoring only sees the complete pools. + const int n_full = 5; // complete pools of real tokens + ggml_tensor * ktail = ggml_view_2d(ctx, cache, D*kpool, 1, (size_t) (D*kpool)*es, + (size_t) (n_full*kpool*D)*es); + ggml_tensor * kall = ggml_concat(ctx, ksel_pre, ggml_cast(ctx, ktail, ksel_pre->type), 1); + + ggml_tensor * mtail = ggml_view_2d(ctx, mt, T*kpool, 1, mt->nb[1], (size_t) n_full*mt->nb[1]); + ggml_tensor * mall = ggml_concat(ctx, msel_pre, ggml_cast(ctx, mtail, msel_pre->type), 1); + mall = ggml_reshape_2d(ctx, mall, T, sel + kpool); + mall = ggml_cont(ctx, ggml_transpose(ctx, mall)); + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, kall); + ggml_build_forward_expand(gf, mall); ggml_build_forward_expand(gf, ksel); ggml_build_forward_expand(gf, vsel_c); ggml_build_forward_expand(gf, msel); @@ -151,6 +168,34 @@ int main() { } check(ok_excl, "unselected pools are absent from the gather"); + // the appended pool must be pool n_full, sitting after the selected ones + const float * ka = (const float *) kall->data; + bool ok_tail = true; + for (int j = 0; j < kpool && ok_tail; j++) { + const int tok = n_full*kpool + j; + for (int c = 0; c < D; c++) + if (ka[(sel + j)*D + c] != 1000.0f*tok + c) { ok_tail = false; break; } + } + check(ok_tail, "tail pool is appended by view, after the selected pools"); + + // and the selected pools must be untouched by the append + bool ok_keep = true; + for (int s2 = 0; s2 < sel && ok_keep; s2++) { + const int tok = pool_ids[s2 / kpool]*kpool + (s2 % kpool); + for (int c = 0; c < D; c++) + if (ka[s2*D + c] != 1000.0f*tok + c) { ok_keep = false; break; } + } + check(ok_keep, "appending the tail does not disturb the selected rows"); + + const float * ma = (const float *) mall->data; + bool ok_mtail = true; + for (int t = 0; t < T && ok_mtail; t++) + for (int j = 0; j < kpool; j++) { + const int tok = n_full*kpool + j; + if (ma[t*(sel + kpool) + sel + j] != 100.0f*tok + t) { ok_mtail = false; break; } + } + check(ok_mtail, "tail mask row stays paired with the tail K rows"); + ggml_free(ctx); printf("\n%s\n", failures ? "FAILED" : "PASS"); return failures ? 1 : 0; -- 2.43.0 From 2a4a41238175cc5d0ee3e591865653e49096c782 Mon Sep 17 00:00:00 2001 From: Patrick Devaney Date: Sun, 30 Aug 2026 22:07:19 -0400 Subject: [PATCH 12/12] convert: carry GLM-5.3's other two EOG tokens into the GGUF generation_config.json lists eos_token_id = [154820 <|endoftext|>, 154827 <|user|>, 154829 <|observation|>], but a GGUF header carries only one eos id, so the other two were dropped and llama.cpp ended up with <|endoftext|> as its only EOG token. The model ends an assistant turn with <|user|>, so nothing stopped generation: it answered, emitted <|user|>, then hallucinated a follow-up question and answered that too. Observed directly on a chart-reading question - correct in the first turn, wrong in the fabricated second - so any harness reading the tail of the output grades the wrong turn. llama.cpp folds eos, eot and eom into its EOG set and neither <|user|> nor <|observation|> is in its name-matching list, so map the extra two onto eot/eom, which is what those fields are for. Verified with --override-kv on an already built file: all three are then listed as EOG and generation stops at the end of the answer. Co-Authored-By: Claude Opus 5 Claude-Session: https://claude.ai/code/session_016B6ZbadovyJKFiN1CwDBLA --- convert_hf_to_gguf.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index 2f37829cf..45a8eadaf 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -6127,6 +6127,28 @@ class Glm5NextModel(TextModel): def set_vocab(self): self._set_vocab_gpt2() + # GLM-5.3 has THREE end-of-generation tokens. generation_config.json lists + # eos_token_id = [154820 <|endoftext|>, 154827 <|user|>, 154829 <|observation|>], but a + # GGUF carries only one eos id, so the other two were being dropped and llama.cpp ended up + # with <|endoftext|> as its only EOG token. + # + # That is not cosmetic. The model ends an assistant turn with <|user|>, so with only + # <|endoftext|> registered nothing stops generation: the model answers, emits <|user|>, + # then hallucinates a follow-up question and answers that too. Observed directly - a + # chart-reading answer was correct in the first turn and wrong in the fabricated second. + # + # llama.cpp folds eos, eot and eom into its EOG set (llama-vocab.cpp), and neither + # <|user|> nor <|observation|> is in its name-matching list, so map the extra two onto + # eot/eom, which is what those fields are for. + gen_cfg = self.dir_model / "generation_config.json" + if gen_cfg.is_file(): + with open(gen_cfg, encoding="utf-8") as f: + eos_ids = json.load(f).get("eos_token_id") + if isinstance(eos_ids, list) and len(eos_ids) > 1: + for field, tid in zip(("eot", "eom"), eos_ids[1:]): + getattr(self.gguf_writer, f"add_{field}_token_id")(int(tid)) + logger.info(f"gguf: {field} token id = {tid}") + def set_gguf_parameters(self): super().set_gguf_parameters() self.gguf_writer.add_vocab_size(self.hparams["vocab_size"]) -- 2.43.0