Instructions to use patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF with libraries, inference providers, notebooks, and local apps. Follow these links to get started.
- Notebooks
- Google Colab
- Kaggle
- Local Apps Settings
- llama.cpp
How to use patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF with llama.cpp:
Install (macOS, Linux)
curl -LsSf https://llama.app/install.sh | sh # Start a local OpenAI-compatible server with a web UI: llama serve -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M # Run inference directly in the terminal: llama cli -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
Install from WinGet (Windows)
winget install llama.cpp # Start a local OpenAI-compatible server with a web UI: llama serve -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M # Run inference directly in the terminal: llama cli -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
Use pre-built binary
# Download pre-built binary from: # https://github.com/ggerganov/llama.cpp/releases # Start a local OpenAI-compatible server with a web UI: ./llama-server -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M # Run inference directly in the terminal: ./llama-cli -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
Build from source code
git clone https://github.com/ggerganov/llama.cpp.git cd llama.cpp cmake -B build cmake --build build -j --target llama-server llama-cli # Start a local OpenAI-compatible server with a web UI: ./build/bin/llama-server -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M # Run inference directly in the terminal: ./build/bin/llama-cli -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
Use Docker
docker model run hf.co/patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
- LM Studio
- Jan
- Ollama
How to use patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF with Ollama:
ollama run hf.co/patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
- Unsloth Desktop
- Pi
How to use patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF with Pi:
Start the llama.cpp server
# Install llama.cpp: brew install llama.cpp # Start a local OpenAI-compatible server: llama serve -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
Configure the model in Pi
# Install Pi: npm install -g @earendil-works/pi-coding-agent # Add to ~/.pi/agent/models.json: { "providers": { "llama-cpp": { "baseUrl": "http://localhost:8080/v1", "api": "openai-completions", "apiKey": "none", "models": [ { "id": "patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M" } ] } } }Run Pi
# Start Pi in your project directory: pi
- Docker Model Runner
How to use patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF with Docker Model Runner:
docker model run hf.co/patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
- Lemonade
How to use patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF with Lemonade:
Pull the model
# Download Lemonade from https://lemonade-server.ai/ lemonade pull patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
Run and chat with the model
lemonade run user.GLM-5.3-Flash-REAP50-GGUF-Q4_K_M
List all available models
lemonade list
- Hermes Agent
How to use patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF with Hermes Agent:
Start the llama.cpp server
# Install llama.cpp: brew install llama.cpp # Start a local OpenAI-compatible server: llama serve -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
Configure Hermes
# Install Hermes: curl -fsSL https://hermes-agent.nousresearch.com/install.sh | bash hermes setup # Point Hermes at the local server: hermes config set model.provider custom hermes config set model.base_url http://127.0.0.1:8080/v1 hermes config set model.default patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
Run Hermes
hermes
- Atomic Chat
- OpenClaw
How to use patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF with OpenClaw:
Start the llama.cpp server
# Install llama.cpp: brew install llama.cpp # Start a local OpenAI-compatible server: llama serve -hf patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M
Configure OpenClaw
# Install OpenClaw: npm install -g openclaw@latest # Register the local server and set it as the default model: openclaw onboard --non-interactive --mode local \ --auth-choice custom-api-key \ --custom-base-url http://127.0.0.1:8080/v1 \ --custom-model-id "patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF:Q4_K_M" \ --custom-provider-id llama-cpp \ --custom-compatibility openai \ --custom-text-input \ --accept-risk \ --skip-health
Run OpenClaw
openclaw agent --local --agent main --message "Hello from Hugging Face"
Download glm5-next-llama.cpp.patch from patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF: direct link, hf CLI and curl.
- Browser
- Download file 300 kB
-
https://huggingface.co/patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF/resolve/main/glm5-next-llama.cpp.patch
- Command line
-
hf download hf://patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF/glm5-next-llama.cpp.patch
-
curl -L -o glm5-next-llama.cpp.patch https://huggingface.co/patrickbdevaney/GLM-5.3-Flash-REAP50-GGUF/resolve/main/glm5-next-llama.cpp.patch
300 kB
| From 922e708702b81a156460ba718a4f6688663fe910 Mon Sep 17 00:00:00 2001 | |
| From: Patrick Devaney <patrickbdevaney@gmail.com> | |
| Date: Sat, 29 Aug 2026 02:11:34 -0400 | |
| Subject: [PATCH 01/15] 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 | |
| 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" | |
| class MODEL_ARCH(IntEnum): | |
| LLAMA_EMBED = auto() | |
| MAINCODER = auto() | |
| KIMI_LINEAR = auto() | |
| + GLM5_NEXT = auto() | |
| class VISION_PROJECTOR_TYPE(IntEnum): | |
| 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() | |
| 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] = { | |
| 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}", | |
| 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 | |
| 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 08:14:02 -0400 | |
| Subject: [PATCH 02/15] 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 | |
| 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 | |
| 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"]) | |
| 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 | |
| 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, | |
| 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 | |
| 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); | |
| 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 | |
| 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 | |
| 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 | |
| #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" | |
| 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; | |
| 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 | |
| +#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<<<grid, block, 0, ctx.stream()>>>( | |
| + (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 | |
| +#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 | |
| static const char * GGML_OP_NAME[GGML_OP_COUNT] = { | |
| "RMS_NORM_BACK", | |
| "GROUP_NORM", | |
| "L2_NORM", | |
| + "MHC_SINKHORN", | |
| "MUL_MAT", | |
| "MUL_MAT_ID", | |
| 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", | |
| 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", | |
| 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"); | |
| 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 | |
| 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" | |
| 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: | |
| 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" | |
| 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() | |
| 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", | |
| 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 | |
| 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) | |
| 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 | |
| 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 | |
| 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 | |
| static const std::map<llm_arch, const char *> 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)" }, | |
| }; | |
| static const std::map<llm_kv, const char *> 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" }, | |
| static const std::map<llm_tensor, const char *> 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_tensor> llm_get_tensor_names(llm_arch arch) { | |
| static std::set<llm_tensor> 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, | |
| static const std::map<llm_tensor, llm_tensor_info> 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}}, | |
| 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 | |
| enum llm_arch { | |
| LLM_ARCH_LLAMA_EMBED, | |
| LLM_ARCH_MAINCODER, | |
| LLM_ARCH_KIMI_LINEAR, | |
| + LLM_ARCH_GLM5_NEXT, | |
| LLM_ARCH_UNKNOWN, | |
| }; | |
| 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, | |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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); | |
| 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); | |
| ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const { | |
| { | |
| llm = std::make_unique<llm_build_mimo2_iswa>(*this, params); | |
| } break; | |
| + case LLM_ARCH_GLM5_NEXT: | |
| + { | |
| + llm = std::make_unique<llm_build_glm5_next>(*this, params); | |
| + } break; | |
| case LLM_ARCH_KIMI_LINEAR: | |
| { | |
| llm = std::make_unique<llm_build_kimi_linear>(*this, params); | |
| 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 | |
| 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 | |
| +#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 | |
| 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 | |
| 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 | |
| +// 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 <algorithm> | |
| +#include <cmath> | |
| +#include <cstdio> | |
| +#include <cstring> | |
| +#include <vector> | |
| + | |
| +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<int32_t> ids(n_tok); | |
| + if (fread(ids.data(), 4, n_tok, f) != (size_t) n_tok) return 1; | |
| + std::vector<float> 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 | |
| +// 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 <cmath> | |
| +#include <cstdio> | |
| +#include <cstdlib> | |
| +#include <cstring> | |
| +#include <algorithm> | |
| +#include <vector> | |
| + | |
| +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<float> 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 | |
| #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 | |
| 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 | |
| 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 { | |
| 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 08:50:53 -0400 | |
| Subject: [PATCH 03/15] 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 | |
| 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 | |
| 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 | |
| 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 | |
| #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 | |
| 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 | |
| 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; | |
| 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 | |
| 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 | |
| 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 11:02:37 -0400 | |
| Subject: [PATCH 04/15] 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 | |
| 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 | |
| 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 | |
| static const std::map<llm_kv, const char *> 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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 | |
| +// 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 <algorithm> | |
| +#include <cmath> | |
| +#include <cstdio> | |
| +#include <cstring> | |
| +#include <vector> | |
| + | |
| +static std::vector<float> rd(FILE * f, size_t n) { | |
| + std::vector<float> 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<float> & 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<float> & 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 11:02:46 -0400 | |
| Subject: [PATCH 05/15] 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 | |
| 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 11:36:25 -0400 | |
| Subject: [PATCH 06/15] 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 | |
| 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) { | |
| 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 | |
| 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 13:45:20 -0400 | |
| Subject: [PATCH 07/15] 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 | |
| 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 | |
| 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 | |
| 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 | |
| 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 14:20:59 -0400 | |
| Subject: [PATCH 08/15] 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 | |
| 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 | |
| +// 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 <algorithm> | |
| +#include <cstdio> | |
| +#include <cstring> | |
| +#include <vector> | |
| + | |
| +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<float> 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<int> 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<int> sp(pools, pools + select_k); | |
| + std::sort(sp.begin(), sp.end()); | |
| + const std::vector<int> 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 18:18:47 -0400 | |
| Subject: [PATCH 09/15] 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 <noreply@anthropic.com> | |
| 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 | |
| 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 | |
| +// 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 <algorithm> | |
| +#include <cmath> | |
| +#include <cstdio> | |
| +#include <cstring> | |
| +#include <random> | |
| +#include <vector> | |
| + | |
| +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<float> 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<float> 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<std::vector<char>> vis(T, std::vector<char>(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<int> sel_pools(select_k); | |
| + { | |
| + std::vector<int> ord(n_pools); for (int i=0;i<n_pools;++i) ord[i]=i; | |
| + std::partial_sort(ord.begin(), ord.begin()+select_k, ord.end(), | |
| + [&](int a,int b){ return pool_scores[a] > pool_scores[b]; }); | |
| + std::copy(ord.begin(), ord.begin()+select_k, sel_pools.begin()); | |
| + } | |
| + std::vector<char> 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<int> want; | |
| + for (int p : sel_pools) for (int j = 0; j < kpool; ++j) want.push_back(p*kpool + j); | |
| + std::vector<int> got((int32_t *) idx->data, (int32_t *) idx->data + sel); | |
| + std::vector<int> 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 21:50:21 -0400 | |
| Subject: [PATCH 10/15] 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 <noreply@anthropic.com> | |
| 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 | |
| #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). | |
| 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) { | |
| 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 | |
| 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 | |
| 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 | |
| +// 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 <cmath> | |
| +#include <cstdio> | |
| +#include <cstdlib> | |
| +#include <cstring> | |
| +#include <vector> | |
| + | |
| +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<int32_t> 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 22:00:19 -0400 | |
| Subject: [PATCH 11/15] 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 <noreply@anthropic.com> | |
| 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 | |
| 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; | |
| 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 | |
| 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; | |
| 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; | |
| 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 | |
| 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 | |
| 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); | |
| 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, | |
| 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); | |
| 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] | |
| 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 | |
| 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 | |
| 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], | |
| 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); | |
| 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 <patrickbdevaney@gmail.com> | |
| Date: Sun, 30 Aug 2026 22:07:19 -0400 | |
| Subject: [PATCH 12/15] 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 <noreply@anthropic.com> | |
| 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 | |
| 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 | |
| From 7c63bd83e251a1ac81f3b087f4a3b87b3e7f5aac Mon Sep 17 00:00:00 2001 | |
| From: Patrick Devaney <patrickbdevaney@gmail.com> | |
| Date: Thu, 3 Sep 2026 14:13:46 -0400 | |
| Subject: [PATCH 13/15] embedding: dump raw fp16 hidden states and the token | |
| ids beside them | |
| Capturing a teacher trace for a draft head means keeping the model's final hidden | |
| state for every token, and the existing output paths are not usable for it: the | |
| JSON writer costs more than the forward pass at 4096 dims, and without the ids the | |
| rows cannot be matched back to tokens after the tokenizer has had its say. | |
| Writes the states as a flat fp16 array with a sidecar of ids and per-sequence | |
| lengths, so a capture can be compared row for row against another run - which is | |
| how the quantisation drift between IQ3_M and IQ4_XS was measured. | |
| Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> | |
| Claude-Session: https://claude.ai/code/session_016B6ZbadovyJKFiN1CwDBLA | |
| --- | |
| common/common.h | 1 + | |
| examples/embedding/embedding.cpp | 50 ++++++++++++++++++++++++++++++++ | |
| 2 files changed, 51 insertions(+) | |
| diff --git a/common/common.h b/common/common.h | |
| index 020b6a721..35761eb84 100644 | |
| --- a/common/common.h | |
| +++ b/common/common.h | |
| enum common_speculative_type { | |
| COMMON_SPECULATIVE_TYPE_NONE, // no speculative decoding | |
| COMMON_SPECULATIVE_TYPE_DRAFT, // draft model | |
| COMMON_SPECULATIVE_TYPE_EAGLE3, // eagle draft model | |
| + COMMON_SPECULATIVE_TYPE_MTP, // native multi-token-prediction draft (GLM-5.3 blk.45) | |
| COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, // simple self-speculative decoding | |
| COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, // self-speculative decoding with n-gram keys only | |
| COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, // self-speculative decoding with n-gram keys and 4 m-gram values | |
| diff --git a/examples/embedding/embedding.cpp b/examples/embedding/embedding.cpp | |
| index f6a20ef9d..b93265cab 100644 | |
| --- a/examples/embedding/embedding.cpp | |
| +++ b/examples/embedding/embedding.cpp | |
| int main(int argc, char ** argv) { | |
| float * out = emb + e * n_embd_out; | |
| batch_decode(ctx, batch, out, s, n_embd_out, params.embd_normalize); | |
| + // Raw fp16 dump for draft-head distillation. JSON of 4096 floats/token costs ~6x the | |
| + // bytes of the fp16 it encodes, which is untenable across a multi-day extraction. | |
| + // Driven by an env var so the shared arg parser stays untouched. | |
| + if (const char * bin_path = getenv("LLAMA_EMBD_BIN")) { | |
| + FILE * bf = fopen(bin_path, "wb"); | |
| + if (bf == NULL) { | |
| + LOG_ERR("%s: could not open LLAMA_EMBD_BIN=%s\n", __func__, bin_path); | |
| + } else { | |
| + const int64_t nrow = (pooling_type == LLAMA_POOLING_TYPE_NONE) ? n_embd_count : n_prompts; | |
| + std::vector<ggml_fp16_t> row(n_embd_out); | |
| + for (int64_t j = 0; j < nrow; j++) { | |
| + ggml_fp32_to_fp16_row(emb + j*n_embd_out, row.data(), n_embd_out); | |
| + fwrite(row.data(), sizeof(ggml_fp16_t), n_embd_out, bf); | |
| + } | |
| + fclose(bf); | |
| + LOG_INF("%s: wrote %lld x %d fp16 to %s\n", __func__, (long long) nrow, n_embd_out, bin_path); | |
| + } | |
| + } | |
| + | |
| + // Companion token dump. The whole point of the trace is that row i is the hidden state | |
| + // OF token i, so the targets have to be the ids this binary actually tokenized - not a | |
| + // re-tokenization done elsewhere that can drift by a token here and there. Writing them | |
| + // from inside the same process makes the alignment true by construction instead of an | |
| + // assertion we hope holds. | |
| + if (const char * bin_path = getenv("LLAMA_EMBD_BIN")) { | |
| + if (pooling_type == LLAMA_POOLING_TYPE_NONE) { | |
| + const std::string ids_path = std::string(bin_path) + ".ids"; | |
| + const std::string lens_path = std::string(bin_path) + ".lens"; | |
| + FILE * idf = fopen(ids_path.c_str(), "wb"); | |
| + FILE * lnf = fopen(lens_path.c_str(), "wb"); | |
| + if (idf == NULL || lnf == NULL) { | |
| + LOG_ERR("%s: could not open %s / %s\n", __func__, ids_path.c_str(), lens_path.c_str()); | |
| + } else { | |
| + int64_t ntok = 0; | |
| + for (int k = 0; k < n_prompts; k++) { | |
| + const int32_t len = (int32_t) inputs[k].size(); | |
| + fwrite(&len, sizeof(int32_t), 1, lnf); | |
| + for (int32_t j = 0; j < len; j++) { | |
| + const int32_t id = (int32_t) inputs[k][j]; | |
| + fwrite(&id, sizeof(int32_t), 1, idf); | |
| + } | |
| + ntok += len; | |
| + } | |
| + LOG_INF("%s: wrote %lld ids over %d seqs to %s\n", __func__, (long long) ntok, n_prompts, ids_path.c_str()); | |
| + } | |
| + if (idf) fclose(idf); | |
| + if (lnf) fclose(lnf); | |
| + } | |
| + } | |
| + | |
| if (params.embd_out.empty()) { | |
| LOG("\n"); | |
| -- | |
| 2.43.0 | |
| From 009abb19733c92fc93e4b1aa4584f3e44f8cbb4a Mon Sep 17 00:00:00 2001 | |
| From: Patrick Devaney <patrickbdevaney@gmail.com> | |
| Date: Thu, 3 Sep 2026 14:13:57 -0400 | |
| Subject: [PATCH 14/15] glm5-next: run the native MTP block as a standalone | |
| draft model | |
| GLM-5.3 ships a real multi-token-prediction module as blk.45, and nothing runs it. | |
| transformers discards it on load, llama.cpp keeps the tensors and never builds a | |
| graph for them, and the EAGLE3 slot in common/speculative.cpp is an empty stub. | |
| Adds LLM_ARCH_GLM5_NEXT_MTP: the same block lifted into a one-layer model that | |
| takes the target's hidden state as an input. Two details are easy to invert and | |
| both are taken from vLLM's glm4_moe_mtp.py, which is the authoritative | |
| implementation: | |
| - enorm normalizes the token EMBEDDING and hnorm the HIDDEN state (e, then h), | |
| concatenated in that order for eh_proj. | |
| - the module's output feeds the next speculative step BEFORE shared_head_norm, | |
| so res->t_embd is set pre-norm here, unlike every other architecture. | |
| The hidden state arrives through llama_set_mtp_hidden() into a buffer owned by the | |
| context, allocated up front rather than lazily because the warm-up decode runs | |
| before any caller can set one and the graph captures the pointer. Position 0's | |
| embedding is masked, and the mask is built in set_input from ubatch->pos rather | |
| than by casting the positions, since I32->F32 casts are not reliable on every | |
| backend. | |
| llama_set_embeddings_outputs_only() is the other half of making this usable: | |
| output_all is cparams.embeddings, so simply enabling embeddings on a generation | |
| context promotes every prompt token to an output and allocates a 620 KB logits row | |
| for each one at GLM-5.3's 154880 vocab. | |
| Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> | |
| Claude-Session: https://claude.ai/code/session_016B6ZbadovyJKFiN1CwDBLA | |
| --- | |
| include/llama.h | 16 ++++ | |
| src/CMakeLists.txt | 1 + | |
| src/llama-arch.cpp | 42 +++++++++ | |
| src/llama-arch.h | 1 + | |
| src/llama-context.cpp | 58 ++++++++++++- | |
| src/llama-context.h | 6 ++ | |
| src/llama-cparams.h | 15 ++++ | |
| src/llama-graph.cpp | 51 +++++++++++ | |
| src/llama-graph.h | 20 +++++ | |
| src/llama-model.cpp | 99 +++++++++++++++++++++ | |
| src/models/glm5-next-mtp.cpp | 163 +++++++++++++++++++++++++++++++++++ | |
| src/models/models.h | 8 ++ | |
| 12 files changed, 477 insertions(+), 3 deletions(-) | |
| create mode 100644 src/models/glm5-next-mtp.cpp | |
| diff --git a/include/llama.h b/include/llama.h | |
| index a940f9d64..66c9e1b8f 100644 | |
| --- a/include/llama.h | |
| +++ b/include/llama.h | |
| extern "C" { | |
| // TODO: rename to avoid confusion with llama_get_embeddings() | |
| LLAMA_API void llama_set_embeddings(struct llama_context * ctx, bool embeddings); | |
| + // Restrict embedding extraction to tokens that are already outputs, instead of promoting | |
| + // every token in the batch to an output (the default, which suits embedding models). A | |
| + // speculative draft that conditions on the target's hidden state needs this: without it, | |
| + // enabling embeddings on a generation context would allocate an n_vocab logits row for | |
| + // every prompt token. | |
| + LLAMA_API void llama_set_embeddings_outputs_only(struct llama_context * ctx, bool value); | |
| + | |
| + // Supply the target model's hidden states to a GLM-5.3 MTP draft context, one n_embd-sized | |
| + // row per token of the batch that is about to be decoded, in batch order. Must be called | |
| + // before each llama_decode() on such a context; other architectures ignore it. | |
| + // The rows are copied, so `data` need not outlive the call. | |
| + LLAMA_API void llama_set_mtp_hidden( | |
| + struct llama_context * ctx, | |
| + const float * data, | |
| + int32_t n_tokens); | |
| + | |
| // Set whether to use causal attention or not | |
| // If set to true, the model will only attend to the past tokens | |
| LLAMA_API void llama_set_causal_attn(struct llama_context * ctx, bool causal_attn); | |
| diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt | |
| index dbac7fdc5..94e4a631d 100644 | |
| --- a/src/CMakeLists.txt | |
| +++ b/src/CMakeLists.txt | |
| add_library(llama | |
| models/gemma4-iswa.cpp | |
| models/glm4-moe.cpp | |
| models/glm5-next.cpp | |
| + models/glm5-next-mtp.cpp | |
| models/glm4.cpp | |
| models/gpt2.cpp | |
| models/gptneox.cpp | |
| diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp | |
| index 15783b5e5..29767db3b 100644 | |
| --- a/src/llama-arch.cpp | |
| +++ b/src/llama-arch.cpp | |
| static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = { | |
| { LLM_ARCH_MAINCODER, "maincoder" }, | |
| { LLM_ARCH_KIMI_LINEAR, "kimi-linear" }, | |
| { LLM_ARCH_GLM5_NEXT, "glm5-next" }, | |
| + { LLM_ARCH_GLM5_NEXT_MTP, "glm5-next-mtp" }, | |
| { LLM_ARCH_UNKNOWN, "(unknown)" }, | |
| }; | |
| static std::set<llm_tensor> llm_get_tensor_names(llm_arch arch) { | |
| LLM_TENSOR_NEXTN_HNORM, | |
| LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, | |
| }; | |
| + // GLM-5.3's MTP module lifted out into a standalone speculative draft. One MLA+DSA | |
| + // layer plus the 144-expert MoE - no KDA, and no hyper-connections, because blk.45 | |
| + // carries no hc_* tensors and uses plain pre-norm residuals. The shared head norm is | |
| + // written as OUTPUT_NORM by the converter, since in a standalone model that is exactly | |
| + // the role it plays. | |
| + case LLM_ARCH_GLM5_NEXT_MTP: | |
| + return { | |
| + LLM_TENSOR_TOKEN_EMBD, | |
| + LLM_TENSOR_OUTPUT_NORM, | |
| + LLM_TENSOR_OUTPUT, | |
| + LLM_TENSOR_ATTN_NORM, | |
| + LLM_TENSOR_FFN_NORM, | |
| + LLM_TENSOR_ATTN_OUT, | |
| + 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, | |
| + 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, | |
| + // The MTP entry point: enorm on the token embedding, hnorm on the incoming | |
| + // hidden state, concatenated and projected back down to n_embd. | |
| + LLM_TENSOR_NEXTN_EH_PROJ, | |
| + LLM_TENSOR_NEXTN_ENORM, | |
| + LLM_TENSOR_NEXTN_HNORM, | |
| + }; | |
| case LLM_ARCH_KIMI_LINEAR: | |
| return { | |
| LLM_TENSOR_TOKEN_EMBD, | |
| diff --git a/src/llama-arch.h b/src/llama-arch.h | |
| index 5cd242dcb..fcaa3ad59 100644 | |
| --- a/src/llama-arch.h | |
| +++ b/src/llama-arch.h | |
| enum llm_arch { | |
| LLM_ARCH_MAINCODER, | |
| LLM_ARCH_KIMI_LINEAR, | |
| LLM_ARCH_GLM5_NEXT, | |
| + LLM_ARCH_GLM5_NEXT_MTP, | |
| LLM_ARCH_UNKNOWN, | |
| }; | |
| diff --git a/src/llama-context.cpp b/src/llama-context.cpp | |
| index a808e3e45..ee084c726 100644 | |
| --- a/src/llama-context.cpp | |
| +++ b/src/llama-context.cpp | |
| llama_context::llama_context( | |
| cparams.yarn_beta_fast = params.yarn_beta_fast >= 0.0f ? params.yarn_beta_fast : hparams.yarn_beta_fast; | |
| cparams.yarn_beta_slow = params.yarn_beta_slow >= 0.0f ? params.yarn_beta_slow : hparams.yarn_beta_slow; | |
| cparams.embeddings = params.embeddings; | |
| + cparams.embd_outputs_only = false; | |
| + cparams.mtp_h_prev = nullptr; | |
| + cparams.mtp_h_prev_n = 0; | |
| cparams.offload_kqv = params.offload_kqv; | |
| cparams.no_perf = params.no_perf; | |
| cparams.pooling_type = params.pooling_type; | |
| llama_context::llama_context( | |
| LLAMA_LOG_INFO("%s: n_ctx = %u\n", __func__, cparams.n_ctx); | |
| LLAMA_LOG_INFO("%s: n_ctx_seq = %u\n", __func__, cparams.n_ctx_seq); | |
| LLAMA_LOG_INFO("%s: n_batch = %u\n", __func__, cparams.n_batch); | |
| + | |
| + // A GLM-5.3 MTP draft reads the target's hidden state from this buffer on every decode. | |
| + // Allocate it up front, at full batch width, for two reasons: the graph captures the | |
| + // pointer when it is built, so it must never move; and the warmup decode happens before | |
| + // any caller has had a chance to supply hidden states, which would otherwise be a crash | |
| + // on model load rather than a working (if uninformative) warmup against zeros. | |
| + if (model.arch == LLM_ARCH_GLM5_NEXT_MTP) { | |
| + mtp_h_prev.resize((size_t) model.hparams.n_embd * cparams.n_batch, 0.0f); | |
| + cparams.mtp_h_prev = mtp_h_prev.data(); | |
| + } | |
| LLAMA_LOG_INFO("%s: n_ubatch = %u\n", __func__, cparams.n_ubatch); | |
| LLAMA_LOG_INFO("%s: causal_attn = %d\n", __func__, cparams.causal_attn); | |
| LLAMA_LOG_INFO("%s: flash_attn = %s\n", __func__, llama_flash_attn_type_name(params.flash_attn_type)); | |
| void llama_context::set_embeddings(bool value) { | |
| //sched_need_reserve = true; | |
| } | |
| +void llama_context::set_embeddings_outputs_only(bool value) { | |
| + LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); | |
| + | |
| + cparams.embd_outputs_only = value; | |
| +} | |
| + | |
| +void llama_context::set_mtp_hidden(const float * data, int32_t n_tokens) { | |
| + const int64_t n_embd = model.hparams.n_embd; | |
| + | |
| + if (model.arch != LLM_ARCH_GLM5_NEXT_MTP) { | |
| + LLAMA_LOG_ERROR("%s: called on a non-MTP model (arch %s)\n", | |
| + __func__, llm_arch_name(model.arch)); | |
| + return; | |
| + } | |
| + | |
| + if (n_tokens <= 0 || (uint32_t) n_tokens > cparams.n_batch) { | |
| + LLAMA_LOG_ERROR("%s: n_tokens = %d out of range (n_batch = %u)\n", | |
| + __func__, n_tokens, cparams.n_batch); | |
| + return; | |
| + } | |
| + | |
| + // Allocated once at full batch width and never resized, because the graph captures this | |
| + // pointer when it is built and would otherwise be left holding a freed buffer after a | |
| + // reallocation. | |
| + GGML_ASSERT(!mtp_h_prev.empty() && "MTP hidden buffer was not allocated at context creation"); | |
| + | |
| + std::copy(data, data + (size_t) n_embd*n_tokens, mtp_h_prev.begin()); | |
| + cparams.mtp_h_prev_n = n_tokens; | |
| +} | |
| + | |
| void llama_context::set_causal_attn(bool value) { | |
| LLAMA_LOG_DEBUG("%s: value = %d\n", __func__, value); | |
| int llama_context::decode(const llama_batch & batch_inp) { | |
| const int64_t n_vocab = vocab.n_tokens(); | |
| const int64_t n_embd = hparams.n_embd_inp(); | |
| - // when computing embeddings, all tokens are output | |
| - const bool output_all = cparams.embeddings; | |
| + // when computing embeddings, all tokens are output - unless the caller only wants hidden | |
| + // states for the tokens it already asked for (see llama_set_embeddings_outputs_only) | |
| + const bool output_all = cparams.embeddings && !cparams.embd_outputs_only; | |
| const bool has_samplers = !sampling.samplers.empty(); | |
| const uint32_t n_seq_max = cparams.kv_unified ? LLAMA_MAX_SEQ : cparams.n_seq_max; | |
| void llama_context::output_reorder() { | |
| // | |
| uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { | |
| - if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE) { | |
| + if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_GLM5_NEXT) { | |
| return std::max<uint32_t>(n_tokens * 40, 32u * model.n_tensors()); | |
| } | |
| uint32_t res = std::max<uint32_t>(1024u, 8u*model.n_tensors()); | |
| void llama_set_embeddings(llama_context * ctx, bool embeddings) { | |
| ctx->set_embeddings(embeddings); | |
| } | |
| +void llama_set_embeddings_outputs_only(llama_context * ctx, bool value) { | |
| + ctx->set_embeddings_outputs_only(value); | |
| +} | |
| + | |
| +void llama_set_mtp_hidden(llama_context * ctx, const float * data, int32_t n_tokens) { | |
| + ctx->set_mtp_hidden(data, n_tokens); | |
| +} | |
| + | |
| void llama_set_causal_attn(llama_context * ctx, bool causal_attn) { | |
| ctx->set_causal_attn(causal_attn); | |
| } | |
| diff --git a/src/llama-context.h b/src/llama-context.h | |
| index e0d0085c1..5ec6072fb 100644 | |
| --- a/src/llama-context.h | |
| +++ b/src/llama-context.h | |
| struct llama_context { | |
| void set_abort_callback(bool (*abort_callback)(void * data), void * abort_callback_data); | |
| void set_embeddings (bool value); | |
| + void set_mtp_hidden (const float * data, int32_t n_tokens); | |
| + void set_embeddings_outputs_only(bool value); | |
| void set_causal_attn(bool value); | |
| void set_warmup(bool value); | |
| private: | |
| // host buffer for the model output (logits and embeddings) | |
| ggml_backend_buffer_ptr buf_output; | |
| + // Target-model hidden states for a GLM-5.3 MTP draft. Sized once from n_embd*n_batch so the | |
| + // pointer handed to the graph never moves; empty for every other architecture. | |
| + std::vector<float> mtp_h_prev; | |
| + | |
| bool has_evaluated_once = false; | |
| // env: LLAMA_GRAPH_REUSE_DISABLE | |
| diff --git a/src/llama-cparams.h b/src/llama-cparams.h | |
| index 9d3594741..4d75f52ab 100644 | |
| --- a/src/llama-cparams.h | |
| +++ b/src/llama-cparams.h | |
| struct llama_cparams { | |
| float yarn_beta_fast; | |
| float yarn_beta_slow; | |
| + // Speculative MTP drafts (LLM_ARCH_GLM5_NEXT_MTP) consume the *target* model's hidden | |
| + // state alongside the token being fed. It cannot ride in the ubatch: the batch carries | |
| + // either tokens or embeddings, never both, and the draft needs its own token lookup. The | |
| + // buffer lives in llama_context and is allocated once, so this pointer is stable for the | |
| + // lifetime of the context and safe to capture at graph-build time. | |
| + const float * mtp_h_prev; // [n_embd, mtp_h_prev_n], row i belongs to batch token i | |
| + uint32_t mtp_h_prev_n; | |
| + | |
| bool embeddings; | |
| + | |
| + // Extract embeddings only for the tokens that are already outputs, instead of promoting | |
| + // every token in the batch to an output. Embedding models want the latter - they pool over | |
| + // the whole sequence - but a speculative MTP draft only needs the hidden state of the token | |
| + // it is about to draft from, and promoting a whole prompt would allocate an n_vocab logits | |
| + // row per token (620 KB each at GLM-5.3's 154880 vocab) that nothing ever reads. | |
| + bool embd_outputs_only; | |
| bool causal_attn; | |
| bool offload_kqv; | |
| bool flash_attn; | |
| diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp | |
| index c3e357a4e..4859563d2 100644 | |
| --- a/src/llama-graph.cpp | |
| +++ b/src/llama-graph.cpp | |
| bool llm_graph_input_embd::can_reuse(const llm_graph_params & params) { | |
| return res; | |
| } | |
| +void llm_graph_input_mtp_hidden::set_input(const llama_ubatch * ubatch) { | |
| + const int64_t n_tokens = ubatch->n_tokens; | |
| + | |
| + // The caller supplies one hidden state per token in the batch it is about to decode. If the | |
| + // batch were split into several ubatches these rows would no longer line up, so the draft | |
| + // context is created with n_ubatch == n_batch and the mismatch is a hard error rather than | |
| + // a silent misalignment - drafting from the wrong hidden state degrades acceptance in a way | |
| + // that looks like a bad head rather than a bug. | |
| + GGML_ASSERT(src != nullptr && "llama_set_mtp_hidden() must be called before llama_decode()"); | |
| + GGML_ASSERT(h_prev->ne[0] == n_embd); | |
| + | |
| + ggml_backend_tensor_set(h_prev, src, 0, n_tokens*n_embd*ggml_element_size(h_prev)); | |
| + | |
| + // The reference zeroes the token embedding at absolute position 0, where there is no | |
| + // preceding token for the module to condition on | |
| + // (glm4_moe_mtp.py: torch.where(positions == 0, 0, inputs_embeds)). | |
| + // | |
| + // It is built here from ubatch->pos rather than derived in the graph from inp_pos: doing it | |
| + // in-graph needs an I32->F32 ggml_cast, which not every backend implements, and this costs | |
| + // one float per token. | |
| + if (emask && ubatch->pos) { | |
| + std::vector<float> keep(n_tokens); | |
| + for (int64_t i = 0; i < n_tokens; ++i) { | |
| + keep[i] = ubatch->pos[i] == 0 ? 0.0f : 1.0f; | |
| + } | |
| + ggml_backend_tensor_set(emask, keep.data(), 0, n_tokens*ggml_element_size(emask)); | |
| + } | |
| +} | |
| + | |
| +bool llm_graph_input_mtp_hidden::can_reuse(const llm_graph_params & params) { | |
| + return h_prev && h_prev->ne[1] == params.ubatch.n_tokens; | |
| +} | |
| + | |
| void llm_graph_input_pos::set_input(const llama_ubatch * ubatch) { | |
| if (ubatch->pos && pos) { | |
| const int64_t n_tokens = ubatch->n_tokens; | |
| ggml_tensor * llm_graph_context::build_inp_embd(ggml_tensor * tok_embd) const { | |
| return cur; | |
| } | |
| +llm_graph_input_mtp_hidden * llm_graph_context::build_inp_mtp_hidden() const { | |
| + auto inp = std::make_unique<llm_graph_input_mtp_hidden>(hparams.n_embd, cparams.mtp_h_prev); | |
| + | |
| + inp->h_prev = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, ubatch.n_tokens); | |
| + ggml_set_input(inp->h_prev); | |
| + cb(inp->h_prev, "inp_mtp_hidden", -1); | |
| + | |
| + inp->emask = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, ubatch.n_tokens); | |
| + ggml_set_input(inp->emask); | |
| + cb(inp->emask, "inp_mtp_emask", -1); | |
| + | |
| + auto * res_inp = inp.get(); | |
| + | |
| + res->add_input(std::move(inp)); | |
| + | |
| + return res_inp; | |
| +} | |
| + | |
| ggml_tensor * llm_graph_context::build_inp_pos() const { | |
| auto inp = std::make_unique<llm_graph_input_pos>(hparams.n_pos_per_embd()); | |
| diff --git a/src/llama-graph.h b/src/llama-graph.h | |
| index b779a70f6..5c5398955 100644 | |
| --- a/src/llama-graph.h | |
| +++ b/src/llama-graph.h | |
| public: | |
| const int64_t n_embd = 0; | |
| }; | |
| +// The previous-token hidden state consumed by a GLM-5.3 MTP draft. Row i is the target model's | |
| +// final hidden state for batch token i, and is supplied out-of-band via llama_set_mtp_hidden() | |
| +// because a llama_batch has no slot for it. | |
| +class llm_graph_input_mtp_hidden : public llm_graph_input_i { | |
| +public: | |
| + llm_graph_input_mtp_hidden(int64_t n_embd, const float * src) : n_embd(n_embd), src(src) {} | |
| + virtual ~llm_graph_input_mtp_hidden() = default; | |
| + | |
| + void set_input(const llama_ubatch * ubatch) override; | |
| + | |
| + bool can_reuse(const llm_graph_params & params) override; | |
| + | |
| + ggml_tensor * h_prev = nullptr; // F32 [n_embd, n_batch] | |
| + ggml_tensor * emask = nullptr; // F32 [1, n_batch] 0 at position 0, else 1 | |
| + | |
| + const int64_t n_embd = 0; | |
| + const float * src = nullptr; // stable, owned by llama_context | |
| +}; | |
| + | |
| class llm_graph_input_pos : public llm_graph_input_i { | |
| public: | |
| llm_graph_input_pos(uint32_t n_pos_per_embd) : n_pos_per_embd(n_pos_per_embd) {} | |
| struct llm_graph_context { | |
| ggml_tensor * build_inp_embd(ggml_tensor * tok_embd) const; | |
| ggml_tensor * build_inp_pos() const; | |
| + llm_graph_input_mtp_hidden * build_inp_mtp_hidden() const; | |
| ggml_tensor * build_inp_attn_scale() const; | |
| ggml_tensor * build_inp_out_ids() const; | |
| ggml_tensor * build_inp_mean() const; | |
| diff --git a/src/llama-model.cpp b/src/llama-model.cpp | |
| index 0a4bc7bf1..8e0378d77 100644 | |
| --- a/src/llama-model.cpp | |
| +++ b/src/llama-model.cpp | |
| void llama_model::load_hparams(llama_model_loader & ml) { | |
| default: type = LLM_TYPE_UNKNOWN; | |
| } | |
| } break; | |
| + case LLM_ARCH_GLM5_NEXT_MTP: | |
| + { | |
| + // The MTP module on its own: one MLA layer and one MoE, no KDA and no | |
| + // hyper-connections. Everything below is the parent's geometry, which the | |
| + // converter copies verbatim rather than restating. | |
| + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); | |
| + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); | |
| + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); | |
| + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); | |
| + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); | |
| + ml.get_key(LLM_KV_SWIGLU_LIMIT, hparams.swiglu_limit, false); | |
| + | |
| + 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); | |
| + | |
| + 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); | |
| + | |
| + // No recurrent layers here, so nothing to mark - but the array is consulted | |
| + // unconditionally elsewhere, so make the absence explicit rather than implied. | |
| + for (uint32_t i = 0; i < hparams.n_layer; ++i) { | |
| + hparams.recurrent_layer_arr[i] = false; | |
| + } | |
| + | |
| + // nextn_predict_layers is deliberately NOT read: in the parent it marks a block | |
| + // to skip, and here that block is the entire model. | |
| + type = hparams.n_layer == 1 ? LLM_TYPE_A13B : LLM_TYPE_UNKNOWN; | |
| + } break; | |
| case LLM_ARCH_KIMI_LINEAR: | |
| { | |
| ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); | |
| 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_MTP: | |
| + { | |
| + // The standalone MTP draft. Shapes are the parent's blk.45 verbatim; the | |
| + // converter renamed shared_head_norm to output_norm because in a model that | |
| + // stands on its own that is precisely what it is. | |
| + 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 n_ff_exp = hparams.n_ff_exp; | |
| + 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(); | |
| + const int64_t idx_h = hparams.indexer_head_size; | |
| + | |
| + for (int i = 0; i < n_layer; ++i) { | |
| + auto & layer = layers[i]; | |
| + | |
| + 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.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); | |
| + | |
| + 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); | |
| + | |
| + 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); | |
| + | |
| + 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); | |
| + } | |
| + } break; | |
| case LLM_ARCH_GLM5_NEXT: | |
| { | |
| tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); | |
| ggml_cgraph * llama_model::build_graph(const llm_graph_params & params) const { | |
| { | |
| llm = std::make_unique<llm_build_glm5_next>(*this, params); | |
| } break; | |
| + case LLM_ARCH_GLM5_NEXT_MTP: | |
| + { | |
| + llm = std::make_unique<llm_build_glm5_next_mtp>(*this, params); | |
| + } break; | |
| case LLM_ARCH_KIMI_LINEAR: | |
| { | |
| llm = std::make_unique<llm_build_kimi_linear>(*this, params); | |
| llama_rope_type llama_model_rope_type(const llama_model * model) { | |
| case LLM_ARCH_NEMOTRON_H: | |
| case LLM_ARCH_NEMOTRON_H_MOE: | |
| case LLM_ARCH_GLM5_NEXT: // NoPE: qk_rope_head_dim == 0 on the MLA path | |
| + case LLM_ARCH_GLM5_NEXT_MTP: // same NoPE MLA block, lifted out of the parent | |
| case LLM_ARCH_KIMI_LINEAR: | |
| return LLAMA_ROPE_TYPE_NONE; | |
| diff --git a/src/models/glm5-next-mtp.cpp b/src/models/glm5-next-mtp.cpp | |
| new file mode 100644 | |
| index 000000000..a6eff3d2f | |
| --- /dev/null | |
| +++ b/src/models/glm5-next-mtp.cpp | |
| +#include "models.h" | |
| + | |
| +// GLM-5.3's Multi-Token Prediction module, run as a standalone speculative draft. | |
| +// | |
| +// The module is one MLA+DSA attention layer and one 144-expert MoE, preceded by the MTP entry | |
| +// point: the incoming token embedding and the *target model's* hidden state are each RMS-normed, | |
| +// concatenated, and projected back down to n_embd by eh_proj. Reference: | |
| +// vllm/model_executor/models/glm4_moe_mtp.py, Glm4MoeMultiTokenPredictorLayer.forward. | |
| +// | |
| +// Two details are easy to get backwards and both change the result: | |
| +// | |
| +// * enorm normalises the token EMBEDDING and hnorm the incoming HIDDEN state - not the other | |
| +// way round - and the concatenation order is [enorm(embed), hnorm(h_prev)], matching the | |
| +// [2*n_embd, n_embd] shape of eh_proj. | |
| +// | |
| +// * res->t_embd is set BEFORE the final norm, unlike every other architecture here. The value | |
| +// the reference feeds back as previous_hidden_states on the next speculative step is the | |
| +// module's raw output; shared_head.norm is applied only on the way to the logits. Exposing | |
| +// the post-norm tensor instead would quietly corrupt drafting at depth > 1. | |
| +// | |
| +// Unlike the 45 transformer layers of the parent, this block has no mHC hyper-connections - it | |
| +// consumes an already-collapsed hidden state, so there are no parallel streams to mix - and | |
| +// therefore uses plain pre-norm residuals. | |
| + | |
| +llm_build_glm5_next_mtp::llm_build_glm5_next_mtp(const llama_model & model, const llm_graph_params & params) : | |
| + llm_graph_context(params), model(model) { | |
| + GGML_ASSERT(n_layer == 1 && "the MTP draft is a single block"); | |
| + | |
| + 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); | |
| + const int64_t n_head = hparams.n_head(); | |
| + | |
| + const int il = 0; | |
| + const auto & layer = model.layers[il]; | |
| + | |
| + ggml_tensor * cur; | |
| + | |
| + ggml_tensor * inp_embd = build_inp_embd(model.tok_embd); | |
| + cb(inp_embd, "model.embed_tokens", -1); | |
| + | |
| + auto * inp_mtp = build_inp_mtp_hidden(); | |
| + | |
| + // NoPE, like the parent's MLA layers: qk_rope_head_dim is 0, so there is no rotary tail. | |
| + auto * inp_attn = build_attn_inp_k(); | |
| + auto * inp_out_ids = build_inp_out_ids(); | |
| + | |
| + // ---------------- MTP entry point ---------------- | |
| + { | |
| + ggml_tensor * e = ggml_mul(ctx0, inp_embd, inp_mtp->emask); | |
| + e = build_norm(e, layer.nextn.enorm, NULL, LLM_NORM_RMS, il); | |
| + cb(e, "mtp_enorm", il); | |
| + | |
| + ggml_tensor * h = build_norm(inp_mtp->h_prev, layer.nextn.hnorm, NULL, LLM_NORM_RMS, il); | |
| + cb(h, "mtp_hnorm", il); | |
| + | |
| + cur = ggml_concat(ctx0, e, h, 0); | |
| + cur = ggml_mul_mat(ctx0, layer.nextn.eh_proj, cur); | |
| + cb(cur, "mtp_eh_proj", il); | |
| + } | |
| + | |
| + ggml_tensor * inpL = cur; | |
| + | |
| + // ---------------- attention ---------------- | |
| + cur = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); | |
| + cb(cur, "attn_norm", il); | |
| + | |
| + { | |
| + 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) { | |
| + // Absorbed MLA: the query is folded through wk_b so attention runs against the | |
| + // compressed latent, and wv_b expands the result on the way out. | |
| + 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, 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, layer.wo, NULL, Qcur, Kcur, Vcur, | |
| + nullptr, nullptr, nullptr, kq_scale_mla, il); | |
| + } | |
| + cb(cur, "mla_out", il); | |
| + } | |
| + | |
| + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); | |
| + cb(ffn_inp, "attn_residual", il); | |
| + | |
| + // ---------------- MoE ---------------- | |
| + cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); | |
| + cb(cur, "ffn_norm", il); | |
| + | |
| + { | |
| + ggml_tensor * moe_out = build_moe_ffn(cur, | |
| + layer.ffn_gate_inp, | |
| + layer.ffn_up_exps, | |
| + layer.ffn_gate_exps, | |
| + layer.ffn_down_exps, | |
| + layer.ffn_exp_probs_b, | |
| + 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); | |
| + | |
| + // Same clamped SwiGLU as the parent's shared expert: the fused SiLU path has no limit, | |
| + // and the clamp only bites once activations exceed it. | |
| + 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 = ggml_add(ctx0, cur, ffn_inp); | |
| + | |
| + if (inp_out_ids) { | |
| + cur = ggml_get_rows(ctx0, cur, inp_out_ids); | |
| + } | |
| + | |
| + // Pre-norm, deliberately: this is the tensor that becomes previous_hidden_states on the next | |
| + // speculative step. See the note at the top of this file. | |
| + cb(cur, "result_mtp_hidden", -1); | |
| + res->t_embd = cur; | |
| + | |
| + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); | |
| + cb(cur, "result_norm", -1); | |
| + | |
| + 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 91bab3fd4..f13fbe5e4 100644 | |
| --- a/src/models/models.h | |
| +++ b/src/models/models.h | |
| struct llm_build_glm5_next : public llm_build_delta_net_base { | |
| const llama_model & model; | |
| }; | |
| +// GLM-5.3's MTP module as a standalone speculative draft: one MLA+DSA layer and one MoE, fed | |
| +// the target model's hidden state through llama_set_mtp_hidden(). No KDA and no mHC. | |
| +struct llm_build_glm5_next_mtp : public llm_graph_context { | |
| + llm_build_glm5_next_mtp(const llama_model & model, const llm_graph_params & params); | |
| + | |
| + 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); | |
| -- | |
| 2.43.0 | |
| From 694a46a9f5b00cd20197bfc12d01871cb84e54b0 Mon Sep 17 00:00:00 2001 | |
| From: Patrick Devaney <patrickbdevaney@gmail.com> | |
| Date: Thu, 3 Sep 2026 14:14:10 -0400 | |
| Subject: [PATCH 15/15] speculative: checkpoint the recurrent state so hybrid | |
| targets can be speculated against | |
| Rejecting a drafted token means removing it from the target's memory, and a model | |
| whose layers carry a rolling recurrent state cannot do that - the state summarises | |
| every token it has seen and has no per-token history to rewind to, so | |
| llama_memory_recurrent::seq_rm refuses any partial removal touching the final | |
| position. That ruled out speculative decoding for the whole hybrid family | |
| regardless of the draft: GLM-5.3 (34 of 45 layers are KDA linear attention), | |
| Mamba, Jamba, Falcon-H1, Granite-4, Qwen3-Next. | |
| It also failed silently. llama_decode returned -1, llama-speculative-simple never | |
| checked, and generation carried on from stale logits while reporting an acceptance | |
| rate for tokens it had never verified - a convincing 1.4-1.5x that was not real. | |
| Only the server called is_compat. The guard is now in the tool as well, before the | |
| prompt decode rather than after, since is_compat clears the context. | |
| The fix checkpoints the recurrent state before the verification batch and, on a | |
| partial rejection, restores it and replays the accepted tokens. | |
| LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY already saves the recurrent half alone, so the | |
| checkpoint is a fixed 145.56 MiB for GLM-5.3 whatever the context length, and the | |
| attention KV needs no checkpoint at all: restoring puts the recurrent cell's final | |
| position back before the removed range, which makes the caller's ordinary seq_rm | |
| legal again. No call site had to change to be correct. | |
| Three things that are not obvious: | |
| - The replay is a decode, and a decode overwrites the target's embeddings, which | |
| is exactly what an MTP draft reads next. The hidden state is now captured | |
| eagerly in set_target_output_idx rather than lazily at draft time. | |
| - With more than one sequence in flight the replay would also overwrite logits | |
| the caller is still sampling other slots from, so restore and replay are split: | |
| restore happens at once, replay waits for common_speculative_flush_rollback(). | |
| - A draft discarded for being shorter than n_min used to leave a checkpoint armed | |
| for tokens that never reached the batch, so n_min is applied before the save. | |
| Measured on GLM-5.3-Flash-REAP50 IQ3_M: the checkpoint costs about 11 ms against a | |
| 335 ms/token decode, and the un-fine-tuned blk.45 draft is accepted 73.7% of the | |
| time at depth 1. | |
| Co-Authored-By: Claude Opus 5 <noreply@anthropic.com> | |
| Claude-Session: https://claude.ai/code/session_016B6ZbadovyJKFiN1CwDBLA | |
| --- | |
| common/speculative.cpp | 534 +++++++++++++++++- | |
| common/speculative.h | 24 +- | |
| .../speculative-simple/speculative-simple.cpp | 23 + | |
| tools/server/server-context.cpp | 12 +- | |
| 4 files changed, 584 insertions(+), 9 deletions(-) | |
| diff --git a/common/speculative.cpp b/common/speculative.cpp | |
| index 3e68c38e4..e12ed5d19 100644 | |
| --- a/common/speculative.cpp | |
| +++ b/common/speculative.cpp | |
| const std::vector<enum common_speculative_type> common_speculative_types = { | |
| COMMON_SPECULATIVE_TYPE_NONE, | |
| COMMON_SPECULATIVE_TYPE_DRAFT, | |
| COMMON_SPECULATIVE_TYPE_EAGLE3, | |
| + COMMON_SPECULATIVE_TYPE_MTP, | |
| COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, | |
| COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K, | |
| COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, | |
| const std::map<std::string, enum common_speculative_type> common_speculative_typ | |
| {"none", COMMON_SPECULATIVE_TYPE_NONE}, | |
| {"draft", COMMON_SPECULATIVE_TYPE_DRAFT}, | |
| {"eagle3", COMMON_SPECULATIVE_TYPE_EAGLE3}, | |
| + {"mtp", COMMON_SPECULATIVE_TYPE_MTP}, | |
| {"ngram_simple", COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE}, | |
| {"ngram_map_k", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K}, | |
| {"ngram_map_k4v", COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V}, | |
| struct common_speculative_state { | |
| llama_tokens & result) = 0; | |
| virtual void accept(uint16_t n_accepted) = 0; | |
| + | |
| + // Hidden-state speculators (MTP, EAGLE) need to know which target output row belongs to the | |
| + // token preceding id_last; it is the count of draft tokens the target just accepted. Nothing | |
| + // else uses it, so the default is to ignore it. | |
| + virtual void set_target_output_idx(int32_t i) { GGML_UNUSED(i); } | |
| }; | |
| struct common_speculative_state_draft : public common_speculative_state { | |
| struct common_speculative_state_eagle3 : public common_speculative_state { | |
| } | |
| }; | |
| + | |
| +// Is this draft model GLM-5.3's native MTP module, repackaged by scripts/make_mtp_draft.py? | |
| +static bool common_speculative_is_mtp(const llama_model * model_dft) { | |
| + if (!model_dft) { | |
| + return false; | |
| + } | |
| + | |
| + char buf[64] = {0}; | |
| + const int32_t n = llama_model_meta_val_str(model_dft, "general.architecture", buf, sizeof(buf)); | |
| + | |
| + return n > 0 && strcmp(buf, "glm5-next-mtp") == 0; | |
| +} | |
| + | |
| +// Speculative decoding with a Multi-Token Prediction head. | |
| +// | |
| +// Unlike a draft model, an MTP module is not a small language model: it cannot predict anything | |
| +// on its own. Every step it consumes the hidden state that the *target* produced for the | |
| +// preceding token, so the two models are coupled at every decode rather than merely agreeing on | |
| +// a vocabulary. That coupling is the whole reason it drafts well from a single layer - the | |
| +// target's 45 layers have already encoded the context into the vector it is handed. | |
| +// | |
| +// Reference: vllm/model_executor/models/glm4_moe_mtp.py. | |
| +// | |
| +// Two consequences shape the code below: | |
| +// | |
| +// * The draft's own KV cache holds only the tokens this speculator has itself decoded. Hidden | |
| +// states for the prompt were never computed - the target emits one only for the position it | |
| +// was asked to output - so there is nothing to seed them with. This costs less than it | |
| +// sounds: GLM-5.3's MLA is NoPE, so a gap in the cache is not a positional inconsistency, | |
| +// merely fewer keys, and the context the module actually relies on arrives through h_prev. | |
| +// | |
| +// * At depth > 1 the module feeds on its OWN output rather than the target's. That output is | |
| +// the pre-shared_head_norm tensor, which is why the draft graph exposes it as t_embd; using | |
| +// the post-norm value would silently degrade every draft past the first. | |
| +struct common_speculative_state_mtp : public common_speculative_state { | |
| + llama_context * ctx_tgt; | |
| + llama_context * ctx_dft; | |
| + | |
| + common_sampler * smpl; | |
| + llama_batch batch; | |
| + | |
| + // Number of cells the draft KV holds that the target has committed to. Everything at or | |
| + // after this position is last round's speculation and is dropped before drafting again. | |
| + int32_t n_dft = 0; | |
| + | |
| + // Which target output row carries the hidden state for the token preceding id_last. The | |
| + // caller knows it (it is the number of draft tokens the target accepted); -1 means "the | |
| + // last row", which is the right answer both on the first call and when every draft was | |
| + // accepted. | |
| + int32_t tgt_out_idx = -1; | |
| + | |
| + std::vector<float> h_buf; | |
| + | |
| + // Whether h_buf holds the hidden state for the current round. It is filled in | |
| + // set_target_output_idx() rather than here in draft(), because the target's embeddings buffer | |
| + // does not survive until drafting: on a hybrid target the accepted tokens are replayed on | |
| + // ctx_tgt in between (see common_spec_rollback), which overwrites it. Reading it late would | |
| + // hand the draft head a hidden state from the replay batch instead of the verification batch. | |
| + bool h_valid = false; | |
| + | |
| + common_speculative_state_mtp( | |
| + enum common_speculative_type type, | |
| + llama_context * ctx_tgt, | |
| + llama_context * ctx_dft) | |
| + : common_speculative_state(type) | |
| + , ctx_tgt(ctx_tgt) | |
| + , ctx_dft(ctx_dft) | |
| + { | |
| + batch = llama_batch_init(llama_n_batch(ctx_dft), 0, 1); | |
| + | |
| + { | |
| + common_params_sampling sparams; | |
| + sparams.no_perf = false; | |
| + sparams.top_k = 10; | |
| + sparams.samplers = { COMMON_SAMPLER_TYPE_TOP_K }; | |
| + | |
| + smpl = common_sampler_init(llama_get_model(ctx_dft), sparams); | |
| + } | |
| + | |
| + // Both contexts must hand back hidden states: the target's feeds depth 1, the draft's | |
| + // feeds every depth after that. llama.cpp extracts embeddings together with logits, so | |
| + // this costs an n_embd-sized row per output and nothing else. | |
| + llama_set_embeddings(ctx_tgt, true); | |
| + llama_set_embeddings(ctx_dft, true); | |
| + | |
| + // Without this, turning embeddings on would promote every prompt token to an output and | |
| + // allocate a full logits row for each. We only ever read the hidden state of tokens the | |
| + // caller already wanted logits for, so ask for exactly that. | |
| + llama_set_embeddings_outputs_only(ctx_tgt, true); | |
| + llama_set_embeddings_outputs_only(ctx_dft, true); | |
| + | |
| + h_buf.resize(llama_model_n_embd(llama_get_model(ctx_dft))); | |
| + } | |
| + | |
| + ~common_speculative_state_mtp() override { | |
| + common_sampler_free(smpl); | |
| + llama_batch_free(batch); | |
| + } | |
| + | |
| + void begin(const llama_tokens & prompt) override { | |
| + GGML_UNUSED(prompt); | |
| + | |
| + // A new generation shares nothing with the previous one. | |
| + llama_memory_clear(llama_get_memory(ctx_dft), false); | |
| + n_dft = 0; | |
| + tgt_out_idx = -1; | |
| + h_valid = false; | |
| + } | |
| + | |
| + void draft( | |
| + const common_params_speculative & params, | |
| + const llama_tokens & prompt_tgt, | |
| + llama_token id_last, | |
| + llama_tokens & result) override { | |
| + GGML_UNUSED(prompt_tgt); | |
| + | |
| + result.clear(); | |
| + result.reserve(params.n_max); | |
| + | |
| + const int n_embd = (int) h_buf.size(); | |
| + | |
| + // Normally set_target_output_idx() has already captured this. It has not on the first | |
| + // draft of a sequence, where the caller has decoded the prompt and never accepted | |
| + // anything, so fall back to reading the target's last output row. | |
| + if (!h_valid) { | |
| + const float * h_tgt = llama_get_embeddings_ith(ctx_tgt, tgt_out_idx); | |
| + if (h_tgt == nullptr) { | |
| + // No hidden state means no draft. Returning empty is not a failure: the caller | |
| + // simply decodes one token the ordinary way, which is strictly better than | |
| + // drafting from a stale or zeroed vector and having the target reject it. | |
| + LOG_DBG("%s: target hidden state unavailable (row %d) - skipping this draft\n", | |
| + __func__, tgt_out_idx); | |
| + return; | |
| + } | |
| + std::copy(h_tgt, h_tgt + n_embd, h_buf.begin()); | |
| + } | |
| + | |
| + // h_buf is overwritten with draft-side hidden states below, so it must be re-captured | |
| + // before the next round. | |
| + h_valid = false; | |
| + | |
| + // Drop last round's speculation, keeping the cells the target committed to. | |
| + auto * mem_dft = llama_get_memory(ctx_dft); | |
| + llama_memory_seq_rm(mem_dft, 0, n_dft, -1); | |
| + | |
| + // Take the write position from the cache rather than from n_dft. The two should agree, | |
| + // but a counter that drifts by one produces "sequence positions must remain consecutive" | |
| + // on the *next* call, at which point the cause is several decodes behind the symptom - | |
| + // and every subsequent draft fails the same way, so a single slip silently turns | |
| + // speculation off for the rest of the generation. Asking the memory module what it | |
| + // actually holds cannot drift. | |
| + const llama_pos pos_max = llama_memory_seq_pos_max(mem_dft, 0); | |
| + | |
| + common_sampler_reset(smpl); | |
| + | |
| + llama_token tok = id_last; | |
| + int32_t pos = pos_max + 1; // pos_max is -1 on an empty cache, so this starts at 0 | |
| + int n_decoded = 0; | |
| + | |
| + for (int i = 0; i < params.n_max; ++i) { | |
| + llama_set_mtp_hidden(ctx_dft, h_buf.data(), 1); | |
| + | |
| + common_batch_clear(batch); | |
| + common_batch_add (batch, tok, pos, { 0 }, true); | |
| + | |
| + if (llama_decode(ctx_dft, batch) != 0) { | |
| + // Leave nothing half-written behind: a failed decode that still advanced our | |
| + // bookkeeping would desynchronise every later call. Start the draft cache over | |
| + // instead - it costs this generation's accumulated MTP context and nothing else. | |
| + LOG_ERR("%s: draft decode failed at depth %d (pos %d) - resetting draft cache\n", | |
| + __func__, i, pos); | |
| + llama_memory_clear(mem_dft, false); | |
| + n_dft = 0; | |
| + return; | |
| + } | |
| + | |
| + n_decoded++; | |
| + | |
| + common_sampler_sample(smpl, ctx_dft, 0, true); | |
| + | |
| + const auto * cur_p = common_sampler_get_candidates(smpl, true); | |
| + const llama_token id = cur_p->data[0].id; | |
| + | |
| + common_sampler_accept(smpl, id, true); | |
| + result.push_back(id); | |
| + | |
| + if ((int) result.size() >= params.n_max) { | |
| + break; | |
| + } | |
| + | |
| + // Drafting is only worth the decode if the module is confident; the target pays for | |
| + // every token it has to reject. | |
| + if (cur_p->data[0].p < params.p_min) { | |
| + break; | |
| + } | |
| + | |
| + // Depth > 1 conditions on the module's own output. See the note above on why this | |
| + // is the pre-norm tensor. | |
| + const float * h_dft = llama_get_embeddings_ith(ctx_dft, -1); | |
| + if (h_dft == nullptr) { | |
| + break; | |
| + } | |
| + std::copy(h_dft, h_dft + n_embd, h_buf.begin()); | |
| + | |
| + tok = id; | |
| + pos++; | |
| + } | |
| + | |
| + // The position of id_last is now committed; everything drafted after it is provisional | |
| + // until accept() tells us how much the target kept. If nothing was decoded there is | |
| + // nothing to commit. | |
| + if (n_decoded > 0) { | |
| + n_dft = pos_max + 2; | |
| + } | |
| + } | |
| + | |
| + void accept(uint16_t n_accepted) override { | |
| + // Draft tokens the target agreed with keep their draft-side KV cells. | |
| + n_dft += n_accepted; | |
| + } | |
| + | |
| + void set_target_output_idx(int32_t i) override { | |
| + tgt_out_idx = i; | |
| + | |
| + const float * h_tgt = llama_get_embeddings_ith(ctx_tgt, i); | |
| + if (h_tgt == nullptr) { | |
| + h_valid = false; | |
| + return; | |
| + } | |
| + | |
| + std::copy(h_tgt, h_tgt + h_buf.size(), h_buf.begin()); | |
| + h_valid = true; | |
| + } | |
| +}; | |
| + | |
| // state of self-speculation (simple implementation, not ngram-map) | |
| struct common_speculative_state_ngram_simple : public common_speculative_state { | |
| common_ngram_simple_config config; | |
| struct common_speculative_state_ngram_cache : public common_speculative_state { | |
| } | |
| }; | |
| + | |
| +// | |
| +// recurrent-state rollback | |
| +// | |
| + | |
| +// Speculative decoding rejects a drafted token by removing it from the target's memory. A model | |
| +// whose layers carry a rolling recurrent state cannot do that: the state summarises every token | |
| +// seen so far and has no per-token history to rewind to, so llama_memory_recurrent::seq_rm refuses | |
| +// any partial removal that touches the final position. That is not a corner case for GLM-5.3 - | |
| +// 34 of its 45 layers are KDA linear attention - and it makes speculation impossible for the whole | |
| +// hybrid family (Mamba, Jamba, Falcon-H1, Granite-4, ...) no matter how good the draft is. | |
| +// | |
| +// The way through is to checkpoint the recurrent state before the verification batch and, when | |
| +// some drafted tokens turn out to be wrong, restore it and replay the ones that were right. Three | |
| +// things keep this cheap: | |
| +// | |
| +// - LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY saves the recurrent half only. Its size is fixed by the | |
| +// model (145 MiB for GLM-5.3-Flash: 9.6 MiB of conv state, 136 MiB of delta-net state), not by | |
| +// the context length, so the cost does not grow as the conversation does. The attention KV | |
| +// needs no checkpoint - it rewinds the ordinary way. | |
| +// - Restoring puts the recurrent cell's final position back at pos_start - 1, which makes the | |
| +// ordinary seq_rm legal again: nothing the cell holds is inside the removed range any more. | |
| +// So the attention half is truncated by the very call the caller was already making, and no | |
| +// call site has to learn about any of this. | |
| +// - The replay is skipped entirely when every drafted token is accepted, which is the case that | |
| +// matters most for throughput. | |
| +// | |
| +// What it costs is one extra forward pass over the accepted tokens on a partial rejection. That is | |
| +// the honest price of the architecture; a future refinement could fold those tokens into the next | |
| +// verification batch instead of paying for a pass of their own. | |
| +struct common_spec_rollback { | |
| + llama_context * ctx = nullptr; | |
| + llama_seq_id seq_id = 0; | |
| + | |
| + std::vector<uint8_t> buf; | |
| + | |
| + llama_pos pos_start = -1; // position of the first token of the verification batch | |
| + llama_tokens fed; // what was fed there: id_last followed by the draft | |
| + bool armed = false; | |
| + | |
| + // Number of tokens waiting to be replayed onto the restored state, or 0 for none. The restore | |
| + // itself is done as soon as the rejection is known, but the replay is a decode, and a decode | |
| + // overwrites the context's output buffer. With more than one sequence in flight the caller is | |
| + // still sampling other sequences out of that buffer at this point, so the decode has to wait | |
| + // until it is asked for - see common_speculative_flush_rollback(). | |
| + int32_t n_pending = 0; | |
| + | |
| + llama_batch batch; | |
| + | |
| + common_spec_rollback(llama_context * ctx, llama_seq_id seq_id, int32_t n_batch) | |
| + : ctx(ctx), seq_id(seq_id), batch(llama_batch_init(n_batch, 0, 1)) {} | |
| + | |
| + ~common_spec_rollback() { | |
| + llama_batch_free(batch); | |
| + } | |
| +}; | |
| + | |
| +// Take the checkpoint. Called with the target's memory holding exactly the committed prefix, i.e. | |
| +// right before the caller builds the batch of id_last + draft. | |
| +static void common_spec_rollback_save( | |
| + common_spec_rollback * rb, | |
| + llama_token id_last, | |
| + const llama_tokens & draft) { | |
| + rb->armed = false; | |
| + | |
| + const size_t n = llama_state_seq_get_size_ext(rb->ctx, rb->seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); | |
| + if (n == 0) { | |
| + LOG_WRN("%s: the target has no recurrent state to checkpoint\n", __func__); | |
| + return; | |
| + } | |
| + | |
| + if (rb->buf.size() < n) { | |
| + rb->buf.resize(n); | |
| + } | |
| + | |
| + if (llama_state_seq_get_data_ext(rb->ctx, rb->buf.data(), rb->buf.size(), rb->seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { | |
| + LOG_WRN("%s: failed to checkpoint the recurrent state\n", __func__); | |
| + return; | |
| + } | |
| + | |
| + // Take the write position from the memory rather than from the caller's token count. The two | |
| + // agree, but only the memory is authoritative about where the next token actually lands. | |
| + rb->pos_start = llama_memory_seq_pos_max(llama_get_memory(rb->ctx), rb->seq_id) + 1; | |
| + | |
| + rb->fed.clear(); | |
| + rb->fed.reserve(1 + draft.size()); | |
| + rb->fed.push_back(id_last); | |
| + rb->fed.insert(rb->fed.end(), draft.begin(), draft.end()); | |
| + | |
| + rb->armed = true; | |
| +} | |
| + | |
| +// Put the sequence back to the accepted prefix. n_accepted is the number of DRAFTED tokens the | |
| +// target agreed with, so the batch keeps n_accepted + 1 of the tokens it was given. Returns false | |
| +// if the sequence could not be restored, in which case speculation must stop - carrying on against | |
| +// a recurrent state that no longer matches the tokens would corrupt the output silently. | |
| +static bool common_spec_rollback_apply(common_spec_rollback * rb, int32_t n_accepted) { | |
| + if (!rb->armed) { | |
| + return true; | |
| + } | |
| + | |
| + rb->armed = false; | |
| + | |
| + const int32_t n_keep = n_accepted + 1; | |
| + | |
| + // Nothing was rejected: the state already covers exactly the tokens that were kept. | |
| + if (n_keep >= (int32_t) rb->fed.size()) { | |
| + return true; | |
| + } | |
| + | |
| + if (llama_state_seq_set_data_ext(rb->ctx, rb->buf.data(), rb->buf.size(), rb->seq_id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { | |
| + LOG_ERR("%s: failed to restore the recurrent state\n", __func__); | |
| + return false; | |
| + } | |
| + | |
| + // Legal now that the recurrent half is back at pos_start - 1; this truncates the attention KV. | |
| + if (!llama_memory_seq_rm(llama_get_memory(rb->ctx), rb->seq_id, rb->pos_start, -1)) { | |
| + LOG_ERR("%s: failed to truncate the memory after restoring the recurrent state\n", __func__); | |
| + return false; | |
| + } | |
| + | |
| + rb->n_pending = n_keep; | |
| + | |
| + LOG_DBG("%s: rolled back to pos %d, %d token(s) queued for replay\n", __func__, rb->pos_start, n_keep); | |
| + | |
| + return true; | |
| +} | |
| + | |
| +// Replay the accepted tokens onto the restored state. Must run before anything else touches the | |
| +// sequence, and only once the caller has finished reading logits out of the context. | |
| +static bool common_spec_rollback_flush(common_spec_rollback * rb) { | |
| + if (rb->n_pending == 0) { | |
| + return true; | |
| + } | |
| + | |
| + const int32_t n_keep = rb->n_pending; | |
| + rb->n_pending = 0; | |
| + | |
| + // Between the restore and here the caller may have thrown the sequence away entirely - the | |
| + // server does exactly that when it releases a child task's slot. Replaying onto a sequence | |
| + // that no longer ends where the checkpoint left it would punch a hole in it, so check that it | |
| + // is still the sequence we rolled back before writing to it. | |
| + const llama_pos pos_cur = llama_memory_seq_pos_max(llama_get_memory(rb->ctx), rb->seq_id) + 1; | |
| + if (pos_cur != rb->pos_start) { | |
| + LOG_DBG("%s: sequence moved to pos %d since the rollback to %d - dropping the replay\n", | |
| + __func__, pos_cur, rb->pos_start); | |
| + return true; | |
| + } | |
| + | |
| + common_batch_clear(rb->batch); | |
| + for (int32_t i = 0; i < n_keep; ++i) { | |
| + // Ask for logits on the last row only. They are discarded - the caller sampled from the | |
| + // verification batch - but a batch that requests no outputs at all is not a shape every | |
| + // path expects, and one row is nothing. | |
| + common_batch_add(rb->batch, rb->fed[i], rb->pos_start + i, { rb->seq_id }, i == n_keep - 1); | |
| + } | |
| + | |
| + if (llama_decode(rb->ctx, rb->batch) != 0) { | |
| + LOG_ERR("%s: failed to replay %d accepted tokens\n", __func__, n_keep); | |
| + return false; | |
| + } | |
| + | |
| + LOG_DBG("%s: replayed %d token(s) from pos %d\n", __func__, n_keep, rb->pos_start); | |
| + | |
| + return true; | |
| +} | |
| + | |
| struct common_speculative { | |
| std::vector<std::unique_ptr<common_speculative_state>> impls; // list of implementations to use and their states | |
| common_speculative_state * curr_impl = nullptr; // current implementation in use (for stats) | |
| + | |
| + // Set only when the target's memory cannot drop rejected tokens on its own. | |
| + std::unique_ptr<common_spec_rollback> rb; | |
| + | |
| + // Latched if a rollback ever fails. Speculation stays off for the rest of the run rather than | |
| + // producing tokens verified against a state that has drifted from the sequence. | |
| + bool broken = false; | |
| }; | |
| static common_ngram_map get_common_ngram_map(const common_speculative_config & config) { | |
| std::string common_speculative_type_to_str(enum common_speculative_type type) { | |
| case COMMON_SPECULATIVE_TYPE_NONE: return "none"; | |
| case COMMON_SPECULATIVE_TYPE_DRAFT: return "draft"; | |
| case COMMON_SPECULATIVE_TYPE_EAGLE3: return "eagle3"; | |
| + case COMMON_SPECULATIVE_TYPE_MTP: return "mtp"; | |
| case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: return "ngram_simple"; | |
| case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K: return "ngram_map_k"; | |
| case COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V: return "ngram_map_k4v"; | |
| bool common_speculative_is_compat(llama_context * ctx_tgt) { | |
| // try to remove the last tokens | |
| if (!llama_memory_seq_rm(mem, 0, 1, -1)) { | |
| - LOG_WRN("%s: the target context does not support partial sequence removal\n", __func__); | |
| - res = false; | |
| - goto done; | |
| + // A recurrent or hybrid target refuses this: its rolling state cannot be rewound by a | |
| + // token. That used to be the end of the matter, but it is recoverable - the state can be | |
| + // checkpointed before the verification batch and the accepted tokens replayed onto it. | |
| + // See common_spec_rollback; common_speculative_init arms it for exactly these models. | |
| + if (!llama_model_is_recurrent(llama_get_model(ctx_tgt)) && !llama_model_is_hybrid(llama_get_model(ctx_tgt))) { | |
| + LOG_WRN("%s: the target context does not support partial sequence removal\n", __func__); | |
| + res = false; | |
| + goto done; | |
| + } | |
| + | |
| + LOG_INF("%s: the target cannot drop rejected tokens directly - speculation will checkpoint " | |
| + "and replay its recurrent state\n", __func__); | |
| } | |
| done: | |
| done: | |
| // | |
| common_speculative * common_speculative_init( | |
| common_params_speculative & params, | |
| - llama_context * ctx_tgt) { | |
| + llama_context * ctx_tgt, | |
| + llama_seq_id seq_id_tgt) { | |
| llama_context * ctx_dft = nullptr; | |
| if (params.model_dft) { | |
| ctx_dft = llama_init_from_model(params.model_dft, params.cparams_dft); | |
| common_speculative * common_speculative_init( | |
| bool has_draft = !params.mparams_dft.path.empty(); | |
| bool has_draft_eagle3 = false; // TODO PR-18039: if params.speculative.eagle3 | |
| + // A GLM-5.3 MTP draft announces itself through its architecture, so -md is all | |
| + // the user has to pass. Treating it as an ordinary draft model would run it | |
| + // without hidden states and produce noise, so this is a redirect, not an option. | |
| + bool has_mtp = has_draft && common_speculative_is_mtp(params.model_dft); | |
| + if (has_mtp) { | |
| + has_draft = false; | |
| + LOG_INF("%s: draft model is a GLM-5.3 MTP module - using hidden-state speculation\n", __func__); | |
| + } | |
| + | |
| bool has_ngram_cache = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_CACHE); | |
| bool has_ngram_simple = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE); | |
| bool has_ngram_map_k = (params.type == COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K); | |
| common_speculative * common_speculative_init( | |
| if (has_draft_eagle3) { | |
| configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_EAGLE3, params)); | |
| } | |
| + if (has_mtp) { | |
| + configs.push_back(common_speculative_config(COMMON_SPECULATIVE_TYPE_MTP, params)); | |
| + } | |
| } | |
| std::vector<std::unique_ptr<common_speculative_state>> impls = {}; | |
| common_speculative * common_speculative_init( | |
| impls.push_back(std::make_unique<common_speculative_state_eagle3>(config.type)); | |
| break; | |
| } | |
| + case COMMON_SPECULATIVE_TYPE_MTP: { | |
| + impls.push_back(std::make_unique<common_speculative_state_mtp>(config.type, | |
| + /* .ctx_tgt = */ ctx_tgt, | |
| + /* .ctx_dft = */ ctx_dft | |
| + )); | |
| + break; | |
| + } | |
| case COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE: { | |
| common_ngram_map ngram_map = get_common_ngram_map(config); | |
| common_speculative * common_speculative_init( | |
| return nullptr; | |
| } | |
| - auto * result = new common_speculative { | |
| - /* .impls = */ std::move(impls) | |
| - }; | |
| + auto * result = new common_speculative(); | |
| + | |
| + result->impls = std::move(impls); | |
| + | |
| + // Arm the recurrent-state rollback for targets whose memory cannot drop rejected tokens. The | |
| + // test is on the model rather than on a decode probe: common_speculative_is_compat() has | |
| + // already paid for the probe, and repeating it here would clear the caller's memory. | |
| + { | |
| + const llama_model * model_tgt = llama_get_model(ctx_tgt); | |
| + | |
| + if (llama_model_is_recurrent(model_tgt) || llama_model_is_hybrid(model_tgt)) { | |
| + result->rb = std::make_unique<common_spec_rollback>(ctx_tgt, seq_id_tgt, (int32_t) llama_n_batch(ctx_tgt)); | |
| + | |
| + LOG_INF("%s: recurrent-state rollback enabled for seq %d (checkpoint %.2f MiB)\n", __func__, | |
| + seq_id_tgt, | |
| + llama_state_seq_get_size_ext(ctx_tgt, seq_id_tgt, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY)/1024.0/1024.0); | |
| + } | |
| + } | |
| return result; | |
| } | |
| llama_tokens common_speculative_draft( | |
| spec->curr_impl = nullptr; // reset current implementation | |
| + if (spec->broken) { | |
| + return result; | |
| + } | |
| + | |
| + // Safety net for a caller that does not flush on its own: the checkpoint below reads the | |
| + // target's write position, and that is only right once the replay has landed. | |
| + common_speculative_flush_rollback(spec); | |
| + | |
| for (auto & impl : spec->impls) { | |
| { | |
| common_time_meas tm(impl->t_draft_us, !impl->gen_perf); | |
| llama_tokens common_speculative_draft( | |
| } | |
| } | |
| + // Drop a draft too short to be worth verifying. Both callers already do this on their side, | |
| + // but deciding it here as well is what keeps the checkpoint below in step with what the target | |
| + // is actually fed - a draft discarded after the checkpoint was taken would leave the rollback | |
| + // expecting tokens that never reached the batch. | |
| + if ((int) result.size() < params.n_min) { | |
| + result.clear(); | |
| + } | |
| + | |
| + // The target's memory holds nothing but the committed prefix at this moment, which is the only | |
| + // point at which the checkpoint is worth taking. | |
| + if (spec->rb) { | |
| + if (result.empty()) { | |
| + // Nothing drafted means nothing can be rejected. Disarm rather than leave the previous | |
| + // round's checkpoint armed, which would make the next rollback restore to a stale | |
| + // position. | |
| + spec->rb->armed = false; | |
| + } else { | |
| + common_spec_rollback_save(spec->rb.get(), id_last, result); | |
| + } | |
| + } | |
| + | |
| return result; | |
| } | |
| +void common_speculative_set_target_output_idx(common_speculative * spec, int32_t i) { | |
| + if (spec == nullptr) { | |
| + return; | |
| + } | |
| + | |
| + // Applies to whichever implementation is live; the others ignore it. | |
| + for (const auto & impl : spec->impls) { | |
| + impl->set_target_output_idx(i); | |
| + } | |
| + | |
| + // i is also the number of drafted tokens the target accepted, which is all the rollback needs | |
| + // to know. It has to run after the implementations, not before: the MTP implementation reads | |
| + // the target's hidden state in the call above, and the replay overwrites it. | |
| + if (spec->rb && !common_spec_rollback_apply(spec->rb.get(), i)) { | |
| + LOG_ERR("%s: recurrent-state rollback failed - disabling speculative decoding\n", __func__); | |
| + spec->broken = true; | |
| + } | |
| +} | |
| + | |
| +void common_speculative_flush_rollback(common_speculative * spec) { | |
| + if (spec == nullptr || spec->rb == nullptr || spec->broken) { | |
| + return; | |
| + } | |
| + | |
| + if (!common_spec_rollback_flush(spec->rb.get())) { | |
| + LOG_ERR("%s: recurrent-state replay failed - disabling speculative decoding\n", __func__); | |
| + spec->broken = true; | |
| + } | |
| +} | |
| + | |
| void common_speculative_accept(common_speculative * spec, uint16_t n_accepted) { | |
| if (n_accepted == 0) { | |
| return; | |
| diff --git a/common/speculative.h b/common/speculative.h | |
| index 876cde3d1..941bdfa1b 100644 | |
| --- a/common/speculative.h | |
| +++ b/common/speculative.h | |
| std::string common_speculative_type_to_str(enum common_speculative_type type); | |
| // note: clears the memory of the context | |
| bool common_speculative_is_compat(llama_context * ctx_tgt); | |
| +// seq_id_tgt is the sequence this speculator drives on the target context. It only matters for | |
| +// targets whose memory cannot drop rejected tokens by itself (recurrent and hybrid models such as | |
| +// GLM-5.3): those are rolled back by checkpointing and replaying the recurrent state of that one | |
| +// sequence. Callers that run a single sequence can leave it at 0. | |
| common_speculative * common_speculative_init( | |
| common_params_speculative & params, | |
| - llama_context * ctx_tgt); | |
| + llama_context * ctx_tgt, | |
| + llama_seq_id seq_id_tgt = 0); | |
| void common_speculative_free(common_speculative * spec); | |
| llama_tokens common_speculative_draft( | |
| const llama_tokens & prompt, | |
| llama_token id_last); | |
| +// Tell a hidden-state speculator (MTP) which target output row holds the hidden state for the | |
| +// token preceding id_last - i.e. the number of draft tokens the target just accepted. Harmless | |
| +// and ignored for every other speculative type. Call it after sampling the target, before the | |
| +// next common_speculative_draft(). | |
| +void common_speculative_set_target_output_idx(common_speculative * spec, int32_t i); | |
| + | |
| // informs the speculative decoder that n_accepted tokens were accepted by the target model | |
| void common_speculative_accept(common_speculative * spec, uint16_t n_accepted); | |
| +// Finish the rollback that common_speculative_set_target_output_idx() started, on targets that | |
| +// need one (recurrent and hybrid models - see common_spec_rollback in speculative.cpp). The | |
| +// rejected tokens are already gone from the target's memory by then; what is left is a decode to | |
| +// replay the accepted ones, and a decode overwrites the context's output buffer. So the caller | |
| +// has to say when it is done reading logits for every sequence sharing that context. | |
| +// | |
| +// Call it once after the whole sample-and-accept pass. It is a no-op for ordinary targets and for | |
| +// sequences with nothing pending, and common_speculative_draft() calls it too, so forgetting it | |
| +// costs correctness only if the sequence is dropped before it ever drafts again. | |
| +void common_speculative_flush_rollback(common_speculative * spec); | |
| + | |
| // print statistics about the speculative decoding | |
| void common_speculative_print_stats(const common_speculative * spec); | |
| diff --git a/examples/speculative-simple/speculative-simple.cpp b/examples/speculative-simple/speculative-simple.cpp | |
| index a03dbce88..e49b09840 100644 | |
| --- a/examples/speculative-simple/speculative-simple.cpp | |
| +++ b/examples/speculative-simple/speculative-simple.cpp | |
| int main(int argc, char ** argv) { | |
| struct common_sampler * smpl = common_sampler_init(model_tgt, params.sampling); | |
| // eval the prompt | |
| + // Not every target can be speculated against. Rejecting a drafted token means removing it from | |
| + // the KV cache, and a recurrent (linear-attention) layer has no way to rewind its state to an | |
| + // earlier position - llama_memory_recurrent::seq_rm refuses a partial removal that touches the | |
| + // final position. For those models common_speculative_init arms a checkpoint of the recurrent | |
| + // state instead, and this check passes; it still fails for a target that can do neither. | |
| + // | |
| + // The check has to be here rather than further down: without it the run does not fail, it | |
| + // lies. llama_decode returns -1, nobody looks, and generation continues from stale logits, | |
| + // reporting an acceptance rate for tokens it never really verified. | |
| + if (!common_speculative_is_compat(ctx_tgt)) { | |
| + LOG_ERR("%s: the target model does not support speculative decoding - its memory module " | |
| + "cannot roll back rejected tokens (typical of hybrid/recurrent architectures)\n", __func__); | |
| + return 1; | |
| + } | |
| + | |
| llama_decode(ctx_tgt, llama_batch_get_one(inp.data(), inp.size() - 1)); | |
| // note: keep the last token separate! | |
| int main(int argc, char ** argv) { | |
| // | |
| const auto ids = common_sampler_sample_and_accept_n(smpl, ctx_tgt, draft); | |
| + // The hidden state an MTP draft needs next is the one at the row of the last accepted | |
| + // token, which is exactly ids.size() - 1. Other speculative types ignore this. | |
| + common_speculative_set_target_output_idx(spec, (int32_t) ids.size() - 1); | |
| + | |
| + // Only one sequence here, so nothing else is waiting to read logits: the rollback of a | |
| + // hybrid target can be finished immediately. (No-op for targets that do not need one.) | |
| + common_speculative_flush_rollback(spec); | |
| + | |
| //LOG_DBG("ids: %s\n", string_from(ctx_tgt, ids).c_str()); | |
| GGML_ASSERT(ids.size() > 0); // there will always be at least one accepted token | |
| diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp | |
| index 5523f23b5..e28fe38a2 100644 | |
| --- a/tools/server/server-context.cpp | |
| +++ b/tools/server/server-context.cpp | |
| private: | |
| // try speculative decoding | |
| if (can_spec) { | |
| - slot.spec = common_speculative_init(params_base.speculative, slot.ctx); | |
| + slot.spec = common_speculative_init(params_base.speculative, slot.ctx, slot.id); | |
| if (slot.spec) { | |
| if (mctx) { | |
| SRV_ERR("%s\n", "speculative decoding is not supported with multimodal"); | |
| private: | |
| slot.n_draft_accepted += ids.size() - 1; | |
| // inform the speculative decoding about the number of accepted tokens | |
| + common_speculative_set_target_output_idx(slot.spec, (int32_t) ids.size() - 1); | |
| common_speculative_accept(slot.spec, ids.size() - 1); | |
| // rollback to the state before sampling the draft tokens | |
| private: | |
| SLT_DBG(slot, "accepted %d/%d draft tokens, new n_tokens = %d\n", (int) ids.size() - 1, (int) n_draft, slot.prompt.n_tokens()); | |
| } | |
| + | |
| + // A target that cannot drop rejected tokens by itself - any recurrent or hybrid model - | |
| + // has had its state restored above but not yet replayed onto, because replaying is a | |
| + // decode and a decode would have overwritten the logits the loop above was still | |
| + // sampling other slots from. Every slot is done reading now, so finish the job here | |
| + // rather than leaving a sequence whose memory is shorter than its token list. | |
| + for (auto & slot : slots) { | |
| + common_speculative_flush_rollback(slot.spec); | |
| + } | |
| } | |
| SRV_DBG("%s", "run slots completed\n"); | |
| -- | |
| 2.43.0 | |