diff --git a/conversion/__init__.py b/conversion/__init__.py index a5632fcc4bb..e48ae475f11 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -104,6 +104,8 @@ "Glm4MoeLiteForCausalLM": "glm", "Glm4vForConditionalGeneration": "glm", "Glm4vMoeForConditionalGeneration": "glm", + "Glm5NextForCausalLM": "glm5next", + "Glm5NextForConditionalGeneration": "glm5next", "GlmForCausalLM": "chatglm", "GlmMoeDsaForCausalLM": "glm", "GlmOcrForConditionalGeneration": "glm", @@ -296,6 +298,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 daae28e92ad..a31ecf01db1 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -2541,6 +2541,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..f1b4df32e43 100644 --- a/examples/embedding/embedding.cpp +++ b/examples/embedding/embedding.cpp @@ -144,6 +144,11 @@ int main(int argc, char ** argv) { return 1; } + 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/ggml/include/ggml-metal.h b/ggml/include/ggml-metal.h index 433838f0d6d..eb7dec7dbea 100644 --- a/ggml/include/ggml-metal.h +++ b/ggml/include/ggml-metal.h @@ -54,6 +54,11 @@ GGML_BACKEND_API bool ggml_backend_metal_supports_family(ggml_backend_t backend, // capture all command buffers committed the next time `ggml_backend_graph_compute` is called GGML_BACKEND_API void ggml_backend_metal_capture_next_compute(ggml_backend_t backend); +// Returns tensor's Metal GPU virtual address, or 0 for a non-Metal buffer. +// This is a diagnostic interface; ordinary callers should not depend on the +// relative placement of separately allocated MTLBuffers. +GGML_BACKEND_API uint64_t ggml_backend_metal_buffer_get_gpu_address(ggml_backend_buffer_t buffer, const struct ggml_tensor * tensor); + GGML_BACKEND_API ggml_backend_reg_t ggml_backend_metal_reg(void); #ifdef __cplusplus diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index 1227ed39a09..2eb98191322 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -467,6 +467,23 @@ enum ggml_status ggml_metal_graph_compute(ggml_metal_t ctx, struct ggml_cgraph * ctx->n_nodes_per_cb = (ctx->n_nodes_1 + ctx->n_cb - 1) / ctx->n_cb; + // fabley (speed conference 008/009): env-gated arm-file — when + // LLAMA_METAL_CAPTURE_ARM_FILE is set and the file exists, consume it + // and capture THIS compute. Production without the env pays nothing. + { + static const char * arm_path = NULL; + static bool arm_checked = false; + if (!arm_checked) { + arm_checked = true; + arm_path = getenv("LLAMA_METAL_CAPTURE_ARM_FILE"); + } + if (arm_path && access(arm_path, F_OK) == 0) { + unlink(arm_path); + ctx->capture_compute = 1; + GGML_LOG_WARN("%s: capture armed via %s\n", __func__, arm_path); + } + } + if (ctx->capture_compute >= 0) { ctx->capture_compute--; } diff --git a/ggml/src/ggml-metal/ggml-metal-device.h b/ggml/src/ggml-metal/ggml-metal-device.h index 003b688dbac..cef4a083ce7 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.h +++ b/ggml/src/ggml-metal/ggml-metal-device.h @@ -329,6 +329,7 @@ ggml_metal_buffer_t ggml_metal_buffer_map (ggml_metal_device_t dev, void * ptr, void ggml_metal_buffer_free (ggml_metal_buffer_t buf); void * ggml_metal_buffer_get_base (ggml_metal_buffer_t buf); +uint64_t ggml_metal_buffer_get_gpu_address(ggml_metal_buffer_t buf, const struct ggml_tensor * tensor); bool ggml_metal_buffer_is_shared(ggml_metal_buffer_t buf); void ggml_metal_buffer_memset_tensor(ggml_metal_buffer_t buf, struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 41ce90dc8a9..3de01d6b5c6 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -2141,6 +2141,16 @@ void ggml_metal_buffer_free(ggml_metal_buffer_t buf) { return buf->all_data; } +uint64_t ggml_metal_buffer_get_gpu_address(ggml_metal_buffer_t buf, const struct ggml_tensor * tensor) { + struct ggml_metal_buffer_id bid = ggml_metal_buffer_get_id(buf, tensor); + if (bid.metal == nil) { + return 0; + } + + id metal = bid.metal; + return metal.gpuAddress + bid.offs; +} + bool ggml_metal_buffer_is_shared(ggml_metal_buffer_t buf) { return buf->is_shared; } diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index 75de0f6dd08..d04bbe56fd2 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -521,7 +521,10 @@ static int ggml_metal_op_encode_impl(ggml_metal_op_t ctx, int idx) { int ggml_metal_op_encode(ggml_metal_op_t ctx, int idx) { if (ctx->use_capture) { - ggml_metal_encoder_debug_group_push(ctx->enc, ggml_op_desc(ctx->node(idx))); + // fabley: include the tensor name so traces attribute per graph stage + char dbg[192]; + snprintf(dbg, sizeof(dbg), "%s|%s", ctx->node(idx)->name, ggml_op_desc(ctx->node(idx))); + ggml_metal_encoder_debug_group_push(ctx->enc, dbg); } int res = ggml_metal_op_encode_impl(ctx, idx); diff --git a/ggml/src/ggml-metal/ggml-metal.cpp b/ggml/src/ggml-metal/ggml-metal.cpp index 9756d47050c..414851b1335 100644 --- a/ggml/src/ggml-metal/ggml-metal.cpp +++ b/ggml/src/ggml-metal/ggml-metal.cpp @@ -182,6 +182,17 @@ static bool ggml_backend_buffer_is_metal(ggml_backend_buffer_t buffer) { buffer->iface.free_buffer == ggml_backend_metal_buffer_private_free_buffer; } +uint64_t ggml_backend_metal_buffer_get_gpu_address( + ggml_backend_buffer_t buffer, + const struct ggml_tensor * tensor) { + if (buffer == nullptr || tensor == nullptr || !ggml_backend_buffer_is_metal(buffer)) { + return 0; + } + + ggml_metal_buffer_t ctx = (ggml_metal_buffer_t) buffer->context; + return ggml_metal_buffer_get_gpu_address(ctx, tensor); +} + // // buffer types // diff --git a/ggml/src/ggml-metal/kernels/mul_mm.metal b/ggml/src/ggml-metal/kernels/mul_mm.metal index ee848eed6d6..df3be003890 100644 --- a/ggml/src/ggml-metal/kernels/mul_mm.metal +++ b/ggml/src/ggml-metal/kernels/mul_mm.metal @@ -134,7 +134,8 @@ kernel void kernel_mul_mm( // Store result tile to output matrix (with batch offset) // cT.store handles bounds checking via tD's extents (M, N) - device float * dstBatch = (device float *)dst + im * N * M; + // int32 im*N*M wraps for KQ [n_kv, ub, 64] past 2^31 elements (glm5next collapse) + device float * dstBatch = (device float *)dst + (uint64_t)im * (uint64_t)N * (uint64_t)M; auto tD = tensor(dstBatch, dextents(M, N), array({1, M})); cT.store(tD.slice(ra, rb)); @@ -318,7 +319,7 @@ kernel void kernel_mul_mm( // if no bounds checks on the output are needed, we can directly write to device memory device float * C = (device float *) dst + (r0 + 32*(sgitg & 1)) + \ - (r1 + 16*(sgitg >> 1)) * args.ne0 + im*args.ne1*args.ne0; + (uint64_t)(r1 + 16*(sgitg >> 1)) * (uint64_t)args.ne0 + (uint64_t)im*(uint64_t)args.ne1*(uint64_t)args.ne0; for (short i = 0; i < 8; i++) { simdgroup_store(mc[i], C + 8*(i%4) + 8*args.ne0*(i/4), args.ne0, 0, false); @@ -337,7 +338,7 @@ kernel void kernel_mul_mm( if (sgitg == 0) { for (int j = tiitg; j < nr1; j += NR1) { - device float * D = (device float *) dst + r0 + (r1 + j)*args.ne0 + im*args.ne1*args.ne0; + device float * D = (device float *) dst + r0 + (uint64_t)(r1 + j)*(uint64_t)args.ne0 + (uint64_t)im*(uint64_t)args.ne1*(uint64_t)args.ne0; device float4 * D4 = (device float4 *) D; threadgroup float * C = temp_str + (j*NR0); diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index c99feb3c795..a67d067043d 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -225,6 +225,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" @@ -382,6 +383,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" @@ -559,6 +561,7 @@ class MODEL_ARCH(IntEnum): GLM4 = auto() GLM4_MOE = auto() GLM_DSA = auto() + GLM5NEXT = auto() BITNET = auto() T5 = auto() T5ENCODER = auto() @@ -1307,6 +1310,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", @@ -3991,6 +3995,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, @@ -5648,6 +5716,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 1f309ad2eaf..1d29d7896b5 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) @@ -1376,6 +1379,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 861acfe181f..b36cc43b386 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2708,6 +2708,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", + ), + }, MODEL_ARCH.QWEN4EXP: { MODEL_TENSOR.HC_ATTN_NORM: ( "model.layers.{bid}.attn_hyper_connection.hc_norm", diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 8922dc12adc..768be87583e 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 5e61f61f7f0..01cd315b306 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -84,6 +84,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GLM4, "glm4" }, { LLM_ARCH_GLM4_MOE, "glm4moe" }, { LLM_ARCH_GLM_DSA, "glm-dsa" }, + { LLM_ARCH_GLM5NEXT, "glm5next" }, { LLM_ARCH_BITNET, "bitnet" }, { LLM_ARCH_T5, "t5" }, { LLM_ARCH_T5ENCODER, "t5encoder" }, @@ -284,6 +285,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" }, @@ -1077,6 +1079,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_QWEN35MOE: case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_MINIMAX_01: return true; default: @@ -1101,6 +1104,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: @@ -1139,6 +1143,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 ca7d55a5fd7..201d7ceff50 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -89,6 +89,7 @@ enum llm_arch { LLM_ARCH_GLM4, LLM_ARCH_GLM4_MOE, LLM_ARCH_GLM_DSA, + LLM_ARCH_GLM5NEXT, LLM_ARCH_BITNET, LLM_ARCH_T5, LLM_ARCH_T5ENCODER, @@ -289,6 +290,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 fb88919f9d6..cbe23640388 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -1,6 +1,10 @@ #include "llama-context.h" +#include +#include + #include "ggml.h" +#include "ggml-metal.h" #include "llama-arch.h" #include "llama-graph.h" #include "llama-impl.h" @@ -17,8 +21,10 @@ #include #include #include +#include #include #include +#include // // llama_context @@ -236,6 +242,15 @@ llama_context::llama_context( cparams.fused_lid = true; cparams.auto_flid = true; + { + // the fused kernel sums heads in a different order, so a near-tied top-k can differ + 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; @@ -1336,6 +1351,12 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll // in order to correctly reuse a graph, it's full topology has to be uniquely determined by these parameters const auto gparams = graph_params(res, ubatch, mctx, gtype); + // R0 instrumentation (codex 072 s3): graph_rebuild_us and set_inputs_us are + // SEPARATE owners — R0 shrinks only the first, E0c/E1 only the second. + static const bool fbl_host_timers = getenv("LLAMA_HOST_TIMERS") != nullptr; + const int64_t fbl_t_rebuild0 = fbl_host_timers ? ggml_time_us() : 0; + bool fbl_reused = false; + if (!graph_reuse_disable && res->can_reuse(gparams)) { //LLAMA_LOG_DEBUG("%s: reusing previous graph\n", __func__); @@ -1347,6 +1368,7 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll } n_reused++; + fbl_reused = true; } else { res->reset(); @@ -1370,19 +1392,180 @@ llm_graph_result * llama_context::process_ubatch(const llama_ubatch & ubatch, ll ret = GGML_STATUS_ALLOC_FAILED; return nullptr; } + + // fabley: alloc-map dump for the wrap-landing investigation (~/conference). + // Env-gated; dumps every tensor's VA pointer + owning buffer for graphs whose + // batched kq ne0 crosses LLAMA_ALLOCDUMP_MIN_NE0, at most every _STRIDE tokens. + { + static const char * fbl_path = getenv("LLAMA_ALLOCDUMP_FILE"); + if (fbl_path) { + static long long fbl_gate = getenv("LLAMA_ALLOCDUMP_MIN_NE0") ? atoll(getenv("LLAMA_ALLOCDUMP_MIN_NE0")) : 0; + static long long fbl_stride = getenv("LLAMA_ALLOCDUMP_STRIDE") ? atoll(getenv("LLAMA_ALLOCDUMP_STRIDE")) : 2048; + static long long fbl_reqtok = getenv("LLAMA_ALLOCDUMP_REQ_TOKENS") ? atoll(getenv("LLAMA_ALLOCDUMP_REQ_TOKENS")) : 0; + static long long fbl_last = -1; + // grokk-reply: reserve (n_tokens=1) and tail batches are the wrong geometry + const bool fbl_tok_ok = fbl_reqtok == 0 || (long long) ubatch.n_tokens == fbl_reqtok; + long long kq_ne0 = -1; + for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { + const ggml_tensor * n = ggml_graph_node(gf, i); + if (strncmp(n->name, "kq-", 3) == 0 && n->ne[2] >= 8 && n->ne[0] > kq_ne0) { + kq_ne0 = n->ne[0]; + } + } + if (fbl_tok_ok && kq_ne0 >= fbl_gate && (fbl_last < 0 || llabs(kq_ne0 - fbl_last) >= fbl_stride)) { + fbl_last = kq_ne0; + if (FILE * f = fopen(fbl_path, "a")) { + fprintf(f, "GRAPH kq_ne0=%lld n_nodes=%d n_tokens=%d\n", + kq_ne0, ggml_graph_n_nodes(gf), (int) ubatch.n_tokens); + std::unordered_set fbl_seen; + auto fbl_dump = [&](const char * tag, int idx, const ggml_tensor * t) { + if (t == nullptr || fbl_seen.count(t)) { + return; + } + fbl_seen.insert(t); + ggml_backend_buffer_t buf = t->buffer; + const uint64_t gpu_data = buf ? ggml_backend_metal_buffer_get_gpu_address(buf, t) : 0; + const uint64_t gpu_base = gpu_data + ? gpu_data - ((uintptr_t) t->data - (uintptr_t) ggml_backend_buffer_get_base(buf)) + : 0; + fprintf(f, "%s|%d|%s|%s|%p|%zu|%lld,%lld,%lld,%lld|%p|%zu|%s|%p|0x%llx|0x%llx\n", + tag, idx, t->name, ggml_op_name(t->op), t->data, ggml_nbytes(t), + (long long) t->ne[0], (long long) t->ne[1], + (long long) t->ne[2], (long long) t->ne[3], + buf ? ggml_backend_buffer_get_base(buf) : nullptr, + buf ? ggml_backend_buffer_get_size(buf) : (size_t) 0, + buf ? ggml_backend_buffer_name(buf) : "-", + (const void *) t->view_src, + (unsigned long long) gpu_data, (unsigned long long) gpu_base); + }; + for (int i = 0; i < ggml_graph_n_nodes(gf); i++) { + ggml_tensor * n = ggml_graph_node(gf, i); + fbl_dump("N", i, n); + if (n->view_src) { + fbl_dump("V", i, n->view_src); + } + for (int j = 0; j < GGML_MAX_SRC; j++) { + fbl_dump("S", i, n->src[j]); + if (n->src[j] && n->src[j]->view_src) { + fbl_dump("V", i, n->src[j]->view_src); + } + } + } + fprintf(f, "ENDGRAPH\n"); + fclose(f); + } + } + } + } } // set the input data for the input tensors { - //const auto t_start_us = ggml_time_us(); + const int64_t fbl_t_set0 = fbl_host_timers ? ggml_time_us() : 0; // FIXME this call causes a crash if any model inputs were not used in the graph and were therefore not allocated res->set_inputs(&ubatch); - //LLAMA_LOG_INFO("graph set inputs time: %.3f ms\n", (ggml_time_us() - t_start_us)/1000.0); + if (fbl_host_timers && ubatch.n_tokens == 1) { + // decode-1 only: prefill chunks would swamp the percentiles + const int64_t now = ggml_time_us(); + fbl_host_rebuild_us.push_back(fbl_reused ? 0 : (fbl_t_set0 - fbl_t_rebuild0)); + fbl_host_setinp_us.push_back(now - fbl_t_set0); + static const int64_t fbl_every = [](){ + const char * e = getenv("LLAMA_HOST_TIMERS"); + const long long v = e ? atoll(e) : 0; + return v > 0 ? v : 64; + }(); + if ((int64_t) fbl_host_rebuild_us.size() % fbl_every == 0) { + auto pct = [](std::vector v, double p) { + std::sort(v.begin(), v.end()); + return v.empty() ? (int64_t) 0 : v[(size_t) (p*(v.size() - 1))]; + }; + // rebuild EVENT LEDGER, chronological (codex 082 s3): the first + // event is the cold graph after restore, never a pad-256 cost + // (grokk 081 s2); later events are the width-cadence oracle. + // Percentiles over one or two events pretend to knowledge that + // does not exist, so raw ordinals+values are printed instead. + const size_t n_tot = fbl_host_rebuild_us.size(); + int64_t sum_rebuild = 0; + std::string events; + int64_t first_us = 0, later_max = 0; + size_t later_n = 0; + char buf[64]; + for (size_t j = 0; j < n_tot; ++j) { + const int64_t v = fbl_host_rebuild_us[j]; + sum_rebuild += v; + if (v > 0) { + if (first_us == 0) { + first_us = v; + } else { + later_n++; + later_max = std::max(later_max, v); + } + if (events.size() < 256) { + snprintf(buf, sizeof(buf), "%s%zu:%lld", + events.empty() ? "" : ",", j + 1, (long long) v); + events += buf; + } + } + } + fprintf(stderr, "HOST_TIMERS n=%zu n_rebuild=%zu n_reuse=%zu | " + "first_rebuild_us=%lld later_n=%zu later_max_us=%lld | " + "rebuild_events=%s | burst_amort_us=%lld | " + "set_inputs_us p50=%lld p95=%lld\n", + n_tot, (size_t) (later_n + (first_us > 0)), n_tot - later_n - (first_us > 0), + (long long) first_us, later_n, (long long) later_max, + events.empty() ? "none" : events.c_str(), + (long long) (n_tot ? sum_rebuild/(int64_t) n_tot : 0), + (long long) pct(fbl_host_setinp_us, 0.5), (long long) pct(fbl_host_setinp_us, 0.95)); + } + } } const auto status = graph_compute(res->get_gf(), ubatch.n_tokens > 1); + + // E0a: report any layer where the gathered mask is not subsumed by the + // slot certificate (debug env; nodes exist only when the env is set) + static const bool fbl_e0a = getenv("LLAMA_GLM5_E0A") != nullptr; + if (fbl_e0a && status == GGML_STATUS_SUCCESS) { + ggml_cgraph * gfc = res->get_gf(); + std::map> layers; + for (int i = 0; i < ggml_graph_n_nodes(gfc); ++i) { + ggml_tensor * t = ggml_graph_node(gfc, i); + if (strncmp(t->name, "e0a_old-", 8) == 0) { + layers[atoi(t->name + 8)].first = t; + } else if (strncmp(t->name, "e0a_cert-", 9) == 0) { + layers[atoi(t->name + 9)].second = t; + } + } + for (auto & kv : layers) { + ggml_tensor * to = kv.second.first; + ggml_tensor * tc = kv.second.second; + if (!to || !tc) continue; + const int64_t n = to->ne[0]; + std::vector vo(n), vc(n); + ggml_backend_tensor_get(to, vo.data(), 0, n*sizeof(float)); + ggml_backend_tensor_get(tc, vc.data(), 0, n*sizeof(float)); + int bad = 0, finite = 0; + for (int64_t j = 0; j < n; ++j) { + const bool fo = vo[j] > -1e30f; + const bool fc = vc[j] > -1e30f; + finite += fc; + if (fo != fc) bad++; + } + if (bad) { + fprintf(stderr, "E0A VIOLATION layer=%d n=%lld class_mismatch=%d finite_cert=%d\n", + kv.first, (long long) n, bad, finite); + } else { + static bool logged = false; + if (!logged) { + fprintf(stderr, "E0A CHECK ACTIVE n=%lld finite=%d (layer %d clean)\n", + (long long) n, finite, kv.first); + logged = true; + } + } + } + } if (status != GGML_STATUS_SUCCESS) { LLAMA_LOG_ERROR("%s: failed to compute graph, compute status: %d\n", __func__, status); ret = status; @@ -2293,8 +2476,9 @@ 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 runs out by ubatch 3840: KDA costs 182 nodes + ~16/token + // per layer, so 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 || @@ -3257,6 +3441,8 @@ llama_perf_context_data llama_context::perf_get_data() const { } void llama_context::perf_reset() { + fbl_host_rebuild_us.clear(); + fbl_host_setinp_us.clear(); t_start_us = ggml_time_us(); t_eval_us = n_eval = 0; t_p_eval_us = n_p_eval = 0; diff --git a/src/llama-context.h b/src/llama-context.h index bf91daa8b56..9b3f4a2073d 100644 --- a/src/llama-context.h +++ b/src/llama-context.h @@ -391,4 +391,8 @@ struct llama_context { mutable int32_t n_eval = 0; // number of eval calls mutable int32_t n_reused = 0; // number of times the previous graph was reused + + // R0 host-interval buckets (env LLAMA_HOST_TIMERS, decode-1 samples only) + std::vector fbl_host_rebuild_us; + std::vector fbl_host_setinp_us; }; diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 8fca8e1bc0e..f786ca9d5e3 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1,5 +1,7 @@ #include "llama-graph.h" +#include + #include "llama-impl.h" #include "llama-model.h" #include "llama-batch.h" @@ -12,6 +14,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" @@ -504,7 +507,13 @@ bool llm_graph_input_attn_kv::can_reuse(const llm_graph_params & params) { void llm_graph_input_attn_k::set_input(const llama_ubatch * ubatch) { mctx->set_input_k_idxs(self_k_idxs, ubatch); - mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + if (self_v_idxs) { + mctx->set_input_v_idxs(self_v_idxs, ubatch); // Patch B mirror + } + + if (self_kq_mask) { + mctx->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + } } bool llm_graph_input_attn_k::can_reuse(const llm_graph_params & params) { @@ -518,7 +527,17 @@ bool llm_graph_input_attn_k::can_reuse_impl(const llm_graph_params & params) { res &= self_k_idxs->ne[0] == params.ubatch.n_tokens; - res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams); + // Patch B (grokk 036): exact width check, mirroring the hybrid-k wrapper + res &= !self_v_idxs || + (int64_t) self_v_idxs->ne[0] == + (int64_t) params.ubatch.n_tokens * mctx->mirror_width_max(); + + if (self_kq_mask) { + res &= can_reuse_kq_mask(self_kq_mask, mctx, params.ubatch, params.cparams); + } else { + // maskless graph: the width witness is the stored build-time n_kv + res &= b_n_kv == (int64_t) mctx->get_n_kv(); + } return res; } @@ -1087,7 +1106,9 @@ void llm_graph_input_mem_hybrid::set_input(const llama_ubatch * ubatch) { mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); - mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + if (inp_attn->self_kq_mask) { + mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + } if (inp_attn->self_k_rot) { mctx->get_attn()->set_input_k_rot(inp_attn->self_k_rot); @@ -1120,7 +1141,11 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) { res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens; //res &= inp_attn->self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there - res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams); + if (inp_attn->self_kq_mask) { + res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams); + } else { + res &= inp_attn->b_n_kv == (int64_t) mctx->get_attn()->get_n_kv(); + } res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs(); @@ -1139,7 +1164,13 @@ bool llm_graph_input_mem_hybrid::can_reuse(const llm_graph_params & params) { void llm_graph_input_mem_hybrid_k::set_input(const llama_ubatch * ubatch) { mctx->get_attn()->set_input_k_idxs(inp_attn->self_k_idxs, ubatch); - mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + if (inp_attn->self_v_idxs) { + mctx->get_attn()->set_input_v_idxs(inp_attn->self_v_idxs, ubatch); // Patch B mirror + } + + if (inp_attn->self_kq_mask) { + mctx->get_attn()->set_input_kq_mask(inp_attn->self_kq_mask, ubatch, cparams.causal_attn); + } const int64_t n_rs = mctx->get_recr()->get_n_rs(); @@ -1163,7 +1194,11 @@ bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) { res &= inp_attn->self_k_idxs->ne[0] == params.ubatch.n_tokens; - res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams); + if (inp_attn->self_kq_mask) { + res &= can_reuse_kq_mask(inp_attn->self_kq_mask, mctx->get_attn(), params.ubatch, params.cparams); + } else { + res &= inp_attn->b_n_kv == (int64_t) mctx->get_attn()->get_n_kv(); + } res &= inp_rs->s_copy->ne[0] == mctx->get_recr()->get_n_rs(); @@ -1173,6 +1208,12 @@ bool llm_graph_input_mem_hybrid_k::can_reuse(const llm_graph_params & params) { res &= inp_rs->head == mctx->get_recr()->get_head(); res &= inp_rs->rs_z == mctx->get_recr()->get_rs_z(); + if (inp_attn->self_v_idxs && + (int64_t) inp_attn->self_v_idxs->ne[0] != + (int64_t) params.ubatch.n_tokens * mctx->get_attn()->mirror_width_max()) { + return false; // Patch B (grokk 033 2.3) + } + return res; } @@ -1419,7 +1460,7 @@ bool llm_graph_result::can_reuse(const llm_graph_params & params) { const bool cur = input->can_reuse(params); if (debug > 1) { - LLAMA_LOG_DEBUG("%s: can_reuse = %d\n", "placeholder", cur); + LLAMA_LOG_DEBUG("%s: can_reuse = %d\n", typeid(*input).name(), cur); } res = res && cur; @@ -1779,7 +1820,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 +2217,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); @@ -2604,12 +2645,67 @@ ggml_tensor * llm_graph_context::build_attn_mha( cur = ggml_reshape_2d(ctx0, cur, cur->ne[0]*cur->ne[1], cur->ne[2]*cur->ne[3]); } else { - ggml_tensor * kq = ggml_mul_mat(ctx0, k, q); - cb(kq, "kq", il); + { // codex 026: complementary branch marker — OFF cells must prove they + // reached the non-FA branch, not merely lack the packed marker + static bool fbl_branch_logged = false; + if (!fbl_branch_logged) { + fbl_branch_logged = true; + fprintf(stderr, "GLM_ATTN_PATH=NON_FA\n"); + } + } + // fabley speed-conf 019 Patch A: packed-MQA decode transform. + // At decode (n_tokens==1) with an MQA cache (one KV head) and many + // query heads, mul_mat broadcasts over ne2 and Metal re-reads the + // whole K cache once per head (measured 6.03 ms/layer @108K). + // Packing heads into ne1 selects mul_mm (r2=1, K tiled once), then + // permutes back so softmax/mask see the original layout. + // Env-gated: LLAMA_PACKED_MQA_DECODE=1. Prefill is never affected. + // mode: off (default) | kq | kqv | both (legacy "1" == both) + static const int fbl_mode = [] { + const char * e = getenv("LLAMA_PACKED_MQA_DECODE"); + if (!e || strcmp(e, "0") == 0 || strcmp(e, "off") == 0) return 0; + if (strcmp(e, "kq") == 0) return 1; + if (strcmp(e, "kqv") == 0) return 2; + if (strcmp(e, "both") == 0 || strcmp(e, "1") == 0) return 3; // "1" = legacy alias + // grokk 028 §3: bare atoi made "2" mean both — fail loud instead + fprintf(stderr, "LLAMA_PACKED_MQA_DECODE=%s unrecognized (want off|kq|kqv|both) -> off\n", e); + return 0; + }(); + const bool fbl_gate = fbl_mode != 0 && + q->ne[1] == 1 && k->ne[2] == 1 && q->ne[2] > 8 && n_stream == 1; + const bool fbl_pack = fbl_gate && (fbl_mode & 1) != 0; // KQ side + const bool fbl_pack_o = fbl_gate && (fbl_mode & 2) != 0; // KQV side + + ggml_tensor * kq = nullptr; + if (fbl_gate) { + static bool fbl_logged = false; + if (!fbl_logged) { + fbl_logged = true; + fprintf(stderr, "packed-MQA decode path ACTIVE (mode=%d)\n", fbl_mode); + } + } + if (fbl_pack) { + // Patch C (view-only): singleton-axis permutes are already contiguous; + // ggml_cont here was a real Metal copy. Keep the logical shape as a view. + ggml_tensor * qp = ggml_permute(ctx0, q, 0, 2, 1, 3); // [d, n_head, 1] view + GGML_ASSERT(ggml_is_contiguous(qp)); + cb(qp, "q_packed", il); + kq = ggml_mul_mat(ctx0, k, qp); // [n_kv, n_head, 1] + ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + cb(kq, "kq", il); // grokk 020: historical name on the mul_mm + kq = ggml_permute(ctx0, kq, 0, 2, 1, 3); // [n_kv, 1, n_head] view (contiguous AND permuted) + GGML_ASSERT(ggml_is_contiguous(kq)); + cb(kq, "kq_unpacked", il); + } else { + kq = ggml_mul_mat(ctx0, k, q); + cb(kq, "kq", il); + } // note: this op tends to require high floating point range // while for some models F16 is enough, for others it is not, so we default to F32 here - ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + if (!fbl_pack) { + ggml_mul_mat_set_prec(kq, GGML_PREC_F32); // packed branch set it on the mul_mm above + } if (arch == LLM_ARCH_GROK) { // need to do the following: @@ -2648,8 +2744,23 @@ ggml_tensor * llm_graph_context::build_attn_mha( cb(v, "v_cont", il); } - ggml_tensor * kqv = ggml_mul_mat(ctx0, v, kq); - cb(kqv, "kqv", il); + ggml_tensor * kqv = nullptr; + if (fbl_pack_o) { + ggml_tensor * kqp = ggml_permute(ctx0, kq, 0, 2, 1, 3); // [n_kv, n_head, 1] view + GGML_ASSERT(ggml_is_contiguous(kqp)); + cb(kqp, "kq_soft_max_packed", il); + kqv = ggml_mul_mat(ctx0, v, kqp); // [d_v, n_head, 1] + // packing moves this product from mul_mv (f32 accumulation) to mul_mm + // (f16 tile accumulation by default) - request f32 like the KQ side + ggml_mul_mat_set_prec(kqv, GGML_PREC_F32); + cb(kqv, "kqv", il); // historical name on the mul_mm + kqv = ggml_permute(ctx0, kqv, 0, 2, 1, 3); // [d_v, 1, n_head] view + GGML_ASSERT(ggml_is_contiguous(kqv)); + cb(kqv, "kqv_unpacked", il); + } else { + kqv = ggml_mul_mat(ctx0, v, kq); + cb(kqv, "kqv", il); + } // for MLA with the absorption optimization, we need to "decompress" from MQA back to MHA if (v_mla) { @@ -2863,7 +2974,8 @@ static std::unique_ptr build_attn_inp_k_impl( const llama_ubatch & ubatch, const llama_hparams & hparams, const llama_cparams & cparams, - const llama_kv_cache_context * mctx_cur) { + const llama_kv_cache_context * mctx_cur, + bool skip_kq_mask = false) { auto inp = std::make_unique(hparams, cparams, mctx_cur); @@ -2872,8 +2984,17 @@ static std::unique_ptr build_attn_inp_k_impl( inp->self_k_idxs = mctx_cur->build_input_k_idxs(ctx0, ubatch); - inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); - inp->self_kq_mask_cnv = inp->self_kq_mask; + if (mctx_cur->get_has_v_mirror()) { + inp->self_v_idxs = mctx_cur->build_input_v_idxs(ctx0, ubatch); // Patch B + } + + // E0c: a maskless (gathered) graph never consumes the dense KQ mask; an + // unconsumed input is never allocated and the fill would fault (058 class) + if (!skip_kq_mask) { + inp->self_kq_mask = build_attn_inp_kq_mask(ctx0, mctx_cur, ubatch, cparams); + inp->self_kq_mask_cnv = inp->self_kq_mask; + } + inp->b_n_kv = mctx_cur->get_n_kv(); } return inp; @@ -2914,13 +3035,22 @@ ggml_tensor * llm_graph_context::build_attn( const auto & k_idxs = inp->get_k_idxs(); ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + + // Patch B (grokk 033 2.1): dense consumer maintains the mirror too + if (mctx_cur->get_has_v_mirror()) { + ggml_tensor * vupd = mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il); + cb(vupd, "v_mirror_upd", il); + ggml_build_forward_expand(gf, vupd); + } } const auto & kq_mask = inp->get_kq_mask(); 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 * v = mctx_cur->get_has_v_mirror() + ? mctx_cur->get_v(ctx0, il) + : 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, kq_mask, sinks, v_mla, kq_scale, il); cb(cur, "kqv_out", il); @@ -3541,17 +3671,358 @@ llm_graph_input_mem_hybrid * llm_graph_context::build_inp_mem_hybrid() const { return (llm_graph_input_mem_hybrid *) res->add_input(std::move(inp)); } -llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k() const { +llm_graph_input_mem_hybrid_k * llm_graph_context::build_inp_mem_hybrid_k(bool maskless) const { const auto * mctx_cur = static_cast(mctx); auto inp_rs = build_rs_inp_impl (ctx0, ubatch, mctx_cur->get_recr()); - auto inp_attn = build_attn_inp_k_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_attn()); + auto inp_attn = build_attn_inp_k_impl(ctx0, ubatch, hparams, cparams, mctx_cur->get_attn(), maskless); auto inp = std::make_unique(cparams, std::move(inp_attn), std::move(inp_rs), mctx_cur); 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, + bool gathered) 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"); + + // R0: build-time identity for can_reuse + inp->b_n_kv_idx = mctx_idx->get_n_kv(); + inp->b_n_kv_attn = mctx_attn->get_n_kv(); + inp->hparams_indexer_top_k = hparams.indexer_top_k; + + if (scoring) { + const int64_t n_kv = mctx_attn->get_n_kv(); + + // must match build_attn_inp_kq_mask; get_n_stream() is the stream RANGE and is wrong + const int64_t n_stream = cparams.kv_unified ? 1 : ubatch.n_seqs_unq; + const int64_t n_tps = ubatch.n_tokens/n_stream; + + // pool maps are per SEQUENCE; sized on the ubatch, not n_seq_max (256 in llama-embedding) + 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 == nullptr || (kq_mask->ne[0] == n_kv && kq_mask->ne[3] == n_stream)); + + // the selection terms below exist only for real queries + GGML_ASSERT(kq_mask == nullptr || + (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 indexer wants f16; built once, shared by every indexer layer + 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"); + } + + // lossless in f16 (only 0.0f and -INFINITY), and f16 + f32 -> f16 adds the KQ mask uncast. + // E0c: the maskless gathered graph consumes neither dense mask — creating + // them unconsumed would leave unallocated inputs for the host fill (058) + if (!gathered) { + 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"); + } + + // D0 gathered decode inputs. ONLY when the gathered graph will consume them: + // an input tensor no node reads is never allocated, so the host fill would + // write through data == nullptr. + // PAD32 (codex 063 / grokk 064): packed mul_mm's fast path needs the compact + // reduction dim % 32 == 0 (bc_inp law); pad the tail slots so + // n_top + n_tail_slots lands on the tile. Dead slots are cell 0 / -inf — + // the exact representation the duplicate-cell-0 adversarial proved. + if (gathered) { + const int64_t n_top = (int64_t) kpool * + llama_kpool_select_k((uint32_t) n_pools, hparams.indexer_top_k, kpool); + const int64_t n_sel_pad = GGML_PAD(n_top + kpool - 1, 32); + const int64_t n_tail = n_sel_pad - n_top; + + inp->tail_cells = ggml_new_tensor_3d(ctx0, GGML_TYPE_I32, n_tail, n_tps, n_stream); + ggml_set_input(inp->tail_cells); + ggml_set_name(inp->tail_cells, "kpool_tail_cells"); + + inp->tail_valid = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_tail, n_tps, n_stream); + ggml_set_input(inp->tail_valid); + ggml_set_name(inp->tail_valid, "kpool_tail_valid"); + } + + // pooled-key cache. n_new_max is an exact bound on the pools one ubatch can + // complete: a sequence's tokens are a contiguous position run, so a run of L tokens + // closes at most L/kpool + 1 pools. it is FIXED for the whole decode phase so the + // graph shape does not depend on how many pools happened to close this step. + // after a position mutation every cached pooled key is stale, so this one graph has + // to be able to re-emit all of them. a shape change here forces a rebuild, which is + // the point: the wide shape is used for exactly one ubatch and then goes away. + const bool rebuild = mctx_attn->get_kv()->get_kpool_dirty(); + const int64_t n_new_max = rebuild ? n_pools : n_tps/kpool + n_ps; + + inp->n_new_max = (uint32_t) n_new_max; + inp->rebuild = rebuild; + + inp->pool_reps = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_pools, n_stream); + ggml_set_input(inp->pool_reps); + ggml_set_name(inp->pool_reps, "kpool_pool_reps"); + + inp->new_pool_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool*n_new_max, n_stream); + ggml_set_input(inp->new_pool_cells); + ggml_set_name(inp->new_pool_cells, "kpool_new_pool_cells"); + + // I64 because ggml_set_rows takes its row indices as I64 + inp->new_pool_reps = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new_max*n_stream); + ggml_set_input(inp->new_pool_reps); + ggml_set_name(inp->new_pool_reps, "kpool_new_pool_reps"); + } + + 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)); + + // Patch B: keep the transposed mirror in sync (same latent rows, O(1)/token) + if (mctx_cur->get_has_v_mirror()) { + ggml_tensor * vupd = mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il); + cb(vupd, "v_mirror_upd", il); + ggml_build_forward_expand(gf, vupd); + } + } + + 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]); + + // ggml_set_rows writes THROUGH, and sel_mask is shared per ubatch: scatter into a copy + 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 bias: scattering -inf would ERASE a zero granted to the + // tail (cand_mask rejects over-budget picks below). f32: CUDA only does SET_ROWS for f32 + 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; cand_mask is candidates UNION tail + mask_top_k = ggml_add(ctx0, mask_top_k, cand_mask); + + // ggml_flash_attn_ext asserts an f16 mask, and ggml_add would yield src0's f32 + 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); + } + + // load bearing: keeps an empty, future or foreign-sequence cell masked whatever top-k said + 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); + // Patch B: with the mirror, V is the persistent transposed copy — build_attn_mha + // sees v_trans and never builds the per-token cont(transpose(v)) (the 4.83 ms v_cont) + ggml_tensor * v = mctx_cur->get_has_v_mirror() + ? mctx_cur->get_v(ctx0, il) + : 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; +} + + +ggml_tensor * llm_graph_context::build_attn_sparse_gathered( + 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 * slot_valid, + ggml_tensor * tail_cells, + ggml_tensor * tail_valid, + ggml_tensor * cand_mask, + uint32_t kpool, + float kq_scale, + int il) const { + // these nodes are added in the same order as build_attn / build_attn_sparse + 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, INCLUDING the V-mirror update: a gathered decode does not read + // the mirror, but the next dense graph (prefill, or env-off) will + { + const auto & k_idxs = inp->get_k_idxs(); + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + + if (mctx_cur->get_has_v_mirror()) { + ggml_tensor * vupd = mctx_cur->cpy_v(ctx0, k_cur, inp->get_v_idxs(), il); + cb(vupd, "v_mirror_upd", il); + ggml_build_forward_expand(gf, vupd); + } + } + + // decode-1, single stream only (the gate in the model guarantees this) + GGML_ASSERT(top_k->ne[1] == 1 && top_k->ne[2] == 1); + GGML_ASSERT(slot_valid->ne[0] == top_k->ne[0]); + GGML_ASSERT(ggml_are_same_shape(tail_cells, tail_valid)); + + // explicit active width — never derived from a mask that no longer exists + const int64_t n_kv = mctx_cur->get_n_kv(); + const int64_t n_sel = top_k->ne[0] + tail_cells->ne[0]; + + // printed at graph CONSTRUCTION (codex 063 s6) — reservation builds decode-1 + // graphs too, so this proves "gathered graph built", never "evaluated". + // Execution proof is a cb hit on gathered_k_rows/kqv_out plus the oracle. + static bool logged_graph = false; + if (!logged_graph) { + fprintf(stderr, "GATHERED_DSA GRAPH n_sel_pad=%lld n_finite_max=%lld n_kv=%lld\n", + (long long) n_sel, (long long) (top_k->ne[0] + kpool - 1), + (long long) n_kv); + logged_graph = true; + } + + // ids: selected cells then the always-selected tail. all entries are in [0, n_kv); + // invalid slots name cell 0 and die to the -inf slot mask below, BY SLOT. + ggml_tensor * ids = ggml_concat(ctx0, + ggml_reshape_1d(ctx0, top_k, top_k->ne[0]), + ggml_reshape_1d(ctx0, tail_cells, tail_cells->ne[0]), 0); + cb(ids, "gathered_ids", il); + + // E0b MASKLESS (codex 066, grokk 067, proven by E0a): the slot certificate + // subsumes cand+causal on this lane. Fail closed on the lane premises — + // outside them the dense masks can encode facts the certificate does not. + GGML_ASSERT(cparams.causal_attn && "maskless gathered decode requires causal attention"); + GGML_ASSERT(!hparams.use_alibi && "maskless gathered decode: no ALiBi"); + GGML_ASSERT(hparams.swa_type == LLAMA_SWA_TYPE_NONE && "maskless gathered decode: no SWA"); + GGML_ASSERT(!ubatch.is_pos_2d() && "maskless gathered decode: scalar positions only"); + + ggml_tensor * mask_g = ggml_concat(ctx0, + ggml_reshape_1d(ctx0, slot_valid, slot_valid->ne[0]), + ggml_reshape_1d(ctx0, tail_valid, tail_valid->ne[0]), 0); + mask_g = ggml_reshape_4d(ctx0, mask_g, n_sel, 1, 1, 1); + cb(mask_g, "gathered_mask", il); + + + // K rows: [512, 1, n_kv, 1] -> dim-1 view -> gather -> cast -> [512, 1, n_sel, 1]. + // the final reshape is load-bearing (grokk 057 s3): build_attn_mha permutes (0,2,1,3) + // and packed KQ keys on post-permute k->ne[2]==1 (n_head_kv). + ggml_tensor * kfull = mctx_cur->get_k(ctx0, il); + GGML_ASSERT(kfull->ne[1] == 1 && "gathered DSA is MQA-only (n_head_kv == 1)"); + ggml_tensor * krows = ggml_view_3d(ctx0, kfull, + kfull->ne[0], kfull->ne[2], kfull->ne[3], + kfull->nb[2], kfull->nb[3], 0); + + // name the actual compute owners (059 P0): the gather and the cast are what the + // per-node cut must price; the final reshape is a free view + ggml_tensor * kg32 = ggml_get_rows(ctx0, krows, ids); // F32 [512, n_sel, 1] + cb(kg32, "gathered_k_rows", il); + ggml_tensor * kg16 = ggml_cast(ctx0, kg32, kfull->type); + cb(kg16, "gathered_k_cast", il); + ggml_tensor * kg = ggml_reshape_4d(ctx0, kg16, kfull->ne[0], 1, n_sel, 1); + cb(kg, "gathered_k", il); + + // V is the same compact latent; v_trans false -> compact v_cont inside mha + ggml_tensor * cur = build_attn_mha(q_cur, kg, kg, kq_b, mask_g, 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..e5afc17f945 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -32,6 +32,9 @@ 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 forward declared only +class llm_graph_input_kpool; + // certain models (typically multi-modal) can produce different types of graphs enum llm_graph_type { LLM_GRAPH_TYPE_DEFAULT, @@ -341,8 +344,9 @@ class llm_graph_input_attn_kv : public llm_graph_input_i { ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch] ggml_tensor * self_v_idxs = nullptr; // I64 [n_batch] or [n_batch*n_embd_v_gqa] - ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream]; null on maskless gathered graphs (E0c) ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] + int64_t b_n_kv = -1; // E0c: width witness when the mask is absent // note: assumes v_rot^2 == I ggml_tensor * self_k_rot = nullptr; @@ -379,12 +383,15 @@ class llm_graph_input_attn_k : public llm_graph_input_i { bool can_reuse_impl(const llm_graph_params & params); ggml_tensor * get_k_idxs() const { return self_k_idxs; } + ggml_tensor * get_v_idxs() const { return self_v_idxs; } ggml_tensor * get_kq_mask() const { return self_kq_mask_cnv; } ggml_tensor * self_k_idxs = nullptr; // I64 [n_batch] + ggml_tensor * self_v_idxs = nullptr; // I64 expanded [n_batch*w] — Patch B mirror only ggml_tensor * self_kq_mask = nullptr; // F32/F16 [n_kv, n_batch/n_stream, 1, n_stream] + int64_t b_n_kv = -1; // E0c: width witness when the mask is absent ggml_tensor * self_kq_mask_cnv = nullptr; // [n_kv, n_batch/n_stream, 1, n_stream] const llama_hparams hparams; @@ -1340,10 +1347,62 @@ struct llm_graph_context { // llm_graph_input_mem_hybrid * build_inp_mem_hybrid() const; - llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k() const; + llm_graph_input_mem_hybrid_k * build_inp_mem_hybrid_k(bool maskless = false) const; llm_graph_input_mem_hybrid_iswa * build_inp_mem_hybrid_iswa() const; + // one pooling map per ubatch (see llama-kv-cache-kpool.h); `scoring` false gives only k_idxs + llm_graph_input_kpool * build_inp_kpool( + const llama_memory_hybrid_context * mctx_cur, + ggml_tensor * kq_mask, + bool scoring, + bool gathered = false) const; + + // build_attn, but masking with `top_k` over `sel_mask`; `cand_mask` drops over-budget picks + 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; + + // D0 gathered DSA (codex 056 / grokk 057): same contract as build_attn_sparse, but + // attention runs over the <= n_select+kpool-1 gathered rows instead of a dense masked + // n_kv. decode-1 / n_stream==1 only. slot_valid carries the expanded pool_bias of the + // selected pools (-INFINITY kills padded/duplicate picks BY SLOT; a gather, unlike the + // dense scatter, is not idempotent). K rows are gathered from the latent cache + // (F32 return, cast back), and V is the same compact tensor: v_trans is false, the + // per-token v_cont is ~2 MiB, and the V-mirror is not read (it IS still updated). + ggml_tensor * build_attn_sparse_gathered( + 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, 1, 1] physical cells + ggml_tensor * slot_valid, // F32 [n_select, 1, 1] 0/-inf per SLOT + ggml_tensor * tail_cells, // I32 [n_tail_slots, 1, 1] (real tail + PAD32 dead slots) + ggml_tensor * tail_valid, // F32 [n_tail_slots, 1, 1] + ggml_tensor * cand_mask, // F16 [n_kv, 1, 1, 1] + uint32_t kpool, + float kq_scale, + int il) const; + // // pooling // diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 1411692a890..dfa3aa83dde 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -261,6 +261,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..5d816955f21 --- /dev/null +++ b/src/llama-kv-cache-kpool.cpp @@ -0,0 +1,656 @@ +#include "llama-kv-cache-kpool.h" + +#include "llama-batch.h" +#include "llama-kv-cache.h" +#include "llama-memory-hybrid.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"); + + return std::min(n_pools, indexer_top_k/kpool); +} + +// sel_mask and cand_mask hold only 0.0f and -INFINITY, so f16 is exact here +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)); +} + +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; + // the 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, + ggml_tensor * tail_cells, + ggml_tensor * tail_valid, + int64_t n_kv_arg, + int64_t n_ns_arg, + ggml_tensor * pool_reps, + ggml_tensor * new_pool_cells, + ggml_tensor * new_pool_reps, + const uint32_t * strm_of, + int64_t kv_size, + bool rebuild, + const llama_ubatch * ubatch, + uint32_t kpool) { + GGML_ASSERT(kv != nullptr); + GGML_ASSERT(kpool > 0); + + GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(pool_bias ->buffer)); + GGML_ASSERT((sel_mask == nullptr) == (cand_mask == nullptr)); + if (sel_mask) { + GGML_ASSERT(ggml_backend_buffer_is_host(sel_mask ->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(cand_mask ->buffer)); + } + + GGML_ASSERT(pool_cells->type == GGML_TYPE_I32); + GGML_ASSERT(pool_bias ->type == GGML_TYPE_F32); + if (sel_mask) { + 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"); + GGML_ASSERT(ggml_is_contiguous(sel_mask)); + GGML_ASSERT(ggml_is_contiguous(cand_mask)); + GGML_ASSERT(sel_mask->ne[0] == n_kv_arg && "explicit n_kv must match the masks"); + } + GGML_ASSERT(ggml_is_contiguous(pool_cells)); + GGML_ASSERT(ggml_is_contiguous(pool_bias)); + + // E0c: n_kv is an explicit argument (grokk 067/070) — never inferred from + // a mask that may not exist + const int64_t n_kv = n_kv_arg; + const int64_t n_ns = sel_mask ? sel_mask->ne[3] : n_ns_arg; + 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 + 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; + const int64_t n_pools = pool_cells->ne[0]/r; + + 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); + if (sel_mask) { + 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; + // maskless gathered decode has no query-pad tensor; n_tps==1 on that lane + const int64_t n_padq = sel_mask ? sel_mask->ne[1] : 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 shared cell has nowhere to put its second pool + 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); + } + + // pooled-key cache inputs travel together or not at all + const bool kcache = pool_reps != nullptr; + + GGML_ASSERT((new_pool_cells != nullptr) == kcache); + GGML_ASSERT((new_pool_reps != nullptr) == kcache); + + int64_t n_new_max = 0; + + if (kcache) { + GGML_ASSERT(ggml_backend_buffer_is_host(pool_reps ->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_cells->buffer)); + GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_reps ->buffer)); + + GGML_ASSERT(pool_reps ->type == GGML_TYPE_I32); + GGML_ASSERT(new_pool_cells->type == GGML_TYPE_I32); + GGML_ASSERT(new_pool_reps ->type == GGML_TYPE_I64); + + GGML_ASSERT(ggml_is_contiguous(pool_reps)); + GGML_ASSERT(ggml_is_contiguous(new_pool_cells)); + GGML_ASSERT(ggml_is_contiguous(new_pool_reps)); + + GGML_ASSERT(pool_reps->ne[0] == n_pools && pool_reps->ne[1] == n_ns); + GGML_ASSERT(new_pool_cells->ne[0] % r == 0 && new_pool_cells->ne[1] == n_ns); + + n_new_max = new_pool_cells->ne[0]/r; + + GGML_ASSERT(new_pool_reps->ne[0] == n_new_max*n_ns); + GGML_ASSERT(strm_of != nullptr && kv_size > 0); + } + + int32_t * dst_pool_reps = kcache ? (int32_t *) pool_reps ->data : nullptr; + int32_t * dst_new_cells = kcache ? (int32_t *) new_pool_cells->data : nullptr; + int64_t * dst_new_reps = kcache ? (int64_t *) new_pool_reps ->data : nullptr; + + 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 = sel_mask ? (char *) sel_mask ->data : nullptr; + char * dst_cand_mask = cand_mask ? (char *) cand_mask ->data : nullptr; + + const bool mask_f16 = sel_mask ? (sel_mask->type == GGML_TYPE_F16) : true; + const size_t mask_ts = sel_mask ? ggml_type_size(sel_mask->type) : 0; + + // -1 marks a cell with no usable pool; host side only, never copied into cell_pool + std::vector pool_of(n_kv); + std::vector filled(n_pools); + std::vector pos_at; + + std::vector run_off(n_ps); + std::vector run_len(n_ps); + + 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 ? dst_sel_mask + s*(n_padq*n_kv)*mask_ts : nullptr; + char * cur_cand_mask = dst_cand_mask ? dst_cand_mask + s*(n_padq*n_kv)*mask_ts : nullptr; + float * cur_pool_bias = dst_pool_bias + s*(n_tps*n_pools); + + 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); + + int32_t * cur_pool_reps = kcache ? dst_pool_reps + s*n_pools : nullptr; + int32_t * cur_new_cells = kcache ? dst_new_cells + s*(r*n_new_max) : nullptr; + int64_t * cur_new_reps = kcache ? dst_new_reps + s*n_new_max : nullptr; + + // count of real entries emitted for this stream; the rest is padding + int64_t n_new = 0; + + // members of any complete pool in this stream, used to pad the fixed-size write. + // recomputing a complete pool is idempotent, so a repeat is always safe. + const int32_t * any_rep_src = nullptr; + + if (kcache) { + // a pool with no rep gathers row 0. that row's pooled third may hold another + // pool's key, but such a pool is always -INFINITY in pool_bias, so the value is + // discarded before it can score. + std::fill(cur_pool_reps, cur_pool_reps + n_pools, 0); + std::fill(cur_new_cells, cur_new_cells + r*n_new_max, 0); + } + + // the token loop writes rows < n_tps in full; only the padding rows need clearing + if (!cur_sel_mask) { + // maskless gathered graph: no dense mask to pad (E0c) + } else 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] one packed run per sequence, sized on the pool range it holds. + // NOT one full-width table per sequence: the indexer scores every slot against + // every query, so that multiplies the score tensor by n_seq_max. + // llama_memory_seq_cp can ask for more slots than exist; then a sequence keeps its + // newest pools, the same cut a large hole already forces. + { + 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); + + 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); + } + + // anchoring at the absolute p/kpool follows vLLM and SGLang, not HF + // (valid_keys.argmax(-1)): it is the only anchor that keeps a pool's identity + // stable from the prefill that built it to the decodes that read it. + 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]++; + } + + // pool_valid = grouped_valid_keys.all(-1): the compressor consumes all r keys + for (int64_t j = 0; j < n_kv; ++j) { + // != rather than <: two cells claiming one position overwrite each other + 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]; + } + } + + if (kcache) { + // the pooled key of a complete pool lives in the row of its LAST member, + // the cell holding pos % r == r-1. that slot is only meaningful once the + // pool is complete: for a partial pool it is still 0 from the fill above, + // and cell 0 is a real cell whose own pooled key we must not overwrite. + for (int64_t p = 0; p < n_run; ++p) { + if (filled[p] == (int32_t) r) { + cur_pool_reps[run_off[ps] + p] = part_pool_cells[p*r + (r - 1)]; + + if (any_rep_src == nullptr) { + any_rep_src = part_pool_cells + p*r; + } + } + } + + // recompute exactly the complete pools this ubatch wrote into. touched[] is + // over the run, so the cost is O(tokens), not O(n_kv). + std::vector touched(n_run, 0); + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + + if (ubatch->seq_id[i][0] != seq_of_pool) { + continue; + } + + const int64_t bo = ubatch->pos[i]/r - b_base; + + if (bo >= 0 && bo < n_run) { + touched[bo] = 1; + } + } + + for (int64_t p = 0; p < n_run; ++p) { + // in rebuild mode every cached pooled key is stale, so re-emit all of + // them, not just the pools this ubatch closed + if ((!touched[p] && !rebuild) || filled[p] != (int32_t) r) { + continue; + } + + // n_new_max = n_tps/kpool + n_ps bounds this while a sequence's tokens in + // one ubatch are a contiguous position run, which llama-batch.cpp + // enforces. dropping a completed pool would silently serve a stale key, + // so fail loudly instead of clamping. + GGML_ASSERT(n_new < n_new_max && "k-pool: more pools completed than the fixed bound"); + + std::copy(part_pool_cells + p*r, part_pool_cells + (p + 1)*r, + cur_new_cells + n_new*r); + + cur_new_reps[n_new] = (int64_t) strm_of[s]*kv_size + part_pool_cells[p*r + (r - 1)]; + + n_new++; + } + } + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + + 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++; + + // index_kpool_always_select_tail, which lands selection 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 the + // query straddles is dropped whole + 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 ? cur_sel_mask + ii*n_kv*mask_ts : nullptr; + char * cur_cand = cur_cand_mask ? cur_cand_mask + ii*n_kv*mask_ts : nullptr; + + if (!sel_mask) { + // masks absent on the maskless gathered graph (E0c) + } else 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; + } + } + + // D0: emit the tail cells [tail_start, q] by slot; unused slots stay + // cell 0 / -INFINITY from the pre-fill below the loop + if (tail_cells) { + // width includes the PAD32 dead slots; the real tail occupies + // only [0, kpool-1). Invariant: every slot t >= kpool-1 stays + // cell 0 / -INFINITY from the pre-fill (grokk 064 s2). + const int64_t n_ts = tail_cells->ne[0]; + int32_t * cur_tcell = (int32_t *) tail_cells->data + (s*n_tps + ii)*n_ts; + float * cur_tvalid = (float *) tail_valid->data + (s*n_tps + ii)*n_ts; + + for (int64_t t = 0; t < n_ts; ++t) { + cur_tcell [t] = 0; + cur_tvalid[t] = -INFINITY; + } + for (int64_t j = 0; j < n_kv; ++j) { + const llama_pos p = pos_at[j]; + if (p >= tail_start && p <= q) { + const int64_t t = p - tail_start; + GGML_ASSERT(t < r - 1); + cur_tcell [t] = (int32_t) j; + cur_tvalid[t] = 0.0f; + } + } + } + + // the query's own sequence run only; every other slot keeps the -INFINITY + // of the fill above, 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; + } + } + } + + // exactly one partition per row, or a query reads another sequence's pools + GGML_ASSERT(n_done == n_tps && "every query must belong to a sequence of the ubatch"); + + if (kcache) { + // the write has a fixed row count, so the unused slots must name a destination + // that is safe to overwrite. two cases, both provably harmless: + // - some complete pool exists: repeat it. recomputing a complete pool yields + // the value already there, so the duplicate write is a no-op in effect. + // - none exists: no cell in this stream is the last member of a complete pool, + // so no pooled third is read (every pool is -INFINITY in pool_bias). cell 0 + // is then free. + for (int64_t p = n_new; p < n_new_max; ++p) { + if (any_rep_src) { + std::copy(any_rep_src, any_rep_src + r, cur_new_cells + p*r); + cur_new_reps[p] = (int64_t) strm_of[s]*kv_size + any_rep_src[r - 1]; + } else { + // cells already 0 from the fill; pool r copies of cell 0 into cell 0 + cur_new_reps[p] = (int64_t) strm_of[s]*kv_size; + } + } + } + } +} + + +bool llm_graph_input_kpool::can_reuse(const llm_graph_params & params) { + const auto * mctx_hyb = static_cast(params.mctx); + + const auto * attn = mctx_hyb->get_attn(); + const auto * idx = mctx_hyb->get_idx(); + + if (attn == nullptr || idx == nullptr) { + return false; + } + + // rebind so a reused graph's set_input reads the new contexts + mctx_attn = attn; + mctx_idx = idx; + + bool res = true; + + res &= k_idxs->ne[0] == (int64_t) params.ubatch.n_tokens; + res &= b_n_kv_idx == (int64_t) idx->get_n_kv(); + + if (pool_cells == nullptr) { + // non-scoring graph: only the key/gate store exists + return res; + } + + const int64_t n_kv = attn->get_n_kv(); + const int64_t n_stream = params.cparams.kv_unified ? 1 : params.ubatch.n_seqs_unq; + + if (n_stream <= 0 || params.ubatch.n_tokens % n_stream != 0 || + (int64_t) params.ubatch.n_seqs_unq % n_stream != 0) { + return false; + } + + const int64_t n_tps = params.ubatch.n_tokens/n_stream; + const int64_t n_ps = (int64_t) params.ubatch.n_seqs_unq/n_stream; + + const int64_t n_pools = llama_kpool_n_pools((uint32_t) n_kv, kpool, (uint32_t) n_ps); + + res &= b_n_kv_attn == n_kv; + + // exhaustive optional-pair contract (codex 077 s4 / grokk 078 s1): + // presence, pairedness and exact shape for every optional tensor + res &= (sel_mask != nullptr) == (cand_mask != nullptr); + if (sel_mask != nullptr) { + res &= sel_mask->ne[0] == n_kv && sel_mask->ne[1] == n_tps && sel_mask->ne[3] == n_stream; + res &= ggml_are_same_shape(sel_mask, cand_mask); + } + res &= pool_cells->ne[0] == (int64_t) kpool*n_pools && pool_cells->ne[1] == n_stream; + res &= pool_bias->ne[0] == n_pools && pool_bias->ne[1] == n_tps && pool_bias->ne[2] == n_stream; + if (pool_reps != nullptr) { + res &= pool_reps->ne[0] == n_pools && pool_reps->ne[1] == n_stream; + res &= new_pool_cells != nullptr && new_pool_reps != nullptr; + res &= new_pool_cells->ne[0] == (int64_t) kpool*(int64_t) n_new_max && + new_pool_cells->ne[1] == n_stream; + res &= new_pool_reps->ne[0] == (int64_t) n_new_max*n_stream; + } else { + res &= new_pool_cells == nullptr && new_pool_reps == nullptr; + } + res &= (tail_cells != nullptr) == (tail_valid != nullptr); + + // policy: rebuild mode and the fixed pooled-key emission bound must match + // what a fresh build would choose (codex 068 s3) + const bool want_rebuild = mctx_attn->get_kv()->get_kpool_dirty(); + const int64_t want_new_max = want_rebuild ? n_pools : n_tps/(int64_t) kpool + n_ps; + + res &= rebuild == want_rebuild; + res &= (int64_t) n_new_max == want_new_max; + + if (tail_cells != nullptr) { + res &= ggml_are_same_shape(tail_cells, tail_valid); + // gathered graph: recompute the PAD32 width — never cache 2080 (grokk 070 s3) + const int64_t n_top = (int64_t) kpool * + llama_kpool_select_k((uint32_t) n_pools, hparams_indexer_top_k, kpool); + const int64_t n_sel_pad = GGML_PAD(n_top + kpool - 1, 32); + + res &= tail_cells->ne[0] == n_sel_pad - n_top; + res &= n_tps == 1 && n_stream == 1; + } + + return res; +} + +void llm_graph_input_kpool::set_input(const llama_ubatch * ubatch) { + // unconditional: the key and gate STORE runs on the dense path too. gating it the + // way the scoring is gated would leave every cell below n_select with no indexer + // state, and the first ubatch to cross n_select would pool cells never written + mctx_idx->set_input_k_idxs(k_idxs, ubatch); + + if (pool_cells == nullptr) { + return; + } + + // the pooled key is written into the INDEXER cache, whose slot layout the attention + // cache defines, so the stream map and cell count come from the indexer side + std::vector strm_of; + + if (pool_reps) { + strm_of.resize(mctx_idx->get_n_stream()); + + for (uint32_t s = 0; s < strm_of.size(); ++s) { + strm_of[s] = mctx_idx->get_strm(s); + } + } + + llama_kv_cache_set_input_kpool( + mctx_attn->get_kv(), + /* cell_pool */ nullptr, pool_cells, /* bias */ nullptr, pool_bias, + sel_mask, cand_mask, + tail_cells, tail_valid, + (int64_t) mctx_attn->get_n_kv(), + (int64_t) (tail_cells ? tail_cells->ne[2] : 1), + pool_reps, new_pool_cells, new_pool_reps, + strm_of.empty() ? nullptr : strm_of.data(), + pool_reps ? (int64_t) mctx_idx->get_kv()->get_size() : 0, + rebuild, + ubatch, kpool); + + // every pool has just been re-emitted, so the cache is consistent again. cleared here + // rather than in build_inp_kpool because a graph that is built but not evaluated must + // not clear it. + if (rebuild) { + mctx_attn->get_kv()->clear_kpool_dirty(); + } +} diff --git a/src/llama-kv-cache-kpool.h b/src/llama-kv-cache-kpool.h new file mode 100644 index 00000000000..224f4881a72 --- /dev/null +++ b/src/llama-kv-cache-kpool.h @@ -0,0 +1,133 @@ +#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. the position -> cell map is built host side because +// find_slot's cell order is arbitrary. no input may hold a negative index: ggml_set_rows +// asserts i1 >= 0 and ggml_get_rows has no sentinel, so unusable entries are clamped into +// range and neutralised by the additive masks instead. + +// pool slots for `n_kv` cells shared by `n_seqs` sequences: n_kv/kpool, exact only while +// the sequences' cells are disjoint, plus 2 per sequence for rebasing. +uint32_t llama_kpool_n_pools(uint32_t n_kv, uint32_t kpool, uint32_t n_seqs = 1); + +// select_k of modular_glm5_next.py, Glm5NextTextIndexer.forward. must run over POOLS, not +// cells: relu ties span pool boundaries, so a cell-level cut takes partial pools. +uint32_t llama_kpool_select_k(uint32_t n_pools, uint32_t indexer_top_k, uint32_t kpool); + +// `kv` must be the ATTENTION (MLA) cache; the indexer cache shares its slot layout. +// cell_pool I32 [n_kv, n_stream] per-cell view, optional, unused here +// pool_cells I32 [kpool*n_pools, n_stream] pool member -> cell, 0 if not resident +// bias F32 [n_kv, n_tps, n_stream] per-cell view, optional, unused here +// pool_bias F32 [n_pools, n_tps, n_stream] pool_valid & pool_visible, -INFINITY +// outside the query's own sequence run; computed, not gathered from `bias` at the +// last member, which an incomplete pool lacks and would inherit cell 0's validity +// sel_mask F16/F32 [n_kv, n_batch, 1, n_stream] 0.0f on the always-selected tail only +// cand_mask F16/F32 [n_kv, n_batch, 1, n_stream] max(bias, sel_mask); bounds the top-k +// spills that a partial seq_rm would otherwise let escape the candidate set +// pool_reps / new_pool_cells / new_pool_reps are optional (nullptr when the pooled-key +// cache is off). an entry is emitted only for a pool with filled == kpool: an incomplete +// pool's last-member slot is 0, and cell 0 is a legitimate cell, so writing it would +// clobber another pool's cached key. +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, + ggml_tensor * tail_cells, // optional, D0 gathered decode + ggml_tensor * tail_valid, // optional, with tail_cells + int64_t n_kv_arg, // explicit active width (E0c) + int64_t n_ns_arg, // explicit stream count when masks are null + ggml_tensor * pool_reps, + ggml_tensor * new_pool_cells, + ggml_tensor * new_pool_reps, + // strm_of[s] is the physical stream behind view s and kv_size the per-stream cell + // count, so a global row is strm_of[s]*kv_size + cell. only read when pool_reps is + // set. the indexer cache shares the attention cache's slot layout, so one cell + // index addresses both. + const uint32_t * strm_of, + int64_t kv_size, + // re-emit every complete pool, not only the ones this ubatch closed. set after a + // position mutation; the graph is built with n_new_max == n_pools to hold them. + bool rebuild, + const llama_ubatch * ubatch, + uint32_t kpool); + +// One pooling map per ubatch; rebuilding it per indexer layer costs O(n_kv * n_tokens) +// host writes and dominates prefill. sharing is valid only while every indexer layer sees +// the same candidate set - true for glm5next (indexer_types all "full"), not for windowed. +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; + + // R0 (codex 068 / grokk 070): exact reuse contract. Rebinds the two memory + // contexts and compares every value that determines tensor shape or graph + // policy against the new params. Without this override the inherited + // can_reuse() == false vetoed reuse of EVERY served GLM graph. + bool can_reuse(const llm_graph_params & params) 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] + + // pooled-key cache. the pooled value of a pool lives in the third head of the row of + // the cell holding its LAST member (pos % kpool == kpool-1), so it is a pure function + // of cell content: independent of sequence and of pool ordinal, which is what lets a + // seq_cp share it and a rebase leave it alone. + ggml_tensor * pool_reps = nullptr; // I32 [n_pools, n_stream] stream-local rep cell + ggml_tensor * new_pool_cells = nullptr; // I32 [kpool*n_new_max, n_stream] members to (re)pool + ggml_tensor * new_pool_reps = nullptr; // I64 [n_new_max*n_stream] GLOBAL dest row + + // n_new_max is fixed for the whole decode phase on purpose. making the shape depend on + // how many pools happened to close this step would flip the graph topology every kpool + // tokens and defeat graph reuse / force CUDA-graph recapture. + uint32_t n_new_max = 0; + + // this graph was built to re-emit every pool after a position mutation, not just the + // ones this ubatch closed. decided at build time because it sets n_new_max above. + bool rebuild = false; + + // exact, since pool_bias only holds 0.0f or -INFINITY. nullptr if 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] + + // D0 gathered DSA (codex 056 / grokk 057): the always-selected incomplete-pool tail as + // explicit cells so a gathered graph can concat them to top_k. slot t holds the cell at + // pos tail_start+t; unused slots are cell 0 with tail_valid -INFINITY (get_rows has no + // sentinel). tail length is (q+1)%kpool in [0, kpool-1]. + // PAD32: physical width is n_sel_pad - n_top (>= kpool-1); only the finite + // tail prefix [0, kpool-1) is ever written with live cells. + ggml_tensor * tail_cells = nullptr; // I32 [n_tail_slots, n_tps, n_stream] + ggml_tensor * tail_valid = nullptr; // F32 [n_tail_slots, n_tps, n_stream] + + // build-time identity for can_reuse (explicit metadata, not inferred from + // mask tensors, so the contract survives E0c making the masks optional) + int64_t b_n_kv_attn = -1; + int64_t b_n_kv_idx = -1; + uint32_t hparams_indexer_top_k = 0; + + 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 8fafcd15304..ef0c862ef6f 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -80,7 +80,8 @@ llama_kv_cache::llama_kv_cache( const layer_filter_cb & filter, const layer_reuse_cb & reuse, const layer_share_cb & share, - const char * name_tag) : + const char * name_tag, + bool v_mirror_opt) : model(model), hparams(hparams), v_trans(v_trans), n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type), other(static_cast(mem_other)), @@ -163,6 +164,21 @@ llama_kv_cache::llama_kv_cache( const bool is_mla = hparams.is_mla(); + // Patch B opt-in. Scope guard (codex 051 s3): mirroring the FULL K slab as V + // is only proven for nope-only GLM5Next, where K carries no rope-key half. + // is_mla alone does not establish that contract; fail closed elsewhere so a + // globally inherited env cannot enable the mirror on an unproven MLA arch. + const bool mirror_supported = model.arch == LLM_ARCH_GLM5NEXT && hparams.n_rot(0) == 0; + + v_mirror = v_mirror_opt && mirror_supported && is_mla && v_trans; + if (v_mirror_opt && !v_mirror) { + fprintf(stderr, "MLA_V_MIRROR requested but DISABLED (arch=%d n_rot=%u is_mla=%d v_trans=%d)\n", + (int) model.arch, hparams.n_rot(0), (int) is_mla, (int) v_trans); + } + if (v_mirror) { + GGML_ASSERT((type_k == GGML_TYPE_F16 || type_k == GGML_TYPE_F32) && "V mirror v0 requires F16/F32 type_k"); + } + for (uint32_t il = 0; il < n_layer; il++) { if (!hparams.has_kv(il)) { LLAMA_LOG_DEBUG("%s: layer %3d: does not have KV cache\n", __func__, il); @@ -229,10 +245,15 @@ llama_kv_cache::llama_kv_cache( } const bool has_k = true; - const bool has_v = !is_mla; + const bool has_v = !is_mla || v_mirror; // Patch B: mirror allocates V despite MLA + // grokk 024/020: the mirror is a layout of K -> width n_embd_k_gqa (512), type_k + const uint32_t n_embd_v_alloc = v_mirror ? n_embd_k_gqa : n_embd_v_gqa; + if (v_mirror && n_embd_k_gqa > mirror_w_max) { + mirror_w_max = n_embd_k_gqa; + } ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr; - ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr; + ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, v_mirror ? type_k : type_v, n_embd_v_alloc, kv_size, n_stream) : nullptr; has_k && ggml_format_name(k, "cache_%sk_l%d", name_tag, il); has_v && ggml_format_name(v, "cache_%sv_l%d", name_tag, il); @@ -242,7 +263,7 @@ llama_kv_cache::llama_kv_cache( for (uint32_t s = 0; s < n_stream; ++s) { k_stream.push_back(has_k ? ggml_view_2d(ctx, k, n_embd_k_gqa, kv_size, k->nb[1], s*k->nb[2]) : nullptr); - v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr); + v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_alloc, kv_size, v->nb[1], s*v->nb[2]) : nullptr); } map_layer_ids[il] = layers.size(); @@ -250,6 +271,15 @@ llama_kv_cache::llama_kv_cache( layers.push_back({ il, k, v, k_stream, v_stream, }); } + if (v_mirror) { + for (const auto & layer : layers) { + if (layer.v) { + GGML_ASSERT(hparams.n_embd_k_gqa(layer.il) == mirror_w_max && "V mirror v0 requires uniform mirrored K width"); + } + } + fprintf(stderr, "MLA_V_MIRROR ACTIVE width=%u type=%s\n", mirror_w_max, ggml_type_name(type_k)); + } + if (reuse) { LLAMA_LOG_DEBUG("%s: reusing layers:\n", __func__); @@ -849,6 +879,7 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co ggml_backend_tensor_copy(layer.k_stream[ssrc], layer.k_stream[sdst]); if (layer.v_stream[ssrc]) { + // mirror copied with the stream -> stays clean (codex 034 s4) ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]); } } @@ -893,9 +924,66 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co } } + if (v_mirror && mirror_dirty) { + LLAMA_LOG_DEBUG("%s: rebuilding derived V mirror from K\n", __func__); + + ggml_backend_sched_reset(sched); + + auto * res = lctx->get_gf_res_reserve(); + + res->reset(); + + auto * gf = build_graph_mirror(res, lctx); + if (!ggml_backend_sched_alloc_graph(sched, gf)) { + LLAMA_LOG_ERROR("%s: failed to allocate compute graph for V-mirror rebuild\n", __func__); + return updated; + } + + res->set_inputs(nullptr); + + if (lctx->graph_compute(gf, false) != GGML_STATUS_SUCCESS) { + LLAMA_LOG_ERROR("%s: failed to compute V-mirror rebuild\n", __func__); + return updated; + } + + mirror_dirty = false; + fprintf(stderr, "MIRROR_REBUILD SUCCESS\n"); + updated = true; + } + return updated; } +ggml_cgraph * llama_kv_cache::build_graph_mirror(llm_graph_result * res, llama_context * lctx) const { + GGML_ASSERT(!other); + GGML_UNUSED(lctx); + + auto * ctx = res->get_ctx(); + auto * gf = res->get_gf(); + + const uint64_t kv_size = get_size(); + + for (const auto & layer : layers) { + if (!layer.v) { + continue; + } + + const uint64_t w = hparams.n_embd_k_gqa(layer.il); + + for (uint32_t st = 0; st < n_stream; ++st) { + ggml_tensor * ksrc = ggml_view_2d(ctx, layer.k, w, kv_size, layer.k->nb[1], st*layer.k->nb[2]); + // mirror layout: row e holds element e of every cell (transposed K) + ggml_tensor * vdst = ggml_view_2d(ctx, layer.v, kv_size, w, + ggml_row_size(layer.v->type, kv_size), + st*ggml_row_size(layer.v->type, kv_size*w)); + + ggml_build_forward_expand(gf, ggml_cpy(ctx, ggml_transpose(ctx, ksrc), vdst)); + } + } + + return gf; +} + llama_kv_cache::slot_info llama_kv_cache::find_slot(const llama_ubatch & ubatch, bool cont) const { if (debug > 0) { @@ -1217,6 +1305,18 @@ bool llama_kv_cache::get_has_shift() const { return result; } +void llama_kv_cache::set_kpool_dirty() { + kpool_dirty = true; +} + +bool llama_kv_cache::get_kpool_dirty() const { + return kpool_dirty; +} + +void llama_kv_cache::clear_kpool_dirty() const { + kpool_dirty = false; +} + ggml_type llama_kv_cache::type_k() const { return layers[0].k->type; } @@ -1307,6 +1407,17 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k ggml_row_size(v->type, n_embd_v_gqa*kv_size)*sinfo.s0); } + if (v_mirror) { + // contracts 4.4/7: mirror view is 512-wide (K slab), one KV head + const uint64_t w = hparams.n_embd_k_gqa(il); + return ggml_view_4d(ctx, v, + n_kv, 1, w, ns, + ggml_row_size(v->type, kv_size*w), // head stride > element stride => v_trans in mha + ggml_row_size(v->type, kv_size), + ggml_row_size(v->type, kv_size*w), + ggml_row_size(v->type, kv_size*w)*sinfo.s0); + } + // note: v->nb[1] > v->nb[2] return ggml_view_4d(ctx, v, n_kv, hparams.n_head_kv(il), hparams.n_embd_head_v(il), ns, @@ -1351,6 +1462,30 @@ ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggm return ggml_set_rows(ctx, k, k_cur, k_idxs); } +ggml_tensor * llama_kv_cache::cpy_k_part(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, + int64_t n_embd, int64_t i_off) const { + const int32_t ikv = map_layer_ids.at(il); + + ggml_tensor * k = layers[ikv].k; + + const int64_t n_embd_gqa = k->ne[0]; + const int64_t kv_size = get_size(); + const int64_t n_stream = k->ne[2]; + + GGML_ASSERT(i_off >= 0 && i_off + n_embd <= n_embd_gqa); + GGML_ASSERT(k_cur->ne[0] == n_embd); + + // merge the streams: k_idxs are global, exactly as in cpy_k + ggml_tensor * k2 = ggml_reshape_2d(ctx, k, n_embd_gqa, kv_size*n_stream); + + // a row-slice view of every cell. ggml_set_rows needs contiguous rows in the DEST, + // which ggml_is_contiguous_rows() grants for a view whose ne[0] slice is contiguous. + ggml_tensor * dst = ggml_view_2d(ctx, k2, n_embd, kv_size*n_stream, + k2->nb[1], ggml_row_size(k2->type, i_off)); + + return ggml_set_rows(ctx, dst, k_cur, k_idxs); +} + ggml_tensor * llama_kv_cache::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const { GGML_UNUSED(sinfo); @@ -1425,7 +1560,7 @@ ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama if (!v_trans) { v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens); } else { - v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens*hparams.n_embd_v_gqa_max()); + v_idxs = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, n_tokens*(v_mirror ? mirror_w_max : hparams.n_embd_v_gqa_max())); } ggml_set_input(v_idxs); @@ -1509,7 +1644,7 @@ void llama_kv_cache::set_input_v_idxs(ggml_tensor * dst, const llama_ubatch * ub // note: the V cache is transposed when not using flash attention const int64_t kv_size = get_size(); - const int64_t n_embd_v_gqa = hparams.n_embd_v_gqa_max(); + const int64_t n_embd_v_gqa = v_mirror ? mirror_w_max : hparams.n_embd_v_gqa_max(); for (uint32_t s = 0; s < sinfo.n_stream(); ++s) { const int64_t offs = sinfo.strm[s]*kv_size*n_embd_v_gqa; @@ -2315,7 +2450,7 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); auto * v = layer.v_stream[cr.strm]; - if (!v) { + if (!v || v_mirror) { // contract 4.5: derived mirror never serialized continue; } @@ -2344,7 +2479,7 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); auto * v = layer.v_stream[cr.strm]; - if (!v) { + if (!v || v_mirror) { // contract 4.5: derived mirror never serialized continue; } @@ -2554,6 +2689,10 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 return false; } + if (v_mirror) { + mirror_dirty = true; // contract 4.6: restored K needs a mirror rebuild before decode + } + // For each layer, read the keys for each cell, one row is one cell, read as one contiguous block for (const auto & layer : layers) { const uint32_t il = layer.il; @@ -2601,7 +2740,7 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); auto * v = layer.v_stream[strm]; - if (!v) { + if (!v || v_mirror) { // contract 4.5: mirror not on the wire continue; } @@ -2644,7 +2783,7 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 const uint32_t n_embd_v_gqa = hparams.n_embd_v_gqa(il); auto * v = layer.v_stream[strm]; - if (!v) { + if (!v || v_mirror) { // contract 4.5: mirror not on the wire continue; } @@ -2726,8 +2865,8 @@ llama_kv_cache_context::llama_kv_cache_context( llama_context * lctx, bool do_shift, stream_copy_info sc_info) : status(LLAMA_MEMORY_STATUS_SUCCESS), kv(kv), lctx(lctx), do_shift(do_shift), sc_info(std::move(sc_info)) { - if (!do_shift && this->sc_info.empty()) { - status = LLAMA_MEMORY_STATUS_NO_UPDATE; + if (!do_shift && this->sc_info.empty() && !kv->get_mirror_dirty()) { + status = LLAMA_MEMORY_STATUS_NO_UPDATE; // contract 4.6: a dirty mirror must vote for update } } @@ -2787,10 +2926,34 @@ ggml_type llama_kv_cache_context::type_v() const { return kv->type_v(); } +uint32_t llama_kv_cache_context::get_n_stream() const { + return sinfos[i_cur].s1 - sinfos[i_cur].s0 + 1; +} + +uint32_t llama_kv_cache_context::get_strm(uint32_t s) const { + const auto & sinfo = sinfos[i_cur]; + + GGML_ASSERT(s < sinfo.strm.size()); + + return sinfo.strm[s]; +} + +const llama_kv_cache * llama_kv_cache_context::get_kv() const { + return kv; +} + ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il) const { return kv->get_k(ctx, il, n_kv, sinfos[i_cur]); } +bool llama_kv_cache_context::get_has_v_mirror() const { + return kv->get_has_v_mirror(); +} + +uint32_t llama_kv_cache_context::mirror_width_max() const { + return kv->mirror_width_max(); +} + ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) const { return kv->get_v(ctx, il, n_kv, sinfos[i_cur]); } @@ -2799,6 +2962,11 @@ ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_ return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]); } +ggml_tensor * llama_kv_cache_context::cpy_k_part(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, + int64_t n_embd, int64_t i_off) const { + return kv->cpy_k_part(ctx, k_cur, k_idxs, il, n_embd, i_off); +} + ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const { return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index c4d8699def1..e3826d7657f 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -114,7 +114,8 @@ class llama_kv_cache : public llama_memory_i { const layer_reuse_cb & reuse, const layer_share_cb & share, // a model can hold more than one cache, so the tensor names have to stay unique - const char * name_tag = ""); + const char * name_tag = "", + bool v_mirror_opt = false); ~llama_kv_cache() = default; @@ -160,6 +161,16 @@ class llama_kv_cache : public llama_memory_i { bool get_has_shift() const; + // GLM-5-Next pooled-key cache. a pool groups cells by absolute position, so anything + // that mutates positions in place (seq_add / seq_div) regroups the pools while every + // cached pooled key still looks complete. set here, consumed and cleared by + // llama_kv_cache_set_input_kpool, which then emits every pool instead of only the new + // ones. sticky by design: a flag that never clears degrades to the pre-cache cost, a + // flag that clears too early is silently wrong. + void set_kpool_dirty(); + bool get_kpool_dirty() const; + void clear_kpool_dirty() const; + ggml_type type_k() const; ggml_type type_v() const; @@ -189,10 +200,22 @@ class llama_kv_cache : public llama_memory_i { ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; + // speed-conf Patch B: derived transposed MLA-V mirror (022/024 contracts) + bool get_has_v_mirror() const { return v_mirror; } + bool get_mirror_dirty() const { return mirror_dirty; } + uint32_t mirror_width_max() const { return mirror_w_max; } + // store k_cur and v_cur in the cache based on the provided head location ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const; ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const; + // write n_embd elements at element offset i_off inside each row, leaving the rest of + // the row untouched. k_idxs are GLOBAL rows (stream-merged), as in cpy_k. + // returns the ggml_set_rows result so a later gather can be chained off it and become + // a real graph edge rather than relying on build order. + ggml_tensor * cpy_k_part(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, + int64_t n_embd, int64_t i_off) const; + // // preparation API // @@ -260,6 +283,13 @@ class llama_kv_cache : public llama_memory_i { }; bool v_trans = true; // the value tensor is transposed + bool v_mirror = false; // derived transposed MLA-V mirror (never serialized; contract 4.5) + bool mirror_dirty = false; // set on state restore; rebuilt from K in update() (contract 4.6) + uint32_t mirror_w_max = 0; // max n_embd_k_gqa over mirrored layers (contract 7 idx sizing) + + // see set_kpool_dirty. mutable because the only consumer runs from set_input, which + // holds the cache by const pointer; nothing else observes it. + mutable bool kpool_dirty = false; const uint32_t n_seq_max = 1; const uint32_t n_stream = 1; @@ -329,6 +359,8 @@ class llama_kv_cache : public llama_memory_i { float freq_scale, uint32_t il) const; + ggml_cgraph * build_graph_mirror(llm_graph_result * res, llama_context * lctx) const; + ggml_cgraph * build_graph_shift( llm_graph_result * res, llama_context * lctx) const; @@ -391,12 +423,24 @@ class llama_kv_cache_context : public llama_memory_context_i { uint32_t get_n_kv() const; + // the stream RANGE s1 - s0 + 1 that get_k/get_v use as `ns`, not n_seqs_unq + uint32_t get_n_stream() const; + + // physical stream backing view index s, i.e. sinfo.strm[s]. a global cache row for a + // cell j of that view is get_strm(s)*kv->get_size() + j, the convention set_input_k_idxs + // uses. needed by the k-pool code, which writes cells that are not ubatch tokens. + uint32_t get_strm(uint32_t s) const; + + const llama_kv_cache * get_kv() const; + ggml_type type_k() const; ggml_type type_v() const; // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il) const; + bool get_has_v_mirror() const; + uint32_t mirror_width_max() const; // store k_cur and v_cur in the cache based on the provided head location // note: the heads in k_cur and v_cur should be laid out contiguously in memory @@ -405,6 +449,8 @@ class llama_kv_cache_context : public llama_memory_context_i { // - v_cur [n_embd_head_v, n_head_v, n_tokens] // - v_idxs [n_tokens] or [n_tokens*n_embd_v_gqa] depending if V cache is transposed ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const; + ggml_tensor * cpy_k_part(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, + int64_t n_embd, int64_t i_off) const; ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const; // create destination indices for each head of the current batch for where it would be written in the KV cache diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6..27de7f0d1c3 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,12 @@ 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, + bool attn_v_mirror) : hparams(model.hparams), + hparams_idx(model.hparams), mem_attn(new llama_kv_cache( model, model.hparams, @@ -49,7 +55,9 @@ llama_memory_hybrid::llama_memory_hybrid( [&](int32_t il) { return !hparams.is_recr(il); } : filter_attn, nullptr, - nullptr + nullptr, + "", + attn_v_mirror // Patch B: mem_attn only (contract 4.1); idx cache stays false )), mem_recr(new llama_memory_recurrent( model, @@ -62,7 +70,28 @@ 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 : [&] { + // a *pooling* indexer (indexer_kpool > 0, glm5next only) needs a second head for + // the compressor gate, or the pool cannot be rebuilt once its tokens leave the + // batch. every other arch leaves indexer_kpool 0 and is unchanged. + // + // the third head is the pooled key of the pool this cell ENDS, i.e. it is written + // only for cells at pos % kpool == kpool-1. caching it turns the per-step pooling + // from O(n_kv) into O(pools completed by this ubatch): see build_indexer. + const uint32_t n_head_idx = model.hparams.indexer_kpool > 0 ? 3 : 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 +144,16 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } + // the indexer takes the attention cache's slot layout rather than finding its + // own: allocated separately the two drift apart when the context is rewritten + // between turns, and the top-k indices 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,42 @@ 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) { + // pools group cells by absolute position, so a shift regroups them; every pooled key + // cached in the indexer rows is stale even though each still looks complete. a shift + // that is a whole number of pools regroups nothing, which is the common context-shift + // case, so it is worth not paying for. + if (mem_idx && hparams.indexer_kpool > 0 && shift % (llama_pos) hparams.indexer_kpool != 0) { + mem_attn->set_kpool_dirty(); + } 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) { + // as seq_add, and a divide regroups for any d != 1 + if (mem_idx && hparams.indexer_kpool > 0 && d != 1) { + mem_attn->set_kpool_dirty(); + } 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 +241,19 @@ 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; skipping them here misselects on restore + if (mem_idx) mem_idx->state_write(io, seq_id, flags); } mem_recr->state_write(io, seq_id, flags); } @@ -197,6 +261,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 +274,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 +293,25 @@ 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, but the pending per-cell delta must + // still be cleared or the two caches disagree about whether a shift is outstanding. + // safe because an indexer only exists for LLAMA_ROPE_TYPE_NONE archs, where + // 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 +320,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 +337,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 + 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 +367,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..c9036023259 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 indexer key cache; absent unless filter_idx */ + const layer_filter_cb & filter_idx = nullptr, + ggml_type type_idx = GGML_TYPE_F16, + bool attn_v_mirror = false); ~llama_memory_hybrid() = default; @@ -82,12 +86,16 @@ 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; + 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 +118,8 @@ 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, + slot_info_vec_t sinfos_idx = {}); ~llama_memory_hybrid_context() = default; @@ -126,6 +135,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 +145,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-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index 57919accf09..402384d0340 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -185,6 +185,12 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos // could be fatal return false; } + // a cache with no resident recurrent layers holds no state that could be partially + // erased, so the restriction does not apply to it. this is the glm5next MTP draft + // context, which runs only the NextN block and filters every KDA layer out. + const bool has_state = std::any_of(s_l.begin(), s_l.end(), + [](const ggml_tensor * t) { return t != nullptr; }); + if (0 <= seq_id) { int32_t & tail_id = cells[seq_id].tail; if (tail_id >= 0) { @@ -200,14 +206,16 @@ bool llama_memory_recurrent::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos cell.pos = p0 - 1; return true; } - return false; + if (has_state) { + return false; + } } // invalidate tails which will be cleared if (p0 <= cell.pos && cell.pos < p1) { tail_id = -1; } } - } else { + } else if (seq_id < 0) { // seq_id is negative, then the range should include everything or nothing if (p0 != p1 && (p0 != 0 || p1 != std::numeric_limits::max())) { //printf("[DEBUG] inside `llama_memory_recurrent::seq_rm`: `seq_id` is negative, so returning false\n"); diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 8860bd3f434..8cd12059ba5 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -298,7 +298,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 fc83658dd7f..42a78b79baa 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -202,6 +202,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: @@ -959,6 +961,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"; @@ -2449,9 +2452,11 @@ 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; - // only the sparse-attention architectures use llama_memory_hybrid_idx - // a null filter_idx means the GGUF has no indexer tensors + // null unless the arch has an indexer cache llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; + ggml_type type_idx = GGML_TYPE_F16; + // qwen4exp uses the dedicated llama_memory_hybrid_idx; glm5next carries its + // indexer in llama_memory_hybrid via filter_idx/type_idx const bool needs_mem_idx = (arch == LLM_ARCH_QWEN4EXP); if (arch == LLM_ARCH_FALCON_H1) { filter_attn = [&](uint32_t) { return true; }; @@ -2463,14 +2468,47 @@ 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_QWEN4EXP || arch == LLM_ARCH_MINIMAX_01) { - filter_attn = [&](uint32_t il) { + } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_MINIMAX_01 || arch == LLM_ARCH_GLM5NEXT) { + // the MTP draft context runs the NextN block and nothing else, so it + // gets a cache for that one layer. the trunk never runs it and so keeps + // the layer range it always had. handing the draft the trunk's cache is + // not just waste: the KDA layers would take cells it can never roll + // back, since a draft context is built with n_rs_seq = 0, and then a + // rejected draft fails seq_rm outright. + const bool mtp_ctx = arch == LLM_ARCH_GLM5NEXT && + cparams.ctx_type == LLAMA_CONTEXT_TYPE_MTP && + hparams.n_layer_all > hparams.n_layer(); + + filter_attn = [&, mtp_ctx](uint32_t il) { + if (mtp_ctx) { + return il >= hparams.n_layer() && il < hparams.n_layer_all; + } return il < hparams.n_layer() && !hparams.is_recr(il); }; - filter_recr = [&](uint32_t il) { - return il < hparams.n_layer() && hparams.is_recr(il); + filter_recr = [&, mtp_ctx](uint32_t il) { + return !mtp_ctx && il < hparams.n_layer() && hparams.is_recr(il); }; + if (arch == LLM_ARCH_GLM5NEXT && hparams.indexer_head_size > 0) { + // unified is fine, the pool map is per SEQUENCE. see [TAG_KPOOL_SEQ_PARTITION] + + // only the DSA layers carry an indexer key cache + filter_idx = [&, mtp_ctx](uint32_t il) { + if (mtp_ctx) { + return il >= hparams.n_layer() && il < hparams.n_layer_all; + } + return il < hparams.n_layer() && !hparams.is_recr(il); + }; + + // the gate cached beside the key feeds a softmax, unlike -ctk q8_0's target + 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 (arch == LLM_ARCH_QWEN4EXP && hparams.indexer_head_size > 0) { // QSA runs on the dense-attention layers only filter_idx = [&](uint32_t il) { @@ -2480,6 +2518,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { + // llama_memory_hybrid_iswa has no indexer cache, so SWA would silently lose it + 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, @@ -2538,7 +2579,16 @@ 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, + /* attn_v_mirror */ [] { + const char * e = getenv("LLAMA_MLA_V_MIRROR"); + if (!e || strcmp(e, "0") == 0 || strcmp(e, "off") == 0) return false; + if (strcmp(e, "1") == 0 || strcmp(e, "on") == 0) return true; + fprintf(stderr, "LLAMA_MLA_V_MIRROR=%s unrecognized (want off|0|on|1) -> off\n", e); + return false; + }()); } } else { llama_kv_cache::layer_filter_cb filter = nullptr; @@ -2812,6 +2862,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 38066538ed1..3cf88459b5c 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -144,6 +144,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 c414caa173f..41513007322 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -348,6 +348,29 @@ 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: quantizing these perturbs pool selection and KDA state retention, errors that + // compound over a sequence, for ~1 GiB. note the compressor tensors spell it with an + // UNDERSCORE and the projections with a DOT, so one "indexer." prefix test is not enough + 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; @@ -479,11 +502,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..23f62f4890d 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2257,6 +2257,12 @@ 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, inflating mixed Chinese-English ~13%. chatglm-bpe + // shares this pre_type but is ChatGLM3 and 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..9284cd2a7d0 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -330,9 +330,9 @@ 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, so KDA's per-key-channel + // decay must broadcast over ne1, not ne0 (for GDA g->ne[0] is 1 and both agree) + 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..83543d726ea --- /dev/null +++ b/src/models/glm5next.cpp @@ -0,0 +1,904 @@ +#include "models.h" + +#include "llama-memory-recurrent.h" +#include "llama-memory-hybrid.h" +#include "llama-kv-cache-kpool.h" + +// ssm_a holds -exp(A_log) (kimi-k3), not +exp(A_log) (bailingmoe3); converter checks + +// positions the indexer keeps; at or below this many the dense path IS the sparse one. +// asserted not measured (invisible to output); the second assert is an independent spelling +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); + // warned not asserted: no output comparison sees it, and an assert breaks test-llama-archs + 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"); + + // no linear_num_heads key: KDA head count is attention.head_count (converter enforces) + 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 + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound); + GGML_ASSERT(hparams.kda_gate_lower_bound < 0.0f); + + 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); + + 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); + + 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); + + // n_embd_out stays n_embd: lm_head sees the stream mean. deepseek4's hc_mult*n_embd + // makes llama-context.cpp overread t_embd, and the assert there sizes the destination + hparams.n_embd_out_impl = 0; + + 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); + + // to n_layer_all, not n_layer(): the NextN block is an attention block, and the loop + // below plus the memory layer filters both read the entry for it + uint32_t n_recr = 0; + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + hparams.is_recr_impl[il] = hparams.n_head_kv(il) == 0; + n_recr += il < hparams.n_layer() ? hparams.is_recr_impl[il] : 0; + } + 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; 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; + + 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 + 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); + } + } +} + +// one conv over concatenated q|k|v: keeps the conv state one block for rollback +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); + + 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 + 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 + + // g = lower_bound * sigmoid(exp(A_log)*(f_b(f_a(x)) + dt_bias)); it scales, not clamps + 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 + ggml_tensor * o = ggml_cont_3d(ctx0, out, head_dim, n_head, n_tokens); + cb(o, "kda_scan_out", il); + + 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); + + // 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 store is NOT gated on the sparse path, the scoring is: gating both leaves cells +// below n_select with no indexer state, which the first ubatch to cross n_select pools. +// * weights_proj runs in fp32; bf16 head-gates flip near-tie pool rankings (vLLM, sglang) +// * k_norm is a LayerNorm WITH BIAS at eps 1e-6, not f_norm_rms_eps (transformers, vLLM) +// * the ReLU between the QK dot and the head weighting is real (modular_glm5_next.py) +// no Hadamard rotation: H is orthogonal so (Hq).(Hk) == q.k; it only helps fp8. +// D0 (speed conference 056-058): gathered sparse decode, strict opt-in +static bool glm5_gathered_dsa_enabled() { + static const bool on = []() { + const char * e = getenv("LLAMA_GLM5_GATHERED_DSA"); + if (e == nullptr) return false; + if (strcmp(e, "on") == 0 || strcmp(e, "1") == 0) return true; + if (strcmp(e, "off") == 0 || strcmp(e, "0") == 0) return false; + fprintf(stderr, "LLAMA_GLM5_GATHERED_DSA=%s unrecognized (want off|0|on|1) -> off\n", e); + return false; + }(); + return on; +} + +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, + ggml_tensor ** out_slot_valid) 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; + + 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); + + // a SECOND, INDEPENDENT projection, not a reuse of the key, and cached beside it + 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 + 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); + // key and gate are the first two heads of a three-head row; the third is the pooled key, + // written later from a different set of cells, so this store must leave it alone + ggml_build_forward_expand(gf, + mctx_idx->cpy_k_part(ctx0, ggml_reshape_2d(ctx0, packed, 2*d_idx, n_tokens), + inp_kp->k_idxs, il, 2*d_idx, 0)); + + if (!scoring) { + return nullptr; + } + + 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] == 3 && + "the pooled indexer cache needs a key head, a gate head and a pooled head"); + GGML_ASSERT(kbuf->nb[1] == (size_t) d_idx*kbuf->nb[0] && "key, gate and pooled must be adjacent in a cell"); + GGML_ASSERT(n_tokens == n_tps*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); + + // pool only what this ubatch closed. the count is fixed (see build_inp_kpool), so the + // decode graph keeps one shape; unused slots repeat a complete pool, whose recompute is + // idempotent. non-resident slots hold 0, not a sentinel; garbage pools die to pool_bias. + const int64_t n_new_max = inp_kp->new_pool_cells->ne[0]/r; + + ggml_tensor * members = ggml_get_rows(ctx0, kg_rows, inp_kp->new_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_new_max, n_stream, + nb_mem, nb_mem*r, members->nb[2], 0); + ggml_tensor * mem_g = ggml_view_4d(ctx0, members, d_idx, r, n_new_max, n_stream, + nb_mem, nb_mem*r, members->nb[2], d_idx*members->nb[0]); + + // r-way softmaxes over the SLOT axis, so it must be dim 0; ape is added PRE-softmax + 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); + + ggml_tensor * pool_new = ggml_sum_rows(ctx0, ggml_mul(ctx0, keys_t, probs)); + pool_new = ggml_reshape_2d(ctx0, pool_new, d_idx, n_new_max*n_stream); + cb(pool_new, "indexer_pool_new", il); + + // store into the third head of each representative cell's row. the write is indexed by + // GLOBAL row (stream-merged) while pool_reps below is stream-local, so the read cannot + // be chained off the write tensor; expand it into the graph first and let build order + // sequence them, exactly as the key/gate store above does. + ggml_build_forward_expand(gf, + mctx_idx->cpy_k_part(ctx0, pool_new, inp_kp->new_pool_reps, il, d_idx, 2*d_idx)); + + ggml_tensor * pooled_rd = ggml_view_3d(ctx0, kbuf, d_idx, n_kv, n_stream, + kbuf->nb[2], kbuf->nb[3], 2*d_idx*kbuf->nb[0]); + + ggml_tensor * pool_k = ggml_get_rows(ctx0, pooled_rd, inp_kp->pool_reps); + pool_k = ggml_reshape_4d(ctx0, pool_k, d_idx, n_pools, 1, n_stream); + cb(pool_k, "indexer_pool_k", il); + + // 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 + // fold in here. GGML_PREC_F32 is not cosmetic: bf16 can swap two near-tied pools + 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); + + ggml_tensor * pool_score = nullptr; + + if (cparams.fused_lid) { + // pool_k stays f32 so the kernel takes its f32 path; f16 wmma would undo the prec + 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 { + ggml_tensor * kq = ggml_mul_mat(ctx0, pool_k, ggml_permute(ctx0, iq, 0, 2, 1, 3)); + + // the ReLU sits BETWEEN the per-head dot and the head weighting; either side differs + kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); + ggml_tensor * score = ggml_relu(ctx0, kq); + cb(score, "indexer_score", il); + + 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 then expand, as in the reference: a cell-level top-k is wrong + // because relu ties span pool boundaries and ggml_top_k splits the pool it lands in + 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); + + // the query axis folds into the gather's row axis, so ONE get_rows serves every query + 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); + + // D0: per-slot validity of the selection. gather pool_bias at the selected pool + // ordinals and expand each pool to its kpool member slots. decode-1 only: pool_bias + // is per-token, and this path is gated to n_tps==1. + if (out_slot_valid != nullptr) { + GGML_ASSERT(n_tps == 1 && n_stream == 1 && "gathered DSA is decode-1/one-stream only"); + + ggml_tensor * pb = ggml_reshape_3d(ctx0, inp_kp->pool_bias, 1, n_pools, n_stream); + ggml_tensor * sb = ggml_get_rows(ctx0, pb, sel_flat); // F32 [1, select_k, 1] + + ggml_tensor * rep = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, r, select_k, n_stream); + sb = ggml_repeat(ctx0, sb, rep); // member-major, matches top_k + + *out_slot_valid = ggml_reshape_3d(ctx0, sb, r*select_k, n_tps, n_stream); + cb(*out_slot_valid, "indexer_slot_valid", il); + } + + 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; +} + +// absorbed form (deepseek2/glm-dsa): q_nope goes through wk_b so q.k is taken against +// the latent the cache holds; the naive form needs a V cache this layout lacks +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; + + GGML_ASSERT(hparams.n_rot() == 0); + + // scale is over the MLA head size, as in the reference, not the absorbed width + 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); + + // D0 gate: the tail inputs exist iff build_inp_kpool's single gathered_decode + // boolean was true — resource presence IS the policy (codex 059 P0) + const bool gathered = inp_kp != nullptr && inp_kp->tail_cells != nullptr; + + ggml_tensor * slot_valid = nullptr; + ggml_tensor * top_k = inp_kp ? build_indexer(layer, inp_kp, cur, qr, scoring, il, + gathered ? &slot_valid : nullptr) : 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); + + q = ggml_permute(ctx0, q, 0, 2, 1, 3); + + q = ggml_mul_mat(ctx0, layer.wk_b, q); + + 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 && gathered) { + GGML_ASSERT(slot_valid != nullptr && inp_kp->tail_cells != nullptr); + cur = build_attn_sparse_gathered(inp_attn, + layer.wo, nullptr, nullptr, + q, k, k, nullptr, nullptr, layer.wv_b, + top_k, slot_valid, inp_kp->tail_cells, inp_kp->tail_valid, + inp_kp->cand_mask, inp_kp->kpool, kq_scale, il); + } else 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 like the experts: one Glm5NextTextMLP serves both + 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 stay unbiased + 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 applies 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(); + + // E0c: the gathered decision is computed BEFORE the hybrid input so a + // maskless graph never creates (or fills) the dense KQ mask (codex 066 s4) + const bool fbl_gathered_decode = glm5_gathered_dsa_enabled() && + cparams.n_ctx > glm5next_n_select(hparams) && n_tokens == 1 && + (cparams.kv_unified ? 1 : (int64_t) ubatch.n_seqs_unq) == 1; + + llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(fbl_gathered_decode); + + // one map for the whole ubatch; nothing in it depends on the layer. gated on n_ctx, + // not n_kv, which grows and would flip the graph topology mid-run + 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); + + // ONE owner for the D0 decision; computed above, before the hybrid input + inp_kp = build_inp_kpool(mctx_hyb, + inp_mem->get_attn()->get_kq_mask(), indexer_scoring, + fbl_gathered_decode); + } + } + + 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 mHC state to the experts + 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]); + } + + // when unmasked nextn embeddings are requested, t_h_nextn must keep all rows, so the + // early output masking has to be skipped (it is applied after the final norm instead) + const bool mask_early = !cparams.embeddings_nextn || cparams.embeddings_nextn_masked; + + if (inp_out_ids && mask_early) { + // 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); + + // post-norm hidden state feeds the NextN/MTP draft head + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (inp_out_ids && !mask_early) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +// NextN draft head. Unlike glm-dsa and deepseek32, whose MTP blocks skip the DSA indexer, +// blk.45 here ships a full indexer and is an ordinary DSA layer, so it runs the same +// build_layer_attn/build_layer_ffn the trunk does. What it does NOT run is the mHC mixer: +// the NextN block keeps a plain residual and has no hc_* tensors (see load_arch_tensors). +llama_model_glm5next::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) + : graph(params) { + GGML_ASSERT(hparams.n_layer_nextn == 1 && "glm5next MTP supports a single NextN block"); + + const int il = hparams.n_layer() + cparams.nextn_layer_offset; + + GGML_ASSERT(cparams.nextn_layer_offset >= 0 && + cparams.nextn_layer_offset < (int) hparams.n_layer_nextn && + "nextn_layer_offset out of range [0, n_layer_nextn)"); + + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && layer.nextn.enorm && layer.nextn.hnorm && + "glm5next MTP block tensors missing; convert without --no-mtp"); + GGML_ASSERT(!layer.hc_attn_fn && "the NextN block has no mHC mixer"); + + auto inp = std::make_unique(hparams.n_embd); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_inp(), n_tokens); + ggml_set_input(inp->embd); + + ggml_tensor * tok_embd; + if (ubatch.token) { + ggml_tensor * tok_embd_w = layer.nextn.embed_tokens ? layer.nextn.embed_tokens : model.tok_embd; + + tok_embd = ggml_get_rows(ctx0, tok_embd_w, inp->tokens); + } else { + tok_embd = inp->embd; + } + cb(tok_embd, "mtp_tok_embd", il); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd, n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + ggml_tensor * h_embd = inp->h; + + res->add_input(std::move(inp)); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + llm_graph_input_mem_hybrid_k * inp_mem = build_inp_mem_hybrid_k(); + + // the NextN block is DSA, so nothing here consumes the recurrent half. s_copy would then + // never enter the graph, never be allocated, and set_input would still read its buffer. + ggml_build_forward_expand(gf, inp_mem->get_recr()->s_copy); + + 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); + + // D0 stays OFF the NextN graph until MTP has its own oracle (grokk 060 s2) + inp_kp = build_inp_kpool(mctx_hyb, + inp_mem->get_attn()->get_kq_mask(), indexer_scoring, + /* gathered */ false); + } + } + + ggml_tensor * h_norm = build_norm(h_embd, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + cb(h_norm, "mtp_hnorm", il); + + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + cb(e_norm, "mtp_enorm", il); + + ggml_tensor * cur = build_lora_mm(layer.nextn.eh_proj, + ggml_concat(ctx0, e_norm, h_norm, /*dim=*/ 0), layer.nextn.eh_proj_s); + cb(cur, "mtp_eh_proj", il); + + // plain pre-norm residual block, the trunk's helpers minus the mHC wrapper + ggml_tensor * residual = cur; + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_attn_norm", il); + + cur = build_layer_attn(model, inp_mem, inp_kp, indexer_scoring, cur, il); + + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, residual); + cb(ffn_inp, "mtp_ffn_inp", il); + + cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "mtp_ffn_norm", il); + + cur = build_layer_ffn(model, cur, il); + cb(cur, "mtp_ffn_out", il); + + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "mtp_post_ffn", il); + + ggml_tensor * head_norm_w = layer.nextn.shared_head_norm ? layer.nextn.shared_head_norm : model.output_norm; + GGML_ASSERT(head_norm_w && "glm5next MTP: no nextn.shared_head_norm and no output_norm"); + cur = build_norm(cur, head_norm_w, nullptr, LLM_NORM_RMS, -1); + + // the post-norm hidden state seeds the next MTP step + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "mtp_shared_head_norm", -1); + + ggml_tensor * head_w = layer.nextn.shared_head_head ? layer.nextn.shared_head_head : model.output; + ggml_tensor * head_s = layer.nextn.shared_head_head ? layer.nextn.shared_head_head_s : model.output_s; + GGML_ASSERT(head_w && "glm5next MTP: no nextn.shared_head_head and no output"); + cur = build_lora_mm(head_w, cur, head_s); + cb(cur, "result_output", -1); + + res->t_logits = cur; + ggml_build_forward_expand(gf, cur); +} + +std::unique_ptr llama_model_glm5next::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } + + return std::make_unique(*this, params); +} diff --git a/src/models/models.h b/src/models/models.h index af60764c2f7..77e0f50ce1e 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1176,8 +1176,10 @@ 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) {} + // method-only mixin; here only so glm5next, deriving from this graph, reaches + // build_delta_net. deepseek4 itself has no recurrent 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( @@ -1297,6 +1299,10 @@ struct llama_model_deepseek4 : public llama_model_base { ggml_tensor * build_hc_sinkhorn( ggml_tensor * comb, int il) const; + + static ggml_tensor * build_hc_mean( + ggml_context * ctx, + ggml_tensor * x); }; struct graph_mtp : public graph { @@ -1334,6 +1340,68 @@ 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 block, collapsed by unweighted mean, not a gated head + struct graph : public llama_model_deepseek4::graph { + graph(const llama_model & model, const llm_graph_params & params); + + // builds nothing: lets graph_mtp reuse the block helpers below without the trunk + graph(const llm_graph_params & params) : llama_model_deepseek4::graph(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, the same function below n_select + 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; + + // always stores the key and gate; when `scoring`, returns the selected CELL indices. + // out_slot_valid (optional, D0): expanded pool_bias of the selected pools, + // F32 [kpool*select_k, n_tps, n_stream], 0/-inf per slot (decode-1 only) + ggml_tensor * build_indexer( + const llama_layer & layer, + llm_graph_input_kpool * inp_kp, + ggml_tensor * cur, + ggml_tensor * qr, + bool scoring, + int il, + ggml_tensor ** out_slot_valid = nullptr) const; + + ggml_tensor * build_layer_ffn( + const llama_model & model, + ggml_tensor * cur, + int il) const; + }; + + // NextN draft head. reuses the trunk's block helpers, so it stays a `graph` + struct graph_mtp : public graph { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_eagle3 : public llama_model_base { llama_model_eagle3(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index d58d90952eb..c1c19ac3484 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -64,6 +64,17 @@ static void set_tensor_data(struct ggml_tensor * tensor, void * userdata) { } } +// depth-sweep overrides (0 = leave defaults). Set only by test_depth_sweep(). +static uint32_t g_depth_sweep_n_ctx = 0; +static uint32_t g_depth_sweep_n_ub = 0; +static std::string g_dump_1tok; +static bool g_skip_1tok = false; +static uint32_t g_sweep_iheads = 0; +static uint32_t g_sweep_iklen = 0; +static uint32_t g_sweep_layers = 0; // glm5next fixture layer-count override +static uint32_t g_sweep_dlead = 0; // leading dense layers (real model: 3 of 45) +static uint32_t g_sweep_heads = 0; // glm5next fixture attention-head override + static void usage(char ** argv) { printf("Usage: %s [-a/--arch arch] [-s/--seed seed] [-v/--verbose]\n", argv[0]); } @@ -118,10 +129,17 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { || arch == LLM_ARCH_KIMI_LINEAR || arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 - || arch == LLM_ARCH_MISTRAL4) { + || arch == LLM_ARCH_MISTRAL4 + || arch == LLM_ARCH_GLM5NEXT) { n_embd = 128; n_head = 1; n_ff = 192; + if (arch == LLM_ARCH_GLM5NEXT && g_sweep_layers > 0) { + n_layer = g_sweep_layers; + } + if (arch == LLM_ARCH_GLM5NEXT && g_sweep_heads > 0) { + n_head = g_sweep_heads; + } } else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { n_layer = 3; } else if (arch == LLM_ARCH_CHAMELEON) { @@ -174,6 +192,20 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { } ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head_per_layer); ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_per_layer); + } else if (arch == LLM_ARCH_GLM5NEXT) { + // head_count doubles as the KDA head count, so it stays uniform; the kv array is what + // marks the recurrent layers, and the loader asserts it holds both a zero and a nonzero + // Under --heads (query-head override), DSA stays MQA: n_head_kv is held at 1 + // (real model: 64 query heads over one latent KV head). KDA layers stay 0. + GGML_ASSERT(n_layer >= 2); + const uint32_t glm_dsa_kv = (g_sweep_heads > 0) ? 1u : n_head_kv; + std::vector n_head_kv_per_layer; + n_head_kv_per_layer.reserve(n_layer); + for (uint32_t il = 0; il < n_layer; il++) { + n_head_kv_per_layer.push_back(il == 1 ? 0 : glm_dsa_kv); + } + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head); + ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, n_head_kv_per_layer); } else { ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head); ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT_KV, arch == LLM_ARCH_DEEPSEEK4 ? uint32_t(1) : n_head_kv); @@ -213,12 +245,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: the cache holds the bare latent, so no rope width is added on top of + // the kv LoRA rank and n_rot has to be an explicit 0, not the head size default + 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)); } 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)); } ms.add_kv(LLM_KV_ATTENTION_CLAMP_KQV, 1.0f); - ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS, 1e-5f); + // glm5next warns on anything but the 1e-6 its indexer k_norm hardcodes + ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_EPS, arch == LLM_ARCH_GLM5NEXT ? 1e-6f : 1e-5f); ms.add_kv(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, 1e-5f); ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_EPS, 1e-5f); ms.add_kv(LLM_KV_ATTENTION_GROUPNORM_GROUPS, uint32_t(8)); @@ -277,6 +318,16 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_SWIGLU_CLAMP_EXP, 10.0f); ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); + } else if (arch == LLM_ARCH_GLM5NEXT) { + ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); // build_hc_pre asserts exactly 4 streams + ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f); + // the only arch that pools indexer keys; top_k must be a whole number of pools, and + // the resulting selection width has to stay under n_ctx or the sparse path goes unused + ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, uint32_t(4)); + // glm5next reads these unconditionally; the if (moe) block below never sets them + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_SCALE, 1.0f); + ms.add_kv(LLM_KV_EXPERT_WEIGHTS_NORM, true); } ms.add_kv(LLM_KV_TOKENIZER_MODEL, "no_vocab"); // ms.add_kv(LLM_KV_DENSE_2_FEAT_OUT, n_embd); @@ -343,9 +394,11 @@ static bool silent_model_load_progress(float /*progress*/, void * /*user_data*/) return true; } + static std::pair get_model_and_ctx( struct gguf_context * gguf_ctx, FILE * file, const size_t seed, const std::vector & devs, - const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false) { + const llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER, bool encode = false, + llama_flash_attn_type fa_type = LLAMA_FLASH_ATTN_TYPE_AUTO) { GGML_ASSERT((gguf_ctx == nullptr) != (file == nullptr)); llama_model_params model_params = llama_model_default_params(); model_params.progress_callback = silent_model_load_progress; @@ -355,11 +408,13 @@ static std::pair get_model_and_ctx( model_params.split_mode = split_mode; llama_context_params ctx_params = llama_context_default_params(); - ctx_params.n_ctx = 0; + ctx_params.flash_attn_type = fa_type; // depth-sweep passes DISABLED (production parity); others keep AUTO + ctx_params.n_ctx = g_depth_sweep_n_ctx; // 0 = model default (original behavior) ctx_params.n_threads = 4; ctx_params.n_threads_batch = 4; if (!encode) { - ctx_params.n_ubatch = 64; + ctx_params.n_ubatch = g_depth_sweep_n_ub ? g_depth_sweep_n_ub : 64; + ctx_params.n_batch = std::max(2048, ctx_params.n_ubatch); } size_t tmp = seed; @@ -432,6 +487,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_GLM4_MOE: case LLM_ARCH_GLM_DSA: + case LLM_ARCH_GLM5NEXT: case LLM_ARCH_EXAONE_MOE: case LLM_ARCH_BAILINGMOE: case LLM_ARCH_BAILINGMOE2: @@ -751,6 +807,156 @@ static int test_backends(const llm_arch target_arch, const size_t seed, const gg return all_ok ? 0 : 1; } +// Long-context CPU-vs-device divergence sweep for the glm5next "@@@@" collapse +// (ggml-org/llama.cpp#27754 / #27752). Runs ONE tiny random-weight glm5next +// model on the CPU backend and on the first non-CPU device with an identical +// token stream, decoding in chunks; after each chunk the last-position logits +// are compared (NMSE + greedy argmax). On the real model the collapse depth +// depends on (depth, n_ctx, n_ubatch); this asks whether that reproduces at +// toy scale, in minutes instead of hours. +static int test_depth_sweep(const size_t seed, const uint32_t max_depth, + const uint32_t n_ctx_arg, const uint32_t n_ub, + const uint32_t n_ub2, const bool b_on_cpu, + const uint32_t topk_override) { + g_depth_sweep_n_ctx = n_ctx_arg; + + gguf_context_ptr gguf_ctx = get_gguf_ctx(LLM_ARCH_GLM5NEXT, /*moe=*/true); + // the fixture bakes context_length=128; the KV/pool structures must be + // sized for the sweep target instead + gguf_set_val_u32(gguf_ctx.get(), "glm5next.context_length", n_ctx_arg); + if (topk_override > 0) { + // real GLM-5.3-Flash: indexer.top_k = 2048, kpool = 4 (fixture: 8/4) + gguf_set_val_u32(gguf_ctx.get(), "glm5next.attention.indexer.top_k", topk_override); + } + if (g_sweep_iheads > 0) { // real: 32 (fixture: 1) + gguf_set_val_u32(gguf_ctx.get(), "glm5next.attention.indexer.head_count", g_sweep_iheads); + } + if (g_sweep_iklen > 0) { // real: 128 (fixture: 64) + gguf_set_val_u32(gguf_ctx.get(), "glm5next.attention.indexer.key_length", g_sweep_iklen); + } + if (g_sweep_dlead > 0) { // real: 3 + gguf_set_val_u32(gguf_ctx.get(), "glm5next.leading_dense_block_count", g_sweep_dlead); + } + + ggml_backend_dev_t dev_gpu = nullptr; + for (size_t i = 0; i < ggml_backend_dev_count(); i++) { + ggml_backend_dev_t dev = ggml_backend_dev_get(i); + if (ggml_backend_dev_buffer_type(dev) != ggml_backend_cpu_buffer_type()) { + dev_gpu = dev; + break; + } + } + if (dev_gpu == nullptr) { + printf("depth-sweep: no non-CPU device found\n"); + return 1; + } + // side A: CPU @ n_ub. side B: (cpu|device) @ (n_ub2 ? n_ub2 : n_ub). + const uint32_t ub_a = n_ub ? n_ub : 64; + const uint32_t ub_b = n_ub2 ? n_ub2 : ub_a; + printf("depth-sweep: glm5next-moe, n_ctx=%u, top_k=%s, max_depth=%u | A=cpu@ub%u B=%s@ub%u\n", + n_ctx_arg, topk_override ? std::to_string(topk_override).c_str() : "fixture(8)", + max_depth, ub_a, b_on_cpu ? "cpu" : ggml_backend_dev_description(dev_gpu), ub_b); + + g_depth_sweep_n_ub = ub_a; + auto mc_cpu = get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, false, LLAMA_FLASH_ATTN_TYPE_DISABLED); + g_depth_sweep_n_ub = ub_b; + auto mc_dev = b_on_cpu ? get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {}, LLAMA_SPLIT_MODE_LAYER, false, LLAMA_FLASH_ATTN_TYPE_DISABLED) + : get_model_and_ctx(gguf_ctx.get(), nullptr, seed, {dev_gpu}, LLAMA_SPLIT_MODE_LAYER, false, LLAMA_FLASH_ATTN_TYPE_DISABLED); + const uint32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(mc_cpu.first.get())); + + const std::vector tokens = get_tokens(max_depth, n_vocab, seed); + const uint32_t chunk = 2048; + + llama_batch batch = llama_batch_init(chunk, 0, 1); + bool diverged = false; + for (uint32_t pos0 = 0; pos0 < max_depth; pos0 += chunk) { + const uint32_t n = std::min(chunk, max_depth - pos0); + common_batch_clear(batch); + for (uint32_t i = 0; i < n; i++) { + common_batch_add(batch, tokens[pos0 + i], pos0 + i, {0}, i == n - 1); + } + if (llama_decode(mc_cpu.second.get(), batch)) { + printf("depth-sweep: CPU decode failed at depth %u\n", pos0 + n); + llama_batch_free(batch); + return 1; + } + if (llama_decode(mc_dev.second.get(), batch)) { + printf("depth-sweep: device decode failed at depth %u\n", pos0 + n); + llama_batch_free(batch); + return 1; + } + const float * lc = llama_get_logits_ith(mc_cpu.second.get(), n - 1); + const float * ld = llama_get_logits_ith(mc_dev.second.get(), n - 1); + double se = 0.0, ref = 0.0; + uint32_t amax_c = 0, amax_d = 0; + for (uint32_t j = 0; j < n_vocab; j++) { + const double d = (double) lc[j] - (double) ld[j]; + se += d * d; + ref += (double) lc[j] * (double) lc[j]; + if (lc[j] > lc[amax_c]) amax_c = j; + if (ld[j] > ld[amax_d]) amax_d = j; + } + const double nmse_val = ref > 0.0 ? se / ref : se; + const bool bad = nmse_val > 1e-3 || amax_c != amax_d; + printf("depth=%7u nmse=%.3e argmax_cpu=%u argmax_dev=%u%s\n", + pos0 + n, nmse_val, amax_c, amax_d, bad ? " <-- DIVERGED" : ""); + fflush(stdout); + if (bad) { + diverged = true; + } + } + // grokk 054 s3 / 060 s2: the chunked sweep never takes the n_tokens==1 gate. + // FOUR successive single-token decodes in one process exercise the packed/ + // gathered branches AND the stateful rolling tail (kpool tail length walks + // an ordered transition; starting right after a chunk boundary the expected + // tail lengths are (pos+1)%kpool = 1, 2, 3, 0 for kpool=4). + if (!diverged && !g_skip_1tok) { + FILE * dump = nullptr; + if (!g_dump_1tok.empty()) { + dump = fopen(g_dump_1tok.c_str(), "wb"); + } + for (uint32_t step = 0; step < 4 && !diverged; step++) { + common_batch_clear(batch); + common_batch_add(batch, tokens[step], max_depth + step, {0}, true); + if (llama_decode(mc_cpu.second.get(), batch) || llama_decode(mc_dev.second.get(), batch)) { + printf("depth-sweep-1tok: decode failed at step %u\n", step); + if (dump) fclose(dump); + llama_batch_free(batch); + return 1; + } + const float * lc = llama_get_logits_ith(mc_cpu.second.get(), 0); + const float * ld = llama_get_logits_ith(mc_dev.second.get(), 0); + double se = 0.0, ref = 0.0; + uint32_t amax_c = 0, amax_d = 0; + for (uint32_t j = 0; j < n_vocab; j++) { + const double d = (double) lc[j] - (double) ld[j]; + se += d * d; + ref += (double) lc[j] * (double) lc[j]; + if (lc[j] > lc[amax_c]) amax_c = j; + if (ld[j] > ld[amax_d]) amax_d = j; + } + const double nmse_val = ref > 0.0 ? se / ref : se; + const bool bad = nmse_val > 1e-3 || amax_c != amax_d; + printf("depth-sweep-1tok: step=%u depth=%u tail_len=%u nmse=%.3e " + "argmax_cpu=%u argmax_dev=%u%s\n", + step, max_depth + step + 1, (max_depth + step + 1) % 4, + nmse_val, amax_c, amax_d, bad ? " <-- DIVERGED" : ""); + if (bad) diverged = true; + if (dump) { + fwrite(ld, sizeof(float), n_vocab, dump); + fwrite(lc, sizeof(float), n_vocab, dump); + } + } + if (dump) { + fclose(dump); + printf("depth-sweep-1tok: dumped 4 steps to %s\n", g_dump_1tok.c_str()); + } + } + llama_batch_free(batch); + printf("depth-sweep: %s\n", diverged ? "DIVERGENCE FOUND" : "no divergence up to max depth"); + return diverged ? 2 : 0; +} + int main(int argc, char ** argv) { // FIXME these tests are disabled in the CI for macOS-latest-cmake-arm64 because they are segfaulting common_init(); @@ -758,6 +964,14 @@ int main(int argc, char ** argv) { llm_arch arch = LLM_ARCH_UNKNOWN; size_t seed = rd(); + uint32_t depth_sweep = 0; + std::string dump_1tok; + bool skip_1tok = false; + uint32_t sweep_ctx = 131072; + uint32_t sweep_ub = 0; + uint32_t sweep_ub2 = 0; + uint32_t sweep_topk = 0; + bool sweep_b_cpu = false; ggml_log_level log_level = GGML_LOG_LEVEL_ERROR; std::string out; @@ -775,6 +989,45 @@ int main(int argc, char ** argv) { return 1; } } + if (strcmp(argv[i], "--depth-sweep") == 0 && i + 1 < argc) { + depth_sweep = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--dump-1tok") == 0 && i + 1 < argc) { + dump_1tok = argv[++i]; + } + if (strcmp(argv[i], "--skip-1tok") == 0) { + skip_1tok = true; + } + if (strcmp(argv[i], "--ctx") == 0 && i + 1 < argc) { + sweep_ctx = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--ub") == 0 && i + 1 < argc) { + sweep_ub = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--ub2") == 0 && i + 1 < argc) { + sweep_ub2 = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--b-cpu") == 0) { + sweep_b_cpu = true; + } + if (strcmp(argv[i], "--topk") == 0 && i + 1 < argc) { + sweep_topk = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--iheads") == 0 && i + 1 < argc) { + g_sweep_iheads = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--iklen") == 0 && i + 1 < argc) { + g_sweep_iklen = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--layers") == 0 && i + 1 < argc) { + g_sweep_layers = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--heads") == 0 && i + 1 < argc) { + g_sweep_heads = std::stoul(argv[++i]); + } + if (strcmp(argv[i], "--dlead") == 0 && i + 1 < argc) { + g_sweep_dlead = std::stoul(argv[++i]); + } if (strcmp(argv[i], "-s") == 0 || strcmp(argv[i], "--seed") == 0) { if (i + 1 < argc) { seed = std::stoull(argv[++i]); @@ -799,6 +1052,12 @@ int main(int argc, char ** argv) { printf("%s: using seed %zu\n", __func__, seed); try { + g_dump_1tok = dump_1tok; + g_skip_1tok = skip_1tok; + if (depth_sweep > 0) { + return test_depth_sweep(seed, depth_sweep, sweep_ctx, sweep_ub, + sweep_ub2, sweep_b_cpu, sweep_topk); + } if (!out.empty()) { return save_models(arch, seed, log_level, out); } 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..a5a87eb0f88 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -684,6 +684,19 @@ ggml_tensor * clip_graph::build_ffn( cur = ggml_sqr(ctx0, cur); cb(cur, "ffn_relu_sqr", il); } break; + case FFN_SILU_CLAMP: + { + // not ggml_swiglu_oai: it clamps the same way but adds one to the up branch + 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 +1094,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 +1743,20 @@ 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 PILImageResampling.BICUBIC, which this only approximates + hparams.image_resize_algo = RESIZE_ALGO_BICUBIC; + get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); + get_f32(KEY_VISION_SWIGLU_LIMIT, hparams.swiglu_limit); + hparams.ffn_op = FFN_SILU_CLAMP; + log_ffn_op = "silu_clamp"; + // the preprocessor's min_pixels/max_pixels, 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 +2574,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 +4033,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 +4060,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 +4142,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 +4768,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 +5917,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..78f6d7622e4 100644 --- a/tools/mtmd/models/glm4v.cpp +++ b/tools/mtmd/models/glm4v.cpp @@ -42,7 +42,10 @@ 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); + // GLM-OCR has none, and build_norm still normalizes on a null weight, so skip the call + 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..5d16cb2c367 --- /dev/null +++ b/tools/mtmd/models/glm5next-vision.cpp @@ -0,0 +1,12 @@ +#include "models.h" + +// the GLM-OCR ViT plus a clamp on the SwiGLU gate and up projections, which the per-block MLP +// and the merger both pick up from hparams.ffn_op +// ref: https://huggingface.co/zai-org/GLM-5.3-Flash/blob/main/modeling_glm5_next.py +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..4627d4f68b2 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1583,3 +1583,105 @@ 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 out, leaving pixel area vs min/max +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 + 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); + + 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) { + // aligning both edges is not monotone in the Qwen sqrt(area / max_pixels) scale, so search + 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) { + // already spending the minimum budget, so shrink only, never upscale 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..41126a98c51 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -239,3 +239,19 @@ 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; }; + +// ref: Glm5NextImageProcessor.{smart_resize,resize}. unlike the Qwen-style smart_resize in +// mtmd_image_preprocessor_dyn_size, an over-budget image is pasted top-left, not centred +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 + }; + + // static 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|>