diff --git a/conversion/__init__.py b/conversion/__init__.py index 8de97e95969..6f23240deb1 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -103,6 +103,8 @@ "Glm4MoeLiteForCausalLM": "glm", "Glm4vForConditionalGeneration": "glm", "Glm4vMoeForConditionalGeneration": "glm", + "Glm5NextForCausalLM": "glm5next", + "Glm5NextForConditionalGeneration": "glm5next", "GlmForCausalLM": "chatglm", "GlmMoeDsaForCausalLM": "glm", "GlmOcrForConditionalGeneration": "glm", @@ -293,6 +295,7 @@ "Gemma4UnifiedForConditionalGeneration": "gemma", "Glm4vForConditionalGeneration": "qwen3vl", "Glm4vMoeForConditionalGeneration": "qwen3vl", + "Glm5NextForConditionalGeneration": "glm5next", "Glm5vForConditionalGeneration": "kimivl", "GlmOcrForConditionalGeneration": "qwen3vl", "GlmasrModel": "ultravox", diff --git a/conversion/base.py b/conversion/base.py index 56547ace009..3ff5634407c 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -2537,6 +2537,20 @@ def set_gguf_parameters(self): if not self.has_vision_encoder and not self.has_audio_encoder: raise ValueError("MmprojModel must have either vision or audio encoder") + def prepare_tensors(self): + super().prepare_tensors() + + # an mmproj has no vocab-only mode, so an empty one is always a silent mapping failure + # rather than a supported output; count and size are both checked because a tensor map + # that produces only zero-sized entries is just as broken as one that produces none + n_tensors = sum(len(t) for t in self.gguf_writer.tensors) + n_bytes = sum(ti.nbytes for t in self.gguf_writer.tensors for ti in t.values()) + if n_tensors == 0 or n_bytes == 0: + raise ValueError( + f"refusing to write an mmproj with no tensor data (n_tensors = {n_tensors}, n_bytes = {n_bytes}); " + "check that the model's vision/audio tensors are named as the tensor map expects" + ) + def write_vocab(self): raise ValueError("MmprojModel does not support vocab writing") diff --git a/conversion/glm5next.py b/conversion/glm5next.py new file mode 100644 index 00000000000..f2db5867765 --- /dev/null +++ b/conversion/glm5next.py @@ -0,0 +1,273 @@ +from __future__ import annotations + +import re + +from typing import Callable, Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import LazyTorchTensor, ModelBase, TextModel, gguf +from .qwen3vl import Glm4VVisionModel + + +@ModelBase.register("Glm5NextForConditionalGeneration", "Glm5NextForCausalLM") +# [TAG_HF_EXAMPLE_MISSING] +class Glm5NextModel(TextModel): + """GLM-5.3-Flash text tower: hybrid KDA + DSA attention, nope-only MLA, mHC + hyper-connections, and a NextN block with its own DSA attention and indexer. + """ + + model_arch = gguf.MODEL_ARCH.GLM5NEXT + supports_mtp_export = True + + _experts: list[dict[str, Tensor]] | None = None + _main_layers: int | None = None + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + nextn_layers = 0 if self.no_mtp else (self.hparams.get("num_nextn_predict_layers", 0) or 0) + self.block_count = self.hparams["num_hidden_layers"] + nextn_layers + self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) + + # two independent spellings of the same partition; disagreement means an + # unexpected config + from_types = {il for il, t in enumerate(self.hparams["layer_types"]) if t == "deepseek_sparse_attention"} + from_list = set(self.hparams["linear_attn_config"]["full_attn_layers"]) + if from_types != from_list: + raise ValueError(f"layer_types picks DSA layers {sorted(from_types)} but full_attn_layers says {sorted(from_list)}") + self._full_attn_layers = from_types + + dense_lead = self.hparams["first_k_dense_replace"] + expected = ["dense"] * dense_lead + ["sparse"] * (self.hparams["num_hidden_layers"] - dense_lead) + if self.hparams["mlp_layer_types"] != expected: + raise ValueError("mlp_layer_types does not match first_k_dense_replace") + + def index_tensors(self, remote_hf_model_id: str | None = None): + # runs before TextModel.__init__ has hoisted text_config to the root + hp = self.hparams.get("text_config", self.hparams) + type(self)._main_layers = hp["num_hidden_layers"] + return super().index_tensors(remote_hf_model_id=remote_hf_model_id) + + def set_vocab(self): + self._set_vocab_glm() + + def is_full_attention(self, bid: int) -> bool: + return bid >= self.hparams["num_hidden_layers"] or bid in self._full_attn_layers + + # -- metadata --------------------------------------------------------- + + def set_gguf_parameters(self): + hp = self.hparams + linear_cfg = hp["linear_attn_config"] + + # checked here, not in the loader: head_count_kv is overwritten below with + # the per-layer 1/0 recurrence marker + if hp["num_attention_heads"] != hp.get("num_key_value_heads"): + raise ValueError("glm5next expects MHA-shaped head counts before MLA absorption") + if hp["qk_rope_head_dim"] != 0 or not hp.get("mla_use_nope"): + raise ValueError("glm5next is nope-only: qk_rope_head_dim must be 0 and mla_use_nope true") + if linear_cfg["num_heads"] != hp["num_attention_heads"]: + raise ValueError("glm5next KDA and full attention are expected to share a head count") + if not hp.get("mhc"): + raise ValueError("glm5next without mHC is not supported") + if hp["index_topk"] % hp["index_kpool"] != 0: + raise ValueError("glm5next index_topk must be a whole number of kpool pools") + + # no GGUF key carries these and the graph cannot express them off, so refuse + # rather than write a silently wrong file + if not hp.get("index_kpool_compress"): + raise ValueError("glm5next without the indexer kpool compressor is not supported") + if not hp.get("index_kpool_always_select_tail"): + raise ValueError("glm5next without always-select-tail kpool is not supported") + if not hp.get("indexer_rope_interleave"): + raise ValueError("glm5next without interleaved indexer rope is not supported") + if set(hp["indexer_types"]) != {"full"}: + raise ValueError("glm5next expects every indexer to be full") + + # drop both: head_dim is 0 in the config, head_count_kv is written as a + # per-layer array below + hp.pop("head_dim", None) + hp.pop("num_key_value_heads", None) + + super().set_gguf_parameters() + + self.gguf_writer.add_vocab_size(hp["vocab_size"]) + + # n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-k3 and bailingmoe3 + self.gguf_writer.add_head_count_kv( + [1 if self.is_full_attention(il) else 0 for il in range(self.block_count)]) + + # --- MLA --- + kv_lora_rank = hp["kv_lora_rank"] + qk_rope_head_dim = hp["qk_rope_head_dim"] + self.gguf_writer.add_q_lora_rank(hp["q_lora_rank"]) + self.gguf_writer.add_kv_lora_rank(kv_lora_rank) + self.gguf_writer.add_rope_dimension_count(qk_rope_head_dim) + self.gguf_writer.add_key_length(kv_lora_rank + qk_rope_head_dim) + self.gguf_writer.add_value_length(kv_lora_rank) + self.gguf_writer.add_key_length_mla(hp["qk_nope_head_dim"] + qk_rope_head_dim) + self.gguf_writer.add_value_length_mla(hp["v_head_dim"]) + + # indexer k_norm is a LayerNorm with bias at a fixed 1e-6, not the model's + # RMS eps. glm-dsa omits this key and runs that norm at eps 0 + self.gguf_writer.add_layer_norm_eps(1e-6) + + # --- KDA --- + self.gguf_writer.add_ssm_conv_kernel(linear_cfg["short_conv_kernel_size"]) + self.gguf_writer.add_kda_head_dim(linear_cfg["head_dim"]) + # not a clamp: scales the sigmoid decay gate. required, a missing key + # silently selects the softplus branch instead + self.gguf_writer.add_kda_gate_lower_bound(linear_cfg["gate_lower_bound"]) + + # --- DSA indexer --- + self.gguf_writer.add_indexer_head_count(hp["index_n_heads"]) + self.gguf_writer.add_indexer_key_length(hp["index_head_dim"]) + self.gguf_writer.add_indexer_top_k(hp["index_topk"]) + self.gguf_writer.add_indexer_kpool(hp["index_kpool"]) + + # --- mHC --- + self.gguf_writer.add_hyper_connection_count(hp["hc_mult"]) + self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hp["hc_sinkhorn_iters"]) + self.gguf_writer.add_hyper_connection_epsilon(hp["hc_eps"]) + + # --- MoE --- + n_ff_exp = hp["moe_intermediate_size"] + self.gguf_writer.add_expert_feed_forward_length(n_ff_exp) + self.gguf_writer.add_expert_shared_feed_forward_length(n_ff_exp * hp["n_shared_experts"]) + self.gguf_writer.add_expert_shared_count(hp["n_shared_experts"]) + self.gguf_writer.add_leading_dense_block_count(hp["first_k_dense_replace"]) + self.gguf_writer.add_expert_weights_scale(hp["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(hp["norm_topk_prob"]) + + # one limit for the whole model. no dense-FFN clamp key exists, so the + # expert arrays are sized for every layer to cover the leading dense ones + swiglu_limit = float(hp["swiglu_limit"]) + self.gguf_writer.add_swiglu_clamp_exp([swiglu_limit] * self.block_count) + self.gguf_writer.add_swiglu_clamp_shexp([swiglu_limit] * self.block_count) + + if not self.no_mtp and (nextn_layers := hp.get("num_nextn_predict_layers", 0)): + self.gguf_writer.add_nextn_predict_layers(nextn_layers) + + # -- tensors ---------------------------------------------------------- + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + if (titem := super().filter_tensors(item)) is None: + return None + name, gen = titem + + assert cls._main_layers is not None + is_mtp = (m := re.match(r"model\.layers\.(\d+)\.", name)) is not None and int(m.group(1)) >= cls._main_layers + + if is_mtp and cls.no_mtp: + return None + if cls.mtp_only and not is_mtp and name not in ( + "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", + ): + return None + + return name, gen + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # --- KDA conv1d: HF [d_inner, 1, d_conv] -> ggml ne [d_conv, 1, d_inner, 1] --- + if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")): + d_inner = data_torch.shape[0] + d_conv = data_torch.shape[-1] + data_torch = data_torch.reshape(1, d_inner, 1, d_conv) + + # ssm_a holds -exp(A_log), the kimi-k3 convention (bailingmoe3 stores + # +exp(A_log)); the wrong sign turns decay into an unchecked growing state + if name.endswith(".A_log"): + # eager: the sign is the point of the check, and A_log is one per head + decay = LazyTorchTensor.to_eager(torch.exp(data_torch.float())) + if not bool(torch.isfinite(decay).all() and (decay > 0).all()): + raise ValueError(f"{name}: exp(A_log) must be finite and positive") + data_torch = -decay + + # dt_bias -> the name SSM_DT's mapping expects + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + + # bare tensors in the checkpoint, but the GGUF names carry .weight + if re.search(r"\.hc_(attn|ffn)_(fn|base|scale)$", name) or name.endswith( + (".index_kpool_compress_gate", ".index_kpool_compress_ape")): + name += ".weight" + + # --- routed experts --- + if ".mlp.experts." in name: + n_experts = self.hparams["n_routed_experts"] + assert bid is not None + + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + self._experts[bid][name] = data_torch + + if len(self._experts[bid]) < n_experts * 3: + return + + for weight_name in ("down_proj", "gate_proj", "up_proj"): + tensors = [] + for expert_id in range(n_experts): + expert_name = f"model.layers.{bid}.mlp.experts.{expert_id}.{weight_name}.weight" + tensors.append(self._experts[bid].pop(expert_name)) + merged_name = f"model.layers.{bid}.mlp.experts.{weight_name}.weight" + yield from super().modify_tensors(torch.stack(tensors, dim=0), merged_name, bid) + return + + # --- MLA absorption --- + if name.endswith(".kv_b_proj.weight"): + assert bid is not None + n_head = self.hparams["num_attention_heads"] + v_head_dim = self.hparams["v_head_dim"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + assert data_torch.shape[0] == n_head * (v_head_dim + qk_nope_head_dim) + kv_b = data_torch.view(n_head, v_head_dim + qk_nope_head_dim, data_torch.shape[-1]) + k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1) + yield from super().modify_tensors(k_b.transpose(1, 2), self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_B, bid), bid) + yield from super().modify_tensors(v_b, self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V_B, bid), bid) + return + + yield from super().modify_tensors(data_torch, name, bid) + + def tensor_force_quant(self, name, new_name, bid, n_dims): + # learned position table, one row per pooled key; pinned for the same + # reason POS_EMBD is in base.py + if self.match_model_tensor_name(new_name, gguf.MODEL_TENSOR.INDEXER_COMPRESSOR_APE, bid): + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts is not None: + experts = [name for layer in self._experts for name in layer] + if experts: + raise ValueError(f"Unprocessed experts: {experts}") + + +@ModelBase.register("Glm5NextForConditionalGeneration") +# [TAG_HF_EXAMPLE_MISSING] +class Glm5NextVisionModel(Glm4VVisionModel): + """The vision tower is the GLM-OCR ViT under a `model.visual.` prefix. + + Every tensor already maps through the GLM-4V entries. The one structural + difference is a clamp on the SwiGLU gate and up projections, applied in the + per-block MLP and again in the merger, so it gets its own projector type + rather than a flag on glm4v. + """ + + clip_projector_type = gguf.VisionProjectorType.GLM5NEXT + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_vision is not None + + # Glm4VVisionModel bypasses Qwen3VLVisionModel entirely, which is also where + # the merge size is written, so no GLM4V-family mmproj carries this key and + # clip.cpp falls back to a hardcoded 2. Write it rather than rely on that + self.gguf_writer.add_vision_spatial_merge_size(int(self.hparams_vision["spatial_merge_size"])) + + self.gguf_writer.add_vision_swiglu_limit(float(self.hparams_vision["swiglu_limit"])) diff --git a/conversion/qwen3vl.py b/conversion/qwen3vl.py index 4fec708c9ff..385f47149fc 100644 --- a/conversion/qwen3vl.py +++ b/conversion/qwen3vl.py @@ -228,10 +228,13 @@ class Qwen3ASRMmprojModel(Qwen3OmniMmprojModel): @ModelBase.register("Glm4vForConditionalGeneration", "Glm4vMoeForConditionalGeneration", "GlmOcrForConditionalGeneration") @ModelBase.example("zai-org/GLM-4.1V-9B-Thinking", "zai-org/GLM-4.5V") class Glm4VVisionModel(Qwen3VLVisionModel): + # subclasses that share this tower but need their own clip graph override this + clip_projector_type = gguf.VisionProjectorType.GLM4V + def set_gguf_parameters(self): MmprojModel.set_gguf_parameters(self) # skip Qwen3VLVisionModel parameters assert self.hparams_vision is not None - self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.GLM4V) + self.gguf_writer.add_clip_projector_type(self.clip_projector_type) hidden_act = str(self.hparams_vision.get("hidden_act", "")).lower() if hidden_act == "gelu": diff --git a/examples/embedding/embedding.cpp b/examples/embedding/embedding.cpp index f6a20ef9d07..a30d271253f 100644 --- a/examples/embedding/embedding.cpp +++ b/examples/embedding/embedding.cpp @@ -144,6 +144,12 @@ int main(int argc, char ** argv) { return 1; } + // a context can fail on its own, and llama_n_ctx below dereferences it + if (ctx == NULL) { + LOG_ERR("%s: unable to create context\n", __func__); + return 1; + } + const llama_vocab * vocab = llama_model_get_vocab(model); const int n_ctx_train = llama_model_n_ctx_train(model); diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index f236a5d2c98..cf28513db38 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -220,6 +220,7 @@ class Indexer: BLOCK_SIZE = "{arch}.attention.indexer.block_size" # MSA LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" # MSA TYPES = "{arch}.attention.indexer.types" + KPOOL = "{arch}.attention.indexer.kpool" # glm5next class HyperConnection: COUNT = "{arch}.hyper_connection.count" @@ -364,6 +365,7 @@ class ClipVision: IMAGE_MEAN = "clip.vision.image_mean" IMAGE_STD = "clip.vision.image_std" SPATIAL_MERGE_SIZE = "clip.vision.spatial_merge_size" + SWIGLU_LIMIT = "clip.vision.swiglu_limit" # glm5next: clamp on the SwiGLU gate/up EXPERT_COUNT_PER_LAYER = "clip.vision.expert_count_per_layer" # dots3note pyramid MoE, 0 = dense layer EXPERT_USED_COUNT = "clip.vision.expert_used_count" USE_GELU = "clip.use_gelu" @@ -540,6 +542,7 @@ class MODEL_ARCH(IntEnum): GLM4 = auto() GLM4_MOE = auto() GLM_DSA = auto() + GLM5NEXT = auto() BITNET = auto() T5 = auto() T5ENCODER = auto() @@ -1263,6 +1266,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GLM4: "glm4", MODEL_ARCH.GLM4_MOE: "glm4moe", MODEL_ARCH.GLM_DSA: "glm-dsa", + MODEL_ARCH.GLM5NEXT: "glm5next", MODEL_ARCH.BITNET: "bitnet", MODEL_ARCH.T5: "t5", MODEL_ARCH.T5ENCODER: "t5encoder", @@ -3871,6 +3875,70 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.GLM5NEXT: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.FFN_NORM, + # mHC, layered on top of the per-layer norms above + MODEL_TENSOR.HC_ATTN_FN, + MODEL_TENSOR.HC_ATTN_BASE, + MODEL_TENSOR.HC_ATTN_SCALE, + MODEL_TENSOR.HC_FFN_FN, + MODEL_TENSOR.HC_FFN_BASE, + MODEL_TENSOR.HC_FFN_SCALE, + # KDA (linear-attention layers) + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.SSM_CONV1D_Q, + MODEL_TENSOR.SSM_CONV1D_K, + MODEL_TENSOR.SSM_CONV1D_V, + MODEL_TENSOR.SSM_F_A, + MODEL_TENSOR.SSM_F_B, + MODEL_TENSOR.SSM_G_A, + MODEL_TENSOR.SSM_G_B, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + # DSA (MLA full-attention layers) + MODEL_TENSOR.ATTN_Q_A, + MODEL_TENSOR.ATTN_Q_B, + MODEL_TENSOR.ATTN_Q_A_NORM, + MODEL_TENSOR.ATTN_KV_A_MQA, + MODEL_TENSOR.ATTN_KV_A_NORM, + MODEL_TENSOR.ATTN_K_B, + MODEL_TENSOR.ATTN_V_B, + MODEL_TENSOR.ATTN_OUT, + # DSA indexer, with the kpool key compressor + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.INDEXER_PROJ, + MODEL_TENSOR.INDEXER_ATTN_K, + MODEL_TENSOR.INDEXER_ATTN_Q_B, + MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE, + MODEL_TENSOR.INDEXER_COMPRESSOR_APE, + # FFN: dense on the leading blocks, MoE elsewhere + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + # NextN/MTP, a full DSA decoder layer with its own indexer + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, + ], MODEL_ARCH.BITNET: [ MODEL_TENSOR.ATTN_Q, MODEL_TENSOR.ATTN_K, @@ -5521,6 +5589,7 @@ class VisionProjectorType: LFM2A = "lfm2a" # audio MUSIC_FLAMINGO = "musicflamingo" # audio GLM4V = "glm4v" + GLM5NEXT = "glm5next" YOUTUVL = "youtuvl" NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index d8a96a27bdd..b60806071fb 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -812,6 +812,9 @@ def add_indexer_block_size(self, block_size: int) -> None: def add_indexer_local_blocks(self, local_blocks: int) -> None: self.add_uint32(Keys.Attention.Indexer.LOCAL_BLOCKS.format(arch=self.arch), local_blocks) + def add_indexer_kpool(self, kpool: int) -> None: + self.add_uint32(Keys.Attention.Indexer.KPOOL.format(arch=self.arch), kpool) + def add_indexer_types(self, value: Sequence[bool]) -> None: key = Keys.Attention.Indexer.TYPES.format(arch=self.arch) self.add_array(key, value) @@ -1327,6 +1330,9 @@ def add_vision_image_std(self, values: Sequence[float]) -> None: def add_vision_spatial_merge_size(self, value: int) -> None: self.add_uint32(Keys.ClipVision.SPATIAL_MERGE_SIZE, value) + def add_vision_swiglu_limit(self, value: float) -> None: + self.add_float32(Keys.ClipVision.SWIGLU_LIMIT, value) + def add_vision_expert_count_per_layer(self, value: Sequence[int]) -> None: self.add_array(Keys.ClipVision.EXPERT_COUNT_PER_LAYER, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index ef580518e97..050d96391ef 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2680,6 +2680,34 @@ class TensorNameMap: "model.layers.{bid}.post_attention_layernorm", ), }, + MODEL_ARCH.GLM5NEXT: { + # the converter appends the .weight the checkpoint omits, so these + # match through the usual try_suffixes path + MODEL_TENSOR.HC_ATTN_FN: ( + "model.layers.{bid}.hc_attn_fn", + ), + MODEL_TENSOR.HC_ATTN_BASE: ( + "model.layers.{bid}.hc_attn_base", + ), + MODEL_TENSOR.HC_ATTN_SCALE: ( + "model.layers.{bid}.hc_attn_scale", + ), + MODEL_TENSOR.HC_FFN_FN: ( + "model.layers.{bid}.hc_ffn_fn", + ), + MODEL_TENSOR.HC_FFN_BASE: ( + "model.layers.{bid}.hc_ffn_base", + ), + MODEL_TENSOR.HC_FFN_SCALE: ( + "model.layers.{bid}.hc_ffn_scale", + ), + MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: ( + "model.layers.{bid}.self_attn.indexer.index_kpool_compress_gate", + ), + MODEL_TENSOR.INDEXER_COMPRESSOR_APE: ( + "model.layers.{bid}.self_attn.indexer.index_kpool_compress_ape", + ), + }, } mapping: dict[str, tuple[MODEL_TENSOR, str]] diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c6df19f2ecf..fceddfdb479 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -26,6 +26,7 @@ add_library(llama llama-kv-cache-iswa.cpp llama-kv-cache-dsa.cpp llama-kv-cache-dsa-iswa.cpp + llama-kv-cache-kpool.cpp llama-kv-cache-msa.cpp llama-kv-cache-dsv4.cpp llama-memory.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index eecf444fcf3..076dbfc2c89 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -83,6 +83,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GLM4, "glm4" }, { LLM_ARCH_GLM4_MOE, "glm4moe" }, { LLM_ARCH_GLM_DSA, "glm-dsa" }, + { LLM_ARCH_GLM5NEXT, "glm5next" }, { LLM_ARCH_BITNET, "bitnet" }, { LLM_ARCH_T5, "t5" }, { LLM_ARCH_T5ENCODER, "t5encoder" }, @@ -283,6 +284,7 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" }, { LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" }, { LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" }, + { LLM_KV_ATTENTION_INDEXER_KPOOL, "%s.attention.indexer.kpool" }, { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" }, { LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" }, { LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" }, @@ -1010,6 +1012,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_MINIMAX_01: return true; default: @@ -1034,6 +1037,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) { case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_LFM2: @@ -1072,6 +1076,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN3TTS: diff --git a/src/llama-arch.h b/src/llama-arch.h index 7159e23bf7a..cadbdfff7ef 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -88,6 +88,7 @@ enum llm_arch { LLM_ARCH_GLM4, LLM_ARCH_GLM4_MOE, LLM_ARCH_GLM_DSA, + LLM_ARCH_GLM5NEXT, LLM_ARCH_BITNET, LLM_ARCH_T5, LLM_ARCH_T5ENCODER, @@ -288,6 +289,7 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, LLM_KV_ATTENTION_INDEXER_TYPES, + LLM_KV_ATTENTION_INDEXER_KPOOL, LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, LLM_KV_ATTENTION_OUTPUT_LORA_RANK, LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0402044da6b..6ed803bc58a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -236,6 +236,16 @@ llama_context::llama_context( cparams.fused_lid = true; cparams.auto_flid = true; + { + // escape hatch: the fused kernel sums the heads in a different order than the + // unfused graph, so a near-tied indexer top-k can come out different + const char * LLAMA_FUSED_LID_DISABLE = getenv("LLAMA_FUSED_LID_DISABLE"); + if (LLAMA_FUSED_LID_DISABLE && atoi(LLAMA_FUSED_LID_DISABLE) != 0) { + cparams.fused_lid = false; + cparams.auto_flid = false; + } + } + cparams.fused_dsv4_hc_pre = true; cparams.fused_dsv4_hc_comb = true; cparams.fused_dsv4_hc_post = true; @@ -2293,8 +2303,10 @@ void llama_context::output_reorder() { uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { uint32_t res; - if (model.arch == LLM_ARCH_KIMI_K3) { - // the n_tokens*40 budget below is exhausted at ubatch 3840 + if (model.arch == LLM_ARCH_KIMI_K3 || model.arch == LLM_ARCH_GLM5NEXT) { + // the n_tokens*40 budget below is exhausted at ubatch 3840 for kimi-k3, and + // earlier for glm5next: each KDA layer costs 182 nodes + ~16/token, so its 34 + // KDA layers alone need 6.2k + 31.9*n_tokens before DSA or the MoE res = std::max(n_tokens * 160, 64u * model.n_tensors()); } else if (model.arch == LLM_ARCH_QWEN3NEXT || model.arch == LLM_ARCH_KIMI_LINEAR || diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0e..1bb06702a8a 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -12,6 +12,7 @@ #include "llama-kv-cache-dsa-iswa.h" #include "llama-kv-cache-msa.h" #include "llama-kv-cache-dsv4.h" +#include "llama-kv-cache-kpool.h" #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" #include "llama-memory-recurrent.h" @@ -1779,7 +1780,7 @@ ggml_tensor * llm_graph_context::build_ffn( tmp = ggml_clamp(ctx0, tmp, -limit, limit); cb(tmp, "ffn_up_clamped", il); - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_clamp(ctx0, cur, -INFINITY, limit); cb(cur, "ffn_gate_clamped", il); cur = ggml_swiglu_split(ctx0, cur, tmp); @@ -2176,7 +2177,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn( up = ggml_clamp(ctx0, up, -limit, limit); cb(up, "ffn_moe_up_clamped", il); - if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { + if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) { cur = ggml_clamp(ctx0, cur, -INFINITY, limit); cb(cur, "ffn_moe_gate_clamped", il); cur = ggml_swiglu_split(ctx0, cur, up); @@ -3552,6 +3553,210 @@ llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k() const return (llm_graph_input_mem_hybrid_k *) res->add_input(std::move(inp)); } +llm_graph_input_kpool * llm_graph_context::build_inp_kpool( + const llama_memory_hybrid_context * mctx_cur, + ggml_tensor * kq_mask, + bool scoring) const { + const auto * mctx_attn = mctx_cur->get_attn(); + const auto * mctx_idx = mctx_cur->get_idx(); + + GGML_ASSERT(mctx_idx != nullptr && "a pooled indexer needs the indexer KV cache"); + + const uint32_t kpool = hparams.indexer_kpool; + GGML_ASSERT(kpool > 0); + + auto inp = std::make_unique(mctx_attn, mctx_idx, kpool); + + inp->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + ggml_set_input(inp->k_idxs); + ggml_set_name(inp->k_idxs, "kpool_k_idxs"); + + if (scoring) { + const int64_t n_kv = mctx_attn->get_n_kv(); + + // spelled exactly as build_attn_inp_kq_mask spells it, so that sel_mask, + // cand_mask and the KQ mask are the same shape and add without a + // broadcast. get_n_stream() is the cache's stream RANGE and is a + // different number as soon as a server has non-contiguous slots busy + const int64_t n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq; + const int64_t n_tps = ubatch.n_tokens/n_stream; + + // one pool map per SEQUENCE. a non-unified cache has one sequence per stream, a + // unified cache puts every sequence of the ubatch in stream 0 and cuts the + // stream's pool table into one run per sequence, so the table needs the rebasing + // slack once per sequence. sized on the ubatch and not on n_seq_max, which + // llama-embedding sets to 256; n_seqs_unq is already part of + // llm_graph_params::allow_reuse, so the shape holds while a graph is reused + const int64_t n_ps = (int64_t) ubatch.n_seqs_unq/n_stream; + + GGML_ASSERT(n_ps >= 1 && (int64_t) ubatch.n_seqs_unq == n_ps*n_stream); + + const int64_t n_pools = llama_kpool_n_pools(n_kv, kpool, n_ps); + + GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[3] == n_stream); + + // this tree's build_attn_inp_kq_mask has no GGML_KQ_MASK_PAD, so the mask + // is exactly n_tps rows. the host side supports n_padq > n_tps, the graph + // below does not: the selection terms only exist for real queries + GGML_ASSERT(kq_mask->ne[1] == n_tps && "the pooled indexer needs an unpadded KQ mask"); + + inp->pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool*n_pools, n_stream); + ggml_set_input(inp->pool_cells); + ggml_set_name(inp->pool_cells, "kpool_pool_cells"); + + inp->pool_bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_pools, n_tps, n_stream); + ggml_set_input(inp->pool_bias); + ggml_set_name(inp->pool_bias, "kpool_pool_bias"); + + // the fused lightning indexer wants an f16 mask. built here, once, because + // every indexer layer shares it + if (cparams.fused_lid) { + inp->pool_bias_f16 = ggml_cast(ctx0, + ggml_reshape_4d(ctx0, inp->pool_bias, n_pools, n_tps, 1, n_stream), + GGML_TYPE_F16); + ggml_set_name(inp->pool_bias_f16, "kpool_pool_bias_f16"); + } + + // f16, not f32. The only two values either mask holds are 0.0f and + // -INFINITY, both exact in f16, so the narrower type is lossless and halves + // two [n_kv, n_tps, n_stream] inputs that live for the whole ubatch: 2 GiB + // each at n_ctx = 1 Mi, n_ubatch = 512. + // + // Every consumer takes f16. Under flash attention this is the KQ mask's own + // type, so build_attn_sparse adds the two with no conversion at all. With + // flash attention off - which GLM-5.3-Flash requires, so it is the path that + // matters here - the KQ mask is f32, and ggml_add gives its result src0's + // type: f16 + f32 -> f16 is a supported bin_bcast on CUDA and on the CPU, + // and ggml_soft_max_ext takes an f16 mask as readily as an f32 one + inp->sel_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, n_kv, n_tps, 1, n_stream); + ggml_set_input(inp->sel_mask); + ggml_set_name(inp->sel_mask, "kpool_sel_mask"); + + inp->cand_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, n_kv, n_tps, 1, n_stream); + ggml_set_input(inp->cand_mask); + ggml_set_name(inp->cand_mask, "kpool_cand_mask"); + } + + return (llm_graph_input_kpool *) res->add_input(std::move(inp)); +} + +ggml_tensor * llm_graph_context::build_attn_sparse( + llm_graph_input_attn_k * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * kq_b, + ggml_tensor * sinks, + ggml_tensor * v_mla, + ggml_tensor * top_k, + ggml_tensor * sel_mask, + ggml_tensor * cand_mask, + float kq_scale, + int il) const { + ggml_build_forward_expand(gf, q_cur); + ggml_build_forward_expand(gf, v_cur); + ggml_build_forward_expand(gf, k_cur); + + const auto * mctx_cur = inp->mctx; + + // store to KV cache + { + const auto & k_idxs = inp->get_k_idxs(); + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + } + + const auto & kq_mask = inp->get_kq_mask(); + + GGML_ASSERT(sel_mask->type == GGML_TYPE_F16 || sel_mask->type == GGML_TYPE_F32); + GGML_ASSERT(sel_mask->type == cand_mask->type); + GGML_ASSERT(ggml_are_same_shape(sel_mask, cand_mask)); + GGML_ASSERT(sel_mask->ne[0] == kq_mask->ne[0] && sel_mask->ne[1] == kq_mask->ne[1] && + sel_mask->ne[3] == kq_mask->ne[3]); + + // The dense DSA path (build_attn on llm_graph_input_attn_k_dsa) opens with + // ggml_fill(kq_mask, -INFINITY). Here the scatter starts from sel_mask + // instead, which already holds 0.0 on the query's own trailing incomplete + // pool: GLM always attends to that tail (index_kpool_always_select_tail), and + // keeping it out of the top-k budget is what lets the budget stay a whole + // number of pools. + // + // ggml_set_rows writes THROUGH to its destination and returns a view of it, + // and sel_mask is one shared per-ubatch input read by every indexer layer. + // Scattering into it directly makes each layer inherit the previous layer's + // unmasked cells: measured on TinySparse, layer 3 stayed inside its budget + // while layer 7, running second, reached 411 cells. The dense path never + // meets this because ggml_fill hands it a fresh tensor every layer. Take a + // private copy per layer. + ggml_tensor * mask_all = ggml_dup(ctx0, sel_mask); + + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + mask_all = ggml_view_4d(ctx0, mask_all, 1, mask_all->ne[0], mask_all->ne[1], mask_all->ne[3], + mask_all->nb[0], mask_all->nb[1], mask_all->nb[2], 0); + + // [n_select, n_tps, n_stream] -> [n_select, n_tps, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[2], 1, + top_k->nb[1], top_k->nb[2], top_k->ne[2]*top_k->nb[2], 0); + + // A constant 0, never the cell's own bias. The scatter must not be able to + // ERASE a zero sel_mask already granted: a tail cell can also be named by the + // top-k (through an over-budget pool whose unfilled slots point at it), and + // scattering that cell's -inf score bias would leave the query attending to + // nothing at all. Rejecting an over-budget selection is done additively, + // below, by cand_mask. + // + // f32 whatever mask_all is: ggml_set_rows converts its values into the + // destination, and the CUDA backend only advertises SET_ROWS for f32 values + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); + + ggml_tensor * mask_top_k = ggml_set_rows(ctx0, mask_all, zeros, top_k_3d); + + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + mask_top_k = ggml_view_4d(ctx0, mask_top_k, mask_top_k->ne[1], mask_top_k->ne[2], 1, mask_top_k->ne[3], + mask_top_k->nb[2], mask_top_k->nb[3], mask_top_k->nb[3], 0); + + // the reference's `selected_valid` gather, additively: a cell the top-k named + // that is not in the reference's candidate set goes back to -inf, and a tail + // cell that sel_mask granted stays at 0 because cand_mask is the UNION of the + // candidates and the tail + mask_top_k = ggml_add(ctx0, mask_top_k, cand_mask); + + // ggml_add gives its result src0's type, so an f16 selection mask absorbs an + // f32 KQ mask and no cast is needed. The one direction that still needs one is + // an f32 mask meeting an f16 KQ mask: the add would yield f32 and + // ggml_flash_attn_ext asserts its mask is f16 + if (mask_top_k->type == GGML_TYPE_F32 && kq_mask->type == GGML_TYPE_F16) { + mask_top_k = ggml_cast(ctx0, mask_top_k, GGML_TYPE_F16); + } + + // and finally re-apply causality, occupancy and padding. load bearing: it is + // what keeps an empty, future or foreign-sequence cell masked no matter what + // the top-k returned + mask_top_k = ggml_add(ctx0, mask_top_k, kq_mask); + cb(mask_top_k, "kpool_kq_mask", il); + + ggml_tensor * q = q_cur; + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); + + ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, mask_top_k, sinks, v_mla, kq_scale, il); + cb(cur, "kqv_out", il); + + if (wo) { + cur = build_lora_mm(wo, cur, wo_s); + } + + if (wo_b) { + cur = ggml_add(ctx0, cur, wo_b); + } + + return cur; +} + llm_graph_input_mem_hybrid_iswa * llm_graph_context::build_inp_mem_hybrid_iswa() const { const auto * mctx_cur = static_cast(mctx); diff --git a/src/llama-graph.h b/src/llama-graph.h index b388e028cb5..c994a1774d5 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -32,6 +32,10 @@ class llama_memory_recurrent_context; class llama_memory_hybrid_context; class llama_memory_hybrid_iswa_context; +// defined in llama-kv-cache-kpool.h, which includes this header, so it can only +// be forward declared here +class llm_graph_input_kpool; + // certain models (typically multi-modal) can produce different types of graphs enum llm_graph_type { LLM_GRAPH_TYPE_DEFAULT, @@ -1344,6 +1348,53 @@ struct llm_graph_context { llm_graph_input_mem_hybrid_iswa * build_inp_mem_hybrid_iswa() const; + // + // pooled (GLM-5-Next lightning) indexer + // + + // one pooling map per ubatch, shared by every indexer layer. see + // llama-kv-cache-kpool.h for what the tensors mean and why they are built + // host side rather than derived in the graph. + // + // `scoring` false allocates only k_idxs: the indexer key and gate store is + // unconditional, the selection is not. An input tensor with no consumer is + // never backed by the allocator, so the rest must not be created either + llm_graph_input_kpool * build_inp_kpool( + const llama_memory_hybrid_context * mctx_cur, + ggml_tensor * kq_mask, + bool scoring) const; + + // sparse (pooled top-k) variant of the llm_graph_input_attn_k build_attn. + // + // Identical to it except for the mask. `top_k` names the cells the pooled + // indexer selected, already expanded from whole pools; they are unmasked on + // top of `sel_mask`, which arrives already holding 0.0f on the query's own + // always-selected trailing pool. `cand_mask` is the reference's candidate + // set and is what makes an over-budget selection harmless: ggml_top_k + // returns a full budget of pool ordinals even when fewer pools carry a + // finite score, which during prefill is the normal state. + // + // Both masks may be f16 or f32, and build_inp_kpool allocates f16: the only + // values either holds are 0.0f and -INFINITY, both exact in f16. The combined + // mask inherits their type, which the KQ mask then adds into whatever its own + // type is + ggml_tensor * build_attn_sparse( + llm_graph_input_attn_k * inp, + ggml_tensor * wo, + ggml_tensor * wo_b, + ggml_tensor * wo_s, + ggml_tensor * q_cur, // [n_embd_head_q, n_head_q, n_tokens] + ggml_tensor * k_cur, // [n_embd_head_k, n_head_k, n_tokens] + ggml_tensor * v_cur, // [n_embd_head_v, n_head_v, n_tokens] + ggml_tensor * kq_b, + ggml_tensor * sinks, // [n_head_q] + ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] + ggml_tensor * top_k, // I32 [n_select, n_tokens/n_stream, n_stream] + ggml_tensor * sel_mask, // F16/F32 [n_kv, n_batch, 1, n_stream] + ggml_tensor * cand_mask, // F16/F32 [n_kv, n_batch, 1, n_stream] + float kq_scale, + int il) const; + // // pooling // diff --git a/src/llama-hparams.h b/src/llama-hparams.h index c3c14292c32..bacef90a66e 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -252,6 +252,8 @@ struct llama_hparams { uint32_t indexer_n_head = 0; uint32_t indexer_head_size = 0; uint32_t indexer_top_k = 0; + // glm5next: the indexer scores pools of this many keys instead of single keys + uint32_t indexer_kpool = 0; // MSA uint32_t indexer_block_size = 0; uint32_t indexer_local_blocks = 0; diff --git a/src/llama-kv-cache-kpool.cpp b/src/llama-kv-cache-kpool.cpp new file mode 100644 index 00000000000..9f214dd16cd --- /dev/null +++ b/src/llama-kv-cache-kpool.cpp @@ -0,0 +1,441 @@ +#include "llama-kv-cache-kpool.h" + +#include "llama-batch.h" +#include "llama-kv-cache.h" +#include "llama-kv-cells.h" + +#include +#include +#include + +uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool, uint32_t n_seqs) { + GGML_ASSERT(kpool > 0); + GGML_ASSERT(n_seqs > 0); + + return n_kv/kpool + 2*n_seqs; +} + +uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t kpool) { + GGML_ASSERT(kpool > 0); + GGML_ASSERT(n_pools > 0); + GGML_ASSERT(indexer_top_k % kpool == 0 && "indexer_top_k must be a whole number of pools"); + + // min(index_topk // index_kpool, n_pools), exactly the reference's select_k + return std::min(n_pools, indexer_top_k/kpool); +} + +// sel_mask and cand_mask hold only 0.0f and -INFINITY, so they can be written in the +// KQ mask's f16 exactly as in f32 +template struct kpool_mask_of; + +template <> struct kpool_mask_of { + static float from(float v) { return v; } +}; + +template <> struct kpool_mask_of { + static ggml_fp16_t from(float v) { return ggml_fp32_to_fp16(v); } +}; + +template +static void kpool_mask_fill(T * dst, int64_t n) { + std::fill(dst, dst + n, kpool_mask_of::from(-INFINITY)); +} + +// one query's row of both masks; the two predicates share their operands, and the +// unsigned compares are what let this vectorise +template +static void kpool_mask_row( + T * cur_sel, + T * cur_cand, + const llama_pos * pos_at, + const int32_t * pool_of, + int64_t n_kv, + llama_pos q, + llama_pos tail_start, + int64_t bo_vis) { + const T v_sel = kpool_mask_of::from(0.0f); + const T v_mask = kpool_mask_of::from(-INFINITY); + + for (int64_t j = 0; j < n_kv; ++j) { + const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; + const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; + const bool tail = pos_at[j] >= tail_start; + + cur_sel [j] = vis && tail ? v_sel : v_mask; + // max(bias, sel_mask): the reference's candidate set, which the + // top-k budget may overrun but must never escape + cur_cand[j] = vis && (pooled || tail) ? v_sel : v_mask; + } +} + +void llama_kv_cache_set_input_kpool( + const llama_kv_cache * kv, + ggml_tensor * cell_pool, + ggml_tensor * pool_cells, + ggml_tensor * bias, + ggml_tensor * pool_bias, + ggml_tensor * sel_mask, + ggml_tensor * cand_mask, + const llama_ubatch * ubatch, + uint32_t kpool) { + GGML_ASSERT(kv != nullptr); + GGML_ASSERT(kpool > 0); + + // the per-CELL view is optional: the pooled graph does not consume it, and an + // input tensor with no consumer is never backed by the allocator + GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(pool_bias ->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(sel_mask ->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(cand_mask ->buffer)); + + // both masks are written through raw strides below, so writing the wrong width + // would overrun the allocation 2x. check rather than trust + GGML_ASSERT(pool_cells->type == GGML_TYPE_I32); + GGML_ASSERT(pool_bias ->type == GGML_TYPE_F32); + GGML_ASSERT((sel_mask->type == GGML_TYPE_F16 || sel_mask->type == GGML_TYPE_F32) && + "sel_mask must be f16 or f32"); + GGML_ASSERT(cand_mask->type == sel_mask->type && "both masks must have the KQ mask's type"); + + // everything below is written through raw strides + GGML_ASSERT(ggml_is_contiguous(pool_cells)); + GGML_ASSERT(ggml_is_contiguous(pool_bias)); + GGML_ASSERT(ggml_is_contiguous(sel_mask)); + GGML_ASSERT(ggml_is_contiguous(cand_mask)); + + const int64_t n_kv = sel_mask->ne[0]; + const int64_t n_ns = sel_mask->ne[3]; // streams in this ubatch + const int64_t r = kpool; + const int64_t n_tokens = ubatch->n_tokens; + + // [TAG_KPOOL_SEQ_PARTITION] + // positions are unambiguous only within one sequence, so one pool map per SEQUENCE, + // not per stream. a non-unified cache gives each stream its own cells array and one + // sequence (n_ps == 1, the layout this file had before), a unified cache gives one + // stream carrying every sequence in the ubatch + GGML_ASSERT(n_ns == 1 || (int64_t) ubatch->n_seqs_unq == n_ns); + + const int64_t n_ps = (int64_t) ubatch->n_seqs_unq/n_ns; // sequences per stream + const int64_t n_pools = pool_cells->ne[0]/r; // pool slots per stream + + GGML_ASSERT(n_ps > 0 && (int64_t) ubatch->n_seqs_unq == n_ns*n_ps); + GGML_ASSERT(pool_cells->ne[0] % r == 0); + GGML_ASSERT(n_pools >= 2*n_ps); + GGML_ASSERT(pool_cells->ne[1] == n_ns); + GGML_ASSERT(sel_mask->ne[2] == 1); + GGML_ASSERT(ggml_are_same_shape(cand_mask, sel_mask)); + GGML_ASSERT(pool_bias->ne[0] == n_pools && pool_bias->ne[2] == n_ns); + GGML_ASSERT(n_tokens % n_ns == 0); + + const int64_t n_tps = n_tokens/n_ns; // tokens per stream + const int64_t n_padq = sel_mask->ne[1]; // KQ mask rows, >= n_tps + + GGML_ASSERT(pool_bias->ne[1] == n_tps); + GGML_ASSERT(n_padq >= n_tps); + + if (cell_pool) { + GGML_ASSERT(ggml_backend_buffer_is_host(cell_pool->buffer)); + GGML_ASSERT(cell_pool->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(cell_pool)); + GGML_ASSERT(cell_pool->ne[0] == n_kv && cell_pool->ne[1] == n_ns); + + // one row per stream, so a cell that two sequences of one stream share has + // nowhere to put its second pool. the graph never asks for this view + GGML_ASSERT(n_ps == 1 && "the per-cell pool view needs one sequence per stream"); + } + + if (bias) { + GGML_ASSERT(ggml_backend_buffer_is_host(bias->buffer)); + GGML_ASSERT(bias->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(bias)); + GGML_ASSERT(bias->ne[0] == n_kv && bias->ne[1] == n_tps && bias->ne[2] == n_ns); + } + + int32_t * dst_cell_pool = cell_pool ? (int32_t *) cell_pool->data : nullptr; + int32_t * dst_pool_cells = (int32_t *) pool_cells->data; + float * dst_bias = bias ? (float *) bias->data : nullptr; + float * dst_pool_bias = (float *) pool_bias ->data; + char * dst_sel_mask = (char *) sel_mask ->data; + char * dst_cand_mask = (char *) cand_mask ->data; + + const bool mask_f16 = sel_mask->type == GGML_TYPE_F16; + const size_t mask_ts = ggml_type_size(sel_mask->type); + + // -1 marks a cell with no usable pool. host side only: never copied into cell_pool, + // where ggml_get_rows would read it as an index + std::vector pool_of(n_kv); + std::vector filled(n_pools); + std::vector pos_at; + + // one contiguous run of pool slots per sequence of the stream + std::vector run_off(n_ps); + std::vector run_len(n_ps); + + // which cells array a sequence of this stream uses; same convention as + // llama_kv_cache::set_input_kq_mask. with one sequence per stream the ubatch's + // unique list and the stream's own sequence are the same thing + auto seq_of = [&](int64_t s, int64_t ps) { + return n_ps == 1 ? ubatch->seq_id[s*n_tps][0] : ubatch->seq_id_unq[ps]; + }; + + for (int64_t s = 0; s < n_ns; ++s) { + int32_t * cur_pool_cells = dst_pool_cells + s*(r*n_pools); + char * cur_sel_mask = dst_sel_mask + s*(n_padq*n_kv)*mask_ts; + char * cur_cand_mask = dst_cand_mask + s*(n_padq*n_kv)*mask_ts; + float * cur_pool_bias = dst_pool_bias + s*(n_tps*n_pools); + + // slots of a pool that is not resident, and slots outside the query's own + // sequence run, are cleared once per stream here. the per-sequence pass below + // only writes what it owns + std::fill(cur_pool_cells, cur_pool_cells + r*n_pools, 0); + std::fill(cur_pool_bias, cur_pool_bias + n_tps*n_pools, -INFINITY); + + // the token loop writes rows < n_tps in full; only the padding rows need clearing + if (mask_f16) { + kpool_mask_fill((ggml_fp16_t *) (cur_sel_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); + kpool_mask_fill((ggml_fp16_t *) (cur_cand_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); + } else { + kpool_mask_fill((float *) (cur_sel_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); + kpool_mask_fill((float *) (cur_cand_mask + n_tps*n_kv*mask_ts), (n_padq - n_tps)*n_kv); + } + + // [TAG_KPOOL_PACK] + // cut the stream's pool table into one run per sequence, sized on the pool range + // that sequence actually holds. the table is n_kv/kpool shared plus 2 per sequence + // for rebasing, which covers it whenever the sequences' cells are disjoint - every + // case but a prefix shared through llama_memory_seq_cp. that one can ask for more + // slots than exist, and then a sequence keeps its newest pools, the same cut a + // large hole already forces. + // + // NOT one full-width table per sequence: the indexer scores every slot against + // every query, so a full-width table would multiply the score tensor by the + // sequence count, and llama-embedding asks for n_seq_max 256 + { + int64_t n_want = 0; + + for (int64_t ps = 0; ps < n_ps; ++ps) { + const llama_seq_id seq = seq_of(s, ps); + const auto & cells = kv->get_cells(seq); + + int64_t b_min = 0; + int64_t b_max = 0; + bool found = false; + + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j) || !cells.seq_has(j, seq)) { + continue; + } + const int64_t b = cells.pos_get(j)/r; + b_min = found ? std::min(b_min, b) : b; + b_max = found ? std::max(b_max, b) : b; + found = true; + } + + run_len[ps] = found ? b_max - b_min + 1 : 0; + n_want += run_len[ps]; + } + + if (n_want > n_pools) { + int64_t rem = n_pools; + + for (int64_t ps = 0; ps < n_ps; ++ps) { + run_len[ps] = std::min(run_len[ps], rem/(n_ps - ps)); + rem -= run_len[ps]; + } + } + + int64_t off = 0; + for (int64_t ps = 0; ps < n_ps; ++ps) { + run_off[ps] = off; + off += run_len[ps]; + } + + GGML_ASSERT(off <= n_pools); + } + + int64_t n_done = 0; + + for (int64_t ps = 0; ps < n_ps; ++ps) { + const llama_seq_id seq_of_pool = seq_of(s, ps); + const auto & cells = kv->get_cells(seq_of_pool); + + const int64_t n_run = run_len[ps]; + + int32_t * cur_cell_pool = dst_cell_pool ? dst_cell_pool + s*n_kv : nullptr; + int32_t * part_pool_cells = cur_pool_cells + run_off[ps]*r; + + std::fill(pool_of.begin(), pool_of.end(), -1); + std::fill(filled.begin(), filled.end(), 0); + + // hoist occupancy and sequence membership out of the O(n_kv * n_tokens) loop + // below; neither depends on the query. -1 means the cell holds nothing this + // sequence may pool or attend to. under a unified cache `cells` is shared with + // every other sequence, and seq_has is what keeps their keys out + pos_at.resize(n_kv); + for (int64_t j = 0; j < n_kv; ++j) { + pos_at[j] = cells.is_empty(j) || !cells.seq_has(j, seq_of_pool) ? -1 : cells.pos_get(j); + } + + // a pool ordinal is the absolute p/kpool, which can far exceed n_kv/kpool, so + // rebase on this sequence's lowest resident pool. grouping is untouched, every + // member shifts together. + // + // anchoring at p/kpool follows vLLM and SGLang, not HF (which pools from the + // first *resident* key, valid_keys.argmax(-1), differing under left padding). + // it is the only anchor that keeps a pool's identity stable between the prefill + // that built it and the decodes that read it. + // + // the window is this sequence's run; positions are contiguous in any real + // batch so the resident range fits. seq_rm can leave a hole large enough that + // it does not, and then the newest pools are the ones worth keeping. + int64_t b_base = 0; + { + int64_t b_min = 0; + int64_t b_max = 0; + bool found = false; + + for (int64_t j = 0; j < n_kv; ++j) { + if (pos_at[j] < 0) { + continue; + } + const int64_t b = pos_at[j]/r; + b_min = found ? std::min(b_min, b) : b; + b_max = found ? std::max(b_max, b) : b; + found = true; + } + + b_base = std::max(b_min, b_max - (n_run - 1)); + } + + for (int64_t j = 0; j < n_kv; ++j) { + if (pos_at[j] < 0) { + continue; + } + + const llama_pos p = pos_at[j]; + const int64_t bo = p/r - b_base; + + if (bo < 0 || bo >= n_run) { + continue; + } + + pool_of[j] = (int32_t) bo; + part_pool_cells[bo*r + (p%r)] = (int32_t) j; + filled[bo]++; + } + + // an incompletely resident pool cannot be pooled: the compressor consumes all r + // member keys, and the reference demands pool_valid = grouped_valid_keys.all(-1). + // such cells are the sequence tail, which sel_mask forces in below regardless of + // score, so they point at pool slot 0 purely to keep the gather in range + for (int64_t j = 0; j < n_kv; ++j) { + // != rather than <: two cells claiming one position overwrite each other in + // pool_cells, so an over-filled pool is not usable either + if (pool_of[j] >= 0 && filled[pool_of[j]] != (int32_t) r) { + pool_of[j] = -1; + } + if (cur_cell_pool) { + cur_cell_pool[j] = pool_of[j] < 0 ? 0 : pool_of[j]; + } + } + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + + // a query is pooled by the sequence it is being written to. with several + // sequences in one stream the other partitions own the rest of the rows + if (ubatch->seq_id[i][0] != seq_of_pool) { + continue; + } + + const llama_pos q = ubatch->pos[i]; + + // q >= 0 is what makes the unsigned range test below a range test + GGML_ASSERT(q >= 0); + + n_done++; + + // the query's own incomplete pool ((q + 1) % r cells, its own token + // included) is always attended to (index_kpool_always_select_tail), which is + // what makes the selection land on pool boundaries + const llama_pos tail_start = (q + 1)/r*r; + + // the reference tests visibility at a pool's LAST member, so a pool + // straddling the query is dropped whole. pools are position-aligned here, so + // that test collapses to b*r < tail_start + const int64_t bo_vis = std::max(0, tail_start/r - b_base); + + float * cur_bias = dst_bias ? dst_bias + i*n_kv : nullptr; + char * cur_sel = cur_sel_mask + ii*n_kv*mask_ts; + char * cur_cand = cur_cand_mask + ii*n_kv*mask_ts; + + // the unsigned compares inside fold "empty or another sequence" (pos_at -1) + // and "no usable pool" (pool_of -1) into the range test + if (mask_f16) { + kpool_mask_row((ggml_fp16_t *) cur_sel, (ggml_fp16_t *) cur_cand, + pos_at.data(), pool_of.data(), n_kv, q, tail_start, bo_vis); + } else { + kpool_mask_row((float *) cur_sel, (float *) cur_cand, + pos_at.data(), pool_of.data(), n_kv, q, tail_start, bo_vis); + } + + if (cur_bias) { + for (int64_t j = 0; j < n_kv; ++j) { + const bool vis = (uint32_t) pos_at [j] <= (uint32_t) q; + const bool pooled = (uint32_t) pool_of[j] < (uint32_t) bo_vis; + + cur_bias[j] = vis && pooled ? 0.0f : -INFINITY; + } + } + + // The same predicate, per POOL, which is where the reference applies + // it: pool_valid (completely resident) & pool_visible (its LAST + // member is visible, so a pool the query straddles is dropped whole). + // Pools are position-aligned here, so pool bo's last member is at + // position (b_base + bo)*r + r - 1 and "last member visible" collapses + // to bo < bo_vis. + // + // NOT gathered from `bias` at the last member cell: an incomplete or + // absent pool has no resident last member, pool_cells points that slot + // at cell 0, and the pool would inherit cell 0's validity. + // + // the query's own sequence run only. every other slot stays at the + // -INFINITY the per-stream fill above left, which is what keeps a foreign + // pool out of the budget + float * q_pool_bias = cur_pool_bias + ii*n_pools + run_off[ps]; + + for (int64_t p = 0; p < n_run; ++p) { + const bool valid = filled[p] == (int32_t) r; + const bool visible = p < bo_vis; + + q_pool_bias[p] = valid && visible ? 0.0f : -INFINITY; + } + } + } + + // every row of sel_mask, cand_mask and pool_bias must have been written by + // exactly one partition, or a query is left reading another sequence's pools + GGML_ASSERT(n_done == n_tps && "every query must belong to a sequence of the ubatch"); + } +} + +void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { + // unconditional: the indexer key and gate STORE runs on the dense path too, + // and k_idxs is what tells cpy_k where to put them. Gating the store the way + // the scoring is gated would leave every cell written below n_select - the + // first 2051 positions of every sequence on the real model - with no indexer + // state, and the first ubatch to cross n_select would pool cells that were + // never written + mctx_idx->set_input_k_idxs(k_idxs, ubatch); + + // the rest exists only when the graph scores. below n_select the indexer + // would select every visible position, so build_inp_kpool does not allocate + // these at all and there is nothing to fill + if (pool_cells == nullptr) { + return; + } + + llama_kv_cache_set_input_kpool( + mctx_attn->get_kv(), + /* cell_pool */ nullptr, pool_cells, /* bias */ nullptr, pool_bias, + sel_mask, cand_mask, ubatch, kpool); +} diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h new file mode 100644 index 00000000000..2ad259c8dcf --- /dev/null +++ b/src/llama-kv-cache-kpool.h @@ -0,0 +1,208 @@ +#pragma once + +#include "ggml.h" +#include "llama.h" +#include "llama-graph.h" + +#include + +struct llama_ubatch; +class llama_kv_cache; +class llama_kv_cache_context; + +// +// GLM-5-Next indexer pooling +// +// GLM's lightning indexer scores pools of `kpool` consecutive *positions*, while its +// top-k budget is counted in tokens (indexer_top_k), i.e. indexer_top_k/kpool whole +// pools. The reference requires indexer_top_k % kpool == 0, so the division is exact. +// +// Pools are defined on positions, but everything the graph indexes is addressed by +// *cell*, and a cell index is whatever llama_kv_cache::find_slot handed out: cells of +// one pool are not adjacent, not ordered, and under a unified cache not even owned by +// the same sequence. The mapping therefore cannot be derived in the graph. It is built +// here, host side, from the cache's cells, and passed in as plain input tensors, as +// qwen4exp's QSA does for its compression blocks. +// +// A position only names a pool inside ONE sequence, so the map is per SEQUENCE. A +// non-unified cache gives every sequence its own stream and its own cells array, and +// the two are the same thing. A unified cache puts every sequence of the ubatch in one +// stream and one cells array, so the stream's pool table is PARTITIONED: each sequence +// gets a contiguous run of slots and rebases inside it. `pool_bias` is -INFINITY outside +// the query's own run, which is what stops a query from selecting a foreign pool. +// +// The runs are packed, not one full-width table per sequence. A full-width table per +// sequence would multiply the pool axis by the sequence count, and the indexer scores +// every pool against every query, so llama-embedding (which asks for n_seq_max 256 with +// a unified cache) reserved 286 GB of compute buffer that way. +// +// Nothing here may emit a negative index: ggml_set_rows asserts i1 >= 0 and +// ggml_get_rows has no sentinel, so unpopulated entries are clamped into range and +// neutralised by the additive masks instead. +// + +// number of pool slots the graph must allocate for `n_kv` cells shared by `n_seqs` +// sequences. +// +// pool ordinals are position-derived, then rebased on each sequence's lowest resident +// pool so sequences not starting at position 0 still land inside the array. rebasing +// can cost one slot at each end, hence 2 per sequence. +// +// n_kv/kpool is the budget the sequences share, which is exact while their cells are +// disjoint: a cell belongs to one pool of one sequence. llama_memory_seq_cp breaks that +// - a shared prefix cell is pooled by every sequence holding it - and then the runs are +// cut to the newest pools rather than the table being grown, see [TAG_KPOOL_PACK]. +uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool, uint32_t n_seqs = 1); + +// how many POOLS ggml_top_k selects. +// +// The reference (modular_glm5_next.py, Glm5NextTextIndexer.forward) scores the pool +// axis and takes select_k = min(index_topk // index_kpool, n_pools), then expands each +// selected pool to its members. This is that number, and top-k over pools rather than +// over cells is not an optimisation, it is the only correct spelling. +// +// The tempting alternative - broadcast a pool's score onto its kpool member cells and +// run one top-k of width indexer_top_k over cells - is WRONG, for a reason a +// set-similarity metric cannot see. Its argument is that a pool's members carry its +// score bit-exactly, so the cut must land on a pool boundary; that holds only if tie +// groups never span pools. They do: F.relu drives most pool scores to exactly 0.0 (on +// TinySparse, 92 cells tie at 0.0 for query 386 at layer 3), so the cut falls inside a +// tie group that straddles pools, and ggml_top_k - explicitly unordered among equals; +// the CPU op deliberately swaps the first two results, the CUDA op declares +// determinism::not_guaranteed - takes an arbitrary 1..kpool-1 members of the pool it +// cuts. Measured on TinySparse at 512 tokens the cell-level form leaves a partial pool +// on 7.51% of query rows at L3 and 5.93% at L7 (70 and 47 partial pools); the +// pool-level form leaves none. A pool-aligned top-k WIDTH does not help, because the +// ties are not aligned to anything. +// +// Selecting whole pools also removes the CPU/CUDA tie-break divergence structurally +// rather than making it unlikely: the backends may disagree about WHICH pools come out +// of a tie group, never about whether a pool is taken whole. +// +// NOT indexer_top_k + kpool - 1 either: that is the width of the reference's *output* +// buffer, tail included. The always-selected tail is forced in via `sel_mask`. +uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t kpool); + +// Fill the host-side inputs of the pooled indexer. +// +// `cell_pool` and `bias` are the per-CELL view of the same information and may be +// nullptr. They are kept because they are the independent spelling +// tests/test-glm5next-memory.cpp checks `pool_bias` and `cand_mask` against - but the +// graph selects at POOL granularity (see llama_kpool_select_k), so it passes nullptr +// for both. An input tensor with no consumer is never backed by the allocator, so +// building them anyway would write through a null buffer. +// +// cell_pool I32 [n_kv, n_stream] OPTIONAL +// cell -> its pool slot, or 0 when it has no usable pool. +// +// pool_cells I32 [kpool*n_pools, n_stream] +// pool slot, member -> the cell holding that position, or 0 when not resident. +// One stream's slots are cut into one contiguous run per sequence. +// Two consumers: the compressor gathers a pool's member keys and gates with it, +// and the top-k expands a selected POOL ordinal back into its kpool member CELLS +// with it - the reference's `pool_indices[batch_idx, selected]`. +// +// bias F32 [n_kv, n_tokens/n_stream, n_stream] OPTIONAL +// additive per-(cell, query) bias on a per-cell score: 0.0f in a complete pool +// whose last member the query can see, else -INFINITY, the trailing incomplete +// pool included. +// +// pool_bias F32 [n_pools, n_tokens/n_stream, n_stream] +// the same predicate per POOL, which is where the reference applies it: 0.0f when +// pool p is completely resident and its LAST member is visible to query q +// (`pool_valid & pool_visible`), else -INFINITY, the query's own trailing pool +// included so no budget is spent on it. Every slot outside the query's own +// sequence run is -INFINITY, so a query never selects another sequence's pool. +// +// Derived here rather than in the graph on purpose. Gathering `bias` at +// each pool's last member looks equivalent and is not: an incomplete or +// entirely absent pool has no resident last member, `pool_cells` points +// that slot at cell 0 to keep the gather in range, and the pool would then +// inherit cell 0's validity and compete for budget with a finite score. +// +// sel_mask F16/F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) +// what the top-k scatter starts from, replacing the ggml_fill(kq_mask, -INFINITY) +// opening the DSA mask build in llm_graph_context::build_attn: 0.0f for the +// query's own incomplete trailing pool, which GLM always attends to +// (index_kpool_always_select_tail), else -INFINITY, padding rows included. +// Forcing the tail in here rather than through the score keeps the budget a whole +// number of pools. +// +// cand_mask F16/F32 [n_kv, n_batch, 1, n_stream] (KQ mask shape, may be padded) +// max(bias, sel_mask): the reference's candidate set, i.e. every cell it could +// return for this query. 0.0f in a complete visible pool OR in the query's own +// tail, else -INFINITY, padding rows included. +// +// ggml_top_k returns `select_k` pool ordinals even when fewer pools carry a +// finite score, and during prefill that is the NORMAL state, not a corner case: +// query q has only ~q/kpool complete visible pools against a budget of +// index_topk/kpool (on TinySparse, 20 of 20 query rows are under budget). The +// spilled ordinals are arbitrary among the -INFINITY ties and expanding them +// unmasks cells. Adding the causal KQ mask kills the spills that are empty, +// foreign or in the future. What it does not kill is a resident, causally visible +// cell in an INCOMPLETE pool below the tail - unreachable while positions are +// contiguous, reachable the moment a partial seq_rm leaves a hole. cand_mask +// kills exactly those, for one store per element in a loop that already computes +// both operands. +// +// (The other spill, an unfilled pool slot pointing at cell 0, is a provable +// no-op: the budget can only overflow once every finitely-scored pool is already +// selected, so cell 0 is either already selected through its own pool, or masked +// for the same reason it would have been anyway.) +// +// `kv` must be the ATTENTION (MLA) cache: its cells define the pools and are what the +// top-k indices are ultimately read against. llama_memory_hybrid gives the indexer +// cache the attention cache's slot layout, so the two agree cell for cell. +void llama_kv_cache_set_input_kpool( + const llama_kv_cache * kv, + ggml_tensor * cell_pool, + ggml_tensor * pool_cells, + ggml_tensor * bias, + ggml_tensor * pool_bias, + ggml_tensor * sel_mask, + ggml_tensor * cand_mask, + const llama_ubatch * ubatch, + uint32_t kpool); + +// One pooling map per ubatch, shared by every indexer layer: all four tensors, and +// k_idxs, depend only on the cells and the ubatch, never on the layer. Rebuilding them +// per layer costs O(n_kv * n_tokens) host writes each time (at 128 Ki cells, 512 tokens +// and 11 DSA layers, ~11 x 67M float stores per ubatch, which dominates prefill). +// +// Sharing `pool_bias`, `sel_mask` and `cand_mask` is correct only while every indexer +// layer sees the same candidate set. That holds for glm5next, whose indexer_types are +// all "full"; a model mixing in windowed indexers would need one map per window. +// +// `k_idxs` is present whenever the model has an indexer cache; the rest only when the +// graph actually scores. The indexer key and gate STORE is not gated on the sparse path +// - below n_select the selection is a no-op but the cells still have to be written, or +// the first ubatch to cross n_select would pool cells that were never filled. +class llm_graph_input_kpool : public llm_graph_input_i { +public: + llm_graph_input_kpool( + const llama_kv_cache_context * mctx_attn, + const llama_kv_cache_context * mctx_idx, + uint32_t kpool) : mctx_attn(mctx_attn), mctx_idx(mctx_idx), kpool(kpool) {} + + ~llm_graph_input_kpool() = default; + + void set_input(const llama_ubatch * ubatch) override; + + ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] + ggml_tensor * pool_cells = nullptr; // I32 [kpool*n_pools, n_stream] + ggml_tensor * pool_bias = nullptr; // F32 [n_pools, n_tps, n_stream] + + // pool_bias in the shape and type the fused lightning indexer wants for its mask. + // A ggml_cast node, not an input: built once per graph rather than once per DSA + // layer, since every indexer layer shares the same mask. The cast is exact - + // pool_bias only ever holds 0.0f or -INFINITY. nullptr when the fused path is off + ggml_tensor * pool_bias_f16 = nullptr; // F16 [n_pools, n_tps, 1, n_stream] + + ggml_tensor * sel_mask = nullptr; // F16 [n_kv, n_batch, 1, n_stream] + ggml_tensor * cand_mask = nullptr; // F16 [n_kv, n_batch, 1, n_stream] + + const llama_kv_cache_context * mctx_attn; + const llama_kv_cache_context * mctx_idx; + + const uint32_t kpool; +}; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index ec0f5a75314..ba9129267a9 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -2585,6 +2585,14 @@ uint32_t llama_kv_cache_context::get_n_kv() const { return n_kv; } +uint32_t llama_kv_cache_context::get_n_stream() const { + return sinfos[i_cur].s1 - sinfos[i_cur].s0 + 1; +} + +const llama_kv_cache * llama_kv_cache_context::get_kv() const { + return kv; +} + ggml_type llama_kv_cache_context::type_k() const { return kv->type_k(); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 6cb6dbd2f98..983352d617c 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -366,6 +366,16 @@ class llama_kv_cache_context : public llama_memory_context_i { uint32_t get_n_kv() const; + // streams covered by the current slot info, matching the `ns` get_k/get_v use for + // their stream dimension. 1 for a unified cache. note: this is the stream RANGE + // s1 - s0 + 1, not n_seqs_unq; they differ when the active sequences are not a + // contiguous run of slots, i.e. exactly when a per-cell input sized from this would + // stop agreeing with a KQ mask + uint32_t get_n_stream() const; + + // the cache this context views, for host-side inputs that resolve cell -> position + const llama_kv_cache * get_kv() const; + ggml_type type_k() const; ggml_type type_v() const; diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6..3d842f99f84 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -4,6 +4,8 @@ #include "llama-model.h" #include "llama-context.h" +#include + // // llama_memory_hybrid // @@ -29,8 +31,11 @@ llama_memory_hybrid::llama_memory_hybrid( bool unified, /* layer filters */ const layer_filter_cb & filter_attn, - const layer_filter_cb & filter_recr) : + const layer_filter_cb & filter_recr, + const layer_filter_cb & filter_idx, + ggml_type type_idx) : hparams(model.hparams), + hparams_idx(model.hparams), mem_attn(new llama_kv_cache( model, model.hparams, @@ -62,7 +67,30 @@ llama_memory_hybrid::llama_memory_hybrid( filter_recr == nullptr ? [&](int32_t il) { return hparams.is_recr(il); } : filter_recr - )) {} + )), + mem_idx(filter_idx == nullptr ? nullptr : [&] { + // MQA with one key head of indexer_head_size, as llama_kv_cache_dsa shapes its + // lightning-indexer cache. n_embd_head_k_full is what n_embd_head_k(il) reads + // for a non-SWA layer, so an MLA model still has is_mla() true here, which + // suppresses the V allocation the indexer does not need. + // + // a *pooling* indexer (indexer_kpool > 0, glm5next only) needs a second head: + // its compressor gate is a projection of the same hidden state, so it must be + // cached alongside the key or the pool cannot be rebuilt once the member tokens + // leave the batch. every other arch leaves indexer_kpool 0 and is unchanged. + const uint32_t n_head_idx = model.hparams.indexer_kpool > 0 ? 2 : 1; + + std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), n_head_idx); + hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; + + LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells, %u x %u per cell, type = %s\n", + __func__, kv_size, n_head_idx, hparams_idx.n_embd_head_k_full, ggml_type_name(type_idx)); + + return new llama_kv_cache( + model, hparams_idx, type_idx, type_idx, v_trans, offload, unified, + kv_size, n_seq_max, n_pad, n_swa, swa_type, + nullptr, filter_idx, nullptr, nullptr); + }()) {} llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { do { @@ -115,8 +143,17 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } + // the indexer is a side buffer addressed by the attention cache's cells, so it + // takes that slot layout rather than finding its own: allocating separately lets + // the two drift apart when the context is rewritten between turns, and the top-k + // indices, read against the attention mask, would then point at the wrong cells + llama_kv_cache::slot_info_vec_t heads_idx; + if (mem_idx) { + heads_idx = heads_attn; + } + return std::make_unique( - this, std::move(heads_attn), std::move(ubatches)); + this, std::move(heads_attn), std::move(ubatches), std::move(heads_idx)); } while(false); return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); @@ -132,11 +169,15 @@ llama_memory_context_ptr llama_memory_hybrid::init_update(llama_context * lctx, bool llama_memory_hybrid::get_can_shift() const { // Shifting is trivially supported for recurrent + if (mem_idx && !mem_idx->get_can_shift()) { + return false; + } return mem_attn->get_can_shift(); } void llama_memory_hybrid::clear(bool data) { mem_attn->clear(data); + if (mem_idx) mem_idx->clear(data); mem_recr->clear(data); } @@ -146,26 +187,31 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 if (!mem_recr->seq_rm(seq_id, p0, p1)) { return false; } + if (mem_idx) mem_idx->seq_rm(seq_id, p0, p1); return mem_attn->seq_rm(seq_id, p0, p1); } void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1); + if (mem_idx) mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1); mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1); } void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { mem_attn->seq_keep(seq_id); + if (mem_idx) mem_idx->seq_keep(seq_id); mem_recr->seq_keep(seq_id); } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { mem_attn->seq_add(seq_id, p0, p1, shift); + if (mem_idx) mem_idx->seq_add(seq_id, p0, p1, shift); mem_recr->seq_add(seq_id, p0, p1, shift); } void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { mem_attn->seq_div(seq_id, p0, p1, d); + if (mem_idx) mem_idx->seq_div(seq_id, p0, p1, d); mem_recr->seq_div(seq_id, p0, p1, d); } @@ -184,12 +230,20 @@ std::map llama_memory_hybrid::memory_breakdo for (const auto & buft_size : mem_recr->memory_breakdown()) { mb[buft_size.first] += buft_size.second; } + if (mem_idx) { + for (const auto & buft_size : mem_idx->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + } return mb; } void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_write(io, seq_id, flags); + // indexer keys are not recomputable from the attention cache, so a restored + // session that skipped them would select the wrong cells + if (mem_idx) mem_idx->state_write(io, seq_id, flags); } mem_recr->state_write(io, seq_id, flags); } @@ -197,6 +251,7 @@ void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_read(io, seq_id, flags); + if (mem_idx) mem_idx->state_read(io, seq_id, flags); } mem_recr->state_read(io, seq_id, flags); } @@ -209,11 +264,16 @@ llama_memory_recurrent * llama_memory_hybrid::get_mem_recr() const { return mem_recr.get(); } +llama_kv_cache * llama_memory_hybrid::get_mem_idx() const { + return mem_idx.get(); +} + llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_status status) : status(status) {} llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_hybrid * mem) : ctx_attn(mem->get_mem_attn()->init_full()), ctx_recr(mem->get_mem_recr()->init_full()), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : mem->get_mem_idx()->init_full()), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } @@ -223,17 +283,26 @@ llama_memory_hybrid_context::llama_memory_hybrid_context( bool optimize) : ctx_attn(mem->get_mem_attn()->init_update(lctx, optimize)), ctx_recr(mem->get_mem_recr()->init_update(lctx, optimize)), + // indexer keys carry no positional encoding, so a shift has nothing to correct in + // them, but the pending per-cell delta must still be cleared or the two caches + // disagree about whether a shift is outstanding. an indexer only exists for + // LLAMA_ROPE_TYPE_NONE archs, which is exactly when llama_kv_cache::update skips the + // K-shift graph and does only that + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : mem->get_mem_idx()->init_update(lctx, optimize)), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } llama_memory_hybrid_context::llama_memory_hybrid_context( llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, - std::vector ubatches) : + std::vector ubatches, + slot_info_vec_t sinfos_idx) : ubatches(std::move(ubatches)), // note: here we copy the ubatches. not sure if this is ideal ctx_attn(new llama_kv_cache_context(mem->get_mem_attn(), std::move(sinfos_attn), this->ubatches)), ctx_recr(new llama_memory_recurrent_context(mem->get_mem_recr(), this->ubatches)), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), this->ubatches)), status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { } @@ -242,6 +311,7 @@ bool llama_memory_hybrid_context::next() { ctx_attn->next(); ctx_recr->next(); + if (ctx_idx) ctx_idx->next(); if (++i_next >= ubatches.size()) { return false; @@ -258,6 +328,17 @@ bool llama_memory_hybrid_context::apply() { res = res & ctx_attn->apply(); res = res & ctx_recr->apply(); + if (ctx_idx) { + res = res & ctx_idx->apply(); + + // a top-k over indexer cells is meaningful only if both caches cover the same + // window. only the batch context has slot infos to compare + if (!ubatches.empty()) { + GGML_ASSERT(get_idx()->get_n_kv() == get_attn()->get_n_kv()); + GGML_ASSERT(get_idx()->get_n_stream() == get_attn()->get_n_stream()); + } + } + return res; } @@ -277,3 +358,7 @@ const llama_kv_cache_context * llama_memory_hybrid_context::get_attn() const { const llama_memory_recurrent_context * llama_memory_hybrid_context::get_recr() const { return static_cast(ctx_recr.get()); } + +const llama_kv_cache_context * llama_memory_hybrid_context::get_idx() const { + return static_cast(ctx_idx.get()); +} diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb7499..92e576332e7 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -39,7 +39,11 @@ class llama_memory_hybrid : public llama_memory_i { bool unified, /* layer filters */ const layer_filter_cb & filter_attn = nullptr, - const layer_filter_cb & filter_recr = nullptr); + const layer_filter_cb & filter_recr = nullptr, + /* optional per-token indexer key cache for sparse-attention + hybrids; absent unless filter_idx is given */ + const layer_filter_cb & filter_idx = nullptr, + ggml_type type_idx = GGML_TYPE_F16); ~llama_memory_hybrid() = default; @@ -82,12 +86,18 @@ class llama_memory_hybrid : public llama_memory_i { llama_kv_cache * get_mem_attn() const; llama_memory_recurrent * get_mem_recr() const; + llama_kv_cache * get_mem_idx() const; // nullptr when the model has no indexer private: const llama_hparams & hparams; + // indexer cache geometry: n_head_kv key heads of indexer_head_size, as + // llama_kv_cache_dsa builds its own + llama_hparams hparams_idx; + const std::unique_ptr mem_attn; const std::unique_ptr mem_recr; + const std::unique_ptr mem_idx; }; class llama_memory_hybrid_context : public llama_memory_context_i { @@ -110,7 +120,10 @@ class llama_memory_hybrid_context : public llama_memory_context_i { llama_memory_hybrid_context( llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, - std::vector ubatches); + std::vector ubatches, + // empty without an indexer. the indexer is addressed by the + // attention cache's cells, so it gets that cache's slots + slot_info_vec_t sinfos_idx = {}); ~llama_memory_hybrid_context() = default; @@ -126,6 +139,7 @@ class llama_memory_hybrid_context : public llama_memory_context_i { const llama_kv_cache_context * get_attn() const; const llama_memory_recurrent_context * get_recr() const; + const llama_kv_cache_context * get_idx() const; // nullptr without an indexer private: // the index of the next ubatch to process @@ -135,6 +149,7 @@ class llama_memory_hybrid_context : public llama_memory_context_i { const llama_memory_context_ptr ctx_attn; const llama_memory_context_ptr ctx_recr; + const llama_memory_context_ptr ctx_idx; // null unless the model has an indexer const llama_memory_status status; }; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 9adaa93f62e..a39a6a56dda 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -291,7 +291,11 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size); add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks); + add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, hparams.indexer_kpool); add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true); + add_kv(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); + add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true); add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count); add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index c34700ff563..0d1bfee0dca 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -201,6 +201,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_deepseek4(params); case LLM_ARCH_GLM_DSA: return new llama_model_glm_dsa(params); + case LLM_ARCH_GLM5NEXT: + return new llama_model_glm5next(params); case LLM_ARCH_MISTRAL4: return new llama_model_mistral4(params); case LLM_ARCH_CHATGLM: @@ -941,6 +943,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_288B_A19B: return "288B.A19B"; case LLM_TYPE_300B_A47B: return "300B.A47B"; case LLM_TYPE_310B_A15B: return "310B.A15B"; + case LLM_TYPE_313B_A17B: return "313B.A17B"; case LLM_TYPE_355B_A32B: return "355B.A32B"; case LLM_TYPE_397B_A17B: return "397B.A17B"; case LLM_TYPE_685B_A37B: return "685B.A37B"; @@ -2431,6 +2434,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; + // null for every arch but the sparse-attention ones, which is what + // keeps the indexer cache from existing + llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; + ggml_type type_idx = GGML_TYPE_F16; if (arch == LLM_ARCH_FALCON_H1) { filter_attn = [&](uint32_t) { return true; }; filter_recr = [&](uint32_t) { return true; }; @@ -2441,16 +2448,41 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_recr = [&](uint32_t il) { return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; - } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01) { + } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01 || arch == LLM_ARCH_GLM5NEXT) { filter_attn = [&](uint32_t il) { return il < hparams.n_layer() && !hparams.is_recr(il); }; filter_recr = [&](uint32_t il) { return il < hparams.n_layer() && hparams.is_recr(il); }; + + if (arch == LLM_ARCH_GLM5NEXT && hparams.indexer_head_size > 0) { + // a unified cache is fine here: the pool map is per SEQUENCE, + // not per stream. see [TAG_KPOOL_SEQ_PARTITION] + + // only the DSA layers carry an indexer key cache + filter_idx = [&](uint32_t il) { + return il < hparams.n_layer() && !hparams.is_recr(il); + }; + + // the gate cached next to the key feeds a softmax, so -ctk + // q8_0 would quantise something far more sensitive than a + // key. keep the indexer float + type_idx = params.type_k; + if (ggml_is_quantized(type_idx)) { + LLAMA_LOG_WARN("%s: indexer key cache stays %s rather than %s: it also holds the compressor gates\n", + __func__, ggml_type_name(GGML_TYPE_F16), ggml_type_name(type_idx)); + type_idx = GGML_TYPE_F16; + } + } } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { + // llama_memory_hybrid_iswa has no indexer cache; glm5next is + // swa_type NONE, but a sparse hybrid with SWA would silently + // lose its indexer + GGML_ASSERT(filter_idx == nullptr && "hybrid-iswa cannot carry an indexer cache"); + // Use hybrid-iswa for hybrid models with SWA res = new llama_memory_hybrid_iswa( /* model */ *this, @@ -2488,7 +2520,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* offload */ cparams.offload_kqv, /* unified */ cparams.kv_unified, /* filter_attn */ std::move(filter_attn), - /* filter_recr */ std::move(filter_recr)); + /* filter_recr */ std::move(filter_recr), + /* filter_idx */ std::move(filter_idx), + /* type_idx */ type_idx); } } else { llama_kv_cache::layer_filter_cb filter = nullptr; @@ -2757,6 +2791,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_NEMOTRON_H: case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_KIMI_K3: return LLAMA_ROPE_TYPE_NONE; diff --git a/src/llama-model.h b/src/llama-model.h index 44bd9675754..1654b54d33c 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -143,6 +143,7 @@ enum llm_type { LLM_TYPE_288B_A19B, // dots3-note LLM_TYPE_300B_A47B, // Ernie MoE big LLM_TYPE_310B_A15B, // /MiMo-V2-Flash + LLM_TYPE_313B_A17B, // GLM-5.3-Flash LLM_TYPE_355B_A32B, // GLM-4.5 LLM_TYPE_397B_A17B, // Qwen3.5 LLM_TYPE_685B_A37B, // DeepSeek V3.2 diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 20252815d5c..97b8f4d22ed 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -350,6 +350,47 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param // do not quantize relative position bias (T5) quantize &= name.find("attn_rel_b.weight") == std::string::npos; + // glm5next: small, precision-sensitive tensors that must stay at source + // precision -- the mHC residual mixers, the lightning indexer (selection gate, + // learned k-pool position table, and the three indexer projections) and the KDA + // recurrence gates. ~1 GiB total on GLM-5.3-Flash, so the size cost is noise + // against a 100-240 GB quant, while quantizing them perturbs *which* pools the + // indexer selects and *how much* state each KDA step retains -- errors that + // compound over a sequence instead of averaging out. + // + // Note both spellings are required. The compressor tensors came in with the + // DeepSeek-V4 merge and are named with an UNDERSCORE (indexer_compressor_ape / + // _gate), while the projections use a DOT (indexer.proj / .attn_k / .attn_q_b). + // A single "indexer." prefix test silently misses the compressor pair. + // + // Deliberately NOT listed: attn_q_a / attn_kv_a_mqa / attn_k_b / attn_v_b. They + // are precision-sensitive too, but our release recipe pins them to q8_0 via + // --tensor-type, and that is the configuration the shipped quants were measured + // in (KLD 0.027 at Q5_K_XL). Forcing them full precision here would change quant + // sizes and invalidate those measurements. + // + // indexer.k_norm.weight needs no entry -- the generic "_norm.weight" rule above + // already excludes it. + if (arch == LLM_ARCH_GLM5NEXT) { + static const char * const glm5next_full_precision[] = { + "hc_attn_fn", + "hc_ffn_fn", + "indexer_compressor_ape", + "indexer_compressor_gate", + "indexer.proj", + "indexer.attn_k", + "indexer.attn_q_b", + "ssm_f_a", + "ssm_f_b", + "ssm_g_a", + "ssm_g_b", + "ssm_beta", + }; + for (const char * pin : glm5next_full_precision) { + quantize &= name.find(pin) == std::string::npos; + } + } + // do not quantize specific multimodal tensors quantize &= name.find(".position_embd") == std::string::npos; quantize &= name.find("sam.pos_embd") == std::string::npos; @@ -475,11 +516,11 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type // MoE tensors -> MXFP4 // other tensors -> Q8_0 // MLA projection tensors are also 3D, so match expert tensor roles explicitly. - const bool is_bailingmoe3_expert = arch == LLM_ARCH_BAILINGMOE3 && - (category == tensor_category::FFN_UP || - category == tensor_category::FFN_GATE || - category == tensor_category::FFN_DOWN); - if (tensor->ne[2] > 1 && (arch != LLM_ARCH_BAILINGMOE3 || is_bailingmoe3_expert)) { + const bool has_3d_mla = arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_GLM5NEXT; + const bool is_expert = category == tensor_category::FFN_UP || + category == tensor_category::FFN_GATE || + category == tensor_category::FFN_DOWN; + if (tensor->ne[2] > 1 && (!has_3d_mla || is_expert)) { new_type = GGML_TYPE_MXFP4; } else { new_type = GGML_TYPE_Q8_0; diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index ff926ceecd1..d37f49a5398 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2257,6 +2257,15 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { tokenizer_pre == "chatglm-bpe") { pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4; special_bos_id = LLAMA_TOKEN_NULL; + // glm4 tokenizer.json sets "ignore_merges": true. without it greedy BPE + // cannot reach some vocab entries: " 王" (Ġçİĭ, 102322) needs (Ġ,çİĭ)=242943 + // but (Ġ,ç)=27944 wins first, so it stops three tokens short. triggered by + // whitespace before CJK, inflating mixed Chinese-English ~13%. + // NOT applied to chatglm-bpe, which shares this pre_type but is a different + // tokenizer (ChatGLM3) not confirmed to declare the flag. + if (tokenizer_pre == "glm4") { + ignore_merges = true; + } } else if ( tokenizer_pre == "viking") { pre_type = LLAMA_VOCAB_PRE_TYPE_VIKING; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index fc816e2aeb4..b1a8e828bbf 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -264,7 +264,7 @@ static constexpr int64_t DSV4_CSA_RATIO = 4; static constexpr int64_t DSV4_HCA_RATIO = 128; // mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens] -static ggml_tensor * dsv4_hc_mean(ggml_context * ctx, ggml_tensor * x) { +ggml_tensor * llama_model_deepseek4::graph::build_hc_mean(ggml_context * ctx, ggml_tensor * x) { const int64_t hc = x->ne[1]; ggml_tensor * acc = ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], 0); @@ -1217,7 +1217,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl( } llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_params & params) : - llm_graph_context(params) { + llm_build_delta_net_base(params) { ggml_tensor * cur; ggml_tensor * inp = build_inp_embd(model.tok_embd); @@ -1234,7 +1234,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p for (int il = 0; il < n_layer; ++il) { if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) { - res->t_layer_inp[il] = dsv4_hc_mean(ctx0, inpL); + res->t_layer_inp[il] = build_hc_mean(ctx0, inpL); cb(res->t_layer_inp[il], "layer_inp", il); ggml_build_forward_expand(gf, res->t_layer_inp[il]); } @@ -1316,7 +1316,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p } if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) { - res->t_layer_inp[n_layer] = dsv4_hc_mean(ctx0, inpL); + res->t_layer_inp[n_layer] = build_hc_mean(ctx0, inpL); cb(res->t_layer_inp[n_layer], "layer_inp", n_layer); ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); } diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ad661264773..4b6aec8046d 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -330,9 +330,12 @@ std::pair llm_build_delta_net_base::build_delta_ne cb(b, "b_in", il); cb(g, "g_in", il); - // GDA: [1, 1, H_v, n_seqs] - // KDA: [1, S_k, H_v, n_seqs] - g = ggml_reshape_4d(ctx0, g, 1, g->ne[0], H_v, n_seqs); + // the state is indexed [key, value]: ne0 is the key axis, as the k/q reductions + // below rely on, so KDA's per-key-channel decay must broadcast over ne1, not ne0. + // for GDA g->ne[0] is 1 and both spellings agree + // GDA: [1, 1, H_v, n_seqs] + // KDA: [S_k, 1, H_v, n_seqs] + g = ggml_reshape_4d(ctx0, g, g->ne[0], 1, H_v, n_seqs); b = ggml_reshape_4d(ctx0, b, 1, 1, H_v, n_seqs); // [S_v, S_v, H_v, n_seqs] diff --git a/src/models/glm5next.cpp b/src/models/glm5next.cpp new file mode 100644 index 00000000000..2cb8d4c20b3 --- /dev/null +++ b/src/models/glm5next.cpp @@ -0,0 +1,873 @@ +#include "models.h" + +#include "llama-memory-recurrent.h" +#include "llama-memory-hybrid.h" +#include "llama-kv-cache-kpool.h" + +// +// GLM-5.3-Flash: hybrid KDA (linear) + DSA (nope-only MLA) attention, mHC +// hyper-connections, and a NextN block that is a full DSA decoder layer. +// +// ssm_a holds -exp(A_log), the kimi-k3 convention, so the decay gate reads +// exp(A_log) back as -ssm_a; bailingmoe3 stores +exp(A_log). indistinguishable at +// load time, so conversion/glm5next.py is the only place the sign is checked. +// + +// positions the indexer keeps: index_topk/index_kpool whole pools plus the +// always-selected tail pool minus one. below this many cached tokens every position is +// selected, so sparse selection is exactly the dense path built here. +// +// asserted, not measured: an off-by-one here is invisible to output comparison (the +// reference's own off-by-one on this width is bit-identical on both fixtures). the +// second assert is the parity harness's independent spelling, so the two must agree +static uint32_t glm5next_n_select(const llama_hparams & hparams) { + GGML_ASSERT(hparams.indexer_kpool > 0); + GGML_ASSERT(hparams.indexer_top_k >= hparams.indexer_kpool); + GGML_ASSERT(hparams.indexer_top_k % hparams.indexer_kpool == 0); + + const uint32_t n_select = hparams.indexer_top_k + hparams.indexer_kpool - 1; + + GGML_ASSERT(n_select > hparams.indexer_top_k); + GGML_ASSERT(n_select == (hparams.indexer_top_k/hparams.indexer_kpool + 1)*hparams.indexer_kpool - 1); + + return n_select; +} + +void llama_model_glm5next::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + // indexer k_norm is a LayerNorm with bias; without this key it runs at eps 0 + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); + // The reference HARDCODES nn.LayerNorm(index_head_dim, eps=1e-6); it does not + // read it from the config, and rms_norm_eps is 1e-5. Warned about rather than + // asserted, for two reasons. No output comparison can see it: seeding eps=1e-5 + // into the reference leaves the logits BIT-IDENTICAL at 512 tokens and moves + // the index jaccard by 1.7e-5 at 2048, against a bf16 floor of 0.63. And an + // assert would abort test-llama-archs, whose synthetic models have no reason + // to carry a converter contract + if (hparams.f_norm_eps <= 0.0f || hparams.f_norm_eps > 2e-6f) { + LLAMA_LOG_WARN("%s: indexer k_norm eps is %g, but the reference hardcodes 1e-6. " + "this is invisible to every output comparison; check the converter\n", + __func__, (double) hparams.f_norm_eps); + } + + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + GGML_ASSERT(hparams.n_lora_q > 0 && "glm5next requires a q LoRA"); + GGML_ASSERT(hparams.n_rot() == 0 && "glm5next MLA is nope-only"); + + // KDA. no GGUF key for linear_num_heads, so the KDA head count is + // attention.head_count, which also sizes the recurrent state via n_embd_r/s(). + // conversion/glm5next.py refuses a checkpoint where the two differ + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + GGML_ASSERT(hparams.ssm_d_conv > 1); + GGML_ASSERT(hparams.n_embd_head_kda > 0); + // required: absent, kimi-k3 selects the softplus branch, a different + // function, not a missing clamp + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); + GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f); + + // DSA indexer + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KPOOL, hparams.indexer_kpool); + GGML_ASSERT(hparams.indexer_kpool > 0); + GGML_ASSERT(hparams.indexer_top_k % hparams.indexer_kpool == 0); + + // below n_select resident positions the indexer selects every visible one, so + // the dense path is not an approximation of the sparse one, it IS it + const uint32_t n_select = glm5next_n_select(hparams); + LLAMA_LOG_INFO("%s: indexer selection width = %u cells (%u pools of %u, plus a %u-wide tail)\n", + __func__, n_select, hparams.indexer_top_k/hparams.indexer_kpool, + hparams.indexer_kpool, hparams.indexer_kpool - 1); + + // mHC + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); + ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); + GGML_ASSERT(hparams.dsv4_hc_mult > 0); + + // trunk residual is hc_mult streams wide (deepseek4); lm_head still sees + // n_embd, the streams are averaged first -- so n_embd_out stays n_embd. + // + // deepseek4 sets this to hc_mult*n_embd, and inheriting that here was a bug: + // our t_embd is build_norm(build_hc_mean(...)), which is [n_embd, n_tokens], + // but n_embd_out() would report 4*n_embd, so llama-context.cpp reads + // n_outputs*n_embd_out floats out of a tensor holding a quarter of that. + // The assert there sizes the DESTINATION, so nothing catches the short SOURCE. + // Only --embeddings / llama_get_embeddings* reach it, which is why plain + // generation never showed it. + // + // deepseek4 needs the wide value to size its MTP `h` input; when our NextN + // graph starts consuming it, give MTP its own width rather than widening this. + hparams.n_embd_out_impl = 0; + + // MoE + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false); + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false); + + if (hparams.n_ff_shexp == 0) { + hparams.n_ff_shexp = hparams.n_ff_exp * std::max(1u, hparams.n_expert_shared); + } + + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + GGML_ASSERT(hparams.n_layer_nextn < hparams.n_layer_all); + + // n_head_kv == 0 marks a KDA (recurrent) layer, as in kimi-k3 and bailingmoe3. + // a scalar head_count_kv would make every layer look like DSA, so require both + uint32_t n_recr = 0; + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0; + n_recr += hparams.is_recr_impl[il]; + } + GGML_ASSERT(n_recr > 0 && n_recr < hparams.n_layer() && "glm5next needs a per-layer attention.head_count_kv array"); + + // every glm5next indexer is full; glm-dsa gates its indexer on this + // predicate and the generic loader only zero-fills the array + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + hparams.is_indexer_full_impl[il] = !hparams.is_recr_impl[il]; + } + + switch (hparams.n_layer()) { + case 45: type = hparams.n_embd == 4096 && hparams.n_expert == 288 ? LLM_TYPE_313B_A17B : LLM_TYPE_UNKNOWN; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_glm5next::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = head_dim * n_head; + const int64_t d_conv = hparams.ssm_d_conv; + + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t v_head_dim = hparams.n_embd_head_v_mla(); + + const int64_t n_embd_indexer = hparams.indexer_head_size; + const int64_t kpool = hparams.indexer_kpool; + + const int64_t hc_dim = (int64_t) hparams.dsv4_hc_mult * n_embd; + const int64_t hc_mix_dim = (2 + (int64_t) hparams.dsv4_hc_mult) * hparams.dsv4_hc_mult; + + // the trunk and the NextN block can be split across two GGUFs in either direction + const bool mtp_only = (n_layer_nextn > 0) && (ml.get_weight("blk.0.attn_norm.weight") == nullptr); + const std::string mtp_probe = "blk." + std::to_string(n_layer) + ".nextn.eh_proj.weight"; + const bool trunk_only = (n_layer_nextn > 0) && (ml.get_weight(mtp_probe.c_str()) == nullptr); + const int trunk_flags = mtp_only ? TENSOR_NOT_REQUIRED : 0; + int mtp_flags = trunk_only ? TENSOR_NOT_REQUIRED : 0; + + if (!ml.load_mtp) { + mtp_flags |= TENSOR_SKIP; + } + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + if (!output) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int il = 0; il < n_layer_all; ++il) { + auto & layer = layers[il]; + const int flags = il < n_layer ? trunk_flags : mtp_flags; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", il), {n_embd}, flags); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", il), {n_embd}, flags); + + // the NextN block keeps the plain residual, so it has no mHC mixer + if (il < n_layer) { + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", il), {hc_dim, hc_mix_dim}, flags); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", il), {hc_mix_dim}, flags); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", il), {3}, flags); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", il), {hc_dim, hc_mix_dim}, flags); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", il), {hc_mix_dim}, flags); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", il), {3}, flags); + } + + if (hparams.is_recr(il)) { + create_tensor_qkv(layer, il, n_embd, d_inner, d_inner, d_inner, flags); + + layer.ssm_q_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_Q, "weight", il), {d_conv, 1, d_inner, 1}, flags); + layer.ssm_k_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_K, "weight", il), {d_conv, 1, d_inner, 1}, flags); + layer.ssm_v_conv = create_tensor(tn(LLM_TENSOR_SSM_CONV1D_V, "weight", il), {d_conv, 1, d_inner, 1}, flags); + + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", il), {n_embd, head_dim}, flags); + layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", il), {head_dim, d_inner}, flags); + layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", il), {n_embd, head_dim}, flags); + layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", il), {head_dim, d_inner}, flags); + + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), {n_embd, n_head}, flags); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, il), {n_head}, flags); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), {d_inner}, flags); + + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), {head_dim}, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), {d_inner, n_embd}, flags); + } else { + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", il), {n_embd, q_lora_rank}, flags); + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", il), {q_lora_rank}, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", il), {q_lora_rank, n_head * qk_head_dim}, flags); + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", il), {n_embd, kv_lora_rank}, flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", il), {kv_lora_rank}, flags); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", il), {qk_head_dim, kv_lora_rank, n_head}, flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", il), {kv_lora_rank, v_head_dim, n_head}, flags); + + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), {n_head * v_head_dim, n_embd}, flags); + + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), {n_embd_indexer}, flags); + layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", il), {n_embd_indexer}, flags); + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", il), {n_embd, hparams.indexer_n_head}, flags); + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", il), {n_embd, n_embd_indexer}, flags); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", il), {q_lora_rank, hparams.indexer_n_head * n_embd_indexer}, flags); + + // key pooling: DeepSeek-V4 doubles the compressor width, GLM-5.3 does not + layer.indexer_comp_wgate = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "weight", il), {n_embd, n_embd_indexer}, flags); + layer.indexer_comp_ape = create_tensor(tn(LLM_TENSOR_INDEXER_COMPRESSOR_APE, "weight", il), {n_embd_indexer, kpool}, flags); + } + + if (il < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", il), {n_embd, n_ff}, flags); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", il), {n_embd, n_ff}, flags); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", il), {n_ff, n_embd}, flags); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), {n_embd, n_expert}, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", il), {n_expert}, flags); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", il), {n_embd, hparams.n_ff_exp, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", il), {n_embd, hparams.n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), {hparams.n_ff_exp, n_embd, n_expert}, flags); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), {n_embd, hparams.n_ff_shexp}, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), {n_embd, hparams.n_ff_shexp}, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), {hparams.n_ff_shexp, n_embd}, flags); + } + + if (il >= n_layer) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), {2 * n_embd, n_embd}, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), {n_embd}, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), {n_embd}, flags); + + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", il), {n_embd}, flags); + // absent in the checkpoint: NextN shares the trunk's embeddings and + // lm_head. only accepted if an export adds them + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", il), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", il), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); + } + } +} + +// +// KDA layer +// +// one depthwise conv over the concatenated q|k|v channels, as in the reference: it +// leaves the conv state one contiguous block, which is what build_conv_state needs to +// snapshot for recurrent-state rollback. three separate convs would be numerically +// identical but would restate the rollback write three times +// +ggml_tensor * llama_model_glm5next::graph::build_kda_layer( + const llama_layer & layer, + llm_graph_input_rs * inp_rs, + ggml_tensor * cur, + int il) { + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_inner = head_dim * n_head; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + const auto * mctx_cur = inp_rs->mctx; + + // f, g and beta read the layer input, NOT the convolved q/k/v + ggml_tensor * inp = cur; + + ggml_tensor * Qcur = ggml_mul_mat(ctx0, layer.wq, inp); + ggml_tensor * Kcur = ggml_mul_mat(ctx0, layer.wk, inp); + ggml_tensor * Vcur = ggml_mul_mat(ctx0, layer.wv, inp); + + ggml_tensor * qkv = ggml_concat(ctx0, ggml_concat(ctx0, Qcur, Kcur, 0), Vcur, 0); + qkv = ggml_reshape_3d(ctx0, qkv, 3*d_inner, n_seq_tokens, n_seqs); + cb(qkv, "kda_qkv", il); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_in = build_conv_state(inp_rs, conv_states_all, qkv, d_conv, 3*d_inner, il); + + // stored separately (kimi-linear, kimi-k3), stacked back into the single kernel + ggml_tensor * conv_w = ggml_concat(ctx0, + ggml_concat(ctx0, + ggml_reshape_2d(ctx0, layer.ssm_q_conv, d_conv, d_inner), + ggml_reshape_2d(ctx0, layer.ssm_k_conv, d_conv, d_inner), 1), + ggml_reshape_2d(ctx0, layer.ssm_v_conv, d_conv, d_inner), 1); + + // SiLU is applied to the conv output, not to the projections + ggml_tensor * conv_out = ggml_silu(ctx0, ggml_ssm_conv(ctx0, conv_in, conv_w)); + cb(conv_out, "kda_conv", il); + + const size_t nb_qkv = ggml_row_size(conv_out->type, 3*d_inner); + const size_t nb_head = ggml_row_size(conv_out->type, head_dim); + + Qcur = ggml_view_4d(ctx0, conv_out, head_dim, n_head, n_seq_tokens, n_seqs, + nb_head, nb_qkv, nb_qkv*n_seq_tokens, 0); + Kcur = ggml_view_4d(ctx0, conv_out, head_dim, n_head, n_seq_tokens, n_seqs, + nb_head, nb_qkv, nb_qkv*n_seq_tokens, ggml_row_size(conv_out->type, d_inner)); + Vcur = ggml_view_4d(ctx0, conv_out, head_dim, n_head, n_seq_tokens, n_seqs, + nb_head, nb_qkv, nb_qkv*n_seq_tokens, ggml_row_size(conv_out->type, 2*d_inner)); + + // 1e-6 is the reference's own constant, not the model's norm eps. ggml_l2_norm + // divides by max(sqrt(sum), eps) where the reference uses sqrt(sum + eps); at + // head_dim 128 the clamp never binds, so close but not bit-exact + Qcur = ggml_l2_norm(ctx0, Qcur, 1e-6f); + Kcur = ggml_l2_norm(ctx0, Kcur, 1e-6f); + cb(Qcur, "kda_q_norm", il); + cb(Kcur, "kda_k_norm", il); + + // the 1/sqrt(head_dim) query scale is applied inside build_delta_net, after this norm + + // forget gate. gate_lower_bound is a multiplicative scale, not a clamp: + // g = lower_bound * sigmoid(exp(A_log) * (f_b(f_a(x)) + dt_bias)) + // ssm_a holds -exp(A_log), so exp(A_log) * y == -(y * ssm_a) + ggml_tensor * g = ggml_mul_mat(ctx0, layer.ssm_f_b, ggml_mul_mat(ctx0, layer.ssm_f_a, inp)); + g = ggml_add(ctx0, g, layer.ssm_dt_b); + g = ggml_reshape_3d(ctx0, g, head_dim, n_head, n_tokens); + g = ggml_mul(ctx0, g, ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head, 1)); + g = ggml_sigmoid(ctx0, ggml_scale(ctx0, g, -1.0f)); + g = ggml_scale(ctx0, g, hparams.kda_gate_lower_bound); + g = ggml_reshape_4d(ctx0, g, head_dim, n_head, n_seq_tokens, n_seqs); + cb(g, "kda_gate", il); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, inp); + beta = ggml_sigmoid(ctx0, ggml_reshape_4d(ctx0, beta, 1, n_head, n_seq_tokens, n_seqs)); + cb(beta, "kda_beta", il); + + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head, n_seqs); + + ggml_tensor * out = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g, beta, state, il); + + // the fallbacks return a permuted view, the fused op a contiguous one; cont + // either way rather than depend on which ran + ggml_tensor * o = ggml_cont_3d(ctx0, out, head_dim, n_head, n_tokens); + cb(o, "kda_scan_out", il); + + // low-rank output gate (kimi-k3 has a single full-rank ssm_g instead) + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.ssm_g_b, ggml_mul_mat(ctx0, layer.ssm_g_a, inp)); + gate = ggml_reshape_3d(ctx0, gate, head_dim, n_head, n_tokens); + + // RMS over head_dim only, one weight shared by every head, then a plain sigmoid + // gate: not the SiLU that FusedRMSNormGated defaults to + ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, gate)); + cb(gated, "kda_normed", il); + + cur = ggml_mul_mat(ctx0, layer.wo, ggml_cont_2d(ctx0, gated, d_inner, n_tokens)); + cb(cur, "kda_out", il); + + return cur; +} + +// +// the lightning indexer, pooled +// +// Writes this layer's indexer key and compressor gate into the indexer cache, and +// then - when the cache is large enough for selection to bind - returns the I32 +// ATTENTION-cache cell indices this layer's queries select, ready for the scatter in +// build_attn_sparse. `qr` is the shared q LoRA residual, i.e. q_a_norm(q_a_proj(x)), +// which the indexer's own wq_b consumes; `cur` is the layer input. +// +// The store is NOT gated on the sparse path and the scoring is. Gating both the same +// way leaves every cell written below n_select - the first 2051 positions of every +// sequence on the real model - with no indexer state at all, and the first ubatch to +// cross n_select then pools cells that were never written. +// +// Three points where the reference implementations disagree, resolved 2-of-3 by +// reading transformers (modular_glm5_next.py Glm5NextTextIndexer), sglang +// (dsa_indexer_kpool.py) and vLLM (glm5next/nvidia/attention.py): +// +// * the weights_proj GEMM runs in fp32. sglang gives it params_dtype=fp32 and feeds +// x.float(); vLLM does the same and says why - bf16 head-gates move logits by +// ~1e-2, enough to flip near-tie pool rankings on long context. transformers is +// the outlier and rounds the activation to bf16. +// * k_norm is a LayerNorm WITH BIAS at eps 1e-6, hardcoded in transformers and in +// vLLM; sglang leaves torch's 1e-5 default. It is NOT f_norm_rms_eps (1e-5 here) +// and NOT ggml's 0 default. The value comes from the GGUF and load_arch_hparams +// warns when it is not 1e-6, because no output comparison can see it. +// * the ReLU between the QK dot and the head weighting is explicit in transformers +// and in sglang's non-pooled indexer; in the pooled path both engines hand it to +// the DeepGEMM MQA-logits kernel, so neither validates it. It is easy to drop and +// is written out here. +// +// No Hadamard rotation: sglang and vLLM rotate q and k by H128 before their fp8 +// kernel, but H is orthogonal so (Hq).(Hk) == q.k. It exists to spread magnitude ahead +// of fp8 quantisation; scoring in f32 here it would only cost accuracy. transformers, +// the semantic reference, has none. +// +ggml_tensor * llama_model_glm5next::graph::build_indexer( + const llama_layer & layer, + llm_graph_input_kpool * inp_kp, + ggml_tensor * cur, + ggml_tensor * qr, + bool scoring, + int il) const { + const int64_t d_idx = hparams.indexer_head_size; + const int64_t n_ihead = hparams.indexer_n_head; + const int64_t r = hparams.indexer_kpool; + + const auto * mctx_idx = inp_kp->mctx_idx; + + // a genuine LayerNorm: weight AND bias. glm-dsa runs this same norm at eps 0 today + GGML_ASSERT(layer.indexer_k_norm_b != nullptr && "the indexer k_norm is a LayerNorm with bias"); + + ggml_tensor * ik = build_norm(ggml_mul_mat(ctx0, layer.indexer_attn_k, cur), + layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il); + cb(ik, "indexer_k", il); + + // the pooling gate is a SECOND, INDEPENDENT projection of the hidden state, not a + // reuse of the indexer key. it has to be cached beside the key: the compressor + // mixes a pool's member keys with a softmax over these gates, and a pool is only + // rebuilt once its members have left the batch + ggml_tensor * gate = ggml_mul_mat(ctx0, layer.indexer_comp_wgate, cur); + cb(gate, "indexer_gate", il); + + // {d_idx, 2, n_tokens}: head 0 is the key, head 1 the gate. llama_memory_hybrid + // allocates the second head exactly for this and only when indexer_kpool > 0 + ggml_tensor * packed = ggml_concat(ctx0, + ggml_reshape_3d(ctx0, ik, d_idx, 1, n_tokens), + ggml_reshape_3d(ctx0, gate, d_idx, 1, n_tokens), 1); + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, packed, inp_kp->k_idxs, il)); + + if (!scoring) { + return nullptr; + } + + // {d_idx, 2, n_kv, n_stream} + ggml_tensor * kbuf = mctx_idx->get_k(ctx0, il); + + const int64_t n_kv = kbuf->ne[2]; + const int64_t n_stream = kbuf->ne[3]; + const int64_t n_tps = n_tokens/n_stream; + const int64_t n_pools = inp_kp->pool_cells->ne[0]/r; + + GGML_ASSERT(kbuf->ne[0] == d_idx && kbuf->ne[1] == 2 && + "the pooled indexer cache needs a key head and a gate head"); + GGML_ASSERT(kbuf->nb[1] == (size_t) d_idx*kbuf->nb[0] && "key and gate must be adjacent in a cell"); + GGML_ASSERT(n_tokens == n_tps*n_stream); + + // one cell's key and gate as a single row, so that the members of a pool are + // gathered once rather than twice: {2*d_idx, n_kv, n_stream} + ggml_tensor * kg_rows = ggml_view_3d(ctx0, kbuf, 2*d_idx, n_kv, n_stream, + kbuf->nb[2], kbuf->nb[3], 0); + + // gather each pool's members. pool_cells names a cell per (pool, slot); slots that + // are not resident hold 0 rather than a negative sentinel, because ggml_get_rows + // has none. The resulting garbage pools are neutralised by pool_bias below, never + // by a NaN: unlike the reference, no -inf ever enters the compressor softmax. + // ggml_get_rows always yields F32, so the pooling runs in F32 even though the + // indexer cache is F16 + ggml_tensor * members = ggml_get_rows(ctx0, kg_rows, inp_kp->pool_cells); + cb(members, "indexer_pool_members", il); + + const size_t nb_mem = members->nb[1]; + + ggml_tensor * mem_k = ggml_view_4d(ctx0, members, d_idx, r, n_pools, n_stream, + nb_mem, nb_mem*r, members->nb[2], 0); + ggml_tensor * mem_g = ggml_view_4d(ctx0, members, d_idx, r, n_pools, n_stream, + nb_mem, nb_mem*r, members->nb[2], d_idx*members->nb[0]); + + // d_idx independent r-way softmaxes over the SLOT axis, so the slot axis has to be + // dim 0. ape is added PRE-softmax and is indexed by LOGICAL SLOT, so it broadcasts + // over pools and streams; pool_cells is built in position order, so slot m is + // position p % kpool and the two agree by construction + ggml_tensor * keys_t = ggml_cont(ctx0, ggml_permute(ctx0, mem_k, 1, 0, 2, 3)); + ggml_tensor * gate_t = ggml_cont(ctx0, ggml_permute(ctx0, mem_g, 1, 0, 2, 3)); + + ggml_tensor * ape = ggml_cont(ctx0, ggml_transpose(ctx0, layer.indexer_comp_ape)); + gate_t = ggml_add(ctx0, gate_t, ggml_reshape_4d(ctx0, ape, r, d_idx, 1, 1)); + + ggml_tensor * probs = ggml_soft_max(ctx0, gate_t); + cb(probs, "indexer_pool_probs", il); + + // per-channel weighted average over the pool's members -> {d_idx, n_pools, 1, n_stream} + ggml_tensor * pool_k = ggml_sum_rows(ctx0, ggml_mul(ctx0, keys_t, probs)); + pool_k = ggml_reshape_4d(ctx0, pool_k, d_idx, n_pools, 1, n_stream); + cb(pool_k, "indexer_pool_k", il); + + // {d_idx, n_ihead, n_tps, n_stream}. no rope: n_rot() is 0 for the whole text tower + ggml_tensor * iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr); + iq = ggml_reshape_4d(ctx0, iq, d_idx, n_ihead, n_tps, n_stream); + cb(iq, "indexer_q", il); + + // sign-unconstrained head weights: no softmax, no abs, no relu. Both scale + // constants - the reference's softmax_scale = d_idx^-0.5 and its n_heads^-0.5 head + // factor - are folded in here, on an {n_ihead, n_tokens} tensor rather than on the + // {n_pools, n_tps, n_ihead} score tensor. relu is positively homogeneous and both + // constants are positive, so this is exactly the same function, and it is what the + // engines and the in-tree glm-dsa both do. + // + // GGML_PREC_F32 is not cosmetic: on vLLM a bf16 head gate moves a logit by ~1e-2, + // which is enough to swap two near-tied pools, and the top-k below is a hard cut + ggml_tensor * w = ggml_mul_mat(ctx0, layer.indexer_proj, cur); + ggml_mul_mat_set_prec(w, GGML_PREC_F32); + w = ggml_reshape_4d(ctx0, w, n_ihead, n_tps, 1, n_stream); + w = ggml_scale(ctx0, w, 1.0f/sqrtf(float(d_idx*n_ihead))); + cb(w, "indexer_weights", il); + + // both paths end with pool_bias added: -INFINITY on every pool the reference's + // `pool_valid & pool_visible` rejects, the query's own trailing pool included, so + // that no budget is spent on it + ggml_tensor * pool_score = nullptr; + + if (cparams.fused_lid) { + // one node for the whole dot product -> relu -> head sum -> mask chain, and the + // mask add comes for free. pool_k stays f32, so the kernel takes its f32 vector + // path rather than the f16 wmma path, which would undo the GGML_PREC_F32 above + ggml_tensor * pool_kf = ggml_reshape_4d(ctx0, pool_k, d_idx, 1, n_pools, n_stream); + + pool_score = ggml_lightning_indexer(ctx0, iq, pool_kf, w, inp_kp->pool_bias_f16); + cb(pool_score, "indexer_pool_score", il); + + res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, pool_score, il}); + + pool_score = ggml_reshape_3d(ctx0, pool_score, n_pools, n_tps, n_stream); + } else { + // {n_pools, n_tps, n_ihead, n_stream}: pool_k is MQA and broadcasts over the heads + ggml_tensor * kq = ggml_mul_mat(ctx0, pool_k, ggml_permute(ctx0, iq, 0, 2, 1, 3)); + + // {n_ihead, n_tps, n_pools, n_stream}, contiguous for the relu and the head sum. + // the ReLU sits BETWEEN the per-head dot product and the head weighting: moving it + // to either side is a different function, because the head weights are sign-free + // and the sum is not a convex combination + kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); + ggml_tensor * score = ggml_relu(ctx0, kq); + cb(score, "indexer_score", il); + + // {1, n_tps, n_pools, n_stream} -> {n_pools, n_tps, n_stream} + pool_score = ggml_sum_rows(ctx0, ggml_mul(ctx0, score, w)); + pool_score = ggml_cont(ctx0, ggml_permute(ctx0, pool_score, 2, 1, 0, 3)); + pool_score = ggml_reshape_3d(ctx0, pool_score, n_pools, n_tps, n_stream); + + pool_score = ggml_add(ctx0, pool_score, inp_kp->pool_bias); + cb(pool_score, "indexer_pool_score", il); + } + + // Top-k over POOLS at index_topk/index_kpool, then expand each selected pool to its + // members. This is the reference's own two-step (topk over the pool axis, then + // selected_indices = pool_indices[batch_idx, selected]) and it is NOT + // interchangeable with a single top-k of width index_topk over member cells. + // + // The cell-level form is the tempting one and the argument for it is false. It says + // a pool's members carry its score bit-exactly, so the cut must land on a pool + // boundary. But F.relu drives most pool scores to exactly 0.0, so tie groups SPAN + // pools, the cut falls inside one, and ggml_top_k - explicitly unordered among + // equals - takes an arbitrary 1..kpool-1 members of the pool it lands in. A + // pool-aligned top-k WIDTH does not save it: the ties are not aligned to anything. + // Measured on TinySparse at 512 tokens, the cell-level form leaves a partial pool + // on 7.51% of query rows at L3 and 5.93% at L7 (70 and 47 partial pools); this form + // leaves none. The index-set jaccard does not separate them - it scored the broken + // form at 0.9958/0.9764 against a bf16 noise floor of 0.9779/0.8385, i.e. ABOVE the + // floor - which is why scripts/glm5next_pool_integrity.py exists + const int64_t select_k = llama_kpool_select_k(n_pools, hparams.indexer_top_k, r); + GGML_ASSERT(select_k > 0 && select_k <= n_pools); + + // {select_k, n_tps, n_stream} of POOL ordinals + ggml_tensor * sel = ggml_cont(ctx0, ggml_top_k(ctx0, pool_score, (int) select_k)); + cb(sel, "indexer_top_k_pools", il); + + // Expand pools to members: gather whole rows of `kpool` cells out of pool_cells. + // The query axis folds into the gather's row axis, which is what lets ONE + // ggml_get_rows serve every query, while the stream axis stays where get_rows wants + // it (src0 dim 2 is indexed by the index tensor's dim 1) + ggml_tensor * pc3 = ggml_reshape_3d(ctx0, inp_kp->pool_cells, r, n_pools, n_stream); + ggml_tensor * sel_flat = ggml_reshape_2d(ctx0, sel, select_k*n_tps, n_stream); + + // {r, select_k*n_tps, n_stream} -> {r*select_k, n_tps, n_stream} + ggml_tensor * top_k = ggml_get_rows(ctx0, pc3, sel_flat); + GGML_ASSERT(top_k->type == GGML_TYPE_I32 && "pool_cells is I32, so the gather stays I32"); + top_k = ggml_reshape_3d(ctx0, top_k, r*select_k, n_tps, n_stream); + cb(top_k, "indexer_top_k", il); + + return top_k; +} + +// +// DSA layer. `scoring` false takes the dense limit: full MLA over the whole cache, no +// kpool, no top-k, which is what sparse selection collapses to below +// glm5next_n_select() resident positions +// +// absorbed form, as in deepseek2/deepseek32/glm-dsa: q_nope is pushed through wk_b so +// q.k is taken against the 512-wide latent the cache actually holds (is_mla() drops the +// V allocation and V becomes a view of K). the naive form would re-expand the latent to +// n_head 256-wide k/v every step and needs a V cache this layout does not have +// +ggml_tensor * llama_model_glm5next::graph::build_dsa_layer( + const llama_layer & layer, + llm_graph_input_attn_k * inp_attn, + llm_graph_input_kpool * inp_kp, + bool scoring, + ggml_tensor * cur, + int il) const { + const int64_t qk_head_dim = hparams.n_embd_head_k_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + + // nope-only: the rope half is zero-width, so no split, no concat and no rope + // anywhere in the text tower + GGML_ASSERT(hparams.n_rot() == 0); + + // scale is over the MLA head size, as in the reference, not over the post-absorption + // width: 1/sqrt(kv_lora_rank) = 1/sqrt(512) would be a different model + const float kq_scale = 1.0f/sqrtf(float(qk_head_dim)); + + ggml_tensor * qr = ggml_mul_mat(ctx0, layer.wq_a, cur); + qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(qr, "dsa_q_a_norm", il); + + // the indexer shares this q LoRA residual with the main MLA path and consumes it + // with its own wq_b, so it is built here rather than being handed the layer input + // twice. it also writes the indexer cache, which happens on the dense path too + ggml_tensor * top_k = inp_kp ? build_indexer(layer, inp_kp, cur, qr, scoring, il) : nullptr; + + ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr); + q = ggml_reshape_3d(ctx0, q, qk_head_dim, n_head, n_tokens); + cb(q, "dsa_q_b", il); + + ggml_tensor * kv = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + kv = build_norm(kv, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + cb(kv, "dsa_kv_a_norm", il); + + // {qk_head_dim, n_tokens, n_head} + q = ggml_permute(ctx0, q, 0, 2, 1, 3); + + // {qk_head_dim, kv_lora_rank, n_head} x {qk_head_dim, n_tokens, n_head} + q = ggml_mul_mat(ctx0, layer.wk_b, q); + + // {kv_lora_rank, n_head, n_tokens}. deepseek2 gets this contiguous for free from the + // concat with the roped half, which does not exist here + q = ggml_cont(ctx0, ggml_permute(ctx0, q, 0, 2, 1, 3)); + cb(q, "dsa_q_absorbed", il); + + // absorbed MLA is MQA: one head of keys, and V is the same latent row as K + ggml_tensor * k = ggml_reshape_3d(ctx0, kv, kv_lora_rank, 1, n_tokens); + cb(k, "dsa_kv_latent", il); + + if (top_k) { + cur = build_attn_sparse(inp_attn, + layer.wo, nullptr, nullptr, + q, k, k, nullptr, nullptr, layer.wv_b, + top_k, inp_kp->sel_mask, inp_kp->cand_mask, kq_scale, il); + } else { + cur = build_attn(inp_attn, + layer.wo, nullptr, nullptr, + q, k, k, nullptr, nullptr, layer.wv_b, kq_scale, il); + } + cb(cur, "dsa_out", il); + + return cur; +} + +ggml_tensor * llama_model_glm5next::graph::build_layer_attn( + const llama_model & model, + llm_graph_input_mem_hybrid_k * inp_mem, + llm_graph_input_kpool * inp_kp, + bool scoring, + ggml_tensor * cur, + int il) { + if (hparams.is_recr(il)) { + return build_kda_layer(model.layers[il], inp_mem->get_recr(), cur, il); + } + + return build_dsa_layer(model.layers[il], inp_mem->get_attn(), inp_kp, scoring, cur, il); +} + +ggml_tensor * llama_model_glm5next::graph::build_layer_ffn( + const llama_model & model, + ggml_tensor * cur, + int il) const { + const auto & layer = model.layers[il]; + + // the leading dense layers clamp the same way the experts do: the reference + // routes both through one Glm5NextTextMLP, so swiglu_limit is not MoE-only + if (il < (int) hparams.n_layer_dense_lead) { + return build_ffn(cur, + layer.ffn_up, nullptr, nullptr, + layer.ffn_gate, nullptr, nullptr, + layer.ffn_down, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + } + + // noaux_tc: exp_probs_b biases top-k selection only, the weights are the + // unbiased sigmoid scores. n_group is 1, so the group mask is a no-op + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, hparams.n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + + ggml_tensor * shexp = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(shexp, "ffn_shexp", il); + + // shared expert unscaled: routed_scaling_factor is applied inside build_moe_ffn, + // after norm_topk_prob, to the routed weights only + return ggml_add(ctx0, moe_out, shexp); +} + +llama_model_glm5next::graph::graph(const llama_model & model, const llm_graph_params & params) : + llama_model_deepseek4::graph(params) { + ggml_tensor * cur; + + ggml_tensor * inp = build_inp_embd(model.tok_embd); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // MLA absorption leaves a K-only cache holding the latent, so the attention half of + // the hybrid memory is the _k variant, as in bailingmoe3 + llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); + + // One pooling map for the whole ubatch. pool_cells, pool_bias, sel_mask and + // cand_mask depend only on the cells and the ubatch, never on the layer, so + // rebuilding them per DSA layer would cost O(n_kv * n_tokens) host writes eleven + // times over - at 128 Ki cells and 512 tokens that dominates prefill on its own. + // + // The map is built whenever the model HAS an indexer cache, because the indexer + // key and gate store runs unconditionally. Only `scoring` is gated: below + // index_topk + index_kpool - 1 resident positions the reference selects every + // visible position, so the dense build_attn is not an approximation there, it is + // the same function. + // + // Gated on n_ctx and not on the ubatch's n_kv, even though n_kv is what actually + // decides whether the budget binds. n_kv grows as the cache fills, so gating on it + // would flip the graph's topology partway through a run. n_ctx is fixed for the + // lifetime of the context, so the topology is decided once. The cost is that a + // context configured larger than n_select runs the indexer even while the cache is + // still short, where it selects every visible pool: wasted work, never a wrong + // answer, and llama_kpool_select_k clamps the budget to the pools that exist. + llm_graph_input_kpool * inp_kp = nullptr; + bool indexer_scoring = false; + { + const auto * mctx_hyb = static_cast(mctx); + + if (mctx_hyb->get_idx() != nullptr) { + indexer_scoring = cparams.n_ctx > glm5next_n_select(hparams); + + inp_kp = build_inp_kpool(mctx_hyb, + inp_mem->get_attn()->get_kq_mask(), indexer_scoring); + } + } + + GGML_ASSERT(ubatch.n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == ubatch.n_seq_tokens * ubatch.n_seqs); + + const int64_t hc = hparams.dsv4_hc_mult; + + // hc_mult exact copies of the embedding: no scaling, no one-hot into stream 0 + ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens); + inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1); + cb(inpL, "hc_init", -1); + + for (int il = 0; il < n_layer; ++il) { + if ((size_t) il < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[il]) { + res->t_layer_inp[il] = build_hc_mean(ctx0, inpL); + cb(res->t_layer_inp[il], "layer_inp", il); + ggml_build_forward_expand(gf, res->t_layer_inp[il]); + } + + ggml_tensor * residual = inpL; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + + cur = build_hc_pre(inpL, + model.layers[il].hc_attn_fn, + model.layers[il].hc_attn_scale, + model.layers[il].hc_attn_base, + &post, &comb, il); + cb(cur, "hc_attn_pre", il); + + cur = build_norm(cur, model.layers[il].attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + cur = build_layer_attn(model, inp_mem, inp_kp, indexer_scoring, cur, il); + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "hc_attn_post", il); + + residual = inpL; + cur = build_hc_pre(inpL, + model.layers[il].hc_ffn_fn, + model.layers[il].hc_ffn_scale, + model.layers[il].hc_ffn_base, + &post, &comb, il); + cb(cur, "hc_ffn_pre", il); + + // expand before the sublayer so op offload does not pull the mHC state + // onto the expert weights' backend, as in deepseek4 + ggml_build_forward_expand(gf, residual); + ggml_build_forward_expand(gf, post); + ggml_build_forward_expand(gf, comb); + + cur = build_norm(cur, model.layers[il].ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_layer_ffn(model, cur, il); + cb(cur, "ffn_out", il); + + inpL = build_hc_post(cur, residual, post, comb, il); + inpL = build_cvec(inpL, il); + cb(inpL, "l_last", il); + } + + if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) { + res->t_layer_inp[n_layer] = build_hc_mean(ctx0, inpL); + cb(res->t_layer_inp[n_layer], "layer_inp", n_layer); + ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); + } + + if (inp_out_ids) { + // flattened: get_rows needs one token's streams to be one contiguous row + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + inpL = ggml_reshape_3d(ctx0, ggml_get_rows(ctx0, flat, inp_out_ids), n_embd, hc, n_outputs); + } + + // no hc_head tensor here: unweighted mean, not DeepSeek-V4's learned gated head + cur = build_hc_mean(ctx0, inpL); + cb(cur, "hc_mean", -1); + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +std::unique_ptr llama_model_glm5next::build_arch_graph(const llm_graph_params & params) const { + // llama_init_from_model accepts an MTP context whenever n_layer_nextn > 0, + // which every glm5next checkpoint has; without this it silently runs the trunk + GGML_ASSERT(params.gtype != LLM_GRAPH_TYPE_DECODER_MTP && "glm5next NextN graph not implemented yet"); + + return std::make_unique(*this, params); +} diff --git a/src/models/models.h b/src/models/models.h index 969429e3b6f..1182fd7207b 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1173,8 +1173,12 @@ struct llama_model_deepseek4 : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; - struct graph : public llm_graph_context { - graph(const llm_graph_params & params) : llm_graph_context(params) {} + // llm_build_delta_net_base is a method-only mixin over llm_graph_context (no data + // members, no new virtuals), so this graph's layout and behaviour are unchanged. + // deepseek4 has no recurrent layers; it is here so glm5next, which derives from + // this graph for the mHC residual, can reach build_delta_net for its KDA layers + struct graph : public llm_build_delta_net_base { + graph(const llm_graph_params & params) : llm_build_delta_net_base(params) {} graph(const llama_model & model, const llm_graph_params & params); ggml_tensor * build_hc_pre( @@ -1294,6 +1298,11 @@ struct llama_model_deepseek4 : public llama_model_base { ggml_tensor * build_hc_sinkhorn( ggml_tensor * comb, int il) const; + + // unweighted mean over the streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens] + static ggml_tensor * build_hc_mean( + ggml_context * ctx, + ggml_tensor * x); }; struct graph_mtp : public graph { @@ -1331,6 +1340,66 @@ struct llama_model_glm_dsa : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +struct llama_model_glm5next : public llama_model_base { + llama_model_glm5next(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + // glm5next's mHC is DeepSeek-V4's hyper-connection block (same wide residual, + // 24-row mixer split, activations, Sinkhorn); only the final collapse differs + // (unweighted mean, not a learned gated head), so derive rather than restate + struct graph : public llama_model_deepseek4::graph { + graph(const llama_model & model, const llm_graph_params & params); + + // not const: the delta-net helpers append to the graph through the base + ggml_tensor * build_layer_attn( + const llama_model & model, + llm_graph_input_mem_hybrid_k * inp_mem, + llm_graph_input_kpool * inp_kp, + bool scoring, + ggml_tensor * cur, + int il); + + ggml_tensor * build_kda_layer( + const llama_layer & layer, + llm_graph_input_rs * inp_rs, + ggml_tensor * cur, + int il); + + // `scoring` false keeps the dense path: at or below + // index_topk + index_kpool - 1 resident positions the indexer selects + // every visible position, so dense is not an approximation there, it is + // the same function without the pooling work. The indexer STORE still + // runs; only the selection is skipped + ggml_tensor * build_dsa_layer( + const llama_layer & layer, + llm_graph_input_attn_k * inp_attn, + llm_graph_input_kpool * inp_kp, + bool scoring, + ggml_tensor * cur, + int il) const; + + // writes this layer's indexer key and compressor gate into the indexer + // cache - always - and then, when `scoring`, returns the I32 attention + // cache CELL indices this layer's queries select, already expanded from + // whole pools: [kpool*select_k, n_tps, n_stream] + ggml_tensor * build_indexer( + const llama_layer & layer, + llm_graph_input_kpool * inp_kp, + ggml_tensor * cur, + ggml_tensor * qr, + bool scoring, + int il) const; + + ggml_tensor * build_layer_ffn( + const llama_model & model, + ggml_tensor * cur, + int il) const; + }; + + 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/CMakeLists.txt b/tests/CMakeLists.txt index b9f9d4b78af..2c793e582e0 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -196,6 +196,8 @@ if (NOT WIN32 OR NOT BUILD_SHARED_LIBS) # llama_build_and_test(test-double-float.cpp) # SLOW llama_build_and_test(test-llama-archs.cpp) + # needs a glm5next GGUF as argv[1], so it is built but not registered + llama_build(test-glm5next-memory.cpp) set(MODEL_DIR "${CMAKE_CURRENT_BINARY_DIR}/test-models/") file(MAKE_DIRECTORY "${MODEL_DIR}") diff --git a/tests/test-glm5next-memory.cpp b/tests/test-glm5next-memory.cpp new file mode 100644 index 00000000000..7c4aa16071b --- /dev/null +++ b/tests/test-glm5next-memory.cpp @@ -0,0 +1,1417 @@ +// The GLM-5-Next hybrid memory: the memory object holds all three halves (KDA +// recurrent + conv state, MLA latent KV, pooled indexer key cache) at the sizes the +// reference implies, one ggml graph reaches every one of them, and the pooled top-k +// cuts on a pool boundary and agrees between CPU and CUDA. The indexer graph itself is +// not built here. +// +// Run as: test-glm5next-memory , as written by +// tests/glm5next_make_tiny_gguf.py + +#include "ggml.h" +#include "ggml-alloc.h" +#include "ggml-backend.h" +#include "llama.h" + +#include "../src/llama-model.h" +#include "../src/llama-hparams.h" +#include "../src/llama-cparams.h" +#include "../src/llama-memory-hybrid.h" +#include "../src/llama-memory-recurrent.h" +#include "../src/llama-kv-cache.h" +#include "../src/llama-kv-cache-kpool.h" +#include "../src/llama-batch.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static int n_fail = 0; + +#define CHECK(cond, ...) \ + do { \ + if (!(cond)) { \ + printf("FAIL %s:%d: ", __func__, __LINE__); \ + printf(__VA_ARGS__); \ + printf("\n"); \ + n_fail++; \ + } else { \ + printf("ok "); \ + printf(__VA_ARGS__); \ + printf("\n"); \ + } \ + } while (0) + +// +// numeric evidence for WHERE the top-k runs, independent of any model +// +// The claim under test is the one that decides the whole design: a top-k over +// CELLS cannot be made pool-aligned by choosing its width, and a top-k over +// POOLS is pool-aligned by construction. +// +// The tempting argument for the cell-level form is that a pool's members carry +// its score bit-exactly, so an intra-pool tie is harmless and a budget that is a +// whole multiple of kpool must cut on a pool boundary. That argument silently +// assumes tie groups never SPAN pools. F.relu drives most pool scores to exactly +// 0.0, so they do, and ggml_top_k - explicitly unordered among equals - then +// takes an arbitrary 1..kpool-1 members of the pool it lands in. This +// reproduces that with no model at all: most pools at exactly 0.0, a minority +// positive, fewer positive pools than the budget. +// + +static std::vector run_top_k( + ggml_backend_t backend, + const std::vector & scores, + int64_t n_kv, + int64_t n_rows, + int64_t width) { + ggml_init_params gparams = { + /* .mem_size */ ggml_tensor_overhead()*8 + ggml_graph_overhead(), + /* .mem_buffer */ nullptr, + /* .no_alloc */ true, + }; + + ggml_context * ctx = ggml_init(gparams); + + ggml_tensor * src = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, n_kv, n_rows); + ggml_set_input(src); + + ggml_tensor * dst = ggml_top_k(ctx, src, (int) width); + ggml_set_output(dst); + + ggml_cgraph * gf = ggml_new_graph(ctx); + ggml_build_forward_expand(gf, dst); + + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + std::vector out(width*n_rows); + + if (buf) { + ggml_backend_tensor_set(src, scores.data(), 0, scores.size()*sizeof(float)); + + if (ggml_backend_graph_compute(backend, gf) == GGML_STATUS_SUCCESS) { + ggml_backend_tensor_get(dst, out.data(), 0, out.size()*sizeof(int32_t)); + } else { + out.clear(); + } + + ggml_backend_buffer_free(buf); + } else { + out.clear(); + } + + ggml_free(ctx); + + return out; +} + +static int64_t n_partial_pools(const std::vector & sel, int64_t off, int64_t width, int64_t kpool) { + std::map cnt; + for (int64_t i = 0; i < width; ++i) { + cnt[sel[off + i]/(int32_t) kpool]++; + } + + int64_t n = 0; + for (const auto & kv : cnt) { + n += kv.second != (int32_t) kpool; + } + + return n; +} + +static void test_top_k_boundary() { + printf("\n--- top-k granularity, standalone ---\n"); + + const int64_t kpool = 4; + const int64_t top_k = 2048; // the reference requires top_k %% kpool == 0 + const int64_t n_kv = 8192; + const int64_t n_rows = 4; // queries + const int64_t n_pools = n_kv/kpool; + + // the shape ReLU actually produces: a minority of pools with a positive score, + // every other pool at EXACTLY 0.0. Fewer positive pools than the budget, which + // during prefill is the normal state and not a corner case + const int64_t n_pos = 300; + + std::vector pool_score(n_pools, 0.0f); + for (int64_t p = 0; p < n_pos; ++p) { + // deterministic, distinct, strictly positive + pool_score[(p*7919) % n_pools] = 1.0f + std::fabs(std::sin(0.7f*(float) p))*1000.0f; + } + + const int64_t select_k = llama_kpool_select_k((uint32_t) n_pools, (uint32_t) top_k, (uint32_t) kpool); + CHECK(select_k == top_k/kpool, "llama_kpool_select_k is index_topk/index_kpool (%d pools)", (int) select_k); + + ggml_backend_t backend_cpu = ggml_backend_cpu_init(); + ggml_backend_t backend_gpu = nullptr; + { + ggml_backend_dev_t dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_GPU); + if (dev) { + backend_gpu = ggml_backend_dev_init(dev, nullptr); + } + } + + printf("%s GPU backend for the top-k comparison: %s\n", + backend_gpu ? "ok " : "note", backend_gpu ? ggml_backend_name(backend_gpu) : "none, CPU only"); + + // t = (q + 1) %% kpool cells of trailing incomplete pool, biased out of the + // budget. swept so that the per-stream tail residue is covered, not just t == 0 + for (int64_t t = 0; t < kpool; ++t) { + const int64_t n_tail = t; + const int64_t n_full = n_kv - n_tail; // cells belonging to complete pools + + // ---- A: the WRONG design. one score per cell, top-k of width index_topk -- + std::vector s_cell((size_t) n_kv*n_rows); + // ---- B: what ships. one score per pool, top-k of width index_topk/kpool -- + std::vector s_pool((size_t) n_pools*n_rows); + + for (int64_t r = 0; r < n_rows; ++r) { + for (int64_t j = 0; j < n_kv; ++j) { + s_cell[r*n_kv + j] = j >= n_full ? -INFINITY : pool_score[j/kpool]; + } + for (int64_t p = 0; p < n_pools; ++p) { + // a pool that the tail bites into is not a candidate at all + s_pool[r*n_pools + p] = (p + 1)*kpool > n_full ? -INFINITY : pool_score[p]; + } + } + + const auto sel_cell = run_top_k(backend_cpu, s_cell, n_kv, n_rows, top_k); + const auto sel_pool = run_top_k(backend_cpu, s_pool, n_pools, n_rows, select_k); + + // 1. the cell-level form, at a budget that IS a whole number of pools, + // still splits pools. this is the measurement the design turns on + int64_t partial_cell = 0; + for (int64_t r = 0; r < n_rows; ++r) { + partial_cell += n_partial_pools(sel_cell, r*top_k, top_k, kpool); + } + CHECK(partial_cell > 0, + "t=%d: a CELL-level top-k at the pool-aligned width %d still leaves %d partial pool(s); " + "a pool-aligned WIDTH does not make a pool-aligned CUT", + (int) t, (int) top_k, (int) partial_cell); + + // 2. the pool-level form: expand each selected pool ordinal to its kpool + // member cells, exactly as the graph does through pool_cells + std::vector expanded((size_t) select_k*kpool*n_rows); + for (int64_t r = 0; r < n_rows; ++r) { + for (int64_t i = 0; i < select_k; ++i) { + const int32_t p = sel_pool[r*select_k + i]; + for (int64_t m = 0; m < kpool; ++m) { + expanded[r*select_k*kpool + i*kpool + m] = (int32_t) (p*kpool + m); + } + } + } + + int64_t partial_pool = 0; + int64_t tail_in_pool = 0; + for (int64_t r = 0; r < n_rows; ++r) { + partial_pool += n_partial_pools(expanded, r*select_k*kpool, select_k*kpool, kpool); + for (int64_t i = 0; i < select_k*kpool; ++i) { + tail_in_pool += expanded[r*select_k*kpool + i] >= n_full; + } + } + CHECK(partial_pool == 0, + "t=%d: a POOL-level top-k of %d pools expands to only whole pools (%d partial)", + (int) t, (int) select_k, (int) partial_pool); + CHECK(tail_in_pool == 0, "t=%d: no tail cell consumes budget (%d did)", (int) t, (int) tail_in_pool); + + // 3. CPU vs CUDA on the same input + if (backend_gpu) { + const auto gpu_cell = run_top_k(backend_gpu, s_cell, n_kv, n_rows, top_k); + const auto gpu_pool = run_top_k(backend_gpu, s_pool, n_pools, n_rows, select_k); + + // the pool SET may legitimately differ between backends when the cut + // falls inside a tie group. What may not differ is pool integrity, + // and that is what is asserted + int64_t partial_gpu = 0; + bool ok_gpu = gpu_pool.size() == sel_pool.size(); + for (int64_t r = 0; ok_gpu && r < n_rows; ++r) { + for (int64_t i = 0; i < select_k; ++i) { + const int32_t p = gpu_pool[r*select_k + i]; + for (int64_t m = 0; m < kpool; ++m) { + expanded[r*select_k*kpool + i*kpool + m] = (int32_t) (p*kpool + m); + } + } + partial_gpu += n_partial_pools(expanded, r*select_k*kpool, select_k*kpool, kpool); + } + CHECK(ok_gpu && partial_gpu == 0, + "t=%d: %s also expands to only whole pools (%d partial)", + (int) t, ggml_backend_name(backend_gpu), (int) partial_gpu); + + bool same_pool = gpu_pool.size() == sel_pool.size(); + for (int64_t r = 0; same_pool && r < n_rows; ++r) { + std::set a(sel_pool.begin() + r*select_k, sel_pool.begin() + (r + 1)*select_k); + std::set b(gpu_pool.begin() + r*select_k, gpu_pool.begin() + (r + 1)*select_k); + same_pool = a == b; + } + int64_t partial_gpu_cell = 0; + for (int64_t r = 0; r < n_rows; ++r) { + partial_gpu_cell += n_partial_pools(gpu_cell, r*top_k, top_k, kpool); + } + // reported, not asserted: which members of a cut pool each backend's + // partial sort happens to keep is not contractual + printf("note t=%d: CPU and %s %s on the selected POOL set; the cell-level form leaves " + "%d partial pool(s) there too\n", + (int) t, ggml_backend_name(backend_gpu), + same_pool ? "agree" : "DISAGREE", (int) partial_gpu_cell); + } + } + + if (backend_gpu) { + ggml_backend_free(backend_gpu); + } + ggml_backend_free(backend_cpu); +} + +// +// the memory object itself +// + +struct kpool_tensors { + ggml_tensor * cell_pool = nullptr; + ggml_tensor * pool_cells = nullptr; + ggml_tensor * bias = nullptr; + ggml_tensor * pool_bias = nullptr; + ggml_tensor * sel_mask = nullptr; + ggml_tensor * cand_mask = nullptr; + + // the same two masks in the type the graph actually allocates. the values are + // only ever 0.0f and -INFINITY, so f16 must reproduce f32 bit for bit + ggml_tensor * sel_mask_f16 = nullptr; + ggml_tensor * cand_mask_f16 = nullptr; +}; + +static bool kpool_mask_f16_matches(const ggml_tensor * f32, const ggml_tensor * f16) { + const float * a = (const float *) f32->data; + const ggml_fp16_t * b = (const ggml_fp16_t *) f16->data; + + for (int64_t i = 0; i < ggml_nelements(f32); ++i) { + if (ggml_fp16_to_fp32(b[i]) != a[i]) { + return false; + } + } + + return true; +} + +// The test asks for all six, including the two the pooled graph does not consume. +// cell_pool and bias are the per-CELL spelling of the same predicate, computed by a +// different loop, and they are what pool_bias and cand_mask are checked against here +static kpool_tensors alloc_kpool_tensors( + ggml_context * ctx, + int64_t n_kv, + int64_t n_tps, + int64_t n_padq, + int64_t n_stream, + int64_t kpool, + int64_t n_pools) { + kpool_tensors t; + + t.cell_pool = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, n_kv, n_stream); + t.pool_cells = ggml_new_tensor_2d(ctx, GGML_TYPE_I32, kpool*n_pools, n_stream); + t.bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_kv, n_tps, n_stream); + t.pool_bias = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_pools, n_tps, n_stream); + t.sel_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_padq, 1, n_stream); + t.cand_mask = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, n_kv, n_padq, 1, n_stream); + + t.sel_mask_f16 = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_padq, 1, n_stream); + t.cand_mask_f16 = ggml_new_tensor_4d(ctx, GGML_TYPE_F16, n_kv, n_padq, 1, n_stream); + + ggml_set_input(t.sel_mask_f16); + ggml_set_input(t.cand_mask_f16); + + ggml_set_input(t.cell_pool); + ggml_set_input(t.pool_cells); + ggml_set_input(t.bias); + ggml_set_input(t.pool_bias); + ggml_set_input(t.sel_mask); + ggml_set_input(t.cand_mask); + + return t; +} + +int main(int argc, char ** argv) { + if (argc < 2) { + printf("usage: %s \n", argv[0]); + return 1; + } + + setvbuf(stdout, nullptr, _IOLBF, 0); + + llama_backend_init(); + + test_top_k_boundary(); + + printf("\n--- model ---\n"); + + llama_model_params mparams = llama_model_default_params(); + mparams.n_gpu_layers = 0; + + llama_model * model = llama_model_load_from_file(argv[1], mparams); + if (model == nullptr) { + printf("FAIL: could not load %s\n", argv[1]); + return 1; + } + + const auto & hparams = model->hparams; + + printf("n_layer = %u (+%u nextn)\n", hparams.n_layer(), hparams.n_layer_nextn); + printf("n_head = %u\n", hparams.n_head()); + printf("n_embd_head_kda = %u\n", hparams.n_embd_head_kda); + printf("ssm_d_conv = %u\n", hparams.ssm_d_conv); + printf("n_lora_kv = %u\n", hparams.n_lora_kv); + printf("indexer_head_size= %u\n", hparams.indexer_head_size); + printf("indexer_kpool = %u\n", hparams.indexer_kpool); + printf("indexer_top_k = %u\n", hparams.indexer_top_k); + printf("n_embd_r() = %u\n", hparams.n_embd_r()); + printf("n_embd_s() = %u\n", hparams.n_embd_s()); + + // ---- sizing ------------------------------------------------------------ + { + const uint32_t d_inner = hparams.n_head()*hparams.n_embd_head_kda; + + CHECK(hparams.n_embd_r() == 3*(hparams.ssm_d_conv - 1)*d_inner, + "n_embd_r == 3 conv states of (d_conv-1)*n_head*head_dim (%u)", hparams.n_embd_r()); + CHECK(hparams.n_embd_s() == hparams.n_embd_head_kda*hparams.n_embd_head_kda*hparams.n_head(), + "n_embd_s == head_dim*head_dim per head (%u)", hparams.n_embd_s()); + CHECK(hparams.indexer_kpool > 0 && hparams.indexer_top_k % hparams.indexer_kpool == 0, + "indexer_top_k (%u) is a whole number of pools of %u", hparams.indexer_top_k, hparams.indexer_kpool); + } + + // ---- a multi-sequence unified cache pools per SEQUENCE ------------------ + // + // [TAG_KPOOL_SEQ_PARTITION]. pools group cells by position, and a position only + // names a pool inside one sequence. a unified cache shares one cells array, so the + // stream's pool table is cut into one run per sequence instead of the cache being + // refused. -kvu with --parallel is reachable (llama-perplexity forces it for + // hellaswag / winogrande / multiple-choice, and llama-embedding turns -np 1 into + // -kvu with n_seq_max 256), so this has to work rather than fail at startup. + { + llama_cparams cp = {}; + cp.n_ctx = 256; + cp.n_ctx_seq = 256; + cp.n_batch = 32; + cp.n_ubatch = 32; + cp.n_seq_max = 2; + cp.kv_unified = true; + cp.causal_attn = true; + + llama_memory_params mp = {}; + mp.type_k = GGML_TYPE_F16; + mp.type_v = GGML_TYPE_F16; + mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; + + llama_memory_i * raw = nullptr; + try { + raw = model->create_memory(mp, cp); + } catch (const std::exception & e) { + printf("note create_memory threw: %s\n", e.what()); + raw = nullptr; + } + CHECK(raw != nullptr, "-kvu with n_seq_max 2 is accepted: the pool map is per sequence"); + + const uint32_t kpool = hparams.indexer_kpool; + + // one shared n_kv/kpool budget plus the rebasing slack per sequence, NOT one + // full-width table each: the indexer scores every slot against every query, so a + // full-width table per sequence multiplies the score tensor by the sequence count + CHECK(llama_kpool_n_pools(256, kpool, 1) == 256/kpool + 2 && + llama_kpool_n_pools(256, kpool, 2) == 256/kpool + 4, + "llama_kpool_n_pools is n_kv/kpool + 2 per sequence (%u slots for 2 sequences)", + llama_kpool_n_pools(256, kpool, 2)); + + auto * m2 = dynamic_cast(raw); + CHECK(m2 != nullptr && m2->get_mem_idx() != nullptr, "-kvu builds an indexer cache too"); + + // drive one ubatch holding BOTH sequences through the map. 16 positions each, so + // every sequence owns several complete pools and the runs have to be told apart + if (m2) { + const int64_t n_tok = 16; + + std::vector tok; + std::vector pos; + std::vector sid; + std::vector sptr; + std::vector nsid; + std::vector lg; + + for (llama_seq_id s = 0; s < 2; ++s) { + for (int64_t i = 0; i < n_tok; ++i) { + tok.push_back(0); + pos.push_back((llama_pos) i); + sid.push_back(s); + nsid.push_back(1); + lg.push_back(1); + } + } + for (size_t i = 0; i < sid.size(); ++i) { + sptr.push_back(&sid[i]); + } + + llama_batch b = {}; + b.n_tokens = (int32_t) tok.size(); + b.token = tok.data(); + b.pos = pos.data(); + b.seq_id = sptr.data(); + b.n_seq_id = nsid.data(); + b.logits = lg.data(); + + llama_batch_allocr ba(hparams.n_pos_per_embd()); + if (ba.init(b, model->vocab, nullptr, hparams.n_embd_inp(), cp.n_seq_max, true)) { + auto c = m2->init_batch(ba, cp.n_ubatch, false); + auto * mc = dynamic_cast(c.get()); + + CHECK(mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, + "both sequences fit one unified ubatch"); + + if (mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { + mc->apply(); + + const llama_ubatch & ub = mc->get_ubatch(); + + const int64_t n_kv = mc->get_attn()->get_n_kv(); + const int64_t n_stream = mc->get_attn()->get_n_stream(); + const int64_t n_tps = ub.n_tokens/n_stream; + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, kpool, ub.n_seqs_unq); + + CHECK(n_stream == 1 && ub.n_seqs_unq == 2, + "a unified cache carries both sequences in one stream (n_seqs_unq %u)", + ub.n_seqs_unq); + + ggml_init_params gp = { ggml_tensor_overhead()*16, nullptr, true }; + ggml_context * ctx = ggml_init(gp); + + kpool_tensors kt = alloc_kpool_tensors(ctx, n_kv, n_tps, n_tps, n_stream, kpool, n_pools); + + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + if (buf) { + // no cell_pool: the per-cell view has one row per stream, and a + // cell two sequences share has nowhere to put its second pool + llama_kv_cache_set_input_kpool(m2->get_mem_attn(), + /* cell_pool */ nullptr, kt.pool_cells, kt.bias, kt.pool_bias, + kt.sel_mask, kt.cand_mask, &ub, kpool); + + const int32_t * pc = (const int32_t *) kt.pool_cells->data; + const float * pb = (const float *) kt.pool_bias->data; + + const auto & cells = m2->get_mem_attn()->get_cells(0); + + // the invariant the partitioning exists for: a pool the query may + // spend budget on holds only that query's own visible cells. this + // is what a shared cells array breaks if the map is per stream + bool own = true; + int64_t n_finite = 0; + std::map seq_mask_of_slot; + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const llama_seq_id sq = ub.seq_id[ii][0]; + const llama_pos q = ub.pos[ii]; + + for (int64_t p = 0; p < n_pools; ++p) { + if (pb[ii*n_pools + p] != 0.0f) { + continue; + } + + n_finite++; + seq_mask_of_slot[p] |= 1 << sq; + + for (int64_t m = 0; m < (int64_t) kpool; ++m) { + const int32_t c = pc[p*kpool + m]; + own &= c >= 0 && c < n_kv && !cells.is_empty(c) && + cells.seq_has(c, sq) && cells.pos_get(c) <= q; + } + } + } + + CHECK(own && n_finite > 0, + "every pool a query may select holds only that query's own visible cells (%d selectable (query, pool) pairs)", + (int) n_finite); + + int mask_all = 0; + bool disjoint = true; + for (const auto & kv : seq_mask_of_slot) { + mask_all |= kv.second; + disjoint &= (kv.second & (kv.second - 1)) == 0; + } + + CHECK(disjoint, "the two sequences get disjoint pool runs (%d slots in use of %d)", + (int) seq_mask_of_slot.size(), (int) n_pools); + CHECK(mask_all == 0x3, "both sequences own selectable pools, so neither run is empty"); + + // the masks are filled by one partition per sequence, each writing + // the rows it owns into the shared stream buffer, so this is where + // a wrong element width would cross partitions + llama_kv_cache_set_input_kpool(m2->get_mem_attn(), + /* cell_pool */ nullptr, kt.pool_cells, kt.bias, kt.pool_bias, + kt.sel_mask_f16, kt.cand_mask_f16, &ub, kpool); + + CHECK(kpool_mask_f16_matches(kt.sel_mask, kt.sel_mask_f16) && + kpool_mask_f16_matches(kt.cand_mask, kt.cand_mask_f16), + "the f16 masks match the f32 ones across two sequences in one stream"); + + ggml_backend_buffer_free(buf); + } + + ggml_backend_free(backend); + ggml_free(ctx); + } + } + } + + delete raw; + + // one sequence in flight stays the simple case: the shared cells array holds + // only its keys, and the map filters on seq_has anyway + cp.n_seq_max = 1; + llama_memory_i * raw1 = nullptr; + try { + raw1 = model->create_memory(mp, cp); + } catch (const std::exception &) { + raw1 = nullptr; + } + CHECK(raw1 != nullptr, "-kvu with a single sequence is still allowed"); + delete raw1; + } + + // ---- the indexer cache keeps its own dtype ------------------------------ + { + llama_cparams cp = {}; + cp.n_ctx = 256; + cp.n_ctx_seq = 256; + cp.n_batch = 32; + cp.n_ubatch = 32; + cp.n_seq_max = 1; + cp.kv_unified = false; + cp.causal_attn = true; + + llama_memory_params mp = {}; + mp.type_k = GGML_TYPE_Q8_0; + mp.type_v = GGML_TYPE_Q8_0; + mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; + + llama_memory_i * raw = model->create_memory(mp, cp); + auto * m = dynamic_cast(raw); + + CHECK(m != nullptr && m->get_mem_idx() != nullptr, "-ctk q8_0 still builds an indexer cache"); + if (m && m->get_mem_idx()) { + CHECK(!ggml_is_quantized(m->get_mem_idx()->type_k()), + "indexer cache stays %s under -ctk q8_0: it also holds the compressor gates", + ggml_type_name(m->get_mem_idx()->type_k())); + CHECK(ggml_is_quantized(m->get_mem_attn()->type_k()), + "the MLA cache still honours -ctk q8_0 (%s)", ggml_type_name(m->get_mem_attn()->type_k())); + } + + delete raw; + } + + // ---- create the memory ------------------------------------------------- + llama_cparams cparams = {}; + cparams.n_ctx = 512; + cparams.n_ctx_seq = 512; + cparams.n_batch = 32; + cparams.n_ubatch = 32; + cparams.n_seq_max = 2; + cparams.n_rs_seq = 0; + cparams.kv_unified = false; // one sequence per stream, the n_ps == 1 layout + cparams.offload_kqv = false; + cparams.flash_attn = false; + cparams.causal_attn = true; + + llama_memory_params mparams_mem = {}; + mparams_mem.type_k = GGML_TYPE_F16; + mparams_mem.type_v = GGML_TYPE_F16; + mparams_mem.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; + mparams_mem.swa_full = false; + + llama_memory_i * mem_raw = model->create_memory(mparams_mem, cparams); + CHECK(mem_raw != nullptr, "create_memory returned a memory object"); + if (!mem_raw) { + return 1; + } + + auto * mem = dynamic_cast(mem_raw); + CHECK(mem != nullptr, "the memory is a llama_memory_hybrid"); + if (!mem) { + return 1; + } + + llama_kv_cache * kv_attn = mem->get_mem_attn(); + llama_memory_recurrent * rs = mem->get_mem_recr(); + llama_kv_cache * kv_idx = mem->get_mem_idx(); + + CHECK(kv_attn != nullptr, "hybrid holds an attention (MLA) cache"); + CHECK(rs != nullptr, "hybrid holds a recurrent (KDA) cache"); + CHECK(kv_idx != nullptr, "hybrid holds an indexer key cache"); + if (!kv_attn || !rs || !kv_idx) { + return 1; + } + + // ---- layer partition --------------------------------------------------- + { + const auto ids_attn = kv_attn->get_layer_ids(); + const auto ids_idx = kv_idx ->get_layer_ids(); + + uint32_t n_dsa = 0; + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + n_dsa += !hparams.is_recr(il); + } + + CHECK(ids_attn.size() == n_dsa, "MLA cache holds the %u DSA trunk layers (%zu)", n_dsa, ids_attn.size()); + CHECK(ids_idx.size() == n_dsa, "indexer cache holds the same %u layers (%zu)", n_dsa, ids_idx.size()); + + bool same = ids_attn.size() == ids_idx.size(); + for (size_t i = 0; same && i < ids_attn.size(); ++i) { + same = ids_attn[i] == ids_idx[i]; + } + CHECK(same, "MLA and indexer caches cover exactly the same layers"); + + for (uint32_t il : ids_attn) { + CHECK(!hparams.is_recr(il), "cached attention layer %u is not recurrent", il); + } + } + + // ---- cache geometry ---------------------------------------------------- + { + const uint32_t il_dsa = kv_attn->get_layer_ids().front(); + + ggml_tensor * k_mla = kv_attn->get_k_storage(il_dsa); + ggml_tensor * k_idx = kv_idx ->get_k_storage(il_dsa); + + CHECK(k_mla != nullptr && k_idx != nullptr, "both caches expose K storage for layer %u", il_dsa); + + // nope-only MLA: the latent row is kv_lora_rank + qk_rope_head_dim + CHECK(k_mla->ne[0] == (int64_t) hparams.n_embd_head_k(il_dsa), + "MLA latent row = %d (n_embd_head_k = %u)", (int) k_mla->ne[0], hparams.n_embd_head_k(il_dsa)); + + // the pooling indexer caches the key AND the compressor gate score + CHECK(k_idx->ne[0] == (int64_t) (2*hparams.indexer_head_size), + "indexer row = %d (expected 2 x indexer_head_size = %u)", + (int) k_idx->ne[0], 2*hparams.indexer_head_size); + + CHECK(k_mla->ne[1] == k_idx->ne[1], "MLA and indexer caches have the same cell count (%d)", (int) k_mla->ne[1]); + CHECK(k_mla->ne[2] == k_idx->ne[2], "MLA and indexer caches have the same stream count (%d)", (int) k_mla->ne[2]); + + // is_mla() stays true for the indexer hparams copy, so no V is allocated + { + size_t bytes = 0; + for (const auto & b : kv_idx->memory_breakdown()) { + bytes += b.second; + } + + const size_t k_only = (size_t) ggml_nbytes(k_idx)*kv_idx->get_layer_ids().size(); + + // the byte count is the observable: K only, no V + CHECK(hparams.is_mla() && bytes == k_only, + "indexer cache allocates K and no V: %zu bytes for %zu layers", + bytes, kv_idx->get_layer_ids().size()); + } + } + + // ---- recurrent geometry ------------------------------------------------ + { + uint32_t n_kda = 0; + for (uint32_t il = 0; il < hparams.n_layer(); ++il) { + n_kda += hparams.is_recr(il); + } + + CHECK(rs->size >= cparams.n_seq_max, "recurrent cache has >= n_seq_max slots (%u)", rs->size); + CHECK(n_kda > 0, "the model has %u KDA layers", n_kda); + } + + // ---- drive a batch through it and reach all three halves in one graph --- + { + const int32_t n_tokens = 10; // pos 9 leaves (9+1) %% kpool = 2 tail cells + + std::vector tokens(n_tokens*cparams.n_seq_max, 0); + std::vector pos; + std::vector seqs; + for (uint32_t s = 0; s < cparams.n_seq_max; ++s) { + for (int32_t i = 0; i < n_tokens; ++i) { + pos.push_back(i); + seqs.push_back((llama_seq_id) s); + } + } + + llama_batch batch = {}; + batch.n_tokens = n_tokens*(int32_t) cparams.n_seq_max; + batch.token = tokens.data(); + batch.pos = pos.data(); + + std::vector seq_ptrs(batch.n_tokens); + for (int32_t i = 0; i < batch.n_tokens; ++i) { + seq_ptrs[i] = &seqs[i]; + } + std::vector n_seq_id(batch.n_tokens, 1); + std::vector logits(batch.n_tokens, 0); + logits.back() = 1; + + batch.seq_id = seq_ptrs.data(); + batch.n_seq_id = n_seq_id.data(); + batch.logits = logits.data(); + + llama_batch_allocr balloc(hparams.n_pos_per_embd()); + const bool ok = balloc.init(batch, model->vocab, nullptr, hparams.n_embd_inp(), + cparams.n_seq_max, true); + CHECK(ok, "batch allocr accepted a %d-token, %u-sequence batch", batch.n_tokens, cparams.n_seq_max); + + auto mctx_ptr = mem->init_batch(balloc, cparams.n_ubatch, false); + CHECK(mctx_ptr != nullptr && mctx_ptr->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, + "init_batch produced a usable memory context"); + + auto * mctx = dynamic_cast(mctx_ptr.get()); + CHECK(mctx != nullptr, "the context is a llama_memory_hybrid_context"); + + if (mctx && mctx->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { + const auto * ctx_attn = mctx->get_attn(); + const auto * ctx_recr = mctx->get_recr(); + const auto * ctx_idx = mctx->get_idx(); + + CHECK(ctx_attn != nullptr, "context exposes the attention half"); + CHECK(ctx_recr != nullptr, "context exposes the recurrent half"); + CHECK(ctx_idx != nullptr, "context exposes the indexer half"); + + // n_kv is only valid once applied; apply() also asserts the caches agree + mctx->apply(); + + // apply() asserts both itself, so reaching this line is the result; printed + // rather than CHECKed so the count stays honest + printf("note indexer and attention caches agree on n_kv (%u) and n_stream (%u)\n", + ctx_idx->get_n_kv(), ctx_idx->get_n_stream()); + CHECK(ctx_idx && ctx_idx->get_kv() == kv_idx, "the indexer context is a view of the indexer cache"); + + ggml_init_params gparams = { + /* .mem_size */ ggml_tensor_overhead()*1024 + ggml_graph_overhead(), + /* .mem_buffer */ nullptr, + /* .no_alloc */ true, + }; + + ggml_context * ctx0 = ggml_init(gparams); + ggml_cgraph * gf = ggml_new_graph(ctx0); + + const uint32_t il_dsa = kv_attn->get_layer_ids().front(); + + uint32_t il_kda = 0; + while (il_kda < hparams.n_layer() && !hparams.is_recr(il_kda)) { + il_kda++; + } + + const int64_t n_kv = ctx_attn->get_n_kv(); + const int64_t n_stream = ctx_attn->get_n_stream(); + const int64_t n_tps = mctx->get_ubatch().n_tokens/n_stream; + + // 1. MLA latent K + ggml_tensor * k_mla = ctx_attn->get_k(ctx0, il_dsa); + CHECK(k_mla != nullptr, "graph reaches the MLA latent cache: [%d, %d, %d, %d]", + (int) k_mla->ne[0], (int) k_mla->ne[1], (int) k_mla->ne[2], (int) k_mla->ne[3]); + + // 2. indexer keys, split into the key half and the gate half + ggml_tensor * k_idx_all = ctx_idx->get_k(ctx0, il_dsa); + CHECK(k_idx_all->ne[1] == 2, "indexer cache view has 2 heads (key | compressor gate)"); + + const int64_t d_idx = hparams.indexer_head_size; + + ggml_tensor * k_idx_v = ggml_view_3d(ctx0, k_idx_all, d_idx, n_kv, n_stream, + k_idx_all->nb[2], k_idx_all->nb[3], 0); + ggml_tensor * g_idx_v = ggml_view_3d(ctx0, k_idx_all, d_idx, n_kv, n_stream, + k_idx_all->nb[2], k_idx_all->nb[3], k_idx_all->nb[1]); + + // 3. KDA recurrent + conv state + ggml_tensor * r_kda = ctx_recr->get_r_l(il_kda); + ggml_tensor * s_kda = ctx_recr->get_s_l(il_kda); + CHECK(r_kda != nullptr && s_kda != nullptr, + "graph reaches the KDA conv state [%d x %d] and recurrent state [%d x %d]", + (int) r_kda->ne[0], (int) r_kda->ne[1], (int) s_kda->ne[0], (int) s_kda->ne[1]); + + // 4. host-side pool map + const int64_t kpool = hparams.indexer_kpool; + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); + const int64_t n_padq = n_tps; // this tree does not pad the KQ mask + + kpool_tensors kt = alloc_kpool_tensors(ctx0, n_kv, n_tps, n_padq, n_stream, kpool, n_pools); + + // shape of the real thing: gather pool members, mix, score, broadcast the + // pool score back onto its member cells, top-k + ggml_tensor * members = ggml_get_rows(ctx0, k_idx_v, kt.pool_cells); + ggml_tensor * gates = ggml_get_rows(ctx0, g_idx_v, kt.pool_cells); + members = ggml_reshape_4d(ctx0, members, d_idx, kpool, n_pools, n_stream); + gates = ggml_reshape_4d(ctx0, gates, d_idx, kpool, n_pools, n_stream); + + // stand-in for softmax(gate + ape) * member, summed over kpool + ggml_tensor * pooled = ggml_mul(ctx0, members, gates); + pooled = ggml_reshape_3d(ctx0, pooled, d_idx*kpool, n_pools, n_stream); + // stand-in for the q . k_pool score: one number per (pool, query) + ggml_tensor * score = ggml_mul_mat(ctx0, + ggml_cont(ctx0, ggml_view_3d(ctx0, pooled, d_idx, n_pools, n_stream, + pooled->nb[1], pooled->nb[2], 0)), + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, d_idx, n_tps, n_stream)); + + // mask the pools the reference rejects, then top-k over POOLS + score = ggml_add(ctx0, ggml_cont(ctx0, score), kt.pool_bias); + + const int64_t select_k = llama_kpool_select_k((uint32_t) n_pools, hparams.indexer_top_k, (uint32_t) kpool); + ggml_tensor * sel = ggml_cont(ctx0, ggml_top_k(ctx0, score, (int) select_k)); + + // expand each selected POOL ordinal into its kpool member CELLS, which is + // the reference's selected_indices = pool_indices[batch_idx, selected]. + // the query axis folds into the gather's row axis so that one get_rows + // serves every query, while the stream axis stays where get_rows wants it + ggml_tensor * pc3 = ggml_reshape_3d(ctx0, kt.pool_cells, kpool, n_pools, n_stream); + ggml_tensor * sel_flat = ggml_reshape_2d(ctx0, sel, select_k*n_tps, n_stream); + + const int64_t width = kpool*select_k; + ggml_tensor * top_k = ggml_reshape_3d(ctx0, + ggml_get_rows(ctx0, pc3, sel_flat), width, n_tps, n_stream); + + CHECK(top_k->type == GGML_TYPE_I32, "the pool -> cell expansion stays I32 (pool_cells is I32)"); + printf("note top-k over %d POOL scores expands to %d I32 CELL indices per query\n", + (int) select_k, (int) width); + + // 5. the scatter the mask is built from, starting at sel_mask rather than + // an all -INFINITY fill, which is what forces the tail in + ggml_tensor * top_k_4d = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream); + ggml_tensor * base = ggml_view_4d(ctx0, kt.sel_mask, 1, n_kv, n_padq, n_stream, + kt.sel_mask->nb[0], kt.sel_mask->nb[1], kt.sel_mask->nb[2], 0); + ggml_tensor * idxs = ggml_view_4d(ctx0, top_k_4d, width, n_tps, n_stream, 1, + top_k_4d->nb[1], top_k_4d->nb[2], n_stream*top_k_4d->nb[3], 0); + ggml_tensor * zeros = ggml_fill(ctx0, + ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, width, n_tps, n_stream), 0.0f); + ggml_tensor * mask = ggml_set_rows(ctx0, base, zeros, idxs); + + ggml_build_forward_expand(gf, k_mla); + ggml_build_forward_expand(gf, r_kda); + ggml_build_forward_expand(gf, s_kda); + ggml_build_forward_expand(gf, mask); + + printf("note one graph reaches all three halves in %d nodes\n", ggml_graph_n_nodes(gf)); + + // 6. fill the pool map on a real allocated buffer + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx0, backend); + CHECK(buf != nullptr, "allocated the graph's input tensors on the CPU backend"); + + if (buf) { + const llama_ubatch & ub = mctx->get_ubatch(); + llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, kt.sel_mask, kt.cand_mask, + &ub, (uint32_t) kpool); + + const int32_t * cp = (const int32_t *) kt.cell_pool->data; + const int32_t * pc = (const int32_t *) kt.pool_cells->data; + const float * bi = (const float *) kt.bias->data; + const float * sm = (const float *) kt.sel_mask->data; + + bool in_range = true; + for (int64_t i = 0; i < ggml_nelements(kt.cell_pool); ++i) { + in_range &= cp[i] >= 0 && cp[i] < n_pools; + } + for (int64_t i = 0; i < ggml_nelements(kt.pool_cells); ++i) { + in_range &= pc[i] >= 0 && pc[i] < n_kv; + } + CHECK(in_range, "every emitted index is non-negative and in range (ggml_set_rows asserts i1 >= 0)"); + + bool finite = true; + for (int64_t i = 0; i < ggml_nelements(kt.bias); ++i) { + finite &= bi[i] == 0.0f || bi[i] == -INFINITY; + } + for (int64_t i = 0; i < ggml_nelements(kt.sel_mask); ++i) { + finite &= sm[i] == 0.0f || sm[i] == -INFINITY; + } + for (int64_t i = 0; i < ggml_nelements(kt.pool_bias); ++i) { + const float * pb = (const float *) kt.pool_bias->data; + finite &= pb[i] == 0.0f || pb[i] == -INFINITY; + } + for (int64_t i = 0; i < ggml_nelements(kt.cand_mask); ++i) { + const float * cm = (const float *) kt.cand_mask->data; + finite &= cm[i] == 0.0f || cm[i] == -INFINITY; + } + CHECK(finite, "bias, pool_bias, sel_mask and cand_mask hold only 0 and -INFINITY: " + "no +1e9 to meet a -inf"); + + // which is why the graph stores them in the KQ mask's f16 and halves + // two [n_kv, n_batch, n_stream] inputs. same fill, half the bytes + { + llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, + kt.sel_mask_f16, kt.cand_mask_f16, &ub, (uint32_t) kpool); + + CHECK(kpool_mask_f16_matches(kt.sel_mask, kt.sel_mask_f16) && + kpool_mask_f16_matches(kt.cand_mask, kt.cand_mask_f16), + "an f16 sel_mask/cand_mask is the exact image of the f32 one"); + } + + // cand_mask is exactly max(bias, sel_mask) lifted to KQ shape, and it + // is what stops an over-budget top-k from escaping the reference's + // candidate set. + // + // ggml_top_k always returns select_k pool ordinals even when fewer + // than select_k pools are finite. During prefill the query at + // position q has only ~q/kpool complete visible pools against a + // budget of index_topk/kpool, so the budget spills into -INFINITY + // pools and picks among them arbitrarily. Adding the causal mask + // kills the expansions that are empty or in the future. What it does + // not kill is a resident, causally visible cell that sits in an + // INCOMPLETE pool below the tail: the reference never selects it + // (pool_valid = grouped_valid_keys.all(-1)), the graph would. + // Unreachable while positions are contiguous, reachable the moment a + // partial seq_rm leaves a hole. + { + const float * cm = (const float *) kt.cand_mask->data; + + bool is_union = true; + int64_t n_spill = 0; // candidate pools strictly fewer than the budget + + const int64_t select_k = + llama_kpool_select_k((uint32_t) n_pools, hparams.indexer_top_k, (uint32_t) kpool); + + for (int64_t s = 0; s < n_stream; ++s) { + for (int64_t ii = 0; ii < n_tps; ++ii) { + const float * rb = bi + (s*n_tps + ii)*n_kv; + const float * rs = sm + (s*n_padq + ii)*n_kv; + const float * rc = cm + (s*n_padq + ii)*n_kv; + + for (int64_t j = 0; j < n_kv; ++j) { + is_union &= rc[j] == std::max(rb[j], rs[j]); + } + + const float * rp = (const float *) kt.pool_bias->data + (s*n_tps + ii)*n_pools; + int64_t n_cand = 0; + for (int64_t p = 0; p < n_pools; ++p) { + n_cand += rp[p] == 0.0f; + } + n_spill += n_cand < select_k; + } + } + + CHECK(is_union, "cand_mask == max(bias, sel_mask): the reference's candidate set"); + + // not a failure: it is the normal prefill state, and the whole + // reason the gate has to exist. Printed so that a future change + // making it zero cannot quietly turn the check above vacuous + printf("note %d of %d (query, stream) rows have fewer candidate pools than the\n" + " top-k budget of %d, so the budget spills and cand_mask is load bearing\n", + (int) n_spill, (int) (n_stream*n_tps), (int) select_k); + } + + // pool_bias is the same predicate as bias, evaluated where the + // reference evaluates it. For a COMPLETE pool the two must agree + // cell for cell; for an incomplete or absent pool bias has no cell + // to speak for it, which is exactly why pool_bias is computed here + // rather than gathered from bias at each pool's last member + { + const float * pb = (const float *) kt.pool_bias->data; + + bool agree = true; + bool exact = true; + + for (int64_t s = 0; s < n_stream; ++s) { + for (int64_t ii = 0; ii < n_tps; ++ii) { + const float * rb = bi + (s*n_tps + ii)*n_kv; + const float * rp = pb + (s*n_tps + ii)*n_pools; + + std::set pools_of_scored_cells; + + for (int64_t j = 0; j < n_kv; ++j) { + if (rb[j] == 0.0f) { + // a scored cell's pool must be a candidate pool + agree &= rp[cp[s*n_kv + j]] == 0.0f; + pools_of_scored_cells.insert(cp[s*n_kv + j]); + } + } + + int64_t n_cand = 0; + for (int64_t p = 0; p < n_pools; ++p) { + n_cand += rp[p] == 0.0f; + } + + // and nothing else may be one. an entirely absent pool + // would pass the first check vacuously; this is what + // catches it, and it is exactly the failure mode of + // gathering pool_bias from bias at each pool's last + // member instead of computing it + exact &= n_cand == (int64_t) pools_of_scored_cells.size(); + } + } + + CHECK(agree, "every cell bias scores lies in a pool pool_bias also accepts"); + CHECK(exact, "and pool_bias accepts no pool that has no scored cell " + "(an absent pool must not become a candidate)"); + } + + // per query: (q+1) %% kpool cells sit in its own incomplete pool and + // must be forced in, the q+1-minus-that below must be scored, and + // nothing else either. both streams, to cover the per-stream strides + { + const llama_ubatch & u = mctx->get_ubatch(); + + for (int64_t s = 0; s < n_stream; ++s) { + bool ok_tail = true; + bool ok_scored = true; + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const llama_pos q = u.pos[s*n_tps + ii]; + const int64_t t = (q + 1) % kpool; + + const float * row_s = sm + (s*n_padq + ii)*n_kv; + const float * row_b = bi + (s*n_tps + ii)*n_kv; + + int64_t n_forced = 0; + int64_t n_scored = 0; + for (int64_t j = 0; j < n_kv; ++j) { + n_forced += row_s[j] == 0.0f; + n_scored += row_b[j] == 0.0f; + } + + ok_tail &= n_forced == t; + ok_scored &= n_scored == (q + 1) - t; + } + + CHECK(ok_tail, "stream %d: every query forces in exactly its (q+1) %% %d tail cells", + (int) s, (int) kpool); + CHECK(ok_scored, "stream %d: and scores exactly the cells below that tail", + (int) s); + } + } + + // while the window starts at position 0, pos/kpool and the reference's + // first-resident-key anchor are the same grouping, so the map must + // reproduce it. incomplete pools carry slot 0 and are masked instead + { + const llama_ubatch & u = mctx->get_ubatch(); + bool agree = true; + for (int64_t s = 0; s < n_stream; ++s) { + const llama_seq_id seq = u.seq_id[s*n_tps][0]; + const auto & cells = kv_attn->get_cells(seq); + + std::map members; + llama_pos p_min = -1; + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j)) { + continue; + } + members[cells.pos_get(j)/kpool]++; + if (p_min < 0 || cells.pos_get(j) < p_min) { + p_min = cells.pos_get(j); + } + } + agree &= p_min == 0; // HF's first_key would also be 0 + + for (int64_t j = 0; j < n_kv && agree; ++j) { + if (cells.is_empty(j)) { + continue; + } + const int64_t b = cells.pos_get(j)/kpool; + agree &= cp[s*n_kv + j] == (int32_t) (members[b] == kpool ? b : 0); + } + } + CHECK(agree, "pool ordinals match the reference anchor while the window starts at position 0"); + } + + { + ggml_status st = ggml_backend_graph_compute(backend, gf); + CHECK(st == GGML_STATUS_SUCCESS, "the pooled-indexer-shaped graph computes (%d)", (int) st); + } + + ggml_backend_buffer_free(buf); + } + + ggml_backend_free(backend); + ggml_free(ctx0); + } + } + + // ---- pool origin across a front eviction -------------------------------- + // + // the reference anchors a pool at the first *resident* key (valid_keys.argmax(-1)); + // this port anchors at pos/kpool, as vLLM and SGLang do. same grouping until the + // window front is dropped by a non-multiple of kpool, when the reference regroups + // every surviving key and this port does not. a cache cannot afford regrouping: the + // pooled key a decode step scores was built during prefill. + { + printf("\n--- pool origin ---\n"); + + // seq 0 holds positions 0..9. drop 0..1, not a multiple of kpool = 4 + const llama_pos n_drop = 2; + mem->seq_rm(0, 0, n_drop); + + std::vector tok(1, 0); + std::vector pos(1, 10); + std::vector sid(1, 0); + std::vector sptr(1, sid.data()); + std::vector nsid(1, 1); + std::vector lg(1, 1); + + llama_batch b = {}; + b.n_tokens = 1; + b.token = tok.data(); + b.pos = pos.data(); + b.seq_id = sptr.data(); + b.n_seq_id = nsid.data(); + b.logits = lg.data(); + + llama_batch_allocr ba(hparams.n_pos_per_embd()); + if (ba.init(b, model->vocab, nullptr, hparams.n_embd_inp(), cparams.n_seq_max, true)) { + auto c = mem->init_batch(ba, cparams.n_ubatch, false); + auto * mc = dynamic_cast(c.get()); + + CHECK(mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, + "one more token fits after the eviction"); + + if (mc && mc->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { + mc->apply(); + + const int64_t n_kv = mc->get_attn()->get_n_kv(); + const int64_t kpool = hparams.indexer_kpool; + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); + + ggml_init_params gp = { ggml_tensor_overhead()*16, nullptr, true }; + ggml_context * ctx = ggml_init(gp); + + kpool_tensors kt = alloc_kpool_tensors(ctx, n_kv, 1, 1, 1, kpool, n_pools); + + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + if (buf) { + llama_kv_cache_set_input_kpool(kv_attn, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, kt.sel_mask, kt.cand_mask, + &mc->get_ubatch(), (uint32_t) kpool); + + const int32_t * cp = (const int32_t *) kt.cell_pool->data; + const auto & cells = kv_attn->get_cells(0); + + // positions 4..7 must stay one pool, or the pooled key built during + // prefill no longer describes the cells it is scored against + std::map slot_of; + for (int64_t j = 0; j < n_kv; ++j) { + if (!cells.is_empty(j) && cells.seq_has(j, 0)) { + slot_of[cells.pos_get(j)] = cp[j]; + } + } + + const bool grouped = + slot_of.count(4) && slot_of.count(5) && slot_of.count(6) && slot_of.count(7) && + slot_of[4] == slot_of[5] && slot_of[5] == slot_of[6] && slot_of[6] == slot_of[7]; + CHECK(grouped, "positions 4..7 stay one pool after the eviction (slot %d)", + grouped ? (int) slot_of[4] : -1); + + // the reference's anchor would make 2..5 the first pool; this port + // deliberately differs, so the difference is pinned rather than fixed + const bool differs = !slot_of.count(2) || slot_of[2] != slot_of[4]; + CHECK(differs, "positions 2 and 4 are NOT pooled together, where the reference anchor would"); + + // the leading remnant 2..3 is an incomplete pool and not in the + // query's tail, the one case where a visible cell is neither + const float * bi = (const float *) kt.bias->data; + const float * sm = (const float *) kt.sel_mask->data; + int64_t n_orphan = 0; + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j) || !cells.seq_has(j, 0)) { + continue; + } + n_orphan += bi[j] == -INFINITY && sm[j] == -INFINITY; + } + CHECK(n_orphan == 2, "the %d-cell leading remnant is neither pooled nor tail (%d)", + (int) n_drop, (int) n_orphan); + + ggml_backend_buffer_free(buf); + } + + ggml_backend_free(backend); + ggml_free(ctx); + } + } + } + + delete mem_raw; + + // ---- the shared input: one fill per ubatch, not one per DSA layer ------- + { + printf("\n--- shared input ---\n"); + + llama_cparams cp = {}; + cp.n_ctx = 16384; + cp.n_ctx_seq = 16384; + cp.n_batch = 512; + cp.n_ubatch = 512; + cp.n_seq_max = 1; + cp.kv_unified = false; + cp.causal_attn = true; + + llama_memory_params mp = {}; + mp.type_k = GGML_TYPE_F16; + mp.type_v = GGML_TYPE_F16; + mp.ctx_type = LLAMA_CONTEXT_TYPE_DEFAULT; + + llama_memory_i * raw = model->create_memory(mp, cp); + auto * m = dynamic_cast(raw); + + if (m) { + llama_kv_cache * kv = m->get_mem_attn(); + + const int64_t n_step = cp.n_ubatch; + const int64_t n_fill = cp.n_ctx_seq - n_step; + const uint32_t n_dsa = (uint32_t) kv->get_layer_ids().size(); + const int64_t kpool = hparams.indexer_kpool; + + std::vector tok(n_step, 0); + std::vector pos(n_step); + std::vector sid(n_step, 0); + std::vector sptr(n_step); + std::vector nsid(n_step, 1); + std::vector lg(n_step, 0); + + for (int64_t i = 0; i < n_step; ++i) { + sptr[i] = &sid[i]; + } + lg.back() = 1; + + llama_memory_context_ptr keep; + + for (int64_t base = 0; base <= n_fill; base += n_step) { + for (int64_t i = 0; i < n_step; ++i) { + pos[i] = (llama_pos) (base + i); + } + + llama_batch b = {}; + b.n_tokens = (int32_t) n_step; + b.token = tok.data(); + b.pos = pos.data(); + b.seq_id = sptr.data(); + b.n_seq_id = nsid.data(); + b.logits = lg.data(); + + llama_batch_allocr ba(hparams.n_pos_per_embd()); + if (!ba.init(b, model->vocab, nullptr, hparams.n_embd_inp(), cp.n_seq_max, true)) { + break; + } + + auto c = m->init_batch(ba, cp.n_ubatch, false); + if (!c || c->get_status() != LLAMA_MEMORY_STATUS_SUCCESS) { + break; + } + c->apply(); + keep = std::move(c); + } + + auto * mctx = dynamic_cast(keep.get()); + CHECK(mctx && mctx->get_status() == LLAMA_MEMORY_STATUS_SUCCESS, + "filled a %d-cell cache in %d-token ubatches", (int) cp.n_ctx_seq, (int) cp.n_ubatch); + if (mctx && mctx->get_status() == LLAMA_MEMORY_STATUS_SUCCESS) { + const int64_t n_kv = mctx->get_attn()->get_n_kv(); + const int64_t n_tps = mctx->get_ubatch().n_tokens; + const int64_t n_padq = GGML_PAD(n_tps, 8) + 8; // exercise the padded-mask path + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, (uint32_t) kpool); + + ggml_init_params gp = { + /* .mem_size */ ggml_tensor_overhead()*16, + /* .mem_buffer */ nullptr, + /* .no_alloc */ true, + }; + + ggml_context * ctx = ggml_init(gp); + kpool_tensors kt = alloc_kpool_tensors(ctx, n_kv, n_tps, n_padq, 1, kpool, n_pools); + + ggml_backend_t backend = ggml_backend_cpu_init(); + ggml_backend_buffer_t buf = ggml_backend_alloc_ctx_tensors(ctx, backend); + + if (buf) { + const llama_ubatch & ub = mctx->get_ubatch(); + + // first pass faults in 60+ MiB of fresh pages; time the steady + // state, which is what a prefill pays per ubatch + double ms = 1e9; + for (int rep = 0; rep < 4; ++rep) { + const auto t0 = std::chrono::steady_clock::now(); + llama_kv_cache_set_input_kpool(kv, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, kt.sel_mask, kt.cand_mask, + &ub, (uint32_t) kpool); + const auto t1 = std::chrono::steady_clock::now(); + + if (rep > 0) { + ms = std::min(ms, std::chrono::duration(t1 - t0).count()); + } + } + + printf("note n_kv = %d, n_tokens = %d, %u DSA layers\n", (int) n_kv, (int) n_tps, n_dsa); + printf("note one fill = %.2f ms; shared = %.2f ms/ubatch, per layer = %.2f ms/ubatch\n", + ms, ms, ms*n_dsa); + printf("note extrapolated to 128Ki x 512 x 11 layers: %.0f ms shared, %.0f ms per layer\n", + ms*(131072.0/(double) n_kv)*(512.0/(double) n_tps), + ms*(131072.0/(double) n_kv)*(512.0/(double) n_tps)*11.0); + + CHECK(n_dsa > 1, "the model has %u indexer layers, so sharing one fill saves %ux the host writes", + n_dsa, n_dsa); + + { + const float * sm = (const float *) kt.sel_mask ->data; + const float * cm = (const float *) kt.cand_mask->data; + + bool pad_masked = true; + for (int64_t ii = n_tps; ii < n_padq; ++ii) { + for (int64_t j = 0; j < n_kv; ++j) { + pad_masked &= sm[ii*n_kv + j] == -INFINITY; + pad_masked &= cm[ii*n_kv + j] == -INFINITY; + } + } + CHECK(pad_masked, "the %d padding rows of a wider sel_mask/cand_mask stay -INFINITY", + (int) (n_padq - n_tps)); + } + + // the padding fill has its own code path per mask type + { + llama_kv_cache_set_input_kpool(kv, kt.cell_pool, kt.pool_cells, kt.bias, kt.pool_bias, + kt.sel_mask_f16, kt.cand_mask_f16, &ub, (uint32_t) kpool); + + CHECK(kpool_mask_f16_matches(kt.sel_mask, kt.sel_mask_f16) && + kpool_mask_f16_matches(kt.cand_mask, kt.cand_mask_f16), + "the f16 masks match the f32 ones on a padded ubatch too"); + } + + // the shared object refills the same tensors deterministically. + // it fills only what the pooled graph consumes - cell_pool and + // bias have no consumer there, so it passes nullptr for both and + // pool_cells is what gets compared + std::vector first((size_t) ggml_nelements(kt.pool_cells)); + memcpy(first.data(), kt.pool_cells->data, first.size()*sizeof(int32_t)); + + llm_graph_input_kpool inp(mctx->get_attn(), mctx->get_idx(), (uint32_t) kpool); + inp.k_idxs = mctx->get_idx()->build_input_k_idxs(ctx, ub); + inp.pool_cells = kt.pool_cells; + inp.pool_bias = kt.pool_bias; + inp.sel_mask = kt.sel_mask; + inp.cand_mask = kt.cand_mask; + + ggml_backend_buffer_t buf2 = ggml_backend_alloc_ctx_tensors(ctx, backend); + if (buf2) { + inp.set_input(&ub); + + CHECK(memcmp(first.data(), kt.pool_cells->data, first.size()*sizeof(int32_t)) == 0, + "llm_graph_input_kpool::set_input rebuilds the identical map, so one\n" + " object can back every indexer layer"); + + ggml_backend_buffer_free(buf2); + } + + ggml_backend_buffer_free(buf); + } + + ggml_backend_free(backend); + ggml_free(ctx); + } else { + printf("note could not fill a %d-cell cache; skipping the cost measurement\n", (int) cp.n_ctx_seq); + } + + keep.reset(); + } + + delete raw; + } + + llama_model_free(model); + llama_backend_free(); + + printf("\n%s: %d failure(s)\n", argv[0], n_fail); + return n_fail == 0 ? 0 : 1; +} diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b8fd66ccae5..ce2febd4209 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -116,9 +116,12 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_KIMI_LINEAR + || arch == LLM_ARCH_GLM5NEXT || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_MISTRAL4) { + // MLA absorbs into MQA, so n_head_kv must be 1: otherwise the per-layer + // head_count_kv array below sizes the latent K row n_head times too wide n_embd = 128; n_head = 1; n_ff = 192; @@ -165,7 +168,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE || arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR || - arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3) { + arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_GLM5NEXT) { GGML_ASSERT(n_layer >= 2); std::vector n_head_per_layer; n_head_per_layer.reserve(n_layer); @@ -213,6 +216,21 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types); } + } else if (arch == LLM_ARCH_GLM5NEXT) { + // nope-only MLA: rope dimension count must be written as 0, not omitted, or + // the generic loader defaults it non-zero and load_arch_hparams rejects it + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(512)); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512)); + ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(0)); + ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192)); + ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128)); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, uint32_t(4)); + // build_hc_pre hard-codes a 4-wide residual + ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1e-6f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } else if (arch == LLM_ARCH_MINIMAX_M3) { // partial rotary: n_rot must not exceed the indexer key length (64) ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64)); @@ -444,6 +462,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_PADDLEOCR: case LLM_ARCH_MIMO2: case LLM_ARCH_KIMI_LINEAR: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_KIMI_K3: case LLM_ARCH_STEP35: case LLM_ARCH_MISTRAL4: diff --git a/tests/test-mtmd-impl.cpp b/tests/test-mtmd-impl.cpp index df18b0a42e2..f6b908344b1 100644 --- a/tests/test-mtmd-impl.cpp +++ b/tests/test-mtmd-impl.cpp @@ -70,6 +70,107 @@ MAKE_TEST(test_image_preprocessor_lfm2) { } } +// GLM-5-Next / GLM-5.3-Flash, hparams as loaded by clip.cpp for PROJECTOR_TYPE_GLM5NEXT +static clip_hparams glm5next_hparams() { + clip_hparams hparams; + hparams.patch_size = 14; + hparams.n_merge = 2; + hparams.set_limit_image_tokens(16, 8000); + return hparams; +} + +// set_limit_image_tokens takes token counts; for a still image the reference's temporal factors cancel, +// leaving factor**2 == 28*28 +MAKE_TEST(test_image_preprocessor_glm5next_budget) { + const clip_hparams hparams = glm5next_hparams(); + + t.assert_equal("image_min_pixels", 16 * 28 * 28, hparams.image_min_pixels); + t.assert_equal("image_max_pixels", 8000 * 28 * 28, hparams.image_max_pixels); +} + +MAKE_TEST(test_image_preprocessor_glm5next_resize) { + const clip_hparams hparams = glm5next_hparams(); + + struct test_case { + clip_image_size input; + clip_image_size canvas; // padded output + clip_image_size content; // resized image inside the canvas + int n_tokens; + }; + + // expected values come from Glm5NextImageProcessor.resize in + // adapt_zips/extract_0826/image_processing_glm5_next.py + const std::vector cases = { + // inside the budget + { { 224, 224 }, { 224, 224 }, { 224, 224 }, 64 }, + { { 448, 448 }, { 448, 448 }, { 448, 448 }, 256 }, + { { 1024, 1024 }, { 1036, 1036 }, { 1024, 1024 }, 1369 }, + // the 16-token floor: 112x112 is exactly 16 tokens + { { 112, 112 }, { 112, 112 }, { 112, 112 }, 16 }, + { { 111, 111 }, { 112, 112 }, { 111, 111 }, 16 }, + { { 113, 113 }, { 140, 140 }, { 113, 113 }, 25 }, + { { 56, 56 }, { 112, 112 }, { 112, 112 }, 16 }, + { { 50, 50 }, { 112, 112 }, { 112, 112 }, 16 }, + { { 28, 28 }, { 112, 112 }, { 112, 112 }, 16 }, + { { 1, 1 }, { 112, 112 }, { 112, 112 }, 16 }, + // the 8000-token ceiling: 2492x2492 still fits, 2520x2520 does not + { { 2492, 2492 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + { { 2493, 2493 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + { { 2504, 2504 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + { { 2520, 2520 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + { { 4000, 4000 }, { 2492, 2492 }, { 2492, 2492 }, 7921 }, + // wide and tall + { { 800, 600 }, { 812, 616 }, { 800, 600 }, 638 }, + { { 600, 800 }, { 616, 812 }, { 600, 800 }, 638 }, + { { 1920, 480 }, { 1932, 504 }, { 1920, 480 }, 1242 }, + { { 480, 1920 }, { 504, 1932 }, { 480, 1920 }, 1242 }, + { { 1920, 1080 }, { 1932, 1092 }, { 1920, 1080 }, 2691 }, + { { 1080, 1920 }, { 1092, 1932 }, { 1080, 1920 }, 2691 }, + // extreme aspect ratios, single-pixel edges survive the resize + { { 4000, 16 }, { 4004, 28 }, { 4000, 16 }, 143 }, + { { 16, 4000 }, { 28, 4004 }, { 16, 4000 }, 143 }, + { { 5000, 1 }, { 5012, 28 }, { 5012, 1 }, 179 }, + { { 1, 5000 }, { 28, 5012 }, { 1, 5012 }, 179 }, + { { 12000, 100 }, { 12012, 112 }, { 12000, 100 }, 1716 }, + { { 100, 12000 }, { 112, 12012 }, { 100, 12000 }, 1716 }, + // over budget, neither edge a multiple of 28 + { { 4007, 3001 }, { 2884, 2156 }, { 2878, 2156 }, 7931 }, + { { 3001, 4007 }, { 2156, 2884 }, { 2156, 2878 }, 7931 }, + { { 3333, 5000 }, { 2044, 3052 }, { 2034, 3052 }, 7957 }, + { { 2729, 2731 }, { 2492, 2492 }, { 2490, 2492 }, 7921 }, + // the 0826 binary search and a Qwen-style smart_resize disagree on all of these, so a regression + // to the old dynamic-size preprocessor cannot pass + { { 4618, 2282 }, { 3528, 1764 }, { 3528, 1743 }, 7938 }, // smart_resize: 3556x1736 + { { 2794, 6096 }, { 1708, 3668 }, { 1681, 3668 }, 7991 }, // smart_resize: 1680x3696 + { { 8858, 1315 }, { 6412, 952 }, { 6412, 951 }, 7786 }, // smart_resize: 6496x952 + { { 1350, 5856 }, { 1204, 5208 }, { 1200, 5208 }, 7998 }, // smart_resize: 1176x5208 + { { 2134, 1472 }, { 2156, 1484 }, { 2134, 1472 }, 4081 }, // smart_resize: 2128x1484 + { { 1021, 4268 }, { 1036, 4284 }, { 1021, 4268 }, 5661 }, // smart_resize: 1008x4256 + }; + + auto fmt = [](const clip_image_size & s) { + return std::to_string(s.width) + "x" + std::to_string(s.height); + }; + + for (const auto & tc : cases) { + const auto geo = mtmd_image_preprocessor_glm5next::get_geometry(hparams, tc.input); + const auto name = " for " + fmt(tc.input); + + t.assert_equal("canvas" + name, fmt(tc.canvas), fmt(geo.canvas)); + t.assert_equal("content" + name, fmt(tc.content), fmt(geo.content)); + + // mirrors clip_n_output_tokens for PROJECTOR_TYPE_GLM5NEXT + const int n_tokens = (geo.canvas.width / (hparams.patch_size * hparams.n_merge)) + * (geo.canvas.height / (hparams.patch_size * hparams.n_merge)); + t.assert_equal("n_tokens" + name, tc.n_tokens, n_tokens); + + t.assert_true("content fits the canvas" + name, + geo.content.width <= geo.canvas.width && geo.content.height <= geo.canvas.height); + t.assert_true("canvas is within the token budget" + name, + geo.canvas.width * geo.canvas.height <= hparams.image_max_pixels); + } +} + // // mtmd temporal merge // diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index e60c9c8787a..4ff8db9e6be 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -38,6 +38,7 @@ add_library(mtmd models/gemma4ua.cpp models/gemma4uv.cpp models/glm4v.cpp + models/glm5next-vision.cpp models/granite-speech.cpp models/granite4-vision.cpp models/hunyuanvl.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index f6045093c63..5c4821f4027 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -63,6 +63,7 @@ #define KEY_PROJ_SAMPLE_WINDOW_SIDE "clip.vision.projector.window_side" #define KEY_PROJ_SPATIAL_OFFSETS "clip.vision.projector.spatial_offsets" #define KEY_SPATIAL_MERGE_SIZE "clip.vision.spatial_merge_size" +#define KEY_VISION_SWIGLU_LIMIT "clip.vision.swiglu_limit" #define KEY_MM_PATCH_MERGE_TYPE "clip.vision.mm_patch_merge_type" #define KEY_IMAGE_GRID_PINPOINTS "clip.vision.image_grid_pinpoints" @@ -482,6 +483,7 @@ enum projector_type { PROJECTOR_TYPE_DEEPSEEKOCR2, PROJECTOR_TYPE_LFM2A, PROJECTOR_TYPE_GLM4V, + PROJECTOR_TYPE_GLM5NEXT, PROJECTOR_TYPE_YOUTUVL, PROJECTOR_TYPE_YASA2, PROJECTOR_TYPE_KIMIK25, @@ -546,6 +548,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_DEEPSEEKOCR2, "deepseekocr2"}, { PROJECTOR_TYPE_LFM2A, "lfm2a"}, { PROJECTOR_TYPE_GLM4V, "glm4v"}, + { PROJECTOR_TYPE_GLM5NEXT, "glm5next"}, { PROJECTOR_TYPE_YOUTUVL, "youtuvl"}, { PROJECTOR_TYPE_YASA2, "yasa2"}, { PROJECTOR_TYPE_KIMIK25, "kimik25"}, diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 060938d86e3..47d28f9018a 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -17,6 +17,7 @@ enum ffn_op_type { FFN_SILU, FFN_GELU_QUICK, FFN_RELU_SQR, + FFN_SILU_CLAMP, }; enum norm_type { @@ -89,6 +90,9 @@ struct clip_hparams { ffn_op_type ffn_op = FFN_GELU; + // clamp applied to the SwiGLU gate/up before the activation (FFN_SILU_CLAMP) + float swiglu_limit = 0.0f; + patch_merge_type mm_patch_merge_type = PATCH_MERGE_FLAT; float eps = 1e-6; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 90de1957586..913f50e59b1 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -684,6 +684,21 @@ ggml_tensor * clip_graph::build_ffn( cur = ggml_sqr(ctx0, cur); cb(cur, "ffn_relu_sqr", il); } break; + case FFN_SILU_CLAMP: + { + // the gate is bounded above only, the up projection on both sides, and both + // before the activation. ggml_swiglu_oai clamps the same way but then adds + // one to the up branch, which is a gpt-oss detail this model does not share + GGML_ASSERT(gate && "FFN_SILU_CLAMP is a gated activation"); + const float limit = hparams.swiglu_limit; + GGML_ASSERT(limit > 0.0f); + tmp = ggml_clamp(ctx0, tmp, -limit, limit); + cb(tmp, "ffn_up_clamped", il); + cur = ggml_clamp(ctx0, cur, -INFINITY, limit); + cb(cur, "ffn_gate_clamped", il); + cur = ggml_swiglu_split(ctx0, cur, tmp); + cb(cur, "ffn_swiglu_limited", il); + } break; } if (down) { @@ -1081,6 +1096,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_GLM5NEXT: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_QWEN3A: { builder = std::make_unique(ctx, img); @@ -1726,6 +1745,23 @@ struct clip_model_loader { hparams.set_limit_image_tokens(8, 4096); hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup } break; + case PROJECTOR_TYPE_GLM5NEXT: + { + hparams.rope_theta = 10000.0f; + hparams.n_merge = 2; + // the reference asks for BICUBIC (resample = PILImageResampling.BICUBIC); + // the bilinear here came from the GLM-4V case this was copied from. neither + // filter matches torchvision's exactly, so this is closer in kind, not exact + hparams.image_resize_algo = RESIZE_ALGO_BICUBIC; + get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); + // drives the clamp in both the per-block MLP and the merger + get_f32(KEY_VISION_SWIGLU_LIMIT, hparams.swiglu_limit); + hparams.ffn_op = FFN_SILU_CLAMP; + log_ffn_op = "silu_clamp"; + // min_pixels/max_pixels of the GLM-5.3-Flash preprocessor, in tokens + hparams.set_limit_image_tokens(16, 8000); + hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup + } break; case PROJECTOR_TYPE_LLAMA4: { hparams.rope_theta = 10000.0f; @@ -2543,6 +2579,7 @@ struct clip_model_loader { } } break; case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5NEXT: { model.mm_fc_w = get_tensor(string_format(TN_MM_PROJECTOR, "weight")); model.mm_ffn_up_w = get_tensor(string_format(TN_MM_UP, "weight")); @@ -4001,6 +4038,7 @@ int clip_n_output_tokens_x(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_EXAONE4_5: case PROJECTOR_TYPE_MIMOVL: case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5NEXT: case PROJECTOR_TYPE_PADDLEOCR: case PROJECTOR_TYPE_HUNYUANVL: case PROJECTOR_TYPE_YOUTUVL: @@ -4027,6 +4065,7 @@ int clip_n_output_tokens_y(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_EXAONE4_5: case PROJECTOR_TYPE_MIMOVL: case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5NEXT: case PROJECTOR_TYPE_PADDLEOCR: case PROJECTOR_TYPE_HUNYUANVL: case PROJECTOR_TYPE_YOUTUVL: @@ -4108,6 +4147,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_MIMOVL: case PROJECTOR_TYPE_MINIMAX_M3: case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5NEXT: case PROJECTOR_TYPE_YOUTUVL: case PROJECTOR_TYPE_MUSE_GLIMMER: { @@ -4733,6 +4773,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { case PROJECTOR_TYPE_QWEN2VL: case PROJECTOR_TYPE_QWEN3VL: case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5NEXT: { const int merge_ratio = hparams.n_merge; const int pw = image_size_width / patch_size; @@ -5881,6 +5922,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { case PROJECTOR_TYPE_GRANITE4_VISION: return ctx->model.qf_proj_blocks.size() * ctx->model.hparams.projection_dim; case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5NEXT: return ctx->model.mm_ffn_down_w->ne[1]; case PROJECTOR_TYPE_MIMO_AUDIO: return ctx->model.mm_2_w->ne[1]; diff --git a/tools/mtmd/models/glm4v.cpp b/tools/mtmd/models/glm4v.cpp index 0e1d596b41b..ad36e7a483d 100644 --- a/tools/mtmd/models/glm4v.cpp +++ b/tools/mtmd/models/glm4v.cpp @@ -42,7 +42,11 @@ ggml_cgraph * clip_graph_glm4v::build() { cb(inp, "patch_bias", -1); // pos-conv norm - inp = build_norm(inp, model.norm_embd_w, model.norm_embd_b, norm_t, eps, -1); + // Note: GLM-OCR does not have a post-conv norm, and build_norm still normalizes when the + // weight is null, so the whole call must be skipped rather than just the affine scale + if (model.norm_embd_w != nullptr) { + inp = build_norm(inp, model.norm_embd_w, model.norm_embd_b, norm_t, eps, -1); + } ggml_tensor * learned_pos_embd = nullptr; // Note: GLM-OCR does not have learned position embeddings diff --git a/tools/mtmd/models/glm5next-vision.cpp b/tools/mtmd/models/glm5next-vision.cpp new file mode 100644 index 00000000000..e1130f1cc5d --- /dev/null +++ b/tools/mtmd/models/glm5next-vision.cpp @@ -0,0 +1,16 @@ +#include "models.h" + +// GLM-5.3-Flash reuses the GLM-OCR ViT unchanged: no learned position embeddings, no +// post-conv norm, q/k norms and biases throughout, and an RMSNorm called post_layernorm. +// ref: https://huggingface.co/zai-org/GLM-5.3-Flash/blob/main/modeling_glm5_next.py +// +// The one difference is the clamp on the SwiGLU gate and up projections. It applies to +// the per-block MLP and to the merger alike, and both read hparams.ffn_op, so selecting +// FFN_SILU_CLAMP for this projector type is enough to cover the pair. +ggml_cgraph * clip_graph_glm5next::build() { + GGML_ASSERT(hparams.ffn_op == FFN_SILU_CLAMP); + GGML_ASSERT(model.norm_embd_w == nullptr); + GGML_ASSERT(model.position_embeddings == nullptr); + + return clip_graph_glm4v::build(); +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 10546fa5dc7..7a7e9e69870 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -178,6 +178,11 @@ struct clip_graph_glm4v : clip_graph { ggml_cgraph * build() override; }; +struct clip_graph_glm5next : clip_graph_glm4v { + clip_graph_glm5next(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph_glm4v(ctx, img) {} + ggml_cgraph * build() override; +}; + struct clip_graph_hunyuanvl : clip_graph { clip_graph_hunyuanvl(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 0dda8770f29..5dc80faae4e 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1583,3 +1583,108 @@ mtmd_image_preproc_out mtmd_image_preprocessor_muse_glimmer::preprocess(const cl output.append(hparams, resized_image, true); return output; } + +// +// mtmd_image_preprocessor_glm5next +// + +// for a still image the reference's temporal_factor cancels on both sides of its budget comparison, +// leaving a plain pixel area against image_min_pixels / image_max_pixels (n_tokens * factor**2) +clip_image_size mtmd_image_preprocessor_glm5next::smart_resize(const clip_hparams & hparams, const clip_image_size & size) { + const int factor = hparams.patch_size * hparams.n_merge; + GGML_ASSERT(factor > 0); + GGML_ASSERT(hparams.image_min_pixels > 0 && hparams.image_max_pixels > 0); + + const int height = size.height; + const int width = size.width; + if (height <= 0 || width <= 0) { + // an empty bitmap is reachable through the public API, same tolerance as calc_size_preserved_ratio + return { 0, 0 }; + } + + const int64_t min_pixels = hparams.image_min_pixels; + const int64_t max_pixels = hparams.image_max_pixels; + + auto align = [factor](int64_t value) { + return (int) (((value + factor - 1) / factor) * factor); + }; + + int aligned_height = align(height); + int aligned_width = align(width); + + // upscale an image that is too small to spend the minimum token budget + if ((int64_t) aligned_height * aligned_width < min_pixels) { + const double scale = std::sqrt((double) min_pixels / ((double) height * (double) width)); + aligned_height = align(std::max(1, (int64_t) std::ceil(height * scale))); + aligned_width = align(std::max(1, (int64_t) std::ceil(width * scale))); + } + + if ((int64_t) aligned_height * aligned_width > max_pixels) { + // binary search the tallest content height whose aligned canvas still fits the budget. the Qwen + // sqrt(area / max_pixels) scale leaves budget unspent: aligning both edges is not monotone in it + int low = 1; + int high = height; + aligned_height = factor; + aligned_width = factor; + while (low <= high) { + const int content_height = (low + high) / 2; + // the reference divides in float64 before flooring, keep it bit-identical + const int content_width = std::max(1, (int) std::floor((double) width * content_height / (double) height)); + const int cand_height = align(content_height); + const int cand_width = align(content_width); + + if ((int64_t) cand_height * cand_width <= max_pixels) { + aligned_height = cand_height; + aligned_width = cand_width; + low = content_height + 1; + } else { + high = content_height - 1; + } + } + } + + return { aligned_width, aligned_height }; +} + +mtmd_image_preprocessor_glm5next::geometry mtmd_image_preprocessor_glm5next::get_geometry(const clip_hparams & hparams, const clip_image_size & size) { + const clip_image_size canvas = smart_resize(hparams, size); + + const int height = size.height; + const int width = size.width; + if (canvas.width == 0 || canvas.height == 0) { + return { canvas, canvas }; + } + + double scale = std::min((double) canvas.height / height, (double) canvas.width / width); + if ((int64_t) height * (int64_t) width >= hparams.image_min_pixels) { + // an image already spending the minimum budget is only shrunk, never upscaled to fill the canvas + scale = std::min(1.0, scale); + } + + geometry geo; + geo.canvas = canvas; + geo.content = { + std::max(1, std::min(canvas.width, (int) std::floor(width * scale))), + std::max(1, std::min(canvas.height, (int) std::floor(height * scale))), + }; + return geo; +} + +mtmd_image_preproc_out mtmd_image_preprocessor_glm5next::preprocess(const clip_image_u8 & img) { + const geometry geo = get_geometry(hparams, img.get_size()); + + clip_image_u8 content; + img_tool::resize(img, content, geo.content, hparams.image_resize_algo, PAD_NONE); + + // the reference pads bottom/right, not centred like img_tool::resize would + clip_image_u8 canvas; + canvas.set_size(geo.canvas, img.is_placeholder()); + if (!img.is_placeholder()) { + img_tool::fill(canvas, hparams.image_pad_color); + img_tool::composite(canvas, content, 0, 0); + } + + mtmd_image_preproc_out output; + output.append(hparams, canvas, true); + return output; +} diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index 732e27379d2..4e0e707b2b8 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -239,3 +239,22 @@ struct mtmd_image_preprocessor_muse_glimmer : mtmd_image_preprocessor { mtmd_image_preprocessor_muse_glimmer(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; }; + +// GLM-5-Next / GLM-5.3-Flash dynamic resize +// ref: Glm5NextImageProcessor.{smart_resize,resize} of the 2026-08-26 adaptation +// +// unlike the Qwen-style smart_resize in mtmd_image_preprocessor_dyn_size, an over-budget image is fitted +// by binary search on the content height, then pasted top-left rather than centred and stretched to fill +struct mtmd_image_preprocessor_glm5next : mtmd_image_preprocessor { + mtmd_image_preprocessor_glm5next(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} + mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; + + struct geometry { + clip_image_size canvas; // aligned, padded output size + clip_image_size content; // resized image, placed at the top-left of the canvas + }; + + // pure arithmetic, exposed so tests can reach it without a clip_ctx + static clip_image_size smart_resize(const clip_hparams & hparams, const clip_image_size & size); + static geometry get_geometry(const clip_hparams & hparams, const clip_image_size & size); +}; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 5b306180d62..1753bd59174 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -860,6 +860,13 @@ struct mtmd_context { img_end = "<|end_of_image|>"; image_preproc = std::make_unique(ctx_v); } break; + case PROJECTOR_TYPE_GLM5NEXT: + { + // glm5next spells video with its own token pair, but video is not supported here + img_beg = "<|begin_of_image|>"; + img_end = "<|end_of_image|>"; + image_preproc = std::make_unique(ctx_v); + } break; case PROJECTOR_TYPE_PADDLEOCR: { // <|IMAGE_START|> ... (image embeddings) ... <|IMAGE_END|>