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..2b2eeb454860 --- /dev/null +++ b/conversion/glm5next.py @@ -0,0 +1,106 @@ +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. + # 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, + **{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 + 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..2b2de3a8a0e3 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 of 45 layers), see 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), 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, + # --- indexer: k-pool compression --- + 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/llama-arch.cpp b/src/llama-arch.cpp index 446de4ae25b3..af08544fa185 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" }, @@ -282,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" }, @@ -1080,6 +1082,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 +1147,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..cdda15c338c6 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, @@ -287,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/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..9d6db15b80f4 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,81 @@ 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, + top_k ? top_k->ne[0] : 0, 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 +3090,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 +3641,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-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 93b468784a33..b92e970ce846 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,150 @@ 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) { + // 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; + } + + 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-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..d50524fe98e3 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,72 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, nullptr); } } 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); }; + 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); + }; + } + + 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 + /* 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) @@ -2537,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, @@ -2837,6 +2910,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..5c9bceacd488 --- /dev/null +++ b/src/models/glm5next.cpp @@ -0,0 +1,1261 @@ +#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. +// +// 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) { + // 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_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); + 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_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); + 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. + // 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); + 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 & ml) { + 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; + + // 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); + 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). + // 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); + 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. 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); + 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); + + // 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 { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } + + 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, 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, kpool); + + 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 &= 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; + } + + 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 * pool_cells = nullptr; // I32 [kpool*n_pools, 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; + + const uint32_t kpool; +}; + +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, (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); + kp->pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_pool, 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->pool_cells); + ggml_set_input(kp->pool_bias); + + 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); + } + + 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 = 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; + + // 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. (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)); + 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); + + // 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); + + // 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); + + // 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); + + // 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); + + // 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); + } + + // 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 + 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); + } + + // 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 && !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); + } + + 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); + + // 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; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + 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 9b87a40d5af9..4f97125d42cc 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1334,6 +1334,63 @@ 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); + }; + + // 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; +}; + 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..dc50d0a6e01e 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 @@ -68,6 +71,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); @@ -79,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; @@ -111,6 +120,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 @@ -143,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; @@ -213,6 +231,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)); @@ -286,20 +311,22 @@ 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})); - 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); @@ -371,10 +398,13 @@ 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, + 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; + model_params.load_mtp = load_mtp; std::vector devs_copy = devs; devs_copy.push_back(nullptr); model_params.devices = devs_copy.data(); @@ -387,6 +417,10 @@ 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; + 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 ? @@ -458,6 +492,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: @@ -617,6 +652,421 @@ 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); + } + + // 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; +} + +// +// 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 { @@ -850,7 +1300,13 @@ 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_mtp(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;