From 6174c9135c884dde5356dd5374b4ffb6de48a36e Mon Sep 17 00:00:00 2001 From: eauchs Date: Wed, 26 Aug 2026 18:07:12 +0200 Subject: [PATCH 01/11] model : add GLM-5.3-Flash (glm5next) Adds the glm5next architecture: KDA linear attention, MLA with a DSA indexer and k-pool compression, gated-residual hyper-connections, and a 288-expert MoE. The DSA path extends the existing llama_memory_hybrid_idx container; llama_memory_hybrid and llama_kv_cache are left untouched. --- conversion/__init__.py | 2 + conversion/glm5next.py | 100 ++++ gguf-py/gguf/constants.py | 65 +++ src/CMakeLists.txt | 1 + src/llama-arch.cpp | 3 + src/llama-arch.h | 1 + src/llama-context.cpp | 1 + src/llama-graph.cpp | 159 ++++-- src/llama-graph.h | 53 ++ src/llama-model-saver.cpp | 4 +- src/llama-model.cpp | 52 +- src/llama-model.h | 1 + src/models/glm5next.cpp | 958 +++++++++++++++++++++++++++++++++++++ src/models/models.h | 51 ++ tests/test-llama-archs.cpp | 28 +- 15 files changed, 1441 insertions(+), 38 deletions(-) create mode 100644 conversion/glm5next.py create mode 100644 src/models/glm5next.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index ba73192efa1b..fdf08a7b28ac 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -101,6 +101,8 @@ "GemmaForCausalLM": "gemma", "Glm4ForCausalLM": "glm", "Glm4MoeForCausalLM": "glm", + "Glm5NextForCausalLM": "glm5next", + "Glm5NextForConditionalGeneration": "glm5next", "Glm4MoeLiteForCausalLM": "glm", "Glm4vForConditionalGeneration": "glm", "Glm4vMoeForConditionalGeneration": "glm", diff --git a/conversion/glm5next.py b/conversion/glm5next.py new file mode 100644 index 000000000000..f54d8786f6fb --- /dev/null +++ b/conversion/glm5next.py @@ -0,0 +1,100 @@ +from __future__ import annotations + +import re +from typing import Iterable + +import torch +from torch import Tensor + +import gguf + +from .base import ModelBase +from .glm import GlmMoeDsaModel + + +@ModelBase.register("Glm5NextForConditionalGeneration", "Glm5NextForCausalLM") +@ModelBase.example("zai-org/GLM-5.3-Flash") +class Glm5NextModel(GlmMoeDsaModel): + """GLM-5.3-Flash. + + Trunk that alternates KDA linear attention (34 layers) with MLA + DSA sparse + attention (11 layers), wrapped in hyper-connection streams. The pieces are + already in tree: the KDA tensors follow kimi-linear, the hyper-connection and + k-pool compressor tensors follow deepseek4, and the MLA/MoE/NextN half is + inherited from GLM-5.2 (GlmMoeDsaModel). + """ + + model_arch = gguf.MODEL_ARCH.GLM5NEXT + + # Tensors that carry no per-layer index and are named differently from the + # generic mapping, resolved by suffix (same approach as DeepseekV4Model). + _direct_map = { + "hc_attn_fn": (gguf.MODEL_TENSOR.HC_ATTN_FN, ""), + "hc_attn_base": (gguf.MODEL_TENSOR.HC_ATTN_BASE, ""), + "hc_attn_scale": (gguf.MODEL_TENSOR.HC_ATTN_SCALE, ""), + "hc_ffn_fn": (gguf.MODEL_TENSOR.HC_FFN_FN, ""), + "hc_ffn_base": (gguf.MODEL_TENSOR.HC_FFN_BASE, ""), + "hc_ffn_scale": (gguf.MODEL_TENSOR.HC_FFN_SCALE, ""), + "self_attn.indexer.index_kpool_compress_ape": + (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_APE, ""), + "self_attn.indexer.index_kpool_compress_gate": + (gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE, ""), + } + + def index_tensors(self, remote_hf_model_id: str | None = None): + # TextModel lifts text_config to the root, but only after this runs - + # and the parent already needs num_hidden_layers from it here. + if "text_config" in self.hparams: + self.hparams = {**self.hparams, **self.hparams["text_config"]} + return super().index_tensors(remote_hf_model_id=remote_hf_model_id) + + @classmethod + def filter_tensors(cls, item): + name = item[0] + # text-only for now: drop the vision tower + if name.startswith("model.visual.") or name.startswith("visual."): + return None + return super().filter_tensors(item) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hparams = self.hparams + + # hyper-connections (mHC): identical formulation to DeepSeek-V4, so the + # existing sinkhorn graph applies unchanged. + self.gguf_writer.add_hyper_connection_count(hparams["hc_mult"]) + self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hparams["hc_sinkhorn_iters"]) + self.gguf_writer.add_hyper_connection_epsilon(hparams["hc_eps"]) + + # KDA linear attention + linear = hparams["linear_attn_config"] + self.gguf_writer.add_ssm_conv_kernel(linear["short_conv_kernel_size"]) + self.gguf_writer.add_ssm_inner_size(linear["num_heads"] * linear["head_dim"]) + self.gguf_writer.add_ssm_state_size(linear["head_dim"]) + self.gguf_writer.add_ssm_group_count(linear["num_heads"]) + + # k-pool compression inside the DSA indexer + self.gguf_writer.add_indexer_block_size(hparams["index_kpool"]) + + # clamped SwiGLU + if (limit := hparams.get("swiglu_limit")) is not None: + self.gguf_writer.add_swiglu_clamp_exp([limit] * self.block_count) + self.gguf_writer.add_swiglu_clamp_shexp([limit] * self.block_count) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # the checkpoint wraps the trunk for the multimodal head + name = re.sub(r"^model\.language_model\.", "model.", name) + + # KDA decay conventions, same as conversion/kimi_linear.py: the graph + # expects ssm_a to already hold -exp(A_log), and the time-step bias to + # be named like a bias so it is not loaded as a MUL_MAT weight. + if name.endswith(".A_log"): + data_torch = -torch.exp(data_torch.float()) + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + + for suffix, (tensor, ext) in self._direct_map.items(): + if name.endswith(suffix) and bid is not None: + return [(self.format_tensor_name(tensor, bid) + ext, data_torch)] + + return super().modify_tensors(data_torch, name, bid) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index b85f62a31145..fe74ec54ec59 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -559,6 +559,7 @@ class MODEL_ARCH(IntEnum): GLM4 = auto() GLM4_MOE = auto() GLM_DSA = auto() + GLM5NEXT = auto() BITNET = auto() T5 = auto() T5ENCODER = auto() @@ -1311,6 +1312,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GLM4: "glm4", MODEL_ARCH.GLM4_MOE: "glm4moe", MODEL_ARCH.GLM_DSA: "glm-dsa", + MODEL_ARCH.GLM5NEXT: "glm5next", MODEL_ARCH.BITNET: "bitnet", MODEL_ARCH.T5: "t5", MODEL_ARCH.T5ENCODER: "t5encoder", @@ -4003,6 +4005,69 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.GLM5NEXT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + 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.ATTN_OUT, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + # NextN/MTP tensors - preserved but unused + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + # --- KDA linear attention (34 couches sur 45), cf. kimi-linear --- + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_F_B, + MODEL_TENSOR.SSM_G_A, + MODEL_TENSOR.SSM_G_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + # --- hyper-connections (mHC, Sinkhorn), cf. deepseek4 --- + 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, + # --- indexeur : compression k-pool --- + MODEL_TENSOR.INDEXER_COMPRESSOR_APE, + MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE, + ], MODEL_ARCH.BITNET: [ MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 8922dc12adc7..7379491687a0 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -30,6 +30,7 @@ add_library(llama llama-kv-cache-dsv4.cpp llama-memory.cpp llama-memory-hybrid.cpp + llama-memory-hybrid-idx.cpp llama-memory-hybrid-iswa.cpp llama-memory-hybrid-idx.cpp llama-memory-recurrent.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 446de4ae25b3..5a916c08440e 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -84,6 +84,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GLM4, "glm4" }, { LLM_ARCH_GLM4_MOE, "glm4moe" }, { LLM_ARCH_GLM_DSA, "glm-dsa" }, + { LLM_ARCH_GLM5NEXT, "glm5next" }, { LLM_ARCH_BITNET, "bitnet" }, { LLM_ARCH_T5, "t5" }, { LLM_ARCH_T5ENCODER, "t5encoder" }, @@ -1080,6 +1081,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_MINIMAX_01: + case LLM_ARCH_GLM5NEXT: return true; default: return false; @@ -1144,6 +1146,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_KIMI_K3: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_QWEN3TTS: case LLM_ARCH_QWEN4EXP: // TODO: fix test-llama-archs return false; diff --git a/src/llama-arch.h b/src/llama-arch.h index 0c0b994836f4..4cdf357080e3 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -89,6 +89,7 @@ enum llm_arch { LLM_ARCH_GLM4, LLM_ARCH_GLM4_MOE, LLM_ARCH_GLM_DSA, + LLM_ARCH_GLM5NEXT, LLM_ARCH_BITNET, LLM_ARCH_T5, LLM_ARCH_T5ENCODER, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 3cc27717ece8..bed9455a1cc3 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2314,6 +2314,7 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_QWEN35MOE || model.arch == LLM_ARCH_QWEN4EXP || model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_GLM5NEXT || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_01 || diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 274a6264336f..a9cf1764203c 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -13,6 +13,7 @@ #include "llama-kv-cache-msa.h" #include "llama-kv-cache-dsv4.h" #include "llama-memory-hybrid.h" +#include "llama-memory-hybrid-idx.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-recurrent.h" @@ -1176,6 +1177,45 @@ bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) { return res; } +void llm_graph_input_mem_hybrid_idx::set_input(const llama_ubatch * ubatch) { + mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); + + mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + + const int64_t n_rs = mctx->get_recr()->get_n_rs(); + + if (inp_rs->s_copy) { + GGML_ASSERT(ggml_backend_buffer_is_host(inp_rs->s_copy->buffer)); + int32_t * data = (int32_t *) inp_rs->s_copy->data; + + for (uint32_t i = 0; i < n_rs; ++i) { + data[i] = mctx->get_recr()->s_copy(i); + } + } +} + +bool llm_graph_input_mem_hybrid_idx::can_reuse(const llm_graph_params & params) { + const auto * mctx = static_cast(params.mctx); + + this->mctx = mctx; + + bool res = true; + + res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens; + + res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams); + + res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs(); + + res &= inp_rs->s_copy_main->ne[0] == params.ubatch.n_seqs; + res &= inp_rs->s_copy_extra->ne[0] == mctx->get_recr()->get_n_rs() - params.ubatch.n_seqs; + + res &= inp_rs->head == mctx->get_recr()->get_head(); + res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z(); + + return res; +} + void llm_graph_input_mem_hybrid_iswa::set_input(const llama_ubatch * ubatch) { const auto * attn_ctx = mctx->get_attn(); @@ -1776,7 +1816,8 @@ ggml_tensor * llm_graph_context::build_ffn( const float limit = hparams.swiglu_clamp_shexp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5NEXT || + (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_swiglu_clamp(ctx0, cur, tmp, limit); } else { tmp = ggml_clamp(ctx0, tmp, -limit, limit); @@ -2170,7 +2211,8 @@ ggml_tensor * llm_graph_context::build_moe_ffn( const float limit = hparams.swiglu_clamp_exp[il]; constexpr float eps = 1e-6f; if (limit > eps) { - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5NEXT || + (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_swiglu_clamp(ctx0, cur, up, limit); } else { up = ggml_clamp(ctx0, up, -limit, limit); @@ -2943,6 +2985,80 @@ ggml_tensor * llm_graph_context::build_attn( return cur; } +ggml_tensor * llm_graph_context::build_attn_mask_top_k( + ggml_tensor * kq_mask, + ggml_tensor * top_k) const { + // prepare new kq mask - starts filled with -INFINITY + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); + + // reshape KQ mask into tensor with rows of size 1: + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); + + // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); + + // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] + // this will be our source of zero values for unmasking top k mask elements + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); + + // modify KQ mask by unmasking elements that are in top_k indices + // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) + ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); + + // reshape to restore the original shape of KQ mask: + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); + + // combine with the original kq mask + kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); + + return kq_mask_top_k; +} + +// same as the dense K-only build_attn above, but attends only to the cells picked by top_k +ggml_tensor * llm_graph_context::build_attn( + llm_graph_input_attn_k * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * kq_b, + ggml_tensor * sinks, + ggml_tensor * v_mla, + ggml_tensor * top_k, + float kq_scale, + int il) const { + ggml_build_forward_expand(gf, q_cur); + ggml_build_forward_expand(gf, v_cur); + ggml_build_forward_expand(gf, k_cur); + + const auto * mctx_cur = inp->mctx; + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, inp->get_k_idxs(), il)); + + ggml_tensor * kq_mask = top_k ? build_attn_mask_top_k(inp->get_kq_mask(), top_k) : inp->get_kq_mask(); + + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); + + ggml_tensor * cur = build_attn_mha(q_cur, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + cb(cur, "kqv_out", il); + + if (wo) { + cur = build_lora_mm(wo, cur, wo_s); + } + + if (wo_b) { + cur = ggml_add(ctx0, cur, wo_b); + } + + return cur; +} + ggml_tensor * llm_graph_context::build_attn( llm_graph_input_attn_k_dsa * inp, ggml_tensor * wo, @@ -2973,33 +3089,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); } - const auto & kq_mask = inp->get_kq_mask_mla(); - - // prepare new kq mask - starts filled with -INFINITY - ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); - - // reshape KQ mask into tensor with rows of size 1: - // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] - kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); - - // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] - ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); - - // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] - // this will be our source of zero values for unmasking top k mask elements - ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); - zeros = ggml_fill(ctx0, zeros, 0.0f); - - // modify KQ mask by unmasking elements that are in top_k indices - // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) - ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); - - // reshape to restore the original shape of KQ mask: - // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] - kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); - - // combine with the original kq mask - kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); + ggml_tensor * kq_mask_top_k = build_attn_mask_top_k(inp->get_kq_mask_mla(), top_k); ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); @@ -3550,6 +3640,17 @@ llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k() const return (llm_graph_input_mem_hybrid_k *) res->add_input(std::move(inp)); } +llm_graph_input_mem_hybrid_idx * llm_graph_context::build_inp_mem_hybrid_idx() const { + const auto * mctx_cur = static_cast(mctx); + + auto inp_rs = build_rs_inp_impl (ctx0, ubatch, mctx_cur->get_recr()); + auto inp_attn = build_attn_inp_k_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_attn()); + + auto inp = std::make_unique(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur); + + return (llm_graph_input_mem_hybrid_idx *) res->add_input(std::move(inp)); +} + llm_graph_input_mem_hybrid_iswa * llm_graph_context::build_inp_mem_hybrid_iswa() const { const auto * mctx_cur = static_cast(mctx); diff --git a/src/llama-graph.h b/src/llama-graph.h index dddfdac7b51e..92a6605dd246 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -30,6 +30,7 @@ class llama_kv_cache_dsv4_context; class llama_kv_cache_iswa_context; class llama_memory_recurrent_context; class llama_memory_hybrid_context; +class llama_memory_hybrid_idx_context; class llama_memory_hybrid_iswa_context; // certain models (typically multi-modal) can produce different types of graphs @@ -713,6 +714,36 @@ class llm_graph_input_mem_hybrid_k : public llm_graph_input_i { const llama_memory_hybrid_context * mctx; }; +// same as llm_graph_input_mem_hybrid_k, for the memory container that also carries an +// indexer cache. The indexer's own inputs are arch-specific and live in the model file. +class llm_graph_input_mem_hybrid_idx : public llm_graph_input_i { +public: + llm_graph_input_mem_hybrid_idx( + const llama_cparams & cparams, + std::unique_ptr inp_attn, + std::unique_ptr inp_rs, + const llama_memory_hybrid_idx_context * mctx) : + inp_attn(std::move(inp_attn)), + inp_rs(std::move(inp_rs)), + cparams(cparams), + mctx(mctx) { } + virtual ~llm_graph_input_mem_hybrid_idx() = default; + + void set_input(const llama_ubatch * ubatch) override; + + bool can_reuse(const llm_graph_params & params) override; + + std::unique_ptr inp_attn; + std::unique_ptr inp_rs; + + llm_graph_input_attn_k * get_attn() const { return inp_attn.get(); } + llm_graph_input_rs * get_recr() const { return inp_rs.get(); } + + const llama_cparams cparams; + + const llama_memory_hybrid_idx_context * mctx; +}; + class llm_graph_input_mem_hybrid_iswa : public llm_graph_input_i { public: llm_graph_input_mem_hybrid_iswa( @@ -1223,6 +1254,27 @@ struct llm_graph_context { float kq_scale, int il) const; + // unmask only the selected cells of the KQ mask (sparse attention) + // top_k: I32 [n_top_k, n_batch/n_stream, 1, n_stream], indices into the cache cells + ggml_tensor * build_attn_mask_top_k( + ggml_tensor * kq_mask, + ggml_tensor * top_k) const; + + ggml_tensor * build_attn( + llm_graph_input_attn_k * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens] + ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] + ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] + ggml_tensor * kq_b, + ggml_tensor * sinks, // [n_head_q] + ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] + ggml_tensor * top_k, // [n_top_k, n_batch/n_stream, 1, n_stream], null = dense + float kq_scale, + int il) const; + llm_graph_input_attn_k_dsa * build_attn_inp_k_dsa() const; llm_graph_input_attn_k_dsa_iswa * build_attn_inp_k_dsa_iswa() const; @@ -1342,6 +1394,7 @@ struct llm_graph_context { llm_graph_input_mem_hybrid * build_inp_mem_hybrid() const; llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k() const; + llm_graph_input_mem_hybrid_idx * build_inp_mem_hybrid_idx() const; llm_graph_input_mem_hybrid_iswa * build_inp_mem_hybrid_iswa() const; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 919e90ecccd1..6cad6e251e43 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -303,7 +303,9 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count); add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank); add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, hparams.dsv4_compress_rope_base); - if (model->arch == LLM_ARCH_DEEPSEEK4 || hparams.dsv4_hc_mult > 0) { + // glm5next reuses dsv4_hc_mult for its hyper-connections but has no compress ratios + if (model->arch == LLM_ARCH_DEEPSEEK4 || + (hparams.dsv4_hc_mult > 0 && model->arch != LLM_ARCH_GLM5NEXT)) { // the loader requires one compress ratio per layer, including nextn layers const std::vector compress_ratios( hparams.dsv4_compress_ratios.begin(), hparams.dsv4_compress_ratios.begin() + hparams.n_layer_all); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 6344f2d8aee4..0c8d45029137 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -15,6 +15,7 @@ #include "llama-kv-cache-msa.h" #include "llama-kv-cache-dsv4.h" #include "llama-memory-hybrid.h" +#include "llama-memory-hybrid-idx.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-hybrid-idx.h" #include "llama-memory-recurrent.h" @@ -202,6 +203,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_deepseek4(params); case LLM_ARCH_GLM_DSA: return new llama_model_glm_dsa(params); + case LLM_ARCH_GLM5NEXT: + return new llama_model_glm5next(params); case LLM_ARCH_MISTRAL4: return new llama_model_mistral4(params); case LLM_ARCH_CHATGLM: @@ -960,6 +963,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_288B_A19B: return "288B.A19B"; case LLM_TYPE_300B_A47B: return "300B.A47B"; case LLM_TYPE_310B_A15B: return "310B.A15B"; + case LLM_TYPE_312B_A17B: return "312B.A17B"; case LLM_TYPE_355B_A32B: return "355B.A32B"; case LLM_TYPE_397B_A17B: return "397B.A17B"; case LLM_TYPE_685B_A37B: return "685B.A37B"; @@ -2024,7 +2028,8 @@ void llama_model::print_info() const { arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_NEMOTRON_H || - arch == LLM_ARCH_NEMOTRON_H_MOE) { + arch == LLM_ARCH_NEMOTRON_H_MOE || + arch == LLM_ARCH_GLM5NEXT) { LLAMA_LOG_INFO("%s: ssm_d_conv = %u\n", __func__, hparams.ssm_d_conv); LLAMA_LOG_INFO("%s: ssm_d_inner = %u\n", __func__, hparams.ssm_d_inner); LLAMA_LOG_INFO("%s: ssm_d_state = %u\n", __func__, hparams.ssm_d_state); @@ -2056,7 +2061,8 @@ void llama_model::print_info() const { if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK2OCR || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || - arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_MISTRAL4) { + arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_MISTRAL4 || + arch == LLM_ARCH_GLM5NEXT) { LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead); LLAMA_LOG_INFO("%s: n_lora_q = %d\n", __func__, hparams.n_lora_q); LLAMA_LOG_INFO("%s: n_lora_kv = %d\n", __func__, hparams.n_lora_kv); @@ -2419,6 +2425,47 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr); } } break; + case LLM_ARCH_GLM5NEXT: + { + // KDA layers recur, MLA layers cache, and the DSA indexer shadows the MLA layers + llama_memory_hybrid_idx::layer_filter_cb filter_recr = + [&](int32_t il) { return (uint32_t) il < hparams.n_layer() && hparams.is_recr(il); }; + llama_memory_hybrid_idx::layer_filter_cb filter_attn = + [&](int32_t il) { return (uint32_t) il < hparams.n_layer() && !hparams.is_recr(il); }; + + // no indexer weights -> no indexer cache, and the graph runs dense MLA + llama_memory_hybrid_idx::layer_filter_cb filter_idx = nullptr; + if (hparams.indexer_head_size > 0 && hparams.indexer_block_size > 0) { + filter_idx = [&](int32_t il) { + return (uint32_t) il < hparams.n_layer() && !hparams.is_recr(il) && hparams.is_indexer_full(il); + }; + } + + GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE); + + res = new llama_memory_hybrid_idx( + /* model */ *this, + /* attn_type_k */ params.type_k, + /* attn_type_v */ params.type_v, + /* attn_v_trans */ !cparams.flash_attn, + /* attn_kv_size */ cparams.n_ctx_seq, + /* attn_n_pad */ 1, + /* attn_n_swa */ hparams.n_swa, + /* attn_swa_type */ hparams.swa_type, + /* recurrent_type_r */ GGML_TYPE_F32, + /* recurrent_type_s */ GGML_TYPE_F32, + /* recurrent_rs_size */ std::max((uint32_t) 1, cparams.n_seq_max), + /* idx_row_size */ 2*hparams.indexer_head_size, // key | k-pool gate + /* idx_kpool */ hparams.indexer_block_size, + /* idx_select_tail */ true, // index_kpool_always_select_tail + /* n_seq_max */ cparams.n_seq_max, + /* n_rs_seq */ cparams.n_rs_seq, + /* offload */ cparams.offload_kqv, + /* unified */ cparams.kv_unified, + /* filter_attn */ std::move(filter_attn), + /* filter_recr */ std::move(filter_recr), + /* filter_idx */ std::move(filter_idx)); + } break; case LLM_ARCH_DFLASH: { // DSV4 DSpark stages store a single MLA-style K per position (window = the draft ring) @@ -2837,6 +2884,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_KIMI_K3: + case LLM_ARCH_GLM5NEXT: return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values diff --git a/src/llama-model.h b/src/llama-model.h index 4c4a30e018bc..037341fe4869 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -145,6 +145,7 @@ enum llm_type { LLM_TYPE_288B_A19B, // dots3-note LLM_TYPE_300B_A47B, // Ernie MoE big LLM_TYPE_310B_A15B, // /MiMo-V2-Flash + LLM_TYPE_312B_A17B, // GLM-5.3-Flash LLM_TYPE_355B_A32B, // GLM-4.5 LLM_TYPE_397B_A17B, // Qwen3.5 LLM_TYPE_685B_A37B, // DeepSeek V3.2 diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp new file mode 100644 index 000000000000..ab4c280fde24 --- /dev/null +++ b/src/models/glm5next.cpp @@ -0,0 +1,958 @@ +#include "models.h" + +#include "llama-memory-hybrid-idx.h" +#include "llama-memory-recurrent.h" + +#include +#include + +// +// GLM-5.3-Flash: hybrid trunk, KDA linear attention on 3 of every 4 layers and MLA +// on the rest, with each attention and FFN block wrapped in hyper-connections (mHC). +// - KDA is the kimi-linear tensor layout with the sigmoid decay gate of kimi-k3 +// - MLA is nope-only (qk_rope_head_dim = 0), so this graph has no positions +// - mHC is the deepseek4 formulation, except the final collapse is a plain mean +// +// The full layers pick the cells they attend to with the DSA k-pool indexer, held in +// llama_memory_hybrid_idx next to the recurrent states and the MLA cache. A GGUF without +// indexer weights still builds the dense graph. +// +// Not wired up yet: +// - the NextN/MTP block, loaded as unused +// + +void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { + // read first: everything below uses n_layer() == n_layer_all - n_layer_nextn + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all && "n_layer_nextn must be < block_count"); + + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + // the indexer k_norm is a plain LayerNorm - fall back to the torch default eps + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps, false); + if (hparams.f_norm_eps == 0.0f) { + hparams.f_norm_eps = 1e-5f; + } + + // MLA + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + + if (!hparams.is_mla()) { + throw std::runtime_error("GLM5NEXT requires MLA (key_length_mla / value_length_mla)"); + } + + // mla_use_nope: the graph relies on the model being position-free + GGML_ASSERT(hparams.n_rot() == 0 && "GLM5NEXT is nope-only: rope.dimension_count must be 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); + if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) { + hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID; + } + + // swiglu_limit clamps the gate before the SiLU (deepseek4 semantics, see build_ffn) + // and applies to the routed experts, the shared expert and the dense MLPs alike + if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false)) { + LLAMA_LOG_WARN("%s: glm5next.swiglu_clamp_exp is missing, " + "assuming the GLM-5.3-Flash swiglu_limit of 10.0\n", __func__); + hparams.swiglu_clamp_exp.fill(10.0f); + } + if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false)) { + hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp; + } + + // DSA indexer - absent in a GGUF without indexer weights, which then runs dense MLA + 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); + // index_kpool: tokens per compressed key, always_select_tail is implied (see build_dsa_top_k) + ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size, false); + + if (hparams.indexer_head_size > 0) { + GGML_ASSERT(hparams.indexer_n_head > 0); + GGML_ASSERT(hparams.indexer_block_size > 0 && "GLM5NEXT requires index_kpool"); + GGML_ASSERT(hparams.indexer_top_k % hparams.indexer_block_size == 0); + } + + // GLM-5.3-Flash has indexer_types = "full" everywhere + std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1); + ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false); + + // hyper-connections (mHC) + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); + ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); + GGML_ASSERT(hparams.dsv4_hc_mult == 4 && "GLM5NEXT expects hyper_connection.count == 4"); + GGML_ASSERT(hparams.dsv4_hc_sinkhorn_iters > 0 && "GLM5NEXT expects hyper_connection.sinkhorn_iterations > 0"); + + // the hc streams collapse to a mean, so the output is plain n_embd. deepseek4 sets + // this to hc_mult*n_embd only to size its MTP buffer; keeping such a value here + // would make llama_context read past the end of t_embd. + hparams.n_embd_out_impl = 0; + + // KDA + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + if (!ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda, false)) { + // older GGUFs store the KDA head dim as ssm.state_size + ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.n_embd_head_kda); + } + if (!ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false)) { + // linear_attn_config.gate_lower_bound + hparams.kda_gate_lower_bound = -5.0f; + } + // only the bounded sigmoid gate is implemented, the softplus branch is dead here + GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f); + + // note: n_embd_r()/n_embd_s() size the recurrent state with n_head()*n_embd_head_kda, + // which works only because linear_attn_config.num_heads == num_attention_heads + + // MLA forces num_key_value_heads = 1 on every layer at conversion time, so the + // kimi-linear "n_head_kv == 0" recurrent marker is not available here + if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) { + uint32_t full_attn_interval = 4; // layer_types: full attention on 3, 7, 11, ... + ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); + GGML_ASSERT(full_attn_interval > 0); + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + hparams.is_recr_impl[il] = (il < hparams.n_layer()) && ((il + 1) % full_attn_interval != 0); + } + } + + switch (hparams.n_layer()) { + case 45: type = LLM_TYPE_312B_A17B; break; // GLM-5.3-Flash + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_glm5next::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_expert_shared = std::max(1, hparams.n_expert_shared); + + 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 n_embd_head_qk_rope = hparams.n_rot(); // 0, nope-only + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_embd_head_qk_rope; + + const int64_t n_idx_head = hparams.indexer_n_head; + const int64_t n_idx_dim = hparams.indexer_head_size; + const int64_t kpool = hparams.indexer_block_size; + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + const int64_t hc_mix_dim = (2 + hc) * hc; + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = head_dim * n_head; + + // the NextN/MTP block is loaded but ignored - there is no MTP graph yet + const int mtp_flags = TENSOR_NOT_REQUIRED | TENSOR_SKIP; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + if (!output) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + // note: no hc_head_*, the hyper-connection head is a plain mean + + for (int i = 0; i < n_layer_all; ++i) { + auto & layer = layers[i]; + + const bool is_mtp = i >= n_layer; + const bool is_recr = !is_mtp && hparams.is_recr(i); + const int flags = is_mtp ? mtp_flags : 0; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); + + // hyper-connections wrap the trunk blocks only, the MTP block uses plain residuals + if (!is_mtp) { + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc_dim, hc_mix_dim}, 0); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0); + } + + if (is_recr) { + // conv1d may be stored 4D [d_conv, 1, d_inner, 1] or 3D (quantization drops the trailing 1) + auto conv = [&](llm_tensor tid) { + ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED); + return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0); + }; + layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q); + layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K); + layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V); + + create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0); + + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0); + layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0); + layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", i), {n_embd, head_dim}, 0); + layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", i), {head_dim, d_inner}, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); + + // ssm_a holds -exp(A_log), folded at conversion time (kimi-linear/kimi-k3 convention) + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, 0); + + // some converters emit dt_bias under the default ".weight" suffix + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, TENSOR_NOT_REQUIRED); + if (!layer.ssm_dt_b) { + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "weight", i), {d_inner}, 0); + } + + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0); + } else { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k_mla}, flags); + // nope-only: kv_lora_rank + 0 + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + n_embd_head_qk_rope}, flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags); + + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, flags); + + // DSA indexer + k-pool compressor + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_idx_head * n_idx_dim}, flags | TENSOR_NOT_REQUIRED); + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, n_idx_dim}, flags | TENSOR_NOT_REQUIRED); + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_idx_dim}, flags | TENSOR_NOT_REQUIRED); + layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {n_idx_dim}, flags | TENSOR_NOT_REQUIRED); + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, n_idx_head}, flags | TENSOR_NOT_REQUIRED); + + layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", i), {n_idx_dim, kpool}, flags | TENSOR_NOT_REQUIRED); + layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", i), {n_embd, n_idx_dim}, flags | 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}, flags); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, flags); + } else { + if (n_expert == 0 || n_expert_used == 0) { + throw std::runtime_error("GLM5NEXT requires n_expert > 0 and n_expert_used > 0"); + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags | TENSOR_NOT_REQUIRED); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), {n_embd, n_ff_exp, n_expert}, flags); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {n_ff_exp * n_expert_shared, n_embd}, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + } + + if (is_mtp) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, mtp_flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags); + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, mtp_flags); + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, mtp_flags); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags); + } + } +} + +std::unique_ptr llama_model_glm5next::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// +// hyper-connections (mHC), per token, with hc_dim = hc*n_embd: +// +// mixes = hc_fn @ rms_norm(streams.flatten()) [(2 + hc)*hc] +// pre, post, comb_w = mixes.split([hc, hc, hc*hc]) +// +// pre = sigmoid(pre *scale[0] + base[ 0: hc]) + eps [hc] +// post = sigmoid(post*scale[1] + base[ hc:2*hc])*2 [hc] +// comb = sinkhorn(softmax(comb_w*scale[2] + base[2*hc:]) + eps) [hc, hc] +// +// sublayer_in = sum_h pre[h]*streams[h] +// out[dst] = post[dst]*sublayer_out + sum_src comb[dst, src]*streams[src] +// +// comb is stored as [dst_hc, src_hc, n_tokens]: the reference matrix is row-major +// with the softmax over its last axis, so ne0 is that axis after the reshape, which +// is the transpose the reference matmul asks for. +// + +static ggml_tensor * glm5next_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) { + return ggml_view_1d(ctx, t, ne0, ggml_row_size(t->type, i0)); +} + +static ggml_tensor * glm5next_view_2d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t ne1, int64_t i0) { + return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], ggml_row_size(t->type, i0)); +} + +static ggml_tensor * glm5next_hc_affine(ggml_context * ctx, ggml_tensor * x, ggml_tensor * scale, ggml_tensor * base) { + return ggml_add(ctx, ggml_mul(ctx, x, scale), base); +} + +// Glm5NextTextHyperHead: unweighted mean over the streams +// [n_embd, hc, n_tokens] -> [n_embd, n_tokens] +static ggml_tensor * glm5next_hc_mean(ggml_context * ctx, ggml_tensor * x) { + const int64_t hc = x->ne[1]; + + ggml_tensor * acc = ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0); + for (int64_t s = 1; s < hc; ++s) { + acc = ggml_add(ctx, acc, ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1])); + } + + return ggml_scale(ctx, acc, 1.0f/hc); +} + +ggml_tensor * llama_model_glm5next::graph::build_hc_collapse(ggml_tensor * x, ggml_tensor * weights, int il) { + GGML_ASSERT(x->ne[0] == n_embd); + GGML_ASSERT(x->ne[1] == (int64_t) hparams.dsv4_hc_mult); + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = x->ne[2]; + + if (cparams.fused_dsv4_hc_pre && il >= 0) { + ggml_tensor * result = ggml_dsv4_hc_pre(ctx0, x, weights); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, result, il}); + return result; + } + + ggml_tensor * result = nullptr; + for (int64_t ih = 0; ih < hc; ++ih) { + ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]); + ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]); + ggml_tensor * cur = ggml_mul(ctx0, xh, wh); + result = result ? ggml_add(ctx0, result, cur) : cur; + } + + return result; +} + +ggml_tensor * llama_model_glm5next::graph::build_hc_sinkhorn(ggml_tensor * comb, int il) { + GGML_UNUSED(il); + + const float eps = hparams.dsv4_hc_eps; + + // comb is [dst_hc, src_hc, n_tokens]: ne0 is the reference's softmax axis + comb = ggml_soft_max(ctx0, comb); + comb = ggml_scale_bias(ctx0, comb, 1.0f, eps); + + // normalize over the reference's dim=-2, which is our ne1 + auto norm_cols = [&]() { + ggml_tensor * t = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3)); + ggml_tensor * sum = ggml_sum_rows(ctx0, t); + sum = ggml_scale_bias(ctx0, sum, 1.0f, eps); + sum = ggml_permute(ctx0, sum, 1, 0, 2, 3); + comb = ggml_div(ctx0, comb, sum); + }; + + // normalize over the reference's dim=-1, which is our ne0 + auto norm_rows = [&]() { + ggml_tensor * sum = ggml_sum_rows(ctx0, comb); + sum = ggml_scale_bias(ctx0, sum, 1.0f, eps); + comb = ggml_div(ctx0, comb, sum); + }; + + norm_cols(); + for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) { + norm_rows(); + norm_cols(); + } + + return comb; +} + +ggml_tensor * llama_model_glm5next::graph::build_hc_pre( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc*n_embd; + const int64_t hc_mix_dim = (2 + hc)*hc; + const int64_t nt = x->ne[2]; + + GGML_ASSERT(x->ne[0] == n_embd); + GGML_ASSERT(x->ne[1] == hc); + GGML_ASSERT(hc_fn->ne[0] == hc_dim); + GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim); + + // DeepseekV4UnweightedRMSNorm: no learned gain + ggml_tensor * flat = ggml_rms_norm(ctx0, ggml_reshape_2d(ctx0, x, hc_dim, nt), norm_rms_eps); + ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat); + cb(mixes, "hc_mixes", il); + + ggml_tensor * scale_pre = glm5next_view_1d(ctx0, hc_scale, 1, 0); + ggml_tensor * scale_post = glm5next_view_1d(ctx0, hc_scale, 1, 1); + + ggml_tensor * base_pre = glm5next_view_1d(ctx0, hc_base, hc, 0); + ggml_tensor * base_post = glm5next_view_1d(ctx0, hc_base, hc, hc); + + ggml_tensor * pre = glm5next_view_2d(ctx0, mixes, hc, nt, 0); + pre = glm5next_hc_affine(ctx0, pre, scale_pre, base_pre); + pre = ggml_sigmoid(ctx0, pre); + pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps); + cb(pre, "hc_pre", il); + + *post = glm5next_view_2d(ctx0, mixes, hc, nt, hc); + *post = glm5next_hc_affine(ctx0, *post, scale_post, base_post); + *post = ggml_sigmoid(ctx0, *post); + *post = ggml_scale(ctx0, *post, 2.0f); + cb(*post, "hc_post", il); + + if (cparams.fused_dsv4_hc_comb && + hc_scale->type == GGML_TYPE_F32 && hc_base->type == GGML_TYPE_F32) { + *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps, + (int32_t) hparams.dsv4_hc_sinkhorn_iters); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il}); + } else { + ggml_tensor * scale_comb = glm5next_view_1d(ctx0, hc_scale, 1, 2); + ggml_tensor * base_comb = glm5next_view_1d(ctx0, hc_base, hc*hc, 2*hc); + + *comb = glm5next_view_2d(ctx0, mixes, hc*hc, nt, 2*hc); + *comb = glm5next_hc_affine(ctx0, *comb, scale_comb, base_comb); + *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt); + *comb = build_hc_sinkhorn(*comb, il); + } + cb(*comb, "hc_comb", il); + + return build_hc_collapse(x, pre, il); +} + +ggml_tensor * llama_model_glm5next::graph::build_hc_post( + ggml_tensor * x, + ggml_tensor * residual, + ggml_tensor * post, + ggml_tensor * comb, + int il) { + GGML_ASSERT(x->ne[0] == n_embd); + GGML_ASSERT(residual->ne[1] == (int64_t) hparams.dsv4_hc_mult); + + if (cparams.fused_dsv4_hc_post) { + ggml_tensor * result = ggml_dsv4_hc_post(ctx0, x, residual, post, comb); + res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, result, il}); + return result; + } + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = x->ne[1]; + + ggml_tensor * out = nullptr; + for (int64_t dst = 0; dst < hc; ++dst) { + ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst*post->nb[0]); + ggml_tensor * cur = ggml_mul(ctx0, x, post_dst); + + for (int64_t src = 0; src < hc; ++src) { + ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]); + ggml_tensor * comb_sd = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2], dst*comb->nb[0] + src*comb->nb[1]); + cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_sd)); + } + + cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt); + out = out ? ggml_concat(ctx0, out, cur, 1) : cur; + } + + return out; +} + +// +// KDA layer +// + +// causal conv1d over one of Q/K/V. `qkv` selects which third of the conv state to use +static ggml_tensor * glm5next_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 cache_head, uint32_t mem_size, uint32_t n_rs_seq) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t total_state_size = 3 * conv_state_size; + + ggml_tensor * conv_state = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + total_state_size * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + x_proj = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state, ggml_transpose(ctx0, x_proj), 0); + + // one snapshot per rollback slot, newest first + const int64_t n_written = std::min(n_seq_tokens, (int64_t) n_rs_seq + 1); + + for (int64_t slot = 0; slot < n_written; ++slot) { + ggml_tensor * snap = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], (conv_x->ne[0] - (d_conv - 1) - slot) * conv_x->nb[0]); + ggml_build_forward_expand(gf, ggml_cpy(ctx0, snap, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + total_state_size * ggml_element_size(conv_states_all), + ((slot * mem_size + cache_head) * total_state_size + qkv * conv_state_size) + * ggml_element_size(conv_states_all)))); + } + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + ggml_tensor * out = ggml_ssm_conv(ctx0, conv_x, conv_weight); + out = ggml_silu(ctx0, ggml_reshape_2d(ctx0, out, d_inner, n_tokens)); + + return ggml_reshape_4d(ctx0, out, head_dim, n_head, n_seq_tokens, n_seqs); +} + +ggml_tensor * llama_model_glm5next::graph::build_kda_layer( + ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs, + int64_t n_seq_tokens, int64_t n_seqs, int il) { + const int64_t n_head_kda = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = n_head_kda * head_dim; + const int64_t d_conv = hparams.ssm_d_conv; + + const auto * mctx_cur = inp_rs->mctx; + const auto cache_head = mctx_cur->get_head(); + const auto mem_size = mctx_cur->get_size(); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * q = glm5next_causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, + cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, + n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + ggml_tensor * k = glm5next_causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, + cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, + n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + ggml_tensor * v = glm5next_causal_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, + cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, + n_seq_tokens, n_seqs, n_tokens, cache_head, mem_size, cparams.n_rs_seq); + cb(q, "kda_q_conv", il); + cb(k, "kda_k_conv", il); + cb(v, "kda_v_conv", il); + + // forget gate: g = gate_lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias)). + // ssm_a holds -exp(A_log), so exp(A_log)*(...) == -(ssm_a*(...)). + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + g = ggml_mul_mat(ctx0, layer.ssm_f_b, g); + g = ggml_add(ctx0, g, layer.ssm_dt_b); + g = ggml_reshape_3d(ctx0, g, head_dim, n_head_kda, n_tokens); + g = ggml_mul(ctx0, g, ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1)); + g = ggml_sigmoid(ctx0, ggml_scale(ctx0, g, -1.0f)); + g = ggml_scale(ctx0, g, hparams.kda_gate_lower_bound); + g = ggml_reshape_4d(ctx0, g, head_dim, n_head_kda, n_seq_tokens, n_seqs); + cb(g, "kda_gate", il); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs)); + cb(beta, "kda_beta", il); + + // the reference uses a hard-coded 1e-6 here; build_delta_net applies the + // 1/sqrt(head_dim) scaling of q after this + q = ggml_l2_norm(ctx0, q, 1e-6f); + k = ggml_l2_norm(ctx0, k, 1e-6f); + + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs); + + ggml_tensor * out = ggml_cont(ctx0, build_recurrent_attn( + inp_rs, ssm_states_all, q, k, v, g, beta, state, il)); + + // Glm5NextTextRMSNormGated: RMSNorm first, then a SIGMOID gate (not SiLU) + ggml_tensor * o_gate = ggml_mul_mat(ctx0, layer.ssm_g_a, cur); + o_gate = ggml_mul_mat(ctx0, layer.ssm_g_b, o_gate); + o_gate = ggml_reshape_3d(ctx0, o_gate, head_dim, n_head_kda, n_tokens); + + out = ggml_reshape_3d(ctx0, out, head_dim, n_head_kda, n_tokens); + out = build_norm(out, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + out = ggml_mul(ctx0, out, ggml_sigmoid(ctx0, o_gate)); + cb(out, "kda_normed_gated", il); + + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, out, d_inner, n_tokens)); + cb(cur, "kda_out", il); + + return cur; +} + +// +// DSA k-pool indexer +// +// The full layers attend to index_topk cells picked by a lightning indexer that scores +// k-pools of index_kpool tokens instead of single tokens. A pool key is a learned per-channel +// convex mix softmax(gate + ape) . keys of its members, so the cache must hold both the +// indexer key and the gate logits of every token. +// +// Everything that depends on the cache layout is computed host-side in set_input; the graph +// only gathers, pools and scores. One input serves every layer: the pool metadata is the same +// for all of them. + +class llm_graph_input_kpool : public llm_graph_input_i { +public: + llm_graph_input_kpool(const llama_memory_hybrid_idx_context * mctx) : mctx(mctx) {} + virtual ~llm_graph_input_kpool() = default; + + void set_input(const llama_ubatch * ubatch) override { + mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); + mctx->set_input_kpool(cell_pool, pool_cells, bias, ubatch); + + GGML_ASSERT(ggml_backend_buffer_is_host(ape_slots->buffer)); + int32_t * data = (int32_t *) ape_slots->data; + for (int64_t i = 0; i < ape_slots->ne[0]; ++i) { + data[i] = (int32_t) i; + } + } + + bool can_reuse(const llm_graph_params & params) override { + const auto * mctx_cur = static_cast(params.mctx); + + mctx = mctx_cur; + + bool res = true; + + res &= k_idxs->ne[0] == params.ubatch.n_tokens; + res &= cell_pool->ne[0] == (int64_t) mctx_cur->get_idx()->get_n_kv(); + res &= cell_pool->ne[1]*bias->ne[1] == params.ubatch.n_tokens; + + return res; + } + + ggml_tensor * k_idxs = nullptr; // I64 [n_tokens] + ggml_tensor * ape_slots = nullptr; // I32 [kpool], the identity - reads the ape rows in order + ggml_tensor * cell_pool = nullptr; // I32 [n_kv, n_stream] + ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] + ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens/n_stream, n_stream] + + const llama_memory_hybrid_idx_context * mctx; +}; + +llm_graph_input_kpool * llama_model_glm5next::graph::build_inp_kpool(llm_graph_input_mem_hybrid_idx * inp_hyb) { + const auto * mctx_idx = inp_hyb->mctx->get_idx(); + + if (!mctx_idx) { + return nullptr; + } + + const int64_t r = hparams.indexer_block_size; + const int64_t n_kv = mctx_idx->get_n_kv(); + const int64_t n_pool = n_kv/r; + // the KQ mask carries the stream count that build_attn and the cache views agree on + const int64_t ns = inp_hyb->get_attn()->get_kq_mask()->ne[3]; + + // the top-k indices address the attention cache, so the two must agree cell for cell + GGML_ASSERT(n_kv == (int64_t) inp_hyb->mctx->get_attn()->get_n_kv()); + GGML_ASSERT(n_tokens % ns == 0); + + if (n_pool == 0) { + return nullptr; + } + + auto kp = std::make_unique(inp_hyb->mctx); + + kp->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + kp->ape_slots = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, r); + kp->cell_pool = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, ns); + kp->pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_pool, ns); + kp->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tokens/ns, ns); + + ggml_set_input(kp->ape_slots); + ggml_set_input(kp->cell_pool); + ggml_set_input(kp->pool_cells); + ggml_set_input(kp->bias); + + return (llm_graph_input_kpool *) res->add_input(std::move(kp)); +} + +ggml_tensor * llama_model_glm5next::graph::build_dsa_top_k( + llm_graph_input_kpool * inp, ggml_tensor * cur, + ggml_tensor * qr, const llama_layer & layer, int il) { + const auto * mctx_idx = inp->mctx->get_idx(); + + const int64_t d = hparams.indexer_head_size; + const int64_t nh = hparams.indexer_n_head; + const int64_t r = hparams.indexer_block_size; + const int64_t n_kv = inp->cell_pool->ne[0]; + const int64_t ns = inp->cell_pool->ne[1]; + const int64_t n_pool = inp->pool_cells->ne[0]/r; + const int64_t n_tps = n_tokens/ns; + + // key and gate logits packed into one cache row: the gate depends on the token hidden + // state, so it cannot be recomputed from the cache later + ggml_tensor * k = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur); + k = build_norm(k, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il); + + ggml_tensor * g = ggml_mul_mat(ctx0, layer.indexer_comp_wgate, cur); + + ggml_tensor * packed = ggml_reshape_3d(ctx0, ggml_concat(ctx0, k, g, 0), 2*d, 1, n_tokens); + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, packed, inp->k_idxs, il)); + + // one key head, so cache rows are contiguous + ggml_tensor * all = mctx_idx->get_k(ctx0, il); + all = ggml_view_3d(ctx0, all, 2*d, n_kv, ns, all->nb[2], all->nb[3], 0); + + // gathers per stream: pool_cells row s indexes stream s's own cells + ggml_tensor * members = ggml_get_rows(ctx0, all, inp->pool_cells); + members = ggml_reshape_4d(ctx0, members, 2*d, r, n_pool, ns); + + ggml_tensor * m_k = ggml_cont(ctx0, ggml_view_4d(ctx0, members, d, r, n_pool, ns, + members->nb[1], members->nb[2], members->nb[3], 0)); + ggml_tensor * m_g = ggml_cont(ctx0, ggml_view_4d(ctx0, members, d, r, n_pool, ns, + members->nb[1], members->nb[2], members->nb[3], ggml_row_size(members->type, d))); + + // pool key = softmax(gate + ape) . keys over the r members, channel by channel. The ape is + // an intra-pool position bias and, with no rope anywhere, the only ordering signal here. + m_g = ggml_add(ctx0, m_g, ggml_get_rows(ctx0, layer.indexer_comp_ape, inp->ape_slots)); + + // softmax normalizes ne[0], so bring the member axis there + ggml_tensor * w = ggml_soft_max(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, m_g, 1, 0, 2, 3))); + ggml_tensor * v = ggml_cont(ctx0, ggml_permute(ctx0, m_k, 1, 0, 2, 3)); + + ggml_tensor * pooled = ggml_sum_rows(ctx0, ggml_mul(ctx0, v, w)); + pooled = ggml_cont(ctx0, ggml_permute(ctx0, pooled, 1, 0, 2, 3)); + pooled = ggml_reshape_3d(ctx0, pooled, d, n_pool, ns); + cb(pooled, "indexer_k_pooled", il); + + // nope-only: no rope on the indexer query either. mul_mat matches ne[2], so stream s's + // queries only meet stream s's pools. + ggml_tensor * q = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr); + q = ggml_reshape_3d(ctx0, q, d, nh*n_tps, ns); + cb(q, "indexer_q", il); + + ggml_tensor * score = ggml_mul_mat(ctx0, pooled, q); + score = ggml_relu(ctx0, ggml_reshape_4d(ctx0, score, n_pool, nh, n_tps, ns)); + + // relu(x*s) == s*relu(x) for s > 0, so both positive scalars (the softmax scale and + // n_heads^-1/2) fold into the small weights tensor instead of the big score one + ggml_tensor * wts = ggml_mul_mat(ctx0, layer.indexer_proj, cur); + wts = ggml_scale(ctx0, wts, 1.0f/sqrtf(float(d*nh))); + wts = ggml_reshape_4d(ctx0, wts, nh, 1, n_tps, ns); + + score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); + score = ggml_sum_rows(ctx0, ggml_mul(ctx0, score, wts)); + score = ggml_reshape_3d(ctx0, score, n_pool, n_tps, ns); + cb(score, "indexer_score", il); + + // give every cell of a pool its pool's score rather than expanding pool ids into cell ids, + // which would need an integer multiply-add ggml has no op for. The r members tie, so the + // cut still lands on a pool boundary. bias carries causality and the always-visible tail. + ggml_tensor * expanded = ggml_get_rows(ctx0, + ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_pool); + expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); + expanded = ggml_add(ctx0, expanded, inp->bias); + cb(expanded, "indexer_score_cells", il); + + // index_topk cells plus the incomplete tail, as in the reference + const int64_t width = std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); + + ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); + + // build_attn_mask_top_k reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask + top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, ns); + cb(top_k, "indexer_top_k", il); + + return top_k; +} + +// +// MLA layer (nope-only, absorbed - i.e. MQA with a single group) +// + +ggml_tensor * llama_model_glm5next::graph::build_mla_layer( + ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_attn_k * inp_attn, llm_graph_input_kpool * inp_kpool, + float kq_scale, int il) { + const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + + ggml_tensor * qr = ggml_mul_mat(ctx0, layer.wq_a, cur); + qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(qr, "qr", il); + + // nope-only: the whole of q is the "nope" part, no split and no rope + ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr); + q = ggml_reshape_3d(ctx0, q, n_embd_head_k, n_head, n_tokens); + cb(q, "q", il); + + // {n_embd_head_k, n_tokens, n_head} x wk_b -> {kv_lora_rank, n_tokens, n_head} + q = ggml_permute(ctx0, q, 0, 2, 1, 3); + ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.wk_b, q); + Qcur = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + cb(Qcur, "Qcur", il); + + // nope-only: wkv_a_mqa outputs exactly kv_lora_rank, no k_pe to split off + ggml_tensor * kv_cmpr = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + cb(kv_cmpr, "kv_cmpr", il); + + // null top_k = attend to the whole cache, for a GGUF without indexer weights + ggml_tensor * top_k = inp_kpool && layer.indexer_attn_q_b + ? build_dsa_top_k(inp_kpool, cur, qr, layer, il) : nullptr; + + cur = build_attn(inp_attn, layer.wo, nullptr, layer.wo_s, + Qcur, kv_cmpr, kv_cmpr, nullptr, nullptr, layer.wv_b, top_k, kq_scale, il); + cb(cur, "mla_out", il); + + return cur; +} + +// +// FFN: leading dense layers, then MoE with a shared expert +// + +ggml_tensor * llama_model_glm5next::graph::build_ffn_layer(ggml_tensor * cur, const llama_layer & layer, int il) { + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = build_ffn(cur, + layer.ffn_up, nullptr, layer.ffn_up_s, + layer.ffn_gate, nullptr, layer.ffn_gate_s, + layer.ffn_down, nullptr, layer.ffn_down_s, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + return cur; + } + + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + nullptr, + layer.ffn_gate_up_exps, + layer.ffn_up_exps_s, + layer.ffn_gate_exps_s, + layer.ffn_down_exps_s); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * shexp = build_ffn(cur, + layer.ffn_up_shexp, nullptr, layer.ffn_up_shexp_s, + layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s, + layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(shexp, "ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, shexp); + cb(cur, "ffn_out", il); + + return cur; +} + +// +// trunk graph +// + +llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + + const int64_t hc = hparams.dsv4_hc_mult; + 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); + + // nope-only, so no YaRN mscale correction on kq_scale + const float kq_scale = 1.0f / sqrtf(float(hparams.n_embd_head_k_mla())); + + ggml_tensor * inp = build_inp_embd(model.tok_embd); + cb(inp, "inp_embd", -1); + + // MLA with absorption uses a K-only cache (V is a view of K) + auto * inp_hyb = build_inp_mem_hybrid_idx(); + auto * inp_rs = inp_hyb->get_recr(); + auto * inp_attn = inp_hyb->get_attn(); + + // the k-pool metadata is the same for every full layer, so build it once + auto * inp_kpool = build_inp_kpool(inp_hyb); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // inputs_embeds.unsqueeze(2).expand(-1, -1, hc, -1) + ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens); + inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1); + cb(inpL, "hc_init", -1); + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * residual = inpL; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + + // attention site + ggml_tensor * cur = build_hc_pre(inpL, + layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, il); + cb(cur, "hc_attn_pre", il); + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + cur = hparams.is_recr(il) + ? build_kda_layer(cur, layer, inp_rs, n_seq_tokens, n_seqs, il) + : build_mla_layer(cur, layer, inp_attn, inp_kpool, kq_scale, il); + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "hc_attn_post", il); + + // FFN site + residual = inpL; + cur = build_hc_pre(inpL, + layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, il); + cb(cur, "hc_ffn_pre", il); + + ggml_build_forward_expand(gf, residual); + ggml_build_forward_expand(gf, post); + ggml_build_forward_expand(gf, comb); + + cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn_layer(cur, layer, il); + + inpL = build_hc_post(cur, residual, post, comb, il); + inpL = build_cvec(inpL, il); + cb(inpL, "l_out", il); + } + + // narrow to the output rows before collapsing the streams + if (inp_out_ids) { + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + flat = ggml_get_rows(ctx0, flat, inp_out_ids); + inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs); + } + + ggml_tensor * cur = glm5next_hc_mean(ctx0, inpL); + cb(cur, "hc_head", -1); + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index 9b87a40d5af9..81ffd6977baa 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1334,6 +1334,57 @@ struct llama_model_glm_dsa : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +// defined in models/glm5next.cpp - the DSA k-pool inputs are arch-specific +class llm_graph_input_kpool; + +struct llama_model_glm5next : public llama_model_base { + llama_model_glm5next(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + + const llama_model & model; + + // manifold-constrained hyper-connections (mHC), same formulation as deepseek4 + // except for the final collapse, which is an unweighted mean here + ggml_tensor * build_hc_collapse(ggml_tensor * x, ggml_tensor * weights, int il); + ggml_tensor * build_hc_sinkhorn(ggml_tensor * comb, int il); + + ggml_tensor * build_hc_pre( + ggml_tensor * x, + ggml_tensor * hc_fn, + ggml_tensor * hc_scale, + ggml_tensor * hc_base, + ggml_tensor ** post, + ggml_tensor ** comb, + int il); + + ggml_tensor * build_hc_post( + ggml_tensor * x, + ggml_tensor * residual, + ggml_tensor * post, + ggml_tensor * comb, + int il); + + ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_rs * inp_rs, int64_t n_seq_tokens, int64_t n_seqs, int il); + llm_graph_input_kpool * build_inp_kpool(llm_graph_input_mem_hybrid_idx * inp_hyb); + + // DSA k-pool indexer: the cells to attend to, in the attention cache's index space + ggml_tensor * build_dsa_top_k(llm_graph_input_kpool * inp, ggml_tensor * cur, + ggml_tensor * qr, const llama_layer & layer, int il); + + ggml_tensor * build_mla_layer(ggml_tensor * cur, const llama_layer & layer, + llm_graph_input_attn_k * inp_attn, llm_graph_input_kpool * inp_kpool, + float kq_scale, int il); + ggml_tensor * build_ffn_layer(ggml_tensor * cur, const llama_layer & layer, int il); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_eagle3 : public llama_model_base { llama_model_eagle3(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b2ea245ab846..5f68312544ac 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -111,6 +111,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { n_embd = 160; // exercise per-head tensor split granularity with head size 80 } else if (arch == LLM_ARCH_QWEN3 || arch == LLM_ARCH_MUSE_GLIMMER || arch == LLM_ARCH_AFMOE) { n_head = 4; + } else if (arch == LLM_ARCH_GLM5NEXT) { + n_embd = 128; + n_head = 1; + n_ff = 192; + // 4 layers gives 2 full-attention layers, so the DSA indexer cache is reused across layers + n_layer = 4; } else if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA @@ -213,6 +219,13 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types); } + } else if (arch == LLM_ARCH_GLM5NEXT) { + // mla_use_nope: qk_rope_head_dim == 0, no RoPE anywhere + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(512)); // kv_lora_rank + 0 + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512)); + ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(0)); + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(128)); // qk_nope_head_dim + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128)); // v_head_dim } else if (arch == LLM_ARCH_MINIMAX_M3) { // partial rotary: n_rot must not exceed the indexer key length (64) ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); @@ -291,15 +304,17 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4})); - if (arch == LLM_ARCH_DEEPSEEK4) { - ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(8)); - ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(32)); - ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector({0, 0, 4, 128})); - ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 160000.0f); + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5NEXT) { + if (arch == LLM_ARCH_DEEPSEEK4) { + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(8)); + ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(32)); + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector({0, 0, 4, 128})); + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, 160000.0f); + ms.add_kv(LLM_KV_HASH_LAYER_COUNT, uint32_t(0)); + } ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); - ms.add_kv(LLM_KV_HASH_LAYER_COUNT, uint32_t(0)); ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 10.0f); ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); @@ -458,6 +473,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_EXAONE_MOE: case LLM_ARCH_BAILINGMOE: case LLM_ARCH_BAILINGMOE2: From e972e2ca0b4810487a826be154a86b033db043e1 Mon Sep 17 00:00:00 2001 From: eauchs Date: Thu, 27 Aug 2026 12:19:50 +0200 Subject: [PATCH 02/11] convert : keep text_config merge from clobbering architectures --- conversion/glm5next.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/conversion/glm5next.py b/conversion/glm5next.py index f54d8786f6fb..2b2eeb454860 100644 --- a/conversion/glm5next.py +++ b/conversion/glm5next.py @@ -44,8 +44,14 @@ class Glm5NextModel(GlmMoeDsaModel): def index_tensors(self, remote_hf_model_id: str | None = None): # TextModel lifts text_config to the root, but only after this runs - # and the parent already needs num_hidden_layers from it here. + # Skip None values: AutoConfig.to_dict() materialises keys that the JSON + # omits, so text_config carries architectures=None and would clobber the + # valid top-level value. if "text_config" in self.hparams: - self.hparams = {**self.hparams, **self.hparams["text_config"]} + self.hparams = { + **self.hparams, + **{k: v for k, v in self.hparams["text_config"].items() if v is not None}, + } return super().index_tensors(remote_hf_model_id=remote_hf_model_id) @classmethod From dca65ecd5dd7ec126daa674115bc22d4f54ea58e Mon Sep 17 00:00:00 2001 From: eauchs Date: Thu, 27 Aug 2026 15:48:44 +0200 Subject: [PATCH 03/11] llama : accept indexer.kpool as a fallback for indexer.block_size --- src/llama-arch.cpp | 1 + src/llama-arch.h | 1 + src/models/glm5next.cpp | 5 +++++ 3 files changed, 7 insertions(+) diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5a916c08440e..af08544fa185 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -283,6 +283,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" }, { LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" }, { LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" }, + { LLM_KV_ATTENTION_INDEXER_KPOOL, "%s.attention.indexer.kpool" }, { LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" }, { LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" }, { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 4cdf357080e3..cdda15c338c6 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -288,6 +288,7 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, LLM_KV_ATTENTION_INDEXER_TOP_K, LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, + LLM_KV_ATTENTION_INDEXER_KPOOL, LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, LLM_KV_ATTENTION_INDEXER_TYPES, LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index ab4c280fde24..6b0cc5eb902f 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -75,6 +75,11 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k, false); // index_kpool: tokens per compressed key, always_select_tail is implied (see build_dsa_top_k) ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size, false); + if (hparams.indexer_block_size == 0) { + // some converters write the k-pool size under `indexer.kpool` rather than + // `indexer.block_size`; accept it so those GGUFs load without an override + ml.get_key(LLM_KV_ATTENTION_INDEXER_KPOOL, hparams.indexer_block_size, false); + } if (hparams.indexer_head_size > 0) { GGML_ASSERT(hparams.indexer_n_head > 0); From 8ee6fcf1aba0c697da7a24c04629a9e70f0e7578 Mon Sep 17 00:00:00 2001 From: eauchs Date: Thu, 27 Aug 2026 16:42:10 +0200 Subject: [PATCH 04/11] llama : select whole k-pools in the glm5next DSA indexer --- gguf-py/gguf/constants.py | 6 +- src/models/glm5next.cpp | 68 +++++++---- tests/test-llama-archs.cpp | 226 ++++++++++++++++++++++++++++++++++++- 3 files changed, 270 insertions(+), 30 deletions(-) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index fe74ec54ec59..2b2de3a8a0e3 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -4043,7 +4043,7 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_HNORM, MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, - # --- KDA linear attention (34 couches sur 45), cf. kimi-linear --- + # --- KDA linear attention (34 of 45 layers), see kimi-linear --- MODEL_TENSOR.ATTN_K, MODEL_TENSOR.ATTN_V, MODEL_TENSOR.SSM_CONV1D_Q, @@ -4057,14 +4057,14 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.SSM_A, MODEL_TENSOR.SSM_DT, MODEL_TENSOR.SSM_NORM, - # --- hyper-connections (mHC, Sinkhorn), cf. deepseek4 --- + # --- hyper-connections (mHC, Sinkhorn), see deepseek4 --- 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, - # --- indexeur : compression k-pool --- + # --- indexer: k-pool compression --- MODEL_TENSOR.INDEXER_COMPRESSOR_APE, MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE, ], diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 6b0cc5eb902f..847899b133e7 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -277,6 +277,10 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader &) { } std::unique_ptr llama_model_glm5next::build_arch_graph(const llm_graph_params & params) const { + // n_layer_nextn > 0 lets llama_init_from_model accept an MTP context, so refuse here + // rather than silently return the trunk graph against the wrong memory + GGML_ASSERT(params.gtype != LLM_GRAPH_TYPE_DECODER_MTP && "GLM5NEXT has no MTP graph yet"); + return std::make_unique(*this, params); } @@ -615,7 +619,7 @@ class llm_graph_input_kpool : public llm_graph_input_i { void set_input(const llama_ubatch * ubatch) override { mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); - mctx->set_input_kpool(cell_pool, pool_cells, bias, ubatch); + mctx->set_input_kpool(pool_cells, pool_bias, tail_cells, ubatch); GGML_ASSERT(ggml_backend_buffer_is_host(ape_slots->buffer)); int32_t * data = (int32_t *) ape_slots->data; @@ -632,17 +636,17 @@ class llm_graph_input_kpool : public llm_graph_input_i { bool res = true; res &= k_idxs->ne[0] == params.ubatch.n_tokens; - res &= cell_pool->ne[0] == (int64_t) mctx_cur->get_idx()->get_n_kv(); - res &= cell_pool->ne[1]*bias->ne[1] == params.ubatch.n_tokens; + res &= pool_bias->ne[0] == (int64_t) mctx_cur->get_idx()->get_n_kv()/mctx_cur->get_kpool(); + res &= pool_bias->ne[1]*pool_bias->ne[2] == params.ubatch.n_tokens; return res; } ggml_tensor * k_idxs = nullptr; // I64 [n_tokens] ggml_tensor * ape_slots = nullptr; // I32 [kpool], the identity - reads the ape rows in order - ggml_tensor * cell_pool = nullptr; // I32 [n_kv, n_stream] ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] - ggml_tensor * bias = nullptr; // F32 [n_kv, n_tokens/n_stream, n_stream] + ggml_tensor * pool_bias = nullptr; // F32 [n_pools, n_tokens/n_stream, n_stream] + ggml_tensor * tail_cells = nullptr; // I32 [kpool-1, n_tokens/n_stream, 1, n_stream], null when kpool == 1 const llama_memory_hybrid_idx_context * mctx; }; @@ -672,14 +676,18 @@ llm_graph_input_kpool * llama_model_glm5next::graph::build_inp_kpool(llm_graph_i kp->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); kp->ape_slots = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, r); - kp->cell_pool = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, ns); kp->pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_pool, ns); - kp->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_kv, n_tokens/ns, ns); + kp->pool_bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_pool, n_tokens/ns, ns); ggml_set_input(kp->ape_slots); - ggml_set_input(kp->cell_pool); ggml_set_input(kp->pool_cells); - ggml_set_input(kp->bias); + ggml_set_input(kp->pool_bias); + + if (r > 1 && inp_hyb->mctx->get_select_tail()) { + // 4d, so it concatenates straight onto the expanded pools in build_dsa_top_k + kp->tail_cells = ggml_new_tensor_4d(ctx0, GGML_TYPE_I32, r - 1, n_tokens/ns, 1, ns); + ggml_set_input(kp->tail_cells); + } return (llm_graph_input_kpool *) res->add_input(std::move(kp)); } @@ -692,8 +700,8 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_top_k( const int64_t d = hparams.indexer_head_size; const int64_t nh = hparams.indexer_n_head; const int64_t r = hparams.indexer_block_size; - const int64_t n_kv = inp->cell_pool->ne[0]; - const int64_t ns = inp->cell_pool->ne[1]; + const int64_t n_kv = mctx_idx->get_n_kv(); + const int64_t ns = inp->pool_cells->ne[1]; const int64_t n_pool = inp->pool_cells->ne[0]/r; const int64_t n_tps = n_tokens/ns; @@ -753,22 +761,36 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_top_k( score = ggml_reshape_3d(ctx0, score, n_pool, n_tps, ns); cb(score, "indexer_score", il); - // give every cell of a pool its pool's score rather than expanding pool ids into cell ids, - // which would need an integer multiply-add ggml has no op for. The r members tie, so the - // cut still lands on a pool boundary. bias carries causality and the always-visible tail. - ggml_tensor * expanded = ggml_get_rows(ctx0, - ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_pool); - expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); - expanded = ggml_add(ctx0, expanded, inp->bias); - cb(expanded, "indexer_score_cells", il); + // the cut is on whole pools, never on single cells. Scoring cells with their pool's score + // and cutting there is not the same thing: relu sends many distinct pools to exactly 0.0 + // and ggml_top_k is unordered among equal keys, so the cut splits pools apart. Diagnosis + // and the reference-free check for it (count partly selected pools) are from PR #27754. + score = ggml_add(ctx0, score, inp->pool_bias); + cb(score, "indexer_score_pools", il); + + // index_topk cells worth of whole pools, i.e. the reference's index_topk/index_kpool + const int64_t n_sel = std::min(n_pool, (int64_t) hparams.indexer_top_k/r); - // index_topk cells plus the incomplete tail, as in the reference - const int64_t width = std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); + ggml_tensor * sel = ggml_top_k(ctx0, score, n_sel); + cb(sel, "indexer_top_k_pools", il); - ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); + // expand each selected pool into its members. get_rows indexes src0 ne[2] with the index + // ne[1], so the stream axis must stay there and n_sel*n_tps folds into one row axis + ggml_tensor * pools = ggml_reshape_3d(ctx0, inp->pool_cells, r, n_pool, ns); + + ggml_tensor * top_k = ggml_get_rows(ctx0, pools, + ggml_reshape_3d(ctx0, sel, n_sel*n_tps, ns, 1)); + + // member j of the i-th selected pool is at i*r + j in both layouts, so this is a reshape + top_k = ggml_reshape_4d(ctx0, top_k, r*n_sel, n_tps, 1, ns); + + // index_kpool_always_select_tail: the trailing incomplete pool has no pool key and can + // never be picked above, so its cells are appended instead of taking pool budget + if (inp->tail_cells) { + top_k = ggml_concat(ctx0, top_k, inp->tail_cells, 0); + } // build_attn_mask_top_k reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask - top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, ns); cb(top_k, "indexer_top_k", il); return top_k; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 5f68312544ac..ebd709ae148d 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -68,6 +68,11 @@ static void usage(char ** argv) { printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-o/--out dir] [-v N] [-h/--help]\n", argv[0]); } +// DSA indexer geometry of the synthetic models, shared with test_dsa_kpool below. +// index_topk/index_kpool = 2 whole pools are selected per query row. +static const uint32_t DSA_INDEXER_TOP_K = 8; +static const uint32_t DSA_INDEXER_KPOOL = 4; + static std::vector get_tokens(const uint32_t n_tokens, const uint32_t n_vocab, const size_t seed){ std::mt19937 gen(seed); std::uniform_int_distribution<> dis(0, n_vocab - 1); @@ -299,8 +304,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(128)); - ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); - ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, DSA_INDEXER_TOP_K); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, DSA_INDEXER_KPOOL); ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4})); @@ -386,7 +391,8 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, - const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { + const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, + ggml_backend_sched_eval_callback cb_eval = nullptr, void * cb_eval_user_data = nullptr) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -402,6 +408,8 @@ static std::pair get_model_and_ctx( if (!encode) { ctx_params.n_ubatch = 64; } + ctx_params.cb_eval = cb_eval; + ctx_params.cb_eval_user_data = cb_eval_user_data; size_t tmp = seed; llama_model_ptr model(gguf_ctx != nullptr ? @@ -633,6 +641,211 @@ static int save_models(const llm_arch target_arch, const size_t seed, const int return 0; } +// +// GLM5NEXT DSA k-pool selection: no pool may be selected in part +// +// The reference indexer picks index_topk/index_kpool *whole* pools and expands each into its +// index_kpool members. Running the top-k over individual cells instead is not equivalent: the +// ReLU drives many distinct pool scores to exactly 0.0 and ggml_top_k is unordered among equal +// keys, so a cell-level cut splits pools apart. Comparing selections against a reference does +// not catch it - a cell-level implementation stays above the bf16-vs-bf16 Jaccard noise floor. +// +// What does catch it is counting *partially* selected pools, which needs no reference at all: +// the count is 0 for a pooled implementation and large for a cell-level one. Both the diagnosis +// and this metric are due to danielhanchen (llama.cpp PR #27754). +// +// A pool is partially selected for a query row when at least one of its index_kpool cells is +// picked and at least one is not. The trailing incomplete pool is excluded by construction: for +// query position q only the pools below tail_start = (q + 1)/r*r are complete, and the cells in +// [tail_start, q] are the always_select_tail cells, which are never counted as pool members. +// +struct dsa_kpool_check { + int64_t r = 0; // index_kpool + + int64_t n_tensor = 0; // indexer_top_k tensors inspected + int64_t n_row = 0; // query rows inspected + int64_t n_pool_whole = 0; // pools selected in full + int64_t n_pool_partial = 0; // pools selected in part - must be 0 + int64_t n_row_partial = 0; // rows holding at least one partial pool + int64_t n_tail_missing = 0; // tail cells not selected - must be 0 + + std::vector buf; + std::vector sel; + + // one query row: `sel` holds the cells picked for query position q + void count_row(const int32_t * row, int64_t width, int64_t q) { + const int64_t tail_start = (q + 1)/r*r; + const int64_t n_pool_vis = tail_start/r; + + sel.assign(q + 1, 0); + for (int64_t j = 0; j < width; j++) { + const int32_t c = row[j]; + if (c >= 0 && c <= q) { + sel[c] = 1; // cells past q are masked anyway and belong to no complete pool + } + } + + // always_select_tail; doubles as a check that cell index == position holds here + for (int64_t c = tail_start; c <= q; c++) { + n_tail_missing += sel[c] == 0; + } + + bool row_partial = false; + for (int64_t b = 0; b < n_pool_vis; b++) { + int64_t cnt = 0; + for (int64_t c = b*r; c < (b + 1)*r; c++) { + cnt += sel[c]; + } + if (cnt == r) { + n_pool_whole++; + } else if (cnt > 0) { + n_pool_partial++; + row_partial = true; + } + } + + n_row_partial += row_partial; + n_row++; + } +}; + +// the metric must not be able to read 0 by accident: a pool-aligned selection scores 0 partial +// pools, the same selection with one cell moved across a pool boundary does not +static void dsa_kpool_check_self_test() { + const std::vector aligned = { 0, 1, 2, 3, 8, 9, 10, 11 }; + const std::vector split = { 0, 1, 2, 4, 8, 9, 10, 11 }; + + dsa_kpool_check ok; + ok.r = 4; + ok.count_row(aligned.data(), aligned.size(), /*q =*/ 15); + GGML_ASSERT(ok.n_pool_whole == 2 && ok.n_pool_partial == 0 && ok.n_tail_missing == 0); + + dsa_kpool_check bad; + bad.r = 4; + bad.count_row(split.data(), split.size(), /*q =*/ 15); + GGML_ASSERT(bad.n_pool_whole == 1 && bad.n_pool_partial == 2); +} + +// exactly "indexer_top_k-": views and backend copies inherit the name with a suffix and +// would otherwise be counted a second time +static bool is_indexer_top_k(const char * name) { + static const char * prefix = "indexer_top_k-"; + const size_t n = strlen(prefix); + if (strncmp(name, prefix, n) != 0 || name[n] == '\0') { + return false; + } + for (const char * p = name + n; *p; p++) { + if (*p < '0' || *p > '9') { + return false; + } + } + return true; +} + +// reads the selection out of every "indexer_top_k-" node, I32 [n_top_k, n_batch, 1, n_stream] +static bool dsa_kpool_eval_cb(struct ggml_tensor * t, bool ask, void * user_data) { + auto & st = *(dsa_kpool_check *) user_data; + + if (ask) { + return is_indexer_top_k(t->name); + } + + GGML_ASSERT(t->type == GGML_TYPE_I32); + + const int64_t width = t->ne[0]; + const int64_t n_tps = t->ne[1]; + const int64_t ns = t->ne[3]; + + st.buf.resize(ggml_nelements(t)); + ggml_backend_tensor_get(t, st.buf.data(), 0, ggml_nbytes(t)); + + for (int64_t s = 0; s < ns; s++) { + for (int64_t i = 0; i < n_tps; i++) { + // one stream, one sequence, one ubatch over a fresh cache: cell index == position + st.count_row(st.buf.data() + (s*n_tps + i)*width, width, /*q =*/ i); + } + } + + st.n_tensor++; + return true; +} + +static int test_dsa_kpool(const size_t seed, const int verbosity) { + struct user_data_t { + struct { + ggml_log_callback callback; + void * user_data; + } log_old; + + int verbosity; + + user_data_t(int verbosity) : verbosity(verbosity) { + llama_log_get(&log_old.callback, &log_old.user_data); + } + }; + user_data_t ud(verbosity); + + llama_log_set([](ggml_log_level level, const char * text, void * user_data) { + const user_data_t * ud = (const user_data_t *) user_data; + int verbosity = common_log_get_verbosity(level); + if (verbosity <= ud->verbosity) { + ud->log_old.callback(level, text, ud->log_old.user_data); + } + }, &ud); + + dsa_kpool_check_self_test(); + + // one ubatch, and enough tokens that index_topk + index_kpool - 1 is a real cut: + // from q = 12 on, the selection can no longer be a union of whole pools by accident + const uint32_t n_tokens = 64; + const std::vector tokens = get_tokens(n_tokens, 128, seed); + + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_GLM5NEXT, true); + + std::vector, std::string>> dev_configs; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + dev_configs.emplace_back(std::vector{dev}, ggml_backend_dev_description(dev)); + } + + bool all_ok = true; + common_log_flush(common_log_main()); + printf("test_dsa_kpool: glm5next, index_topk=%u index_kpool=%u, %u tokens\n", + DSA_INDEXER_TOP_K, DSA_INDEXER_KPOOL, n_tokens); + + for (const auto & dc : dev_configs) { + dsa_kpool_check st; + st.r = DSA_INDEXER_KPOOL; + + auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.first, + LLAMA_SPLIT_MODE_LAYER, false, dsa_kpool_eval_cb, &st); + + llama_batch batch = llama_batch_init(n_tokens, 0, 1); + for (uint32_t pos = 0; pos < n_tokens; pos++) { + common_batch_add(batch, tokens[pos], pos, {0}, true); + } + batch.n_tokens = n_tokens; + const int32_t rc = llama_decode(model_and_ctx.second.get(), batch); + llama_batch_free(batch); + if (rc != 0) { + throw std::runtime_error("failed to decode batch"); + } + + // a silent miss (no indexer_top_k node, or the tail not selected) is a failure too + const bool ok = st.n_tensor > 0 && st.n_tail_missing == 0 && st.n_pool_partial == 0; + all_ok &= ok; + + printf("test_dsa_kpool: %-32s rows %5" PRId64 ", whole pools %6" PRId64 ", " + "partial pools %6" PRId64 " in %4" PRId64 " rows, tail misses %4" PRId64 " %s\n", + dc.second.c_str(), st.n_row, st.n_pool_whole, st.n_pool_partial, st.n_row_partial, + st.n_tail_missing, ok ? "\033[1;32mOK\033[0m" : "\033[1;31mFAIL\033[0m"); + fflush(stdout); + } + + llama_log_set(ud.log_old.callback, ud.log_old.user_data); + return all_ok ? 0 : 1; +} + static int test_backends(const llm_arch target_arch, const size_t seed, const int verbosity) { struct user_data_t { struct { @@ -866,7 +1079,12 @@ int main(int argc, char ** argv) { if (!out.empty()) { return save_models(arch, seed, verbosity, out); } - return test_backends(arch, seed, verbosity); + int ret = 0; + if (arch == LLM_ARCH_UNKNOWN || arch == LLM_ARCH_GLM5NEXT) { + ret |= test_dsa_kpool(seed, verbosity); + } + ret |= test_backends(arch, seed, verbosity); + return ret; } catch (const std::exception & err) { fprintf(stderr, "encountered runtime error: %s\n", err.what()); return -1; From 38d16c762542dafd6d1f8d9bfc5c26347df61f26 Mon Sep 17 00:00:00 2001 From: eauchs Date: Thu, 27 Aug 2026 20:25:14 +0200 Subject: [PATCH 05/11] llama : add the glm5next MTP (NextN) draft graph --- src/llama-model.cpp | 31 ++++- src/models/glm5next.cpp | 238 ++++++++++++++++++++++++++++++++++--- src/models/models.h | 6 + tests/test-llama-archs.cpp | 179 +++++++++++++++++++++++++++- 4 files changed, 434 insertions(+), 20 deletions(-) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 0c8d45029137..907225a502f5 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2427,6 +2427,35 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } break; case LLM_ARCH_GLM5NEXT: { + GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE); + + if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && hparams.n_layer_nextn > 0) { + // The NextN/MTP draft head is one dense-MLA block: no KDA state and no + // DSA indexer, so a plain attention cache over the nextn layer(s) is + // enough - same pattern as GLM_DSA / DEEPSEEK32. + llama_kv_cache::layer_filter_cb filter = + [&](uint32_t il) { return il >= hparams.n_layer(); }; + + res = new llama_kv_cache( + *this, + hparams, + params.type_k, + params.type_v, + !cparams.flash_attn, + cparams.offload_kqv, + cparams.kv_unified, + cparams.n_ctx_seq, + cparams.n_seq_max, + 1, + hparams.n_swa, + hparams.swa_type, + nullptr, + filter, + nullptr, + nullptr); + break; + } + // KDA layers recur, MLA layers cache, and the DSA indexer shadows the MLA layers llama_memory_hybrid_idx::layer_filter_cb filter_recr = [&](int32_t il) { return (uint32_t) il < hparams.n_layer() && hparams.is_recr(il); }; @@ -2441,8 +2470,6 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, }; } - GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE); - res = new llama_memory_hybrid_idx( /* model */ *this, /* attn_type_k */ params.type_k, diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 847899b133e7..08ebc5905f5c 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -17,8 +17,10 @@ // llama_memory_hybrid_idx next to the recurrent states and the MLA cache. A GGUF without // indexer weights still builds the dense graph. // -// Not wired up yet: -// - the NextN/MTP block, loaded as unused +// The NextN/MTP block appended after the trunk is a plain pre-norm decoder block, with +// neither hyper-connections nor KDA. It has indexer weights, but the reference shares the +// trunk index for the MTP step (index_share_for_mtp_iteration), which a separate MTP +// context cannot see, so the draft head attends densely. See graph_mtp below. // void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { @@ -120,8 +122,11 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { // which works only because linear_attn_config.num_heads == num_attention_heads // MLA forces num_key_value_heads = 1 on every layer at conversion time, so the - // kimi-linear "n_head_kv == 0" recurrent marker is not available here - if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) { + // kimi-linear "n_head_kv == 0" recurrent marker is not available here. + // the per-layer arrays cover the trunk only, like the other glm5next ones - the + // NextN block is never recurrent, so leave its entry at 0 + std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), 0); + if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer(), false)) { uint32_t full_attn_interval = 4; // layer_types: full attention on 3, 7, 11, ... ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); GGML_ASSERT(full_attn_interval > 0); @@ -136,7 +141,7 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { } } -void llama_model_glm5next::load_arch_tensors(llama_model_loader &) { +void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { LLAMA_LOAD_LOCALS; const int64_t q_lora_rank = hparams.n_lora_q; @@ -161,8 +166,15 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader &) { const int64_t d_conv = hparams.ssm_d_conv; const int64_t d_inner = head_dim * n_head; - // the NextN/MTP block is loaded but ignored - there is no MTP graph yet - const int mtp_flags = TENSOR_NOT_REQUIRED | TENSOR_SKIP; + // a GGUF that declares nextn layers but ships the trunk alone still loads + const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; + const bool trunk_only = (n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); + + // the MTP block is materialized only for an MTP context (--spec-type draft-mtp) + int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; + if (!ml.load_mtp) { + mtp_flags |= TENSOR_SKIP; + } tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); @@ -233,7 +245,8 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader &) { layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, flags); layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, flags); - // DSA indexer + k-pool compressor + // DSA indexer + k-pool compressor. The MTP block ships one too, but the draft + // head runs dense (see graph_mtp), so it stays loaded and unused there layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_idx_head * n_idx_dim}, flags | TENSOR_NOT_REQUIRED); layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, n_idx_dim}, flags | TENSOR_NOT_REQUIRED); layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_idx_dim}, flags | TENSOR_NOT_REQUIRED); @@ -269,17 +282,19 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader &) { layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2*n_embd, n_embd}, mtp_flags); layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, mtp_flags); layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, mtp_flags); - layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, mtp_flags); - layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, mtp_flags); - layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags); + + // GLM-5.3-Flash ties these to the trunk embeddings / LM head, so they are absent + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, mtp_flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, mtp_flags | TENSOR_NOT_REQUIRED); } } } std::unique_ptr llama_model_glm5next::build_arch_graph(const llm_graph_params & params) const { - // n_layer_nextn > 0 lets llama_init_from_model accept an MTP context, so refuse here - // rather than silently return the trunk graph against the wrong memory - GGML_ASSERT(params.gtype != LLM_GRAPH_TYPE_DECODER_MTP && "GLM5NEXT has no MTP graph yet"); + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } return std::make_unique(*this, params); } @@ -963,8 +978,11 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa cb(inpL, "l_out", il); } + // unmasked nextn embeddings need every row, so narrow after the final norm instead + const bool narrow_late = cparams.embeddings_nextn && !cparams.embeddings_nextn_masked; + // narrow to the output rows before collapsing the streams - if (inp_out_ids) { + if (inp_out_ids && !narrow_late) { ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); flat = ggml_get_rows(ctx0, flat, inp_out_ids); inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs); @@ -974,6 +992,15 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa cb(cur, "hc_head", -1); cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + + // post-norm hidden state feeds the NextN/MTP draft head + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (inp_out_ids && narrow_late) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "result_norm", -1); res->t_embd = cur; @@ -983,3 +1010,184 @@ llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_pa ggml_build_forward_expand(gf, cur); } + +// +// NextN/MTP draft head +// +// enorm(embed) + hnorm(prev_hidden) -> concat(e, h) -> eh_proj -> one plain pre-norm +// decoder block (nope-only MLA + sigmoid-gated MoE with a shared expert, built the same +// way as build_mla_layer/build_ffn_layer build the trunk) -> shared_head_norm -> LM head. +// +// Differences from a trunk layer: +// - no hyper-connections: the block has no hc_* tensors, it uses plain residuals +// - dense attention: the reference shares the trunk index for the MTP step, which this +// separate context cannot see. This costs acceptance rate, never correctness. +// + +llama_model_glm5next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) + : llm_graph_context(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "GLM5NEXT MTP supports a single NextN block"); + GGML_ASSERT(hparams.is_mla() && "GLM5NEXT MTP requires MLA"); + GGML_ASSERT(hparams.n_rot() == 0 && "GLM5NEXT MTP is nope-only"); + + GGML_ASSERT(cparams.nextn_layer_offset >= 0 && + cparams.nextn_layer_offset < (int) hparams.n_layer_nextn && + "nextn_layer_offset out of range [0, n_layer_nextn)"); + + const int il = hparams.n_layer() + cparams.nextn_layer_offset; + + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj (load with --spec-type draft-mtp)"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + GGML_ASSERT(layer.ffn_gate_inp && "MTP block missing ffn_gate_inp"); + + const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + + // nope-only, so no YaRN mscale correction - must match the trunk graph + const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k)); + + // TODO: extract in a common llm_graph_context::build_inp_embd_h() + auto inp = std::make_unique(hparams.n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * tok_embd; + if (ubatch.token) { + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + + tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); + } else { + tok_embd = inp->embd; + } + cb(tok_embd, "mtp_tok_embd", il); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + ggml_tensor * h_embd = inp->h; + + res->add_input(std::move(inp)); + + // no build_inp_pos(): glm5next is position-free + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // MLA with the absorption optimization uses a K-only cache (V is a view of K) + auto * inp_attn = build_attn_inp_k(); + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * concat = ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0); + cb(concat, "mtp_concat", il); + + ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, concat, layer.nextn.eh_proj_s); + cb(cur, "mtp_eh_proj", il); + + ggml_tensor * inpSA = cur; + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + { + ggml_tensor * qr = ggml_mul_mat(ctx0, layer.wq_a, cur); + qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(qr, "mtp_qr", il); + + // nope-only: the whole of q is the "nope" part, no split and no rope + ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr); + q = ggml_reshape_3d(ctx0, q, n_embd_head_k, n_head, n_tokens); + cb(q, "mtp_q", il); + + // {n_embd_head_k, n_tokens, n_head} x wk_b -> {kv_lora_rank, n_tokens, n_head} + q = ggml_permute(ctx0, q, 0, 2, 1, 3); + ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.wk_b, q); + Qcur = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); + cb(Qcur, "mtp_Qcur", il); + + // nope-only: wkv_a_mqa outputs exactly kv_lora_rank, no k_pe to split off + ggml_tensor * kv_cmpr = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + cb(kv_cmpr, "mtp_kv_cmpr", il); + + cur = build_attn(inp_attn, layer.wo, nullptr, layer.wo_s, + Qcur, kv_cmpr, kv_cmpr, nullptr, nullptr, layer.wv_b, kq_scale, il); + cb(cur, "mtp_attn_out", il); + } + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "mtp_ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_ffn_norm", il); + + // the NextN block is always past n_layer_dense_lead, so there is no dense-MLP branch + { + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il, + nullptr, + layer.ffn_gate_up_exps, + layer.ffn_up_exps_s, + layer.ffn_gate_exps_s, + layer.ffn_down_exps_s); + cb(moe_out, "mtp_ffn_moe_out", il); + + ggml_tensor * shexp = build_ffn(cur, + layer.ffn_up_shexp, nullptr, layer.ffn_up_shexp_s, + layer.ffn_gate_shexp, nullptr, layer.ffn_gate_shexp_s, + layer.ffn_down_shexp, nullptr, layer.ffn_down_shexp_s, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(shexp, "mtp_ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, shexp); + cb(cur, "mtp_ffn_out", il); + } + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "mtp_post_ffn", il); + + ggml_tensor * head_norm_w = layer.nextn.shared_head_norm + ? layer.nextn.shared_head_norm + : model.output_norm; + GGML_ASSERT(head_norm_w && "GLM5NEXT MTP: missing both nextn.shared_head_norm and output_norm"); + + // the post-norm hidden state would seed a chained head, unused at n_layer_nextn == 1 + cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1); + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "mtp_shared_head_norm", -1); + + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; + GGML_ASSERT(head_w && "GLM5NEXT MTP: missing LM head (nextn.shared_head_head or model.output)"); + + cur = build_lora_mm(head_w, cur, head_s); + cb(cur, "result_output", -1); + + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} diff --git a/src/models/models.h b/src/models/models.h index 81ffd6977baa..4f97125d42cc 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1382,6 +1382,12 @@ struct llama_model_glm5next : public llama_model_base { ggml_tensor * build_ffn_layer(ggml_tensor * cur, const llama_layer & layer, int il); }; + // NextN/MTP draft head: the block appended at index n_layer(). It is dense MLA with + // no hyper-connections, so it runs on a plain attention cache holding only that block + struct graph_mtp : public llm_graph_context { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index ebd709ae148d..ac7dc64747b4 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -11,6 +11,9 @@ #include "../src/llama-arch.h" #include "../src/llama-model-saver.h" +// nextn/MTP accessors are still staging API +#include "../src/llama-ext.h" + #include #include #include @@ -84,7 +87,8 @@ static std::vector get_tokens(const uint32_t n_tokens, const uint32 return ret; } -static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { +// nextn appends one NextN/MTP block after the trunk, leaving the trunk itself unchanged +static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe, const bool nextn = false) { gguf_context_ptr ret(gguf_init_empty()); llama_model_saver ms(arch, ret.get()); const uint32_t n_ctx = 256; @@ -154,8 +158,11 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_CONTEXT_LENGTH, n_ctx); ms.add_kv(LLM_KV_EMBEDDING_LENGTH, n_embd); ms.add_kv(LLM_KV_FEATURES_LENGTH, n_embd); - ms.add_kv(LLM_KV_BLOCK_COUNT, n_layer); + ms.add_kv(LLM_KV_BLOCK_COUNT, nextn ? n_layer + 1 : n_layer); ms.add_kv(LLM_KV_LEADING_DENSE_BLOCK_COUNT, uint32_t(1)); + if (nextn) { + ms.add_kv(LLM_KV_NEXTN_PREDICT_LAYERS, uint32_t(1)); + } if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { std::vector n_ff_per_layer; @@ -392,10 +399,12 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, - ggml_backend_sched_eval_callback cb_eval = nullptr, void * cb_eval_user_data = nullptr) { + ggml_backend_sched_eval_callback cb_eval = nullptr, void * cb_eval_user_data = nullptr, + bool load_mtp = false) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; + model_params.load_mtp = load_mtp; std::vector devs_copy = devs; devs_copy.push_back(nullptr); model_params.devices = devs_copy.data(); @@ -846,6 +855,169 @@ static int test_dsa_kpool(const size_t seed, const int verbosity) { return all_ok ? 0 : 1; } +// +// GLM5NEXT NextN/MTP draft head +// +// The trunk exports its post-norm hidden state as h_nextn, the draft head consumes it next +// to a token id and emits logits. This is a smoke test: it runs the MTP loader flags, the +// MTP memory branch and the graph, and checks that the logits are finite. It cannot check +// the numbers - there is no MTP reference to compare against. +// +// It also saves the model to a GGUF and reloads it, twice: once with load_mtp so the draft +// head has to find blk..nextn.* by name in a real file, and once without, which is +// the default path where the NextN block is skipped and only the trunk runs. +// +static int test_mtp(const size_t seed, const int verbosity) { + struct user_data_t { + struct { + ggml_log_callback callback; + void * user_data; + } log_old; + + int verbosity; + + user_data_t(int verbosity) : verbosity(verbosity) { + llama_log_get(&log_old.callback, &log_old.user_data); + } + }; + user_data_t ud(verbosity); + + llama_log_set([](ggml_log_level level, const char * text, void * user_data) { + const user_data_t * ud = (const user_data_t *) user_data; + int verbosity = common_log_get_verbosity(level); + if (verbosity <= ud->verbosity) { + ud->log_old.callback(level, text, ud->log_old.user_data); + } + }, &ud); + + const uint32_t n_tokens = 16; + const std::vector tokens = get_tokens(n_tokens, 128, seed); + + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_GLM5NEXT, true, /*nextn =*/ true); + + std::vector, std::string>> dev_configs; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + dev_configs.emplace_back(std::vector{dev}, ggml_backend_dev_description(dev)); + } + + bool all_ok = true; + common_log_flush(common_log_main()); + printf("test_mtp: glm5next, %u tokens\n", n_tokens); + + // decode the trunk, hand its hidden state to the draft head, return the draft logits + auto run_mtp = [&](llama_model * model, llama_context * ctx_tgt) { + const uint32_t n_embd = llama_model_n_embd(model); + const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + + llama_context_params ctx_params = llama_context_default_params(); + ctx_params.n_ctx = 0; + ctx_params.n_threads = 4; + ctx_params.n_threads_batch = 4; + ctx_params.n_ubatch = 64; + ctx_params.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + + llama_context_ptr ctx_dft(llama_init_from_model(model, ctx_params)); + if (!ctx_dft) { + throw std::runtime_error("failed to create MTP context"); + } + + // unmasked, so the trunk keeps every row and narrows after the final norm + llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false); + + llama_batch batch_tgt = llama_batch_init(n_tokens, 0, 1); + for (uint32_t pos = 0; pos < n_tokens; pos++) { + common_batch_add(batch_tgt, tokens[pos], pos, {0}, true); + } + const int32_t rc_tgt = llama_decode(ctx_tgt, batch_tgt); + llama_batch_free(batch_tgt); + if (rc_tgt != 0) { + throw std::runtime_error("failed to decode trunk batch"); + } + + const float * h = llama_get_embeddings_nextn_ith(ctx_tgt, n_tokens - 1); + if (!h) { + throw std::runtime_error("trunk did not export h_nextn"); + } + + // the draft head reads a token id and a hidden state, so the batch needs both + llama_batch batch_dft = llama_batch_init(1, n_embd, 1); + batch_dft.token = (llama_token *) malloc(sizeof(llama_token)); + common_batch_add(batch_dft, tokens[n_tokens - 1], n_tokens - 1, {0}, true); + memcpy(batch_dft.embd, h, n_embd*sizeof(float)); + + const int32_t rc_dft = llama_decode(ctx_dft.get(), batch_dft); + free(batch_dft.token); + batch_dft.token = nullptr; + llama_batch_free(batch_dft); + if (rc_dft != 0) { + throw std::runtime_error("failed to decode MTP batch"); + } + + const float * logits = llama_get_logits_ith(ctx_dft.get(), 0); + if (!logits) { + throw std::runtime_error("the MTP graph produced no logits"); + } + return std::vector(logits, logits + n_vocab); + }; + + for (const auto & dc : dev_configs) { + auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.first, + LLAMA_SPLIT_MODE_LAYER, false, nullptr, nullptr, /*load_mtp =*/ true); + + const std::vector logits = run_mtp(model_and_ctx.first.get(), model_and_ctx.second.get()); + + bool ok = true; + for (size_t i = 0; ok && i < logits.size(); i++) { + ok = std::isfinite(logits[i]); + } + + // save the model, then reload it from the file so the NextN block has to be found + // by name rather than synthesized. Skipped where tmpfile() is unavailable. + std::string status_roundtrip = "\033[1;33mSKIP\033[0m"; + FILE * file = tmpfile(); + if (file != nullptr && llama_model_saver_supports_arch(LLM_ARCH_GLM5NEXT)) { + llama_model_saver ms = llama_model_saver(model_and_ctx.first.get()); + ms.add_kv_from_model(); + ms.add_tensors_from_model(); + ms.save(file); + + rewind(file); + auto rt_mtp = get_model_and_ctx(nullptr, file, seed, dc.first, + LLAMA_SPLIT_MODE_LAYER, false, nullptr, nullptr, /*load_mtp =*/ true); + const std::vector logits_rt = run_mtp(rt_mtp.first.get(), rt_mtp.second.get()); + + status_roundtrip = "\033[1;32mOK\033[0m"; + GGML_ASSERT(logits_rt.size() == logits.size()); + for (size_t i = 0; i < logits_rt.size(); i++) { + if (logits_rt[i] != logits[i]) { + ok = false; + status_roundtrip = "\033[1;31mFAIL\033[0m"; + break; + } + } + + // note: reloading the same file without load_mtp is NOT tested here. The + // synthetic model materializes the optional NVFP4 sidecar scales for the + // NextN block, which a real GGUF does not carry, and TENSOR_SKIP leaves + // them unaccounted for. Same for the other NextN archs, so it is not + // specific to this graph. + } + if (file != nullptr) { + fclose(file); + } + + all_ok &= ok; + + printf("test_mtp: %-32s draft %s, reload %s\n", dc.second.c_str(), + ok ? "\033[1;32mOK\033[0m" : "\033[1;31mFAIL\033[0m", status_roundtrip.c_str()); + fflush(stdout); + } + + llama_log_set(ud.log_old.callback, ud.log_old.user_data); + return all_ok ? 0 : 1; +} + static int test_backends(const llm_arch target_arch, const size_t seed, const int verbosity) { struct user_data_t { struct { @@ -1082,6 +1254,7 @@ int main(int argc, char ** argv) { int ret = 0; if (arch == LLM_ARCH_UNKNOWN || arch == LLM_ARCH_GLM5NEXT) { ret |= test_dsa_kpool(seed, verbosity); + ret |= test_mtp(seed, verbosity); } ret |= test_backends(arch, seed, verbosity); return ret; From 347f2e37a306f0fcc8ba45f0792bd6635f1d095d Mon Sep 17 00:00:00 2001 From: eauchs Date: Fri, 28 Aug 2026 12:07:35 +0200 Subject: [PATCH 06/11] llama : extend llama_memory_hybrid_idx for the glm5next DSA path Rather than adding a dedicated container, glm5next reuses llama_memory_hybrid_idx, which gains an idx_row_size constructor parameter and a set_input_kpool() input for k-pool selection. --- src/CMakeLists.txt | 1 - src/llama-memory-hybrid-idx.cpp | 152 +++++++++++++++++++++++++++++++- src/llama-memory-hybrid-idx.h | 12 +++ src/llama-model.cpp | 3 +- src/models/glm5next.cpp | 13 +-- 5 files changed, 171 insertions(+), 10 deletions(-) diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 7379491687a0..8922dc12adc7 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -30,7 +30,6 @@ add_library(llama llama-kv-cache-dsv4.cpp llama-memory.cpp llama-memory-hybrid.cpp - llama-memory-hybrid-idx.cpp llama-memory-hybrid-iswa.cpp llama-memory-hybrid-idx.cpp llama-memory-recurrent.cpp diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 93b468784a33..ca8c4c8aca5d 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -30,6 +30,8 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( ggml_type type_r, ggml_type type_s, uint32_t rs_size, + /* indexer */ + uint32_t idx_row_size, /* common */ uint32_t n_seq_max, uint32_t n_rs_seq, @@ -47,9 +49,11 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx( filter_attn, filter_recr), hparams_idx(model.hparams), mem_idx(filter_idx == nullptr ? nullptr : [&] { - // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own + // MQA with a single key head, as llama_kv_cache_dsa shapes its own. The row is the + // indexer key, unless the architecture packs more into it (glm5next caches the k-pool + // gate alongside the key, since the gate depends on the token's hidden state) std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1); - hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; + hparams_idx.n_embd_head_k_full = idx_row_size > 0 ? idx_row_size : model.hparams.indexer_head_size; // the cached indexer keys are raw, rotation happens after pooling at read time, so a // K-shift must not rotate them while the stream copies in the same update still apply @@ -677,3 +681,147 @@ void llama_memory_hybrid_idx_context::set_input_qsa( mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); } + +void llama_memory_hybrid_idx_context::set_input_kpool( + ggml_tensor * pool_cells, + ggml_tensor * pool_bias, + ggml_tensor * tail_cells, + const llama_ubatch * ubatch, + uint32_t ratio) const { + GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(pool_bias->buffer)); + GGML_ASSERT(pool_cells->type == GGML_TYPE_I32); + GGML_ASSERT(pool_bias->type == GGML_TYPE_F32); + + const int64_t r = ratio; + const int64_t n_kv = get_attn()->get_n_kv(); + + GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); + const int64_t n_ns = pool_cells->ne[1]; // streams in this ubatch + const int64_t n_pool = pool_bias->ne[0]; + + GGML_ASSERT(r > 0 && n_pool > 0); + GGML_ASSERT(pool_cells->ne[0] == r*n_pool); + GGML_ASSERT(ubatch->n_tokens % n_ns == 0); + + const int64_t n_tps = ubatch->n_tokens/n_ns; + + int32_t * dst_pool_cells = (int32_t *) pool_cells->data; + float * dst_pool_bias = (float *) pool_bias->data; + int32_t * dst_tail_cells = nullptr; + + if (tail_cells) { + GGML_ASSERT(ggml_backend_buffer_is_host(tail_cells->buffer)); + GGML_ASSERT(tail_cells->type == GGML_TYPE_I32); + GGML_ASSERT(tail_cells->ne[0] == r - 1); + + dst_tail_cells = (int32_t *) tail_cells->data; + } + + // one pass per stream: cell j is a different token in each + std::vector filled(n_pool); + std::vector cell_of_pos; + + for (int64_t s = 0; s < n_ns; ++s) { + const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; + + const auto & cells = mem->get_mem_attn()->get_cells(seq_of_stream); + + int32_t * cur_pool_cells = dst_pool_cells + s*(r*n_pool); + + // pool b covers token positions [b*r, (b+1)*r) + std::fill(filled.begin(), filled.end(), 0); + std::fill(cur_pool_cells, cur_pool_cells + r*n_pool, 0); + + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j)) { + continue; + } + + const llama_pos p = cells.pos_get(j); + const int64_t b = p/r; + + if (b >= n_pool) { + continue; + } + + cur_pool_cells[b*r + (p%r)] = (int32_t) j; + filled[b]++; + } + + // an incomplete pool has no pool key: pool_bias below never lets a query pick it, and + // cell 0 in pool_cells only keeps the gather in range + for (int64_t b = 0; b < n_pool; ++b) { + if (filled[b] < (int32_t) r) { + std::fill(cur_pool_cells + b*r, cur_pool_cells + (b + 1)*r, 0); + } + } + + llama_pos max_pos = -1; + int32_t pad_masked = -1; + + if (dst_tail_cells) { + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j) || !cells.seq_has(j, seq_of_stream)) { + pad_masked = (int32_t) j; // the KQ mask rejects it for every query below + continue; + } + + max_pos = std::max(max_pos, cells.pos_get(j)); + } + + cell_of_pos.assign(max_pos + 1, -1); + + for (int64_t j = 0; j < n_kv; ++j) { + if (!cells.is_empty(j) && cells.seq_has(j, seq_of_stream)) { + cell_of_pos[cells.pos_get(j)] = (int32_t) j; + } + } + } + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + const llama_pos q = ubatch->pos[i]; + + // a query may pick a pool only if the pool is complete and its last member is + // visible, which is the reference's pool_valid & pool_visible + float * cur_bias = dst_pool_bias + i*n_pool; + + for (int64_t b = 0; b < n_pool; ++b) { + const bool ok = filled[b] == (int32_t) r && (llama_pos) ((b + 1)*r - 1) <= q; + + cur_bias[b] = ok ? 0.0f : -INFINITY; + } + + if (!dst_tail_cells) { + continue; + } + + // positions [tail_start, q] are the trailing incomplete pool. A pickable pool ends + // at or before tail_start - 1, so the tail never overlaps one and nothing is + // counted twice. + const llama_pos tail_start = (q + 1)/(llama_pos) r*(llama_pos) r; + + // the reference pads the tail with -1; set_rows has no -1 but tolerates duplicates, + // so pad with a cell the KQ mask rejects anyway. Only a full cache that holds this + // sequence alone and ends exactly at q has none - the query then pads with itself. + int32_t pad = pad_masked; + + if (pad < 0 && max_pos > q) { + pad = cell_of_pos[max_pos]; + } + + if (pad < 0) { + pad = q <= max_pos && cell_of_pos[q] >= 0 ? cell_of_pos[q] : 0; + } + + int32_t * cur_tail = dst_tail_cells + i*(r - 1); + + for (int64_t t = 0; t < r - 1; ++t) { + const llama_pos p = tail_start + (llama_pos) t; + + cur_tail[t] = (p <= q && p <= max_pos && cell_of_pos[p] >= 0) ? cell_of_pos[p] : pad; + } + } + } +} diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 705189e7eb58..fc6ed71eac25 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -28,6 +28,8 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { ggml_type type_r, ggml_type type_s, uint32_t rs_size, + /* indexer */ + uint32_t idx_row_size, // floats cached per token; 0 means indexer_head_size /* common */ uint32_t n_seq_max, uint32_t n_rs_seq, @@ -145,6 +147,16 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const; + // k-pool selection (glm5next DSA) over the same cells. Where set_input_qsa biases attention + // scores, this biases the top-k that picks the pools, so a pool is offered to a query only + // when it is complete and its last member is already visible: + // pool_cells I32 [ratio*n_pools, ns] cells making up each pool + // pool_bias F32 [n_pools, n_tokens/ns, ns] 0 where selectable, -inf otherwise + // tail_cells I32 [ratio-1, n_tokens/ns, 1, ns] optional: the cells of each query's own + // incomplete trailing pool, which has no pool key and so is expanded directly + void set_input_kpool(ggml_tensor * pool_cells, ggml_tensor * pool_bias, ggml_tensor * tail_cells, + const llama_ubatch * ubatch, uint32_t ratio) const; + private: const llama_memory_hybrid_idx * mem = nullptr; diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 907225a502f5..d50524fe98e3 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2483,8 +2483,6 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* recurrent_type_s */ GGML_TYPE_F32, /* recurrent_rs_size */ std::max((uint32_t) 1, cparams.n_seq_max), /* idx_row_size */ 2*hparams.indexer_head_size, // key | k-pool gate - /* idx_kpool */ hparams.indexer_block_size, - /* idx_select_tail */ true, // index_kpool_always_select_tail /* n_seq_max */ cparams.n_seq_max, /* n_rs_seq */ cparams.n_rs_seq, /* offload */ cparams.offload_kqv, @@ -2611,6 +2609,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* recurrent_type_k */ GGML_TYPE_F32, /* recurrent_type_v */ GGML_TYPE_F32, /* recurrent_kv_size */ std::max((uint32_t) 1, cparams.n_seq_max), + /* idx_row_size */ 0, // the indexer key alone /* n_seq_max */ cparams.n_seq_max, /* n_rs_seq */ cparams.n_rs_seq, /* offload */ cparams.offload_kqv, diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 08ebc5905f5c..7fdc82c998ad 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -629,12 +629,13 @@ ggml_tensor * llama_model_glm5next::graph::build_kda_layer( class llm_graph_input_kpool : public llm_graph_input_i { public: - llm_graph_input_kpool(const llama_memory_hybrid_idx_context * mctx) : mctx(mctx) {} + llm_graph_input_kpool(const llama_memory_hybrid_idx_context * mctx, uint32_t kpool) : + mctx(mctx), kpool(kpool) {} virtual ~llm_graph_input_kpool() = default; void set_input(const llama_ubatch * ubatch) override { mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); - mctx->set_input_kpool(pool_cells, pool_bias, tail_cells, ubatch); + mctx->set_input_kpool(pool_cells, pool_bias, tail_cells, ubatch, kpool); GGML_ASSERT(ggml_backend_buffer_is_host(ape_slots->buffer)); int32_t * data = (int32_t *) ape_slots->data; @@ -651,7 +652,7 @@ class llm_graph_input_kpool : public llm_graph_input_i { bool res = true; res &= k_idxs->ne[0] == params.ubatch.n_tokens; - res &= pool_bias->ne[0] == (int64_t) mctx_cur->get_idx()->get_n_kv()/mctx_cur->get_kpool(); + res &= pool_bias->ne[0] == (int64_t) (mctx_cur->get_idx()->get_n_kv()/kpool); res &= pool_bias->ne[1]*pool_bias->ne[2] == params.ubatch.n_tokens; return res; @@ -664,6 +665,8 @@ class llm_graph_input_kpool : public llm_graph_input_i { ggml_tensor * tail_cells = nullptr; // I32 [kpool-1, n_tokens/n_stream, 1, n_stream], null when kpool == 1 const llama_memory_hybrid_idx_context * mctx; + + const uint32_t kpool; }; llm_graph_input_kpool * llama_model_glm5next::graph::build_inp_kpool(llm_graph_input_mem_hybrid_idx * inp_hyb) { @@ -687,7 +690,7 @@ llm_graph_input_kpool * llama_model_glm5next::graph::build_inp_kpool(llm_graph_i return nullptr; } - auto kp = std::make_unique(inp_hyb->mctx); + auto kp = std::make_unique(inp_hyb->mctx, (uint32_t) r); kp->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); kp->ape_slots = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, r); @@ -698,7 +701,7 @@ llm_graph_input_kpool * llama_model_glm5next::graph::build_inp_kpool(llm_graph_i ggml_set_input(kp->pool_cells); ggml_set_input(kp->pool_bias); - if (r > 1 && inp_hyb->mctx->get_select_tail()) { + if (r > 1) { // 4d, so it concatenates straight onto the expanded pools in build_dsa_top_k kp->tail_cells = ggml_new_tensor_4d(ctx0, GGML_TYPE_I32, r - 1, n_tokens/ns, 1, ns); ggml_set_input(kp->tail_cells); From dba3221199c516b588c913cebaad4e3fec364bdc Mon Sep 17 00:00:00 2001 From: Okoku <37441948+matteoscalabrini@users.noreply.github.com> Date: Sat, 29 Aug 2026 12:54:21 +0200 Subject: [PATCH 07/11] llama : chunk the glm5next DSA indexer over tokens The indexer scores [n_pool, n_head, n_tokens] and then reduces over heads, and that reduction needs the head axis in ne[0] - so the score tensor is materialised twice, once by the mul_mat and once by the permute+cont: 2 * n_pool * n_head * n_tokens * 4 B per device It is per device because the DSA layers are spread across the trunk, so every device in a layer split owns some of them. At n_ctx 262144 with kpool 4 and 32 heads that is 16 MiB per token, so -ub 4096 asks for about 70 GiB on a single device (measured: cudaMalloc fails at 69826 MiB on a 24.5 GiB card). That caps ubatch. With experts on CPU the ubatch is also what amortises the expert stream over PCIe during prefill, so the cap costs throughput as well as memory. Scoring a token depends only on its own query row and on `pooled`, which is shared across the batch, and no reduction in this path runs across tokens. The token loop can therefore be split into chunks with identical results, and ggml-alloc reuses one buffer across the chunks - bounding the scratch by the chunk size rather than by the ubatch. The chunk is derived from a fixed scratch target, so short contexts come out unchunked and take the previous code path unchanged. Measured on 5x RTX 3090, GLM-5.3-Flash UD-Q4_K_XL, f16 KV, -ot exps=CPU: ub 512 unchunked 8788 MiB/device scratch 75 tok/s prefill at 10k ub 4096 chunked ~2 GiB/device scratch 272 tok/s prefill at 10k Greedy output is byte-identical to the unchunked graph at matched ubatch, with --parallel 1, with --parallel 2, and with a ragged final chunk. Needle-in-a- haystack retrieval passes at 28,788 and 243,314 tokens. --- src/models/glm5next.cpp | 111 ++++++++++++++++++++++++++++++---------- 1 file changed, 85 insertions(+), 26 deletions(-) diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 7fdc82c998ad..93b4802972e6 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -765,47 +765,106 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_top_k( q = ggml_reshape_3d(ctx0, q, d, nh*n_tps, ns); cb(q, "indexer_q", il); - ggml_tensor * score = ggml_mul_mat(ctx0, pooled, q); - score = ggml_relu(ctx0, ggml_reshape_4d(ctx0, score, n_pool, nh, n_tps, ns)); - // relu(x*s) == s*relu(x) for s > 0, so both positive scalars (the softmax scale and // n_heads^-1/2) fold into the small weights tensor instead of the big score one ggml_tensor * wts = ggml_mul_mat(ctx0, layer.indexer_proj, cur); wts = ggml_scale(ctx0, wts, 1.0f/sqrtf(float(d*nh))); wts = ggml_reshape_4d(ctx0, wts, nh, 1, n_tps, ns); - score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); - score = ggml_sum_rows(ctx0, ggml_mul(ctx0, score, wts)); - score = ggml_reshape_3d(ctx0, score, n_pool, n_tps, ns); - cb(score, "indexer_score", il); - - // the cut is on whole pools, never on single cells. Scoring cells with their pool's score - // and cutting there is not the same thing: relu sends many distinct pools to exactly 0.0 - // and ggml_top_k is unordered among equal keys, so the cut splits pools apart. Diagnosis - // and the reference-free check for it (count partly selected pools) are from PR #27754. - score = ggml_add(ctx0, score, inp->pool_bias); - cb(score, "indexer_score_pools", il); - // index_topk cells worth of whole pools, i.e. the reference's index_topk/index_kpool const int64_t n_sel = std::min(n_pool, (int64_t) hparams.indexer_top_k/r); - ggml_tensor * sel = ggml_top_k(ctx0, score, n_sel); - cb(sel, "indexer_top_k_pools", il); - // expand each selected pool into its members. get_rows indexes src0 ne[2] with the index // ne[1], so the stream axis must stay there and n_sel*n_tps folds into one row axis ggml_tensor * pools = ggml_reshape_3d(ctx0, inp->pool_cells, r, n_pool, ns); - ggml_tensor * top_k = ggml_get_rows(ctx0, pools, - ggml_reshape_3d(ctx0, sel, n_sel*n_tps, ns, 1)); + // The scores are [n_pool, nh, n_tokens] and the head reduction needs the head axis in + // ne[0], so the tensor is materialised twice - once by the mul_mat and once by the + // permute+cont below. That is 2*n_pool*nh*n_tokens*4 B PER DEVICE (the DSA layers are + // spread across the trunk, so every device in a layer split pays it): 16 MiB per token + // at n_ctx 262144 with kpool 4 and 32 heads. It therefore caps ubatch - -ub 4096 there + // asks for ~70 GiB on a single device. + // + // Scoring a token depends only on its own query and on `pooled`, which is shared across + // the batch, and no reduction in this path runs across tokens. So the token loop can be + // split into chunks with identical results, and ggml-alloc reuses one buffer across the + // chunks - bounding the scratch by the chunk size instead of by the ubatch. + // + // The chunk is sized to keep that scratch near a fixed target: small enough to leave room + // on ~8 GB devices, large enough that the extra kernel launches stay amortised. Short + // contexts, where the tensor is small anyway, come out unchunked and pay nothing. + // One step covers nc*ns tokens, so ns belongs in the divisor. + constexpr int64_t idx_scratch_target = 2ll*1024*1024*1024; + + const int64_t idx_bytes_per_step = 2*n_pool*nh*ns*(int64_t) sizeof(float); + const int64_t idx_chunk = std::clamp(idx_scratch_target/idx_bytes_per_step, 1, n_tps); + + ggml_tensor * top_k = nullptr; + + for (int64_t t0 = 0; t0 < n_tps; t0 += idx_chunk) { + const int64_t nc = std::min(idx_chunk, n_tps - t0); + + ggml_tensor * q_c = q; + ggml_tensor * wts_c = wts; + ggml_tensor * bias_c = inp->pool_bias; + ggml_tensor * tail_c = inp->tail_cells; + + if (nc != n_tps) { + // q packs (head, token) in ne[1] with the head fastest, so a token range is the + // row range [nh*t0, nh*(t0 + nc)). Strides are kept, so this is also correct for + // n_stream > 1, where a token slice is not contiguous across streams. + q_c = ggml_view_3d(ctx0, q, d, nh*nc, ns, + q->nb[1], q->nb[2], (size_t) (nh*t0)*q->nb[1]); + wts_c = ggml_view_4d(ctx0, wts, nh, 1, nc, ns, + wts->nb[1], wts->nb[2], wts->nb[3], (size_t) t0*wts->nb[2]); + bias_c = ggml_view_3d(ctx0, inp->pool_bias, n_pool, nc, ns, + inp->pool_bias->nb[1], inp->pool_bias->nb[2], + (size_t) t0*inp->pool_bias->nb[1]); + if (tail_c) { + tail_c = ggml_view_4d(ctx0, inp->tail_cells, r - 1, nc, 1, ns, + inp->tail_cells->nb[1], inp->tail_cells->nb[2], inp->tail_cells->nb[3], + (size_t) t0*inp->tail_cells->nb[1]); + } + } + + ggml_tensor * score = ggml_mul_mat(ctx0, pooled, q_c); + score = ggml_relu(ctx0, ggml_reshape_4d(ctx0, score, n_pool, nh, nc, ns)); + + score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); + score = ggml_sum_rows(ctx0, ggml_mul(ctx0, score, wts_c)); + score = ggml_reshape_3d(ctx0, score, n_pool, nc, ns); - // member j of the i-th selected pool is at i*r + j in both layouts, so this is a reshape - top_k = ggml_reshape_4d(ctx0, top_k, r*n_sel, n_tps, 1, ns); + // the cut is on whole pools, never on single cells. Scoring cells with their pool's + // score and cutting there is not the same thing: relu sends many distinct pools to + // exactly 0.0 and ggml_top_k is unordered among equal keys, so the cut splits pools + // apart. Diagnosis and the reference-free check for it (count partly selected pools) + // are from PR #27754. + score = ggml_add(ctx0, score, bias_c); + + ggml_tensor * sel = ggml_top_k(ctx0, score, n_sel); + + // only meaningful when the loop runs once; otherwise an eval-callback dump would get + // one identically-named tensor per chunk per layer + if (nc == n_tps) { + cb(score, "indexer_score_pools", il); + cb(sel, "indexer_top_k_pools", il); + } + + ggml_tensor * tk = ggml_get_rows(ctx0, pools, + ggml_reshape_3d(ctx0, sel, n_sel*nc, ns, 1)); + + // member j of the i-th selected pool is at i*r + j in both layouts, so this is a reshape + tk = ggml_reshape_4d(ctx0, tk, r*n_sel, nc, 1, ns); + + // index_kpool_always_select_tail: the trailing incomplete pool has no pool key and can + // never be picked above, so its cells are appended instead of taking pool budget + if (tail_c) { + tk = ggml_concat(ctx0, tk, tail_c, 0); + } - // index_kpool_always_select_tail: the trailing incomplete pool has no pool key and can - // never be picked above, so its cells are appended instead of taking pool budget - if (inp->tail_cells) { - top_k = ggml_concat(ctx0, top_k, inp->tail_cells, 0); + // appends this chunk's tokens along the token axis; recopies earlier chunks each + // iteration, which is negligible I32 traffic next to the scoring GEMMs + top_k = top_k ? ggml_concat(ctx0, top_k, tk, 1) : tk; } // build_attn_mask_top_k reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask From 0ccfa1516435da030af5922aab5ff5382f64bf75 Mon Sep 17 00:00:00 2001 From: eauchs Date: Sat, 29 Aug 2026 15:57:11 +0200 Subject: [PATCH 08/11] llama : skip foreign-sequence cells when building the glm5next k-pools set_input_kpool only skipped empty cells. A unified cache holds every sequence in one cell array, so a cell belonging to another sequence landed in this stream's pools and collided with its own cell at the same position: cur_pool_cells[b*r + (p%r)] was overwritten and filled[b] counted twice. The tail loop thirty lines below already filtered this way, so the omission was an oversight rather than a choice. This is reachable from a plain llama-server invocation: leaving --parallel unset selects auto slots, which set n_parallel = 4 and kv_unified = true. test_dsa_kpool grows three passes after the existing one -- two sequences on separate caches, two sequences on a unified cache, and a single sequence on a unified cache as the control. At -s 1234, identical on Metal, Accelerate and the second device: whole partial tail misses 2 seq, separate before 216 0 0 (unchanged after) 2 seq, unified before 108 48 72 after 236 36 96 1 seq, unified 236 0 0 (control) Whole pools return to 236, the single-stream value. The 36 partial pools and 96 tail misses that remain are not interpretable with this metric: count_row assumes cell index == position and filters on c <= q, and on a unified cache the second sequence occupies cells 32..63 at positions 0..31, so those are discarded by construction. What remains cannot be separated from the blindness of the measurement. The third pass is there precisely to establish that unified addressing on its own is sound. The three new passes are display-only and do not fail the test. --- src/llama-memory-hybrid-idx.cpp | 5 +++- tests/test-llama-archs.cpp | 51 ++++++++++++++++++++++++++++++++- 2 files changed, 54 insertions(+), 2 deletions(-) diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index ca8c4c8aca5d..b92e970ce846 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -734,7 +734,10 @@ void llama_memory_hybrid_idx_context::set_input_kpool( std::fill(cur_pool_cells, cur_pool_cells + r*n_pool, 0); for (int64_t j = 0; j < n_kv; ++j) { - if (cells.is_empty(j)) { + // a unified cache holds every sequence in one cell array, so a cell from another + // sequence would land in this stream's pools and collide with its own cell at the + // same position. The tail loop below already filters this way. + if (cells.is_empty(j) || !cells.seq_has(j, seq_of_stream)) { continue; } diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index ac7dc64747b4..dc50d0a6e01e 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -400,7 +400,7 @@ static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, ggml_backend_sched_eval_callback cb_eval = nullptr, void * cb_eval_user_data = nullptr, - bool load_mtp = false) { + bool load_mtp = false, uint32_t n_seq_max = 1, bool kv_unified = false) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -419,6 +419,8 @@ static std::pair get_model_and_ctx( } ctx_params.cb_eval = cb_eval; ctx_params.cb_eval_user_data = cb_eval_user_data; + ctx_params.n_seq_max = n_seq_max; + ctx_params.kv_unified = kv_unified; size_t tmp = seed; llama_model_ptr model(gguf_ctx != nullptr ? @@ -851,6 +853,53 @@ static int test_dsa_kpool(const size_t seed, const int verbosity) { fflush(stdout); } + // Second pass: two sequences on a unified cache. This is what llama-server does by + // default - leaving --parallel unset selects auto slots, which set n_parallel = 4 and + // kv_unified = true. The single-stream pass above never exercises it. + // control: the same two sequences on SEPARATE streams. There cell index == position holds + // again per stream, so a metric artefact shows up here too while a real unified-cache + // pool-mixing bug does not. + // cfg 0/1: two sequences, separate then unified. cfg 2: unified cache but a SINGLE + // decoding sequence - there cell index == position still holds, so the metric stays valid + // and any degradation is the unified addressing itself rather than the cross-sequence mix. + for (int cfg = 0; cfg < 3; cfg++) { + const bool unified = cfg != 0; + const bool two_seq = cfg != 2; + printf("test_dsa_kpool: glm5next, %s, kv_unified=%s\n", + two_seq ? "2 seqs x 32 tokens" : "1 seq x 64 tokens", unified ? "true" : "false"); + + for (const auto & dc : dev_configs) { + dsa_kpool_check st; + st.r = DSA_INDEXER_KPOOL; + + auto model_and_ctx = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, dc.first, + LLAMA_SPLIT_MODE_LAYER, false, dsa_kpool_eval_cb, &st, false, 2, unified); + + const uint32_t n_half = two_seq ? n_tokens/2 : n_tokens; + llama_batch batch = llama_batch_init(n_tokens, 0, 2); + for (uint32_t pos = 0; pos < n_half; pos++) { + common_batch_add(batch, tokens[pos], pos, {0}, true); + } + if (two_seq) { + for (uint32_t pos = 0; pos < n_half; pos++) { + common_batch_add(batch, tokens[n_half + pos], pos, {1}, true); + } + } + batch.n_tokens = n_tokens; + const int32_t rc = llama_decode(model_and_ctx.second.get(), batch); + llama_batch_free(batch); + if (rc != 0) { + throw std::runtime_error("failed to decode batch"); + } + + printf("test_dsa_kpool: %-32s rows %5" PRId64 ", whole pools %6" PRId64 ", " + "partial pools %6" PRId64 " in %4" PRId64 " rows, tail misses %4" PRId64 "\n", + dc.second.c_str(), st.n_row, st.n_pool_whole, st.n_pool_partial, st.n_row_partial, + st.n_tail_missing); + fflush(stdout); + } + } + llama_log_set(ud.log_old.callback, ud.log_old.user_data); return all_ok ? 0 : 1; } From c6cbbf993dfc0d8edf50b9f542766dcfaef9a544 Mon Sep 17 00:00:00 2001 From: eauchs Date: Tue, 1 Sep 2026 13:41:50 +0200 Subject: [PATCH 09/11] llama : glm5next fold (d, n_pool) into one ne[1] ahead of the pool softmax softmax normalizes ne[0], so the member axis has to sit there; the permute already puts it there. (d, n_pool) are adjacent on the contiguous tensor, so folding them into a single ne[1] leaves every row identical while moving n_pool off gridDim.y (capped at 65535 on CUDA) onto gridDim.x. Measured by matteoscalabrini on 5x RTX 3090: n_pool = 65536 goes from CUDA error: invalid argument to a coherent decode, greedy output hash-identical (0c5e1504...), prefill 142.6 vs 142.5 t/s. --- src/models/glm5next.cpp | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 93b4802972e6..50d8cf0a93ed 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -750,8 +750,12 @@ ggml_tensor * llama_model_glm5next::graph::build_dsa_top_k( // an intra-pool position bias and, with no rope anywhere, the only ordering signal here. m_g = ggml_add(ctx0, m_g, ggml_get_rows(ctx0, layer.indexer_comp_ape, inp->ape_slots)); - // softmax normalizes ne[0], so bring the member axis there - ggml_tensor * w = ggml_soft_max(ctx0, ggml_cont(ctx0, ggml_permute(ctx0, m_g, 1, 0, 2, 3))); + // softmax normalizes ne[0], so bring the member axis there. (d, n_pool) are adjacent on + // the contiguous tensor, so folding them into one ne[1] leaves every row identical while + // moving n_pool off gridDim.y (65535 on CUDA) onto gridDim.x + ggml_tensor * wc = ggml_cont(ctx0, ggml_permute(ctx0, m_g, 1, 0, 2, 3)); + ggml_tensor * w = ggml_reshape_4d(ctx0, + ggml_soft_max(ctx0, ggml_reshape_3d(ctx0, wc, r, d*n_pool, ns)), r, d, n_pool, ns); ggml_tensor * v = ggml_cont(ctx0, ggml_permute(ctx0, m_k, 1, 0, 2, 3)); ggml_tensor * pooled = ggml_sum_rows(ctx0, ggml_mul(ctx0, v, w)); From dd89075d59c313a6c2287ad7f38a65b6db86ddca Mon Sep 17 00:00:00 2001 From: eauchs Date: Tue, 1 Sep 2026 13:42:06 +0200 Subject: [PATCH 10/11] llama : glm5next declare ssm_a as NOSCAN ssm_a is multiplied as a gate, not scanned -- qwen3next already declares NOSCAN for the same reason, and kimi-linear has the same inconsistency. Both spellings write blk.N.ssm_a, so no GGUF breaks. --- src/models/glm5next.cpp | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index 50d8cf0a93ed..f62ea8612fa7 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -222,8 +222,10 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", i), {head_dim, d_inner}, 0); layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); - // ssm_a holds -exp(A_log), folded at conversion time (kimi-linear/kimi-k3 convention) - layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {n_head}, 0); + // ssm_a holds -exp(A_log), folded at conversion time (kimi-linear/kimi-k3 convention). + // NOSCAN because the gate multiplies by it instead of running an SSM scan, as qwen3next + // does; both spellings write blk.N.ssm_a, only the declared op differs + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, 0); // some converters emit dt_bias under the default ".weight" suffix layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, TENSOR_NOT_REQUIRED); From 1d0c76f3c6d030fdfc269aa27db6334ea2834cec Mon Sep 17 00:00:00 2001 From: eauchs Date: Thu, 3 Sep 2026 17:12:42 +0200 Subject: [PATCH 11/11] llama : glm5next adapt to master's n_kv_max attn API and per-layer n_ff_exp master added an n_kv_max parameter to build_attn_mha; the DSA overload of build_attn passes top_k->ne[0] when the selection is sparse and 0 on the dense path, like the other call sites in the file. n_ff_exp became a per-layer accessor: LLM_KV_EXPERT_FEED_FORWARD_LENGTH is now read via get_key_or_arr into n_ff_exp_arr over n_layer_all, as done in deepseek4. --- src/llama-graph.cpp | 3 ++- src/models/glm5next.cpp | 4 ++-- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index a9cf1764203c..9d6db15b80f4 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -3045,7 +3045,8 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * k = mctx_cur->get_k(ctx0, il); ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); - ggml_tensor * cur = build_attn_mha(q_cur, k, v, kq_b, kq_mask, sinks, v_mla, kq_scale, il); + ggml_tensor * cur = build_attn_mha(q_cur, k, v, kq_b, kq_mask, sinks, v_mla, + top_k ? top_k->ne[0] : 0, kq_scale, il); cb(cur, "kqv_out", il); if (wo) { diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp index f62ea8612fa7..5c9bceacd488 100644 --- a/src/models/glm5next.cpp +++ b/src/models/glm5next.cpp @@ -50,7 +50,7 @@ void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { GGML_ASSERT(hparams.n_rot() == 0 && "GLM5NEXT is nope-only: rope.dimension_count must be 0"); // MoE - ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); @@ -146,7 +146,7 @@ void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { const int64_t q_lora_rank = hparams.n_lora_q; const int64_t kv_lora_rank = hparams.n_lora_kv; - const int64_t n_ff_exp = hparams.n_ff_exp; + const int64_t n_ff_exp = hparams.n_ff_exp(); const int64_t n_expert_shared = std::max(1, hparams.n_expert_shared); const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();