Commit 649dcb103 for llama.cpp
commit 649dcb1036fd47d01a83b56c5e35913e65c277ba
Author: timkhronos <timkhronos@gmail.com>
Date: Wed Sep 30 08:20:32 2026 +0200
add GLM-5.3-Flash (GLM5-Next) support (#27773)
* Rebase GLM-Next support onto master, and migrate to llama-memory-hybrid-idx
* Add initial MTP support
* Merge branch optimizations. Reduce allocated compute buffer size, speed up long context decode, fla, and slight MTP improvements.
* Review driven changes, remove env vars, protect tensors
* Strip MTP for initial PR
* Clean up after mtp strip
* Clean up after mtp strip
* Update speculative.cpp
* Update llama-context.h
* Clean up after mtp strip
* Fix tokenizer ignore merges
* Improve quantization protection selection
* Refactor mhc helpers, graph base
* Lint Fixes
* Apply suggestions from code review
Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
* Skip glm5-next in model saver, fix CRLF
* Skip glm5-next in sweep
* Remove T4 fallback
* Review cleanup
* Review suggestions
* Defer separate MTP gguf handling to MTP PR, drop filter
* Repad n_head_kv
* kpool init apply
* Order by descending score
* Drop guard
* read kpool from hparams, clarify kpool cache flags, remove kpool_build_state(nullptr)
* Add glm5-next support to model saver and add arch test fixture
* Review cleanup
* Kpool pooled caching clarify
* Add multi stream support
* Finish Rebase
* Sparse FA fir DSA prefill
* Const
* Update llama-model.cpp to fix rebase error
* gguf-py : merge tensor map entries for HC tensors
* model : use build_gdn_l2_norm in GLM5_NEXT implementation
* chore : remove trailing whitespace
* model : use new OP precision setting API in GLM5_NEXT implementation
* mtmd : use ggml_swiglu_clamp in GLM5V and apply the image token limit
The two clamps around swiglu_split are what ggml_swiglu_clamp already does,
so the clamp bounds collapse back to one value. GLM5V also never called
set_limit_image_tokens(), so --image-max-tokens had no effect.
Assisted-by: Claude Opus 5
(cherry picked from commit 46d18e12d422be4cc04a70e4a9a9e0168bb3d5b7)
* llama : keep the GLM5-Next k-pool layout across ubatches
The layout was rebuilt from a full cell scan on every ubatch. Pools are fixed
by the positions relative to the sequence's first one, so the layout now lives
on the memory and a ubatch only appends to it.
A sequence edit no longer stales every pooled key either, only the ones at or
after the edited position, which makes a tail seq_rm free. The pooling subgraph
is built unconditionally so the graph shape no longer changes every kpool
tokens, and the pool axis is folded into rows before soft_max, which otherwise
exceeds the CUDA gridDim.y limit past n_kv 262144.
Assisted-by: Claude Opus 5
(cherry picked from commit 5d1c40b93e17fddbf73b785efe43e0d02ccb3977)
* model : write the GLM5-Next recurrent rollback checkpoints
The conv state and the delta net state were only written to the live row, so a
rollback restored whatever the checkpoint rows happened to hold. Take the same
route as kimi-k3: build_recurrent_attn for the state, and write all K_rs conv
groups. That also drops a state view that assumed contiguous rows.
Enroll the arch in test-recurrent-state-rollback, which catches this under its
garbage-filled cache pass.
Assisted-by: Claude Opus 5
(cherry picked from commit 5ace37e86d5d448e83ef5dde5632c748185b18cd)
* llama: fix PR #27773 test-save-load-state restore failure
Clear the attention and indexer cache data after a failed hybrid state restore so restored NaNs cannot affect a later sequence.
Assisted-by: Codex
* llama: fix PR #27773 gpu-rocm graph reallocation
Reserve the full GLM5-Next pool capacity and dirty pool count. The gpu-rocm Test step aborts when n_new grows while the graph node count stays fixed; CUDA, Vulkan, Metal, and WebGPU checks report the same error.
Assisted-by: Codex
* llama : fix GLM5-Next k-pool layout staleness after edits and shared teardown
Two defects in the cross-ubatch k-pool layout added by the k-pool commit:
1. Wrong results. An edited sequence only rebuilt its pool layout when its cell
count changed, so if the first ubatch after an edit added back exactly as many
cells as were removed, the stale position-to-cell list survived. With a unified
cache and more than one sequence, where another sequence takes the freed cells,
the reused layout points at the wrong cells (CPU: large logit drift, CUDA: NaN).
Rebuild whenever the sequence is stale, not only on a size mismatch.
2. Slowdown. "shared" mode was assumed to end only with an edit that forces a
rebuild, but sharing also ends when the other sequence is removed. The survivor
kept shared = true, pinning cache_safe off and re-pooling every pool on every
ubatch (server trigger: n>1 completions with -kvu, via the seq_cp in
copy_state_to). In seq_rm, if the layout has shared cells, stale every sequence
so one rebuild re-derives sharing and cache_safe returns to 1.
Assisted-by: Claude Opus 5
* llama : fix build_attn_mha stream stride for non-contiguous q
build_attn_mha split the batch into streams with a stream stride of
q->nb[3]/n_stream. That only equals one stream's span, (ne[2]/n_stream)*nb[2],
when q is contiguous. GLM5-Next is nope-only, so it does not concat a rope part
and passes the permuted q_absorbed straight in, where nb[3] != ne[2]*nb[2]; the
stride was then n_head times too large and every stream s >= 1 read another
head's queries. Split-KV (-np N without --kv-unified) multi-stream prefill was
wrong for every stream past the first. Unified KV and decode were unaffected
(n_stream == 1, and decode takes the gather path). Other MLA models concat rope
so q is contiguous and the computed value is unchanged for them.
Compute the stride from the token dimension, which is identical for a
contiguous q.
Assisted-by: Claude Opus 5
* llama : re-derive GLM5-Next k-pool sharing on state_read/state_drop
The shared-cell teardown added to seq_rm (stale every sequence when the layout
has shared cells, so a survivor does not keep shared = true and pin cache_safe
off) was missing from the other paths that can free shared cells: state_read
and state_drop staled only the one sequence. Apply the same re-derivation there
and correct the comment that claimed sharing ends only via an edit or seq_rm.
Assisted-by: Claude Opus 5
* quant : drop duplicate GLM5-Next hc_ filter
The hc_ name filter was listed twice in the GLM5_NEXT protection block.
Assisted-by: Claude Opus 5
* glm5-next: scope K-pool cache access to indexed operations
* glm5-next: keep K-pool access in hybrid index memory
* glm5-next: keep mHC graph builders model-local
* glm5-next: mark only touched pools per ubatch
---------
Co-authored-by: Sigbjørn Skjæret <sigbjorn.skjaeret@huggingface.co>
Co-authored-by: Stanisław Szymczyk <sszymczy@gmail.com>
Co-authored-by: Piotr Wilkin <ilintar@gmail.com>
diff --git a/conversion/__init__.py b/conversion/__init__.py
index 85db1d643..777b2f982 100644
--- a/conversion/__init__.py
+++ b/conversion/__init__.py
@@ -106,6 +106,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"Glm4MoeLiteForCausalLM": "glm",
"Glm4vForConditionalGeneration": "glm",
"Glm4vMoeForConditionalGeneration": "glm",
+ "Glm5NextForConditionalGeneration": "glm",
"GlmForCausalLM": "chatglm",
"GlmMoeDsaForCausalLM": "glm",
"GlmOcrForConditionalGeneration": "glm",
@@ -305,6 +306,7 @@ MMPROJ_MODEL_MAP: dict[str, str] = {
"Glm4vForConditionalGeneration": "qwen3vl",
"BailingMoeV3VLForConditionalGeneration": "bailingmoe3",
"Glm4vMoeForConditionalGeneration": "qwen3vl",
+ "Glm5NextForConditionalGeneration": "qwen3vl",
"Glm5vForConditionalGeneration": "kimivl",
"GlmOcrForConditionalGeneration": "qwen3vl",
"GlmasrModel": "ultravox",
diff --git a/conversion/glm.py b/conversion/glm.py
index 7544f850c..10f3ef113 100644
--- a/conversion/glm.py
+++ b/conversion/glm.py
@@ -402,3 +402,197 @@ class SolarOpenModel(Glm4MoeModel):
special_vocab._set_special_token("unk", tokenizer.get_added_vocab()["<unk>"]) # ty: ignore[unresolved-attribute]
special_vocab._set_special_token("bos", tokenizer.get_added_vocab()["<|startoftext|>"]) # ty: ignore[unresolved-attribute]
special_vocab.add_to_gguf(self.gguf_writer)
+
+
+@ModelBase.register("Glm5NextForConditionalGeneration")
+@ModelBase.example("zai-org/GLM-5.3-Flash")
+class Glm5NextModel(TextModel):
+
+ model_arch = gguf.MODEL_ARCH.GLM5_NEXT
+ supports_mtp_export = True
+
+ _experts: list[dict[str, Tensor]] | None = None
+ _n_main_layers: int | None = None
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+
+ self.n_nextn_layers = self.hparams.get("num_nextn_predict_layers", 0)
+ self.skip_mtp = self.no_mtp or self.n_nextn_layers == 0
+
+ if not self.skip_mtp:
+ self.block_count += self.n_nextn_layers
+ self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count)
+
+ self.hparams.pop("head_dim", None)
+
+ def set_vocab(self):
+ # requires transformers >= 5, tokpre hash-resolves to glm4
+ return self._set_vocab_glm()
+
+ def index_tensors(self, remote_hf_model_id: str | None = None):
+ hp = self.hparams.get("text_config", self.hparams)
+ type(self)._n_main_layers = hp["num_hidden_layers"]
+ return super().index_tensors(remote_hf_model_id=remote_hf_model_id)
+
+ @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._n_main_layers is not None
+ m = re.match(r"model\.layers\.(\d+)\.", name)
+ is_mtp = m is not None and int(m.group(1)) >= cls._n_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 set_gguf_parameters(self):
+ super().set_gguf_parameters()
+ hp = self.hparams
+
+ layer_types = hp["layer_types"]
+ n_kv_heads = [0 if t == "linear_attention" else 1 for t in layer_types]
+ assert len(n_kv_heads) == hp["num_hidden_layers"]
+ # Pad to block_count
+ n_kv_heads += [1] * (self.block_count - len(n_kv_heads))
+ self.gguf_writer.add_head_count_kv(n_kv_heads)
+ self.gguf_writer.add_vocab_size(hp["vocab_size"])
+ self.gguf_writer.add_layer_norm_eps(1e-6)
+
+ if not self.skip_mtp:
+ self.gguf_writer.add_nextn_predict_layers(self.n_nextn_layers)
+
+ # KDA
+ lin = hp["linear_attn_config"]
+ assert lin["num_heads"] == hp["num_attention_heads"]
+ self.gguf_writer.add_ssm_conv_kernel(lin["short_conv_kernel_size"])
+ self.gguf_writer.add_kda_head_dim(lin["head_dim"])
+ if (lb := lin.get("gate_lower_bound")) is not None:
+ self.gguf_writer.add_kda_gate_lower_bound(lb)
+
+ # MLA (nope only)
+ assert hp.get("mla_use_nope") and hp["qk_rope_head_dim"] == 0, "expected nope-only MLA"
+ kv_lora_rank = hp["kv_lora_rank"]
+ qk_rope = 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)
+ self.gguf_writer.add_key_length(kv_lora_rank + qk_rope)
+ self.gguf_writer.add_value_length(kv_lora_rank)
+ self.gguf_writer.add_key_length_mla(hp["qk_nope_head_dim"] + qk_rope)
+ self.gguf_writer.add_value_length_mla(hp["v_head_dim"])
+
+ # DSA indexer with k-pool compression
+ 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"])
+ self.gguf_writer.add_indexer_kpool_select_tail(hp.get("index_kpool_always_select_tail", True))
+ if (indexer_types := hp.get("indexer_types")) is not None:
+ self.gguf_writer.add_indexer_types([t == "full" for t in indexer_types])
+
+ # mHC
+ assert hp.get("mhc", True)
+ 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
+ self.gguf_writer.add_leading_dense_block_count(hp["first_k_dense_replace"])
+ self.gguf_writer.add_expert_feed_forward_length(hp["moe_intermediate_size"])
+ self.gguf_writer.add_expert_shared_count(hp["n_shared_experts"])
+ self.gguf_writer.add_expert_weights_scale(hp["routed_scaling_factor"])
+ self.gguf_writer.add_expert_weights_norm(hp["norm_topk_prob"])
+ if (limit := hp.get("swiglu_limit")) is not None:
+ self.gguf_writer.add_swiglu_clamp_exp([limit] * self.block_count)
+ self.gguf_writer.add_swiglu_clamp_shexp([limit] * self.block_count)
+
+ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+ if name == "lm_head.weight" and self.hparams.get("tie_word_embeddings", False):
+ return
+
+ # 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 w_name in ("down_proj", "gate_proj", "up_proj"):
+ datas: list[Tensor] = []
+ for xid in range(n_experts):
+ ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight"
+ datas.append(self._experts[bid].pop(ename))
+ merged = f"model.layers.{bid}.mlp.experts.{w_name}.weight"
+ yield from super().modify_tensors(torch.stack(datas, dim=0), merged, bid)
+ return
+
+ # MLA absorption
+ if name.endswith("kv_b_proj.weight"):
+ 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), name.replace("kv_b_proj", "k_b_proj"), bid)
+ yield from super().modify_tensors(v_b, name.replace("kv_b_proj", "v_b_proj"), bid)
+ return
+
+ # KDA conv1d
+ if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")):
+ if data_torch.ndim == 3:
+ d_inner, _, d_conv = data_torch.shape
+ elif data_torch.ndim == 2:
+ d_inner, d_conv = data_torch.shape
+ else:
+ raise ValueError(f"unexpected conv1d rank {data_torch.ndim} for {name}")
+ data_torch = data_torch.reshape(1, d_inner, 1, d_conv)
+
+ if name.endswith(".A_log"):
+ n_head = self.hparams["num_attention_heads"]
+ data_torch = -torch.exp(data_torch.float().flatten()[:n_head])
+
+ if name.endswith(".dt_bias"):
+ name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias"
+
+ if re.search(r"\.(hc_(?:attn|ffn)_(?:fn|base|scale)|index_kpool_compress_(?:ape|gate))$", name):
+ yield self.map_tensor_name(name) + ".weight", data_torch
+ return
+
+ yield from super().modify_tensors(data_torch, name, bid)
+
+ def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool:
+ # keep the small mHC / gating parameters exact
+ exact_keys = ("hc_attn_", "hc_ffn_", "indexer_compressor_", "ssm_a", "ssm_dt", "exp_probs_b")
+ if new_name.startswith(("blk.", "output_hc")) and any(k in new_name for k in exact_keys):
+ return gguf.GGMLQuantizationType.F32
+ return super().tensor_force_quant(name, new_name, bid, n_dims)
+
+ def prepare_metadata(self, vocab_only: bool):
+ from_dir = self.fname_out.is_dir()
+ super().prepare_metadata(vocab_only=vocab_only)
+ if not self.mtp_only or not from_dir:
+ return
+ output_type: str = self.ftype.name.partition("_")[2]
+ fname_default: str = gguf.naming_convention(
+ self.metadata.name, self.metadata.basename, self.metadata.finetune,
+ self.metadata.version, size_label=None, output_type=output_type, model_type=None)
+ self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf"
+
+ def prepare_tensors(self):
+ super().prepare_tensors()
+ if self._experts is not None:
+ leftover = [k for d in self._experts for k in d.keys()]
+ if leftover:
+ raise ValueError(f"Unprocessed experts: {leftover}")
diff --git a/conversion/qwen3vl.py b/conversion/qwen3vl.py
index 4fec708c9..11ce68515 100644
--- a/conversion/qwen3vl.py
+++ b/conversion/qwen3vl.py
@@ -228,10 +228,12 @@ 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):
+ 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.projector_type)
hidden_act = str(self.hparams_vision.get("hidden_act", "")).lower()
if hidden_act == "gelu":
@@ -249,6 +251,32 @@ class Glm4VVisionModel(Qwen3VLVisionModel):
yield from super().modify_tensors(data_torch, name, bid)
+@ModelBase.register("Glm5NextForConditionalGeneration")
+@ModelBase.example("zai-org/GLM-5.3-Flash")
+class Glm5NextVisionModel(Glm4VVisionModel):
+ # GLM-5.3-Flash vision tower. glm4v layout with per-head qk-norm, no post-conv norm and no learned position embeddings.
+ # Images are placed on a ceil aligned canvas with padding.
+
+ projector_type = gguf.VisionProjectorType.GLM5V
+
+ def set_gguf_parameters(self):
+ super().set_gguf_parameters()
+ assert self.hparams_vision is not None
+ self.gguf_writer.add_vision_spatial_merge_size(int(self.hparams_vision.get("spatial_merge_size", 2)))
+ if (limit := self.hparams_vision.get("swiglu_limit")) is not None:
+ self.gguf_writer.add_vision_swiglu_clamp(float(limit))
+
+ # image token budget from the processor, stored as single-frame pixel counts
+ pc = self.preprocessor_config
+ patch = int(pc.get("patch_size", 14))
+ merge = int(pc.get("merge_size", 2))
+ pixels_per_token = (patch * merge) ** 2
+ if (min_tok := pc.get("min_image_tokens")) is not None:
+ self.gguf_writer.add_vision_min_pixels(int(min_tok) * pixels_per_token)
+ if (max_tok := pc.get("max_image_tokens")) is not None:
+ self.gguf_writer.add_vision_max_pixels(int(max_tok) * pixels_per_token)
+
+
@ModelBase.register("Qwen3VLForConditionalGeneration")
@ModelBase.example("Qwen/Qwen3-VL-4B-Instruct")
class Qwen3VLTextModel(Qwen3Model):
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index e193e9995..2daedfda9 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -230,6 +230,8 @@ class Keys:
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" # GLM5-Next
+ KPOOL_SELECT_TAIL = "{arch}.attention.indexer.kpool_select_tail" # GLM5-Next
class HyperConnection:
COUNT = "{arch}.hyper_connection.count"
@@ -396,6 +398,7 @@ class Keys:
IMAGE_MEAN = "clip.vision.image_mean"
IMAGE_STD = "clip.vision.image_std"
SPATIAL_MERGE_SIZE = "clip.vision.spatial_merge_size"
+ SWIGLU_CLAMP = "clip.vision.swiglu_clamp"
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"
@@ -575,6 +578,7 @@ class MODEL_ARCH(IntEnum):
GLM4 = auto()
GLM4_MOE = auto()
GLM_DSA = auto()
+ GLM5_NEXT = auto()
BITNET = auto()
T5 = auto()
T5ENCODER = auto()
@@ -909,6 +913,8 @@ class MODEL_TENSOR(IntEnum):
INDEXER_COMPRESSOR_WGATE = auto()
INDEXER_COMPRESSOR_APE = auto()
INDEXER_COMPRESSOR_NORM = auto()
+ INDEXER_KPOOL_GATE = auto()
+ INDEXER_KPOOL_APE = auto()
# vision
V_MMPROJ = auto()
V_MMPROJ_FC = auto()
@@ -1333,6 +1339,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
MODEL_ARCH.GLM4: "glm4",
MODEL_ARCH.GLM4_MOE: "glm4moe",
MODEL_ARCH.GLM_DSA: "glm-dsa",
+ MODEL_ARCH.GLM5_NEXT: "glm5-next",
MODEL_ARCH.BITNET: "bitnet",
MODEL_ARCH.T5: "t5",
MODEL_ARCH.T5ENCODER: "t5encoder",
@@ -1666,6 +1673,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: "blk.{bid}.indexer_compressor_gate",
MODEL_TENSOR.INDEXER_COMPRESSOR_APE: "blk.{bid}.indexer_compressor_ape",
MODEL_TENSOR.INDEXER_COMPRESSOR_NORM: "blk.{bid}.indexer_compressor_norm",
+ MODEL_TENSOR.INDEXER_KPOOL_GATE: "blk.{bid}.indexer_compressor_gate",
+ MODEL_TENSOR.INDEXER_KPOOL_APE: "blk.{bid}.indexer_compressor_ape",
# vision
MODEL_TENSOR.V_MMPROJ: "mm.{bid}",
MODEL_TENSOR.V_MMPROJ_FC: "mm.model.fc",
@@ -4103,6 +4112,70 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
],
+ MODEL_ARCH.GLM5_NEXT: [
+ MODEL_TENSOR.TOKEN_EMBD,
+ MODEL_TENSOR.OUTPUT_NORM,
+ MODEL_TENSOR.OUTPUT,
+ MODEL_TENSOR.ATTN_NORM,
+ # mHC
+ 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.ATTN_OUT,
+ MODEL_TENSOR.SSM_CONV1D_Q,
+ MODEL_TENSOR.SSM_CONV1D_K,
+ MODEL_TENSOR.SSM_CONV1D_V,
+ MODEL_TENSOR.SSM_F_A,
+ MODEL_TENSOR.SSM_F_B,
+ MODEL_TENSOR.SSM_BETA,
+ MODEL_TENSOR.SSM_A,
+ MODEL_TENSOR.SSM_G_A,
+ MODEL_TENSOR.SSM_G_B,
+ MODEL_TENSOR.SSM_DT,
+ MODEL_TENSOR.SSM_NORM,
+ # MLA (nope) + DSA layers
+ MODEL_TENSOR.ATTN_Q_A,
+ MODEL_TENSOR.ATTN_Q_B,
+ MODEL_TENSOR.ATTN_KV_A_MQA,
+ MODEL_TENSOR.ATTN_KV_B,
+ MODEL_TENSOR.ATTN_K_B,
+ MODEL_TENSOR.ATTN_V_B,
+ MODEL_TENSOR.ATTN_Q_A_NORM,
+ MODEL_TENSOR.ATTN_KV_A_NORM,
+ MODEL_TENSOR.INDEXER_K_NORM,
+ MODEL_TENSOR.INDEXER_PROJ,
+ MODEL_TENSOR.INDEXER_ATTN_K,
+ MODEL_TENSOR.INDEXER_ATTN_Q_B,
+ MODEL_TENSOR.INDEXER_KPOOL_GATE,
+ MODEL_TENSOR.INDEXER_KPOOL_APE,
+ # FFN
+ MODEL_TENSOR.FFN_NORM,
+ MODEL_TENSOR.FFN_GATE,
+ MODEL_TENSOR.FFN_DOWN,
+ MODEL_TENSOR.FFN_UP,
+ MODEL_TENSOR.FFN_GATE_INP,
+ MODEL_TENSOR.FFN_GATE_EXP,
+ MODEL_TENSOR.FFN_DOWN_EXP,
+ MODEL_TENSOR.FFN_UP_EXP,
+ MODEL_TENSOR.FFN_GATE_SHEXP,
+ MODEL_TENSOR.FFN_DOWN_SHEXP,
+ MODEL_TENSOR.FFN_UP_SHEXP,
+ MODEL_TENSOR.FFN_EXP_PROBS_B,
+ # NextN/MTP tensors - preserved but unused
+ MODEL_TENSOR.NEXTN_EH_PROJ,
+ MODEL_TENSOR.NEXTN_EMBED_TOKENS,
+ MODEL_TENSOR.NEXTN_ENORM,
+ MODEL_TENSOR.NEXTN_HNORM,
+ MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD,
+ MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM,
+ ],
MODEL_ARCH.BITNET: [
MODEL_TENSOR.ATTN_Q,
MODEL_TENSOR.ATTN_K,
@@ -5862,6 +5935,7 @@ class VisionProjectorType:
LFM2A = "lfm2a" # audio
MUSIC_FLAMINGO = "musicflamingo" # audio
GLM4V = "glm4v"
+ GLM5V = "glm5v"
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 543d9b744..342f6eb4d 100644
--- a/gguf-py/gguf/gguf_writer.py
+++ b/gguf-py/gguf/gguf_writer.py
@@ -830,6 +830,12 @@ class GGUFWriter:
key = Keys.Attention.Indexer.TYPES.format(arch=self.arch)
self.add_array(key, value)
+ def add_indexer_kpool(self, value: int) -> None:
+ self.add_uint32(Keys.Attention.Indexer.KPOOL.format(arch=self.arch), value)
+
+ def add_indexer_kpool_select_tail(self, value: bool) -> None:
+ self.add_bool(Keys.Attention.Indexer.KPOOL_SELECT_TAIL.format(arch=self.arch), value)
+
def add_max_alibi_bias(self, bias: float) -> None:
self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias)
@@ -1411,6 +1417,9 @@ class GGUFWriter:
def add_vision_image_std(self, values: Sequence[float]) -> None:
self.add_array(Keys.ClipVision.IMAGE_STD, values)
+ def add_vision_swiglu_clamp(self, value: float) -> None:
+ self.add_float32(Keys.ClipVision.SWIGLU_CLAMP, value)
+
def add_vision_spatial_merge_size(self, value: int) -> None:
self.add_uint32(Keys.ClipVision.SPATIAL_MERGE_SIZE, value)
diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py
index d2dfeece5..16df24911 100644
--- a/gguf-py/gguf/tensor_mapping.py
+++ b/gguf-py/gguf/tensor_mapping.py
@@ -1319,6 +1319,14 @@ class TensorNameMap:
"model.layers.{bid}.self_attn.indexer.wq_b", # DSA
),
+ MODEL_TENSOR.INDEXER_KPOOL_GATE: (
+ "model.layers.{bid}.self_attn.indexer.index_kpool_compress_gate", # glm5-next
+ ),
+
+ MODEL_TENSOR.INDEXER_KPOOL_APE: (
+ "model.layers.{bid}.self_attn.indexer.index_kpool_compress_ape", # glm5-next
+ ),
+
MODEL_TENSOR.INDEXER_Q_PROJ: (
"model.layers.{bid}.self_attn.index_q_proj", # MSA
),
@@ -1333,26 +1341,32 @@ class TensorNameMap:
MODEL_TENSOR.HC_ATTN_FN: (
"model.layers.{bid}.hc_attn_layer.hc_pre.hc_fn", # hy-v4
+ "model.layers.{bid}.hc_attn_fn", # glm5-next
),
MODEL_TENSOR.HC_ATTN_BASE: (
"model.layers.{bid}.hc_attn_layer.hc_pre.hc_base", # hy-v4
+ "model.layers.{bid}.hc_attn_base", # glm5-next
),
MODEL_TENSOR.HC_ATTN_SCALE: (
"model.layers.{bid}.hc_attn_layer.hc_pre.hc_scale", # hy-v4
+ "model.layers.{bid}.hc_attn_scale", # glm5-next
),
MODEL_TENSOR.HC_FFN_FN: (
"model.layers.{bid}.hc_mlp_layer.hc_pre.hc_fn", # hy-v4
+ "model.layers.{bid}.hc_ffn_fn", # glm5-next
),
MODEL_TENSOR.HC_FFN_BASE: (
"model.layers.{bid}.hc_mlp_layer.hc_pre.hc_base", # hy-v4
+ "model.layers.{bid}.hc_ffn_base", # glm5-next
),
MODEL_TENSOR.HC_FFN_SCALE: (
"model.layers.{bid}.hc_mlp_layer.hc_pre.hc_scale", # hy-v4
+ "model.layers.{bid}.hc_ffn_scale", # glm5-next
),
MODEL_TENSOR.HC_HEAD_FN: (
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 2c5a228f3..f72164ddb 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -153,6 +153,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_MAINCODER, "maincoder" },
{ LLM_ARCH_KIMI_LINEAR, "kimi-linear" },
{ LLM_ARCH_KIMI_K3, "kimi-k3" },
+ { LLM_ARCH_GLM5_NEXT, "glm5-next" },
{ LLM_ARCH_TALKIE, "talkie" },
{ LLM_ARCH_MELLUM, "mellum" },
{ LLM_ARCH_NANBEIGE, "nanbeige" },
@@ -294,6 +295,8 @@ static const std::map<llm_kv, const char *> 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_INDEXER_KPOOL_SELECT_TAIL, "%s.attention.indexer.kpool_select_tail" },
{ 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" },
@@ -692,6 +695,8 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_NORM, "blk.%d.indexer_compressor_norm" },
+ { LLM_TENSOR_INDEXER_KPOOL_GATE, "blk.%d.indexer_compressor_gate" },
+ { LLM_TENSOR_INDEXER_KPOOL_APE, "blk.%d.indexer_compressor_ape" },
{ LLM_TENSOR_FFN_GATE_TID2EID, "blk.%d.ffn_gate_tid2eid" },
{ LLM_TENSOR_MASKED_EMBD_CENTROIDS, "masked_embd_centroids" },
{ LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" },
@@ -966,6 +971,8 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
{LLM_TENSOR_INDEXER_COMPRESSOR_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+ {LLM_TENSOR_INDEXER_KPOOL_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_INDEXER_KPOOL_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}},
{LLM_TENSOR_FFN_GATE_TID2EID, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
{LLM_TENSOR_NEXTN_PROJ_PRE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_NEXTN_PROJ_POST, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
@@ -1089,6 +1096,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) {
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_BAILINGMOE3:
case LLM_ARCH_KIMI_K3:
+ case LLM_ARCH_GLM5_NEXT:
case LLM_ARCH_QWEN35:
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_QWEN4EXP:
@@ -1119,6 +1127,7 @@ bool llm_arch_supports_rs_rollback(const llm_arch & arch) {
case LLM_ARCH_QWEN35MOE:
case LLM_ARCH_QWEN4EXP:
case LLM_ARCH_DEEPSEEK4:
+ case LLM_ARCH_GLM5_NEXT:
case LLM_ARCH_NEMOTRON_H:
case LLM_ARCH_NEMOTRON_H_MOE:
case LLM_ARCH_LFM2:
@@ -1160,6 +1169,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_BAILINGMOE3:
case LLM_ARCH_KIMI_K3:
+ case LLM_ARCH_GLM5_NEXT:
case LLM_ARCH_QWEN3TTS:
case LLM_ARCH_QWEN4EXP: // TODO: fix test-llama-archs
return false;
diff --git a/src/llama-arch.h b/src/llama-arch.h
index 3b8bd6428..3b437840d 100644
--- a/src/llama-arch.h
+++ b/src/llama-arch.h
@@ -153,6 +153,7 @@ enum llm_arch {
LLM_ARCH_MAINCODER,
LLM_ARCH_KIMI_LINEAR,
LLM_ARCH_KIMI_K3,
+ LLM_ARCH_GLM5_NEXT,
LLM_ARCH_TALKIE,
LLM_ARCH_MELLUM,
LLM_ARCH_EAGLE3,
@@ -299,6 +300,8 @@ 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_INDEXER_KPOOL_SELECT_TAIL,
LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT,
LLM_KV_ATTENTION_OUTPUT_LORA_RANK,
LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
@@ -690,6 +693,8 @@ enum llm_tensor {
LLM_TENSOR_INDEXER_COMPRESSOR_WGATE,
LLM_TENSOR_INDEXER_COMPRESSOR_APE,
LLM_TENSOR_INDEXER_COMPRESSOR_NORM,
+ LLM_TENSOR_INDEXER_KPOOL_GATE, // glm5-next: k-pool gate scores
+ LLM_TENSOR_INDEXER_KPOOL_APE, // glm5-next: k-pool position bias
LLM_TENSOR_FFN_GATE_TID2EID,
LLM_TENSOR_NEXTN_PROJ_PRE,
LLM_TENSOR_NEXTN_PROJ_POST,
diff --git a/src/llama-context.cpp b/src/llama-context.cpp
index 439beabdf..87bb3ccae 100644
--- a/src/llama-context.cpp
+++ b/src/llama-context.cpp
@@ -2362,7 +2362,7 @@ 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) {
+ if (model.arch == LLM_ARCH_KIMI_K3 || model.arch == LLM_ARCH_GLM5_NEXT) {
// the n_tokens*40 budget below is exhausted at ubatch 3840
res = std::max<uint32_t>(n_tokens * 160, 64u * model.n_tensors());
} else if (model.arch == LLM_ARCH_HRM_TEXT) {
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index abf3069bb..d398ee1d7 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -1840,7 +1840,7 @@ ggml_tensor * llm_graph_context::build_ffn(
const float limit = hparams.swiglu_clamp_shexp[il];
constexpr float eps = 1e-6f;
if (limit > eps) {
- if (arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
+ if (arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0)) {
cur = ggml_swiglu_clamp(ctx0, cur, tmp, limit);
} else {
tmp = ggml_clamp(ctx0, tmp, -limit, limit);
@@ -2234,7 +2234,7 @@ ggml_tensor * llm_graph_context::build_moe_ffn(
const float limit = hparams.swiglu_clamp_exp[il];
constexpr float eps = 1e-6f;
if (limit > eps) {
- if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
+ if (arch == LLM_ARCH_MAPLE || arch == LLM_ARCH_DEEPSEEK4 || arch == LLM_ARCH_GLM5_NEXT || (arch == LLM_ARCH_DFLASH && hparams.dsv4_hc_mult > 0) || arch == LLM_ARCH_HY_V4) {
cur = ggml_swiglu_clamp(ctx0, cur, up, limit);
} else {
up = ggml_clamp(ctx0, up, -limit, limit);
@@ -2624,7 +2624,10 @@ ggml_tensor * llm_graph_context::build_attn_mha(
// split the batch into streams if needed
const auto n_stream = k->ne[3];
- q = ggml_view_4d(ctx0, q, q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, q->nb[1], q->nb[2], q->nb[3]/n_stream, 0);
+ // the stream dim steps over one stream's worth of dim-2 (tokens): (ne[2]/n_stream) rows of stride nb[2].
+ // q->nb[3]/n_stream only equals that for a contiguous q; nope-only MLA (glm5-next) passes a permuted
+ // q_absorbed where nb[3] != ne[2]*nb[2], so using nb[3] read another head's queries for streams s >= 1.
+ q = ggml_view_4d(ctx0, q, q->ne[0], q->ne[1], q->ne[2]/n_stream, n_stream, q->nb[1], q->nb[2], q->nb[2]*(q->ne[2]/n_stream), 0);
q = ggml_permute(ctx0, q, 0, 2, 1, 3);
k = ggml_permute(ctx0, k, 0, 2, 1, 3);
diff --git a/src/llama-hparams.h b/src/llama-hparams.h
index d8fbfbdd8..8248add7d 100644
--- a/src/llama-hparams.h
+++ b/src/llama-hparams.h
@@ -282,6 +282,8 @@ struct llama_hparams {
uint32_t indexer_n_head = 0;
uint32_t indexer_head_size = 0;
uint32_t indexer_top_k = 0;
+ uint32_t indexer_kpool = 0; // k-pool size
+ bool indexer_kpool_select_tail = true;
// MSA
uint32_t indexer_block_size = 0;
uint32_t indexer_local_blocks = 0;
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 1c93f908b..1d0554fa1 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -1202,6 +1202,10 @@ uint32_t llama_kv_cache::get_size() const {
return cells.size();
}
+uint32_t llama_kv_cache::get_n_seq_max() const {
+ return n_seq_max;
+}
+
uint32_t llama_kv_cache::get_n_stream() const {
return n_stream;
}
@@ -1247,6 +1251,12 @@ const llama_kv_cells & llama_kv_cache::get_cells(llama_seq_id seq_id) const {
return v_cells[seq_to_stream[seq_id]];
}
+uint32_t llama_kv_cache::get_stream(llama_seq_id seq_id) const {
+ GGML_ASSERT(seq_id >= 0 && (size_t) seq_id < seq_to_stream.size());
+
+ return seq_to_stream[seq_id];
+}
+
uint32_t llama_kv_cache::get_n_kv(const slot_info & sinfo) const {
uint32_t result = 0;
diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h
index 5051f4343..52cc8cd32 100644
--- a/src/llama-kv-cache.h
+++ b/src/llama-kv-cache.h
@@ -156,6 +156,7 @@ public:
//
uint32_t get_size() const;
+ uint32_t get_n_seq_max() const;
uint32_t get_n_stream() const;
bool get_has_shift() const;
@@ -168,6 +169,9 @@ public:
const llama_kv_cells & get_cells(llama_seq_id seq_id) const;
+ // The stream holding seq_id's cells.
+ uint32_t get_stream(llama_seq_id seq_id) const;
+
// state_read, plus the cells the restored tokens were placed in
// a cache that mirrors another one (the qwen4exp indexer) must not search for its own cells: two searches agree only by luck
// sinfos_out: if set, filled with the layout used; a stream with no cells leaves an empty entry
diff --git a/src/llama-kv-cells.h b/src/llama-kv-cells.h
index 5d567a6ed..9c0d51e4e 100644
--- a/src/llama-kv-cells.h
+++ b/src/llama-kv-cells.h
@@ -318,6 +318,14 @@ public:
return seq[i].test(seq_id);
}
+ // the (pos, cell) pairs of sequence seq_id, ordered by position
+ const std::set<std::pair<llama_pos, uint32_t>> & seq_pos_get(llama_seq_id seq_id) const {
+ assert(seq_id >= 0);
+ assert(seq_id < LLAMA_MAX_SEQ);
+
+ return seq_pos[seq_id];
+ }
+
// the token of the cell of sequence seq_id at the largest position <= p
// when several cells share that position, the one with the highest index wins
// return LLAMA_TOKEN_NULL if the sequence has no cell at or before p
diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp
index 2adfa7a69..32b042253 100644
--- a/src/llama-memory-hybrid-idx.cpp
+++ b/src/llama-memory-hybrid-idx.cpp
@@ -1,5 +1,9 @@
#include "llama-memory-hybrid-idx.h"
+#include <algorithm>
+#include <cmath>
+#include <type_traits>
+
#include "llama-impl.h"
#include "llama-batch.h"
#include "llama-io.h"
@@ -49,7 +53,8 @@ llama_memory_hybrid_idx::llama_memory_hybrid_idx(
mem_idx(filter_idx == nullptr ? nullptr : [&] {
// MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own
std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1);
- hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size;
+ // The glm5 next indexer caches key, gate and pooled values per token
+ hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size * (model.hparams.indexer_kpool > 0 ? 3 : 1);
// the cached indexer keys are raw, rotation happens after pooling at read time, so a
// K-shift must not rotate them while the stream copies in the same update still apply
@@ -145,9 +150,37 @@ void llama_memory_hybrid_idx::clear(bool data) {
if (mem_idx) {
mem_idx->clear(data);
+ mem_idx_stale_set(-1, 0);
}
}
+// A pooled key is only valid while the grouping that produced it holds. Grouping is sequence relative,
+// so an edit at p0 leaves every pool that ends before p0 alone.
+void llama_memory_hybrid_idx::mem_idx_stale_set(llama_seq_id seq_id, llama_pos p0) {
+ p0 = std::max<llama_pos>(p0, 0);
+
+ if (seq_id < 0) {
+ for (auto & p : mem_idx_stale) {
+ p = std::min(p, p0);
+ }
+
+ return;
+ }
+
+ GGML_ASSERT(seq_id < (llama_seq_id) LLAMA_MAX_SEQ);
+
+ mem_idx_stale[seq_id] = std::min(mem_idx_stale[seq_id], p0);
+}
+
+// An edit at or below the first position moves pos_min, which regroups the whole sequence.
+llama_pos llama_memory_hybrid_idx::mem_idx_stale_pos(llama_seq_id seq_id, llama_pos p0) const {
+ if (seq_id < 0 || p0 <= mem_idx->seq_pos_min(seq_id)) {
+ return 0;
+ }
+
+ return p0;
+}
+
bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) {
// same order as llama_memory_hybrid::seq_rm: the recurrent cache can refuse, so try it first
if (!get_mem_recr()->seq_rm(seq_id, p0, p1)) {
@@ -155,7 +188,15 @@ bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_po
}
if (mem_idx) {
+ const llama_pos stale = mem_idx_stale_pos(seq_id, p0);
mem_idx->seq_rm(seq_id, p0, p1);
+ mem_idx_stale_set(seq_id, stale);
+
+ // removing a sequence can free cells another sequence shared, but only this one is marked stale, so the
+ // survivor would keep shared = true and pin cache_safe off forever; stale every sequence to re-derive it
+ if (kpool_layout_shared()) {
+ mem_idx_stale_set(-1, 0);
+ }
}
return get_mem_attn()->seq_rm(seq_id, p0, p1);
@@ -166,6 +207,9 @@ void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_i
if (mem_idx) {
mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1);
+ // the copy shares cells, which cannot hold two groupings, so both sides drop their cached keys
+ mem_idx_stale_set(seq_id_src, 0);
+ mem_idx_stale_set(seq_id_dst, 0);
}
}
@@ -174,6 +218,8 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) {
if (mem_idx) {
mem_idx->seq_keep(seq_id);
+ // cells shared with the dropped sequences become exclusive again, their keys were never cached
+ mem_idx_stale_set(-1, 0);
}
}
@@ -181,7 +227,10 @@ void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_p
llama_memory_hybrid::seq_add(seq_id, p0, p1, shift);
if (mem_idx) {
+ // a negative shift moves the cells below p0, so they regroup as well
+ const llama_pos stale = mem_idx_stale_pos(seq_id, shift < 0 ? p0 + shift : p0);
mem_idx->seq_add(seq_id, p0, p1, shift);
+ mem_idx_stale_set(seq_id, stale);
}
}
@@ -190,6 +239,7 @@ void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_p
if (mem_idx) {
mem_idx->seq_div(seq_id, p0, p1, d);
+ mem_idx_stale_set(seq_id, 0);
}
}
@@ -238,6 +288,12 @@ void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_
if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) {
if (mem_idx) {
mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn);
+ // the restore rewrites the cells behind the pool layout's back
+ mem_idx_stale_set(seq_id, 0);
+ // it can also change which cells are shared; re-derive sharing for every sequence, as seq_rm does
+ if (kpool_layout_shared()) {
+ mem_idx_stale_set(-1, 0);
+ }
}
}
@@ -258,11 +314,16 @@ void llama_memory_hybrid_idx::state_drop(llama_seq_id seq_id) {
return;
}
- get_mem_attn()->seq_rm(seq_id, -1, -1);
+ get_mem_attn()->state_clear(seq_id);
get_mem_recr()->seq_rm(seq_id, -1, -1);
if (mem_idx) {
- mem_idx->seq_rm(seq_id, -1, -1);
+ mem_idx->state_clear(seq_id);
+ mem_idx_stale_set(seq_id, 0);
+ // clearing this sequence can end a sharing the survivor would otherwise keep flagged (see seq_rm)
+ if (kpool_layout_shared()) {
+ mem_idx_stale_set(-1, 0);
+ }
}
}
@@ -608,6 +669,153 @@ static std::vector<uint32_t> llama_memory_hybrid_idx_ns(const llama_kv_cache::sl
return res;
}
+// Which cells of a sequence make up which pool, for the whole cache.
+struct llama_memory_hybrid_idx::kpool_layout {
+ struct seq {
+ llama_pos pos_min = 0;
+ uint32_t strm = 0; // Stream holding this sequence's cells
+ std::vector<std::pair<llama_pos, uint32_t>> cells; // Position and stream local cell pairs, sorted by position.
+ std::vector<uint32_t> pools;
+
+ // Where the pool scan stopped, so an append resumes instead of starting over.
+ size_t j_next = 0;
+
+ // Whether any cell also carries another sequence, which rules out caching this sequence's pooled keys.
+ bool shared = false;
+ };
+
+ std::array<seq, LLAMA_MAX_SEQ> seqs;
+
+ uint32_t n_pool_real = 0;
+ bool cache_safe = true;
+};
+
+// Which pools of the layout the current ubatch must re-pool, in the layout's pool order.
+struct llama_memory_hybrid_idx_context::kpool_state {
+ std::vector<uint32_t> is_new;
+ uint32_t generation = 0;
+
+ uint32_t n_pool_real = 0;
+ uint32_t n_new = 0;
+ bool cache_safe = true;
+};
+
+namespace {
+
+// The last padded pool is always unused.
+uint32_t kpool_pad(uint32_t n_pool) {
+ return std::max<uint32_t>(64u, GGML_PAD(n_pool + 1, 64u));
+}
+
+}
+
+llama_memory_hybrid_idx::~llama_memory_hybrid_idx() = default;
+
+const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_layout_get() const {
+ GGML_ASSERT(kpool_lay != nullptr);
+
+ return *kpool_lay;
+}
+
+bool llama_memory_hybrid_idx::kpool_layout_shared() const {
+ return kpool_lay && !kpool_lay->cache_safe;
+}
+
+// Pools are fixed by the positions relative to the sequence's first one, so the layout survives a plain
+// append. A sequence edit can regroup them, and mem_idx_stale tells us it happened.
+const llama_memory_hybrid_idx::kpool_layout & llama_memory_hybrid_idx::kpool_layout_update() {
+ GGML_ASSERT(mem_idx != nullptr);
+
+ if (!kpool_lay) {
+ kpool_lay = std::make_unique<kpool_layout>();
+ }
+
+ auto & lay = *kpool_lay;
+
+ const uint32_t kpool = get_kpool();
+ const uint32_t n_stream_kv = mem_idx->get_n_stream();
+ const bool unified = n_stream_kv == 1;
+
+ lay.n_pool_real = 0;
+ lay.cache_safe = true;
+
+ for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
+ auto & sq = lay.seqs[s];
+
+ // a non unified cache gives each sequence its own stream, with stream local cell indices
+ if (!unified && s >= (llama_seq_id) n_stream_kv) {
+ sq = kpool_layout::seq();
+ continue;
+ }
+
+ const auto & cells = mem_idx->get_cells(unified ? 0 : s);
+ const auto & sp = cells.seq_pos_get(s);
+
+ sq.strm = unified ? 0 : mem_idx->get_stream(s);
+
+ size_t n_kept = 0;
+ if (mem_idx_stale[s] == POS_CLEAN && !sq.cells.empty() && !sp.empty() &&
+ sq.pos_min == sp.begin()->first) {
+ n_kept = sq.cells.size();
+ for (auto it = sp.upper_bound(sq.cells.back()); it != sp.end(); ++it) {
+ sq.cells.push_back(*it);
+ }
+ }
+
+ // the appended tail accounts for every cell only if nothing before it was dropped, but an edit can
+ // regroup a sequence without changing its cell count, so a stale sequence must rebuild regardless
+ if (sq.cells.size() != sp.size() || mem_idx_stale[s] != POS_CLEAN) {
+ sq.cells.assign(sp.begin(), sp.end());
+ sq.pools.clear();
+ sq.j_next = 0;
+ sq.shared = false;
+ sq.pos_min = sp.empty() ? 0 : sp.begin()->first;
+ n_kept = 0;
+ }
+
+ // sharing starts with a seq_cp; it ends with an edit, or a seq_rm/state_drop/state_read that frees the
+ // shared cells - each stales every sequence so the rebuild above re-derives it, so once set it holds
+ // until then and the rescan can be skipped
+ if (unified && !sq.shared) {
+ for (size_t j = n_kept; j < sq.cells.size(); ++j) {
+ if (cells.seq_count(sq.cells[j].second) > 1) {
+ sq.shared = true;
+ break;
+ }
+ }
+ }
+
+ // Pools start at the first valid token
+ size_t j = sq.j_next;
+ while (j + kpool <= sq.cells.size()) {
+ const llama_pos p0 = sq.cells[j].first;
+ if ((p0 - sq.pos_min) % (llama_pos) kpool != 0) {
+ ++j;
+ continue;
+ }
+ bool ok = true;
+ for (uint32_t k = 1; k < kpool; ++k) {
+ if (sq.cells[j + k].first != p0 + (llama_pos) k) {
+ ok = false;
+ break;
+ }
+ }
+ if (ok) {
+ sq.pools.push_back((uint32_t) j);
+ j += kpool;
+ } else {
+ ++j;
+ }
+ }
+ sq.j_next = j;
+
+ lay.n_pool_real += (uint32_t) sq.pools.size();
+ lay.cache_safe = lay.cache_safe && !sq.shared;
+ }
+
+ return lay;
+}
+
llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_status status) :
llama_memory_hybrid_context(status) {}
@@ -619,7 +827,19 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hy
ns_ubatch(mem->get_mem_idx() == nullptr ?
std::vector<uint32_t>() : std::vector<uint32_t>{ mem->get_mem_idx()->get_n_stream() }),
ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :
- new llama_kv_cache_context(mem->get_mem_idx())) {}
+ new llama_kv_cache_context(mem->get_mem_idx())) {
+ if (kpool_track()) {
+ mem->kpool_layout_update();
+ auto st = kpool_build_sizes();
+ const auto * idx = mem->get_mem_idx();
+ const uint64_t n_pool_max = uint64_t(idx->get_size() / mem->get_kpool()) * idx->get_n_seq_max();
+ GGML_ASSERT(n_pool_max <= UINT32_MAX - 64);
+ st.n_pool_real = std::max(st.n_pool_real, uint32_t(n_pool_max));
+ st.n_new = st.n_pool_real;
+ kpool_st = std::make_unique<kpool_state>(std::move(st));
+ i_kpool = 0;
+ }
+}
llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(
llama_memory_hybrid_idx * mem,
@@ -641,9 +861,19 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(
mem(mem),
ns_ubatch(llama_memory_hybrid_idx_ns(sinfos_idx)),
ctx_idx(mem->get_mem_idx() == nullptr ? nullptr :
- new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {}
+ new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {
+ // Sequence edits force the touched positions to re-pool.
+ mem_idx_stale_batch = mem->mem_idx_stale_get();
+}
+
+llama_memory_hybrid_idx_context::~llama_memory_hybrid_idx_context() = default;
bool llama_memory_hybrid_idx_context::next() {
+ // Clear only after a successful ubatch.
+ if (i_cur == 0 && mem != nullptr) {
+ mem->mem_idx_stale_clear();
+ }
+
if (ctx_idx) {
ctx_idx->next();
}
@@ -660,9 +890,24 @@ bool llama_memory_hybrid_idx_context::apply() {
res = res & ctx_idx->apply();
}
+ // Extend the pool layout with this ubatch's cells, then pick what it must re-pool.
+ if (res && kpool_track()) {
+ mem->kpool_layout_update();
+ if (!kpool_st) {
+ kpool_st = std::make_unique<kpool_state>();
+ }
+ kpool_build_state(get_ubatch());
+ i_kpool = i_cur;
+ }
+
return res;
}
+bool llama_memory_hybrid_idx_context::kpool_track() const {
+ // Derived from mem instead of being cached.
+ return mem != nullptr && mem->get_mem_idx() != nullptr && mem->get_kpool() > 0 && !ns_ubatch.empty();
+}
+
const llama_kv_cache_context * llama_memory_hybrid_idx_context::get_idx() const {
return static_cast<const llama_kv_cache_context *>(ctx_idx.get());
}
@@ -686,3 +931,357 @@ void llama_memory_hybrid_idx_context::set_input_qsa(
mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias, causal_attn);
}
+
+llama_memory_hybrid_idx_context::kpool_access::kpool_access(ggml_context * ctx, ggml_tensor * k, int64_t n_embd) : ctx(ctx) {
+ GGML_ASSERT(k->ne[0] == 3*n_embd);
+
+ const int64_t n_cells = k->ne[1]*k->ne[2];
+
+ // Pool indices can refer to other streams. Revisit these full-storage views if that changes:
+ // https://github.com/ggml-org/llama.cpp/pull/27773#discussion_r4130905603
+ key_gate = ggml_view_2d(ctx, k, 2*n_embd, n_cells, k->nb[1], 0);
+ pooled = ggml_view_2d(ctx, k, n_embd, n_cells, k->nb[1], ggml_row_size(k->type, 2*n_embd));
+}
+
+ggml_tensor * llama_memory_hybrid_idx_context::kpool_access::gather_key_gate(ggml_tensor * idxs) const {
+ return ggml_get_rows(ctx, key_gate, idxs);
+}
+
+ggml_tensor * llama_memory_hybrid_idx_context::kpool_access::scatter_pooled(ggml_tensor * values, ggml_tensor * idxs) const {
+ return ggml_set_rows(ctx, pooled, values, idxs);
+}
+
+ggml_tensor * llama_memory_hybrid_idx_context::kpool_access::gather_pooled(ggml_tensor * idxs) const {
+ return ggml_get_rows(ctx, pooled, idxs);
+}
+
+llama_memory_hybrid_idx_context::kpool_access llama_memory_hybrid_idx_context::get_kpool_access(
+ ggml_context * ctx, int32_t il, int64_t n_embd) const {
+ GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr);
+
+ return kpool_access(ctx, mem->get_mem_idx()->get_k_storage(il), n_embd);
+}
+
+ggml_tensor * llama_memory_hybrid_idx_context::gather_mla_rows(
+ ggml_context * ctx, ggml_tensor * idxs, int64_t n_rows, int64_t n_embd, int32_t il) const {
+ GGML_ASSERT(mem != nullptr);
+ ggml_tensor * k = mem->get_mem_attn()->get_k_storage(il);
+ GGML_ASSERT(k->ne[0] == n_embd);
+
+ ggml_tensor * rows = ggml_view_2d(ctx, k, k->ne[0], k->ne[1]*k->ne[2], k->nb[1], 0);
+ return ggml_get_rows(ctx, rows, ggml_reshape_1d(ctx, idxs, n_rows));
+}
+
+// k-pool DSA indexer (glm5-next)
+
+// Sizes only, used by the full cache context so get_n_kpool() works during graph reserve.
+llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kpool_build_sizes() const {
+ const auto & lay = mem->kpool_layout_get();
+
+ kpool_state st;
+ st.n_pool_real = lay.n_pool_real;
+ st.cache_safe = lay.cache_safe;
+
+ return st;
+}
+
+// Which pools this ubatch must re-pool.
+// Pool cache lifecycle:
+// 1. cpy_k writes each token's key | gate into its idx cache row, pooled slot are zeroed.
+// 2. This marks the pools the ubatch touches or completes as new, during decode that's one pool every kpool tokens, zero elsewise.
+// 3. The graph pools only the new pools and set_rows each result into the pooled slot of the pool's last member row.
+// 4. All pools are gathered in one get_rows via pool_cells, fresh ones just written, older ones from whatever batch last wrote them.
+// A seq_* edit regroups the pools from the edited position on, so it stales them and the first ubatch of the next batch
+// rebuilds them from the still-valid key | gate rows, rewriting the (possibly different) rep rows.
+// Orphaned pooled slots are never cleared, a slot is only ever read through pool_cells, which follows the current grouping.
+void llama_memory_hybrid_idx_context::kpool_build_state(const llama_ubatch & ubatch) {
+ const auto & lay = mem->kpool_layout_get();
+ auto & st = *kpool_st;
+
+ st.n_pool_real = lay.n_pool_real;
+ st.cache_safe = lay.cache_safe;
+ st.n_new = 0;
+ if (++st.generation == 0) {
+ std::fill(st.is_new.begin(), st.is_new.end(), 0);
+ st.generation = 1;
+ }
+ st.is_new.resize(lay.n_pool_real, 0);
+
+ auto mark = [&](uint32_t ip) {
+ if (st.is_new[ip] != st.generation) {
+ st.is_new[ip] = st.generation;
+ ++st.n_new;
+ }
+ };
+
+ const uint32_t kpool = mem->get_kpool();
+ std::array<uint32_t, LLAMA_MAX_SEQ> pool_start;
+ uint32_t ip = 0;
+ for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
+ const auto & sq = lay.seqs[s];
+ pool_start[s] = ip;
+ ip += (uint32_t) sq.pools.size();
+
+ if (!st.cache_safe) {
+ continue;
+ }
+
+ // A sequence edit invalidates only pools ending after the edited position.
+ const llama_pos stale_from = i_cur == 0 ?
+ mem_idx_stale_batch[s] : llama_memory_hybrid_idx::POS_CLEAN;
+ if (stale_from == llama_memory_hybrid_idx::POS_CLEAN) {
+ continue;
+ }
+
+ auto first = std::lower_bound(sq.pools.begin(), sq.pools.end(), stale_from,
+ [&](uint32_t j, llama_pos p) { return sq.cells[j].first + (llama_pos) kpool <= p; });
+ for (auto it = first; it != sq.pools.end(); ++it) {
+ mark(pool_start[s] + (uint32_t) (it - sq.pools.begin()));
+ }
+ }
+ GGML_ASSERT(ip == st.is_new.size());
+
+ if (!st.cache_safe) {
+ std::fill(st.is_new.begin(), st.is_new.end(), st.generation);
+ st.n_new = st.n_pool_real;
+ return;
+ }
+
+ for (uint32_t i = 0; i < ubatch.n_tokens; ++i) {
+ const llama_pos p = ubatch.pos[i];
+ for (int32_t k = 0; k < ubatch.n_seq_id[i]; ++k) {
+ const llama_seq_id s = ubatch.seq_id[i][k];
+ const auto & sq = lay.seqs[s];
+ auto it = std::upper_bound(sq.pools.begin(), sq.pools.end(), p,
+ [&](llama_pos pos, uint32_t j) { return pos < sq.cells[j].first; });
+ if (it == sq.pools.begin()) {
+ continue;
+ }
+ --it;
+ if (p < sq.cells[*it].first + (llama_pos) kpool) {
+ mark(pool_start[s] + (uint32_t) (it - sq.pools.begin()));
+ }
+ }
+ }
+}
+
+const llama_memory_hybrid_idx_context::kpool_state & llama_memory_hybrid_idx_context::kpool_cur() const {
+ GGML_ASSERT(kpool_st != nullptr && i_kpool == i_cur && "k-pool state read before apply()");
+
+ return *kpool_st;
+}
+
+uint32_t llama_memory_hybrid_idx_context::get_n_kpool() const {
+ return kpool_pad(kpool_cur().n_pool_real);
+}
+
+uint32_t llama_memory_hybrid_idx_context::get_n_kpool_new() const {
+ return kpool_cur().n_new;
+}
+
+bool llama_memory_hybrid_idx_context::get_kpool_cache_safe() const {
+ return kpool_cur().cache_safe;
+}
+
+void llama_memory_hybrid_idx_context::set_input_kpool(ggml_tensor * pool_cells, ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs,
+ ggml_tensor * gather_mask, bool gather, ggml_tensor * new_pool_idxs, ggml_tensor * new_pool_rep,
+ const llama_ubatch * ubatch) const {
+ GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr);
+ GGML_ASSERT(ggml_backend_buffer_is_host(pool_cells->buffer));
+ GGML_ASSERT(ggml_backend_buffer_is_host(pool_idxs->buffer));
+ GGML_ASSERT(ggml_backend_buffer_is_host(pool_mask->buffer));
+ GGML_ASSERT(ggml_backend_buffer_is_host(tail_idxs->buffer));
+
+ const uint32_t kpool = mem->get_kpool();
+ const uint32_t n_kv = get_idx()->get_n_kv();
+
+ const auto & st = kpool_cur();
+ const auto & lay = mem->kpool_layout_get();
+
+ const uint32_t n_tokens = ubatch->n_tokens;
+ const uint32_t n_pool = (uint32_t) pool_cells->ne[0];
+ const uint32_t n_new = st.n_new;
+ // the graph always pools at least one entry, see build_inp_kpool
+ const uint32_t n_new_g = std::max(n_new, 1u);
+
+ GGML_ASSERT(n_pool == kpool_pad(st.n_pool_real));
+ GGML_ASSERT(st.is_new.size() == st.n_pool_real);
+ GGML_ASSERT(pool_mask->ne[0] == (int64_t) n_pool && pool_mask->ne[1] == (int64_t) n_tokens);
+ GGML_ASSERT(tail_idxs->ne[0] == (int64_t) kpool - 1 && tail_idxs->ne[1] == (int64_t) n_tokens);
+ GGML_ASSERT(pool_idxs->ne[0] == (int64_t) kpool && pool_idxs->ne[1] == (int64_t) n_pool);
+ GGML_ASSERT(st.cache_safe == (new_pool_rep != nullptr));
+ GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_idxs->buffer));
+ GGML_ASSERT(new_pool_idxs->ne[0] == (int64_t) kpool && new_pool_idxs->ne[1] == (int64_t) n_new_g);
+ if (new_pool_rep != nullptr) {
+ GGML_ASSERT(ggml_backend_buffer_is_host(new_pool_rep->buffer));
+ GGML_ASSERT(new_pool_rep->ne[0] == (int64_t) n_new_g);
+ }
+
+ const uint32_t kv_size = mem->get_mem_idx()->get_size();
+ const uint32_t n_stream_kv = mem->get_mem_idx()->get_n_stream();
+
+ auto gcell = [&](const llama_memory_hybrid_idx::kpool_layout::seq & sq, uint32_t cell) {
+ return (int64_t) sq.strm*kv_size + cell;
+ };
+
+ // Sequences present in this ubatch, pools of absent sequences must fall on the scatter sentinel row.
+ std::vector<uint8_t> seq_in_ub(LLAMA_MAX_SEQ, 0);
+ for (uint32_t i = 0; i < n_tokens; ++i) {
+ for (int32_t k = 0; k < ubatch->n_seq_id[i]; ++k) {
+ seq_in_ub[ubatch->seq_id[i][k]] = 1;
+ }
+ }
+
+ // Use the first ubatch cell for padded gathers.
+ int64_t dummy_cell = 0;
+ {
+ const llama_seq_id s = ubatch->seq_id[0][0];
+ const auto & sq = lay.seqs[s];
+ auto it = std::lower_bound(sq.cells.begin(), sq.cells.end(), std::make_pair(ubatch->pos[0], 0u));
+ GGML_ASSERT(it != sq.cells.end() && it->first == ubatch->pos[0]);
+ dummy_cell = gcell(sq, it->second);
+ }
+
+ // Gather maps padding to a real cell and masks it separately.
+ const int32_t sentinel = gather ? (int32_t) dummy_cell : (int32_t) n_kv;
+
+ float * gm = nullptr;
+ uint32_t n_sel = 0;
+ uint32_t n_top = 0; // Pools per token in the selection.
+ if (gather_mask != nullptr) {
+ GGML_ASSERT(ggml_backend_buffer_is_host(gather_mask->buffer));
+ GGML_ASSERT(gather_mask->type == GGML_TYPE_F32);
+ GGML_ASSERT(gather_mask->ne[3] == (int64_t) n_tokens && gather_mask->ne[1] == 1 && gather_mask->ne[2] == 1);
+ n_sel = (uint32_t) gather_mask->ne[0];
+ // The tail slots, when selected, are the n_sel % kpool != 0 remainder.
+ n_top = n_sel / kpool;
+ GGML_ASSERT(n_sel % kpool == 0 || n_sel % kpool == kpool - 1);
+ gm = (float *) gather_mask->data;
+ }
+
+ // pools are laid out per sequence
+ std::vector<uint32_t> seq_pool_start(LLAMA_MAX_SEQ, 0);
+ std::vector<llama_pos> pool_end;
+ pool_end.reserve(n_pool);
+
+ int32_t * pcell = (int32_t *) pool_cells->data;
+ int32_t * pidx = (int32_t *) pool_idxs->data;
+ int32_t * nidx = (int32_t *) new_pool_idxs->data;
+ int64_t * nrep = new_pool_rep != nullptr ? (int64_t *) new_pool_rep->data : nullptr;
+
+ uint32_t i_new = 0;
+ for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) {
+ const auto & sq = lay.seqs[s];
+ seq_pool_start[s] = (uint32_t) pool_end.size();
+
+ const bool inert = !gather && n_stream_kv > 1 && !seq_in_ub[s];
+
+ for (size_t pi = 0; pi < sq.pools.size(); ++pi) {
+ const uint32_t j = sq.pools[pi];
+ const uint32_t ip = (uint32_t) pool_end.size();
+ GGML_ASSERT(ip + 1 < n_pool);
+
+ // The pooled key lives in the last member's row.
+ const uint32_t rep = sq.cells[j + kpool - 1].second;
+ pcell[ip] = (int32_t) gcell(sq, rep);
+
+ for (uint32_t k = 0; k < kpool; ++k) {
+ pidx[(size_t) ip*kpool + k] = inert ? sentinel :
+ (int32_t) (gather ? gcell(sq, sq.cells[j + k].second) : (int64_t) sq.cells[j + k].second);
+ }
+
+ if (st.is_new[ip] == st.generation) {
+ GGML_ASSERT(i_new < n_new);
+ for (uint32_t k = 0; k < kpool; ++k) {
+ nidx[(size_t) i_new*kpool + k] = (int32_t) gcell(sq, sq.cells[j + k].second);
+ }
+ if (nrep != nullptr) {
+ nrep[i_new] = gcell(sq, rep);
+ }
+ ++i_new;
+ }
+
+ pool_end.push_back(sq.cells[j + kpool - 1].first);
+ }
+ }
+ GGML_ASSERT(i_new == n_new);
+
+ // A ubatch that completes no pool re-pools the cell of its first token. That cell cannot belong to
+ // a complete pool here, else the pool would be marked new, so the write never touches a cached key.
+ if (n_new == 0) {
+ for (uint32_t k = 0; k < kpool; ++k) {
+ nidx[k] = (int32_t) dummy_cell;
+ }
+ if (nrep != nullptr) {
+ nrep[0] = dummy_cell;
+ }
+ }
+
+ const uint32_t n_pool_real = (uint32_t) pool_end.size();
+ for (uint32_t ip = n_pool_real; ip < n_pool; ++ip) {
+ pcell[ip] = (int32_t) dummy_cell; // pool_cells always addresses the K storage
+ for (uint32_t k = 0; k < kpool; ++k) {
+ pidx[(size_t) ip*kpool + k] = sentinel;
+ }
+ }
+
+ // a pool is visible when it belongs to the token's sequence and ends at or before it
+ auto fill_mask = [&](auto * data) {
+ using T = std::remove_pointer_t<decltype(data)>;
+ const T keep = llama_cast<T>(0.0f);
+ const T drop = llama_cast<T>(-INFINITY);
+
+ for (uint32_t i = 0; i < n_tokens; ++i) {
+ const llama_seq_id s = ubatch->seq_id[i][0];
+ const llama_pos p = ubatch->pos[i];
+
+ T * row = data + (size_t) i*n_pool;
+ std::fill(row, row + n_pool, drop);
+
+ const uint32_t p0 = seq_pool_start[s];
+ const uint32_t p1 = p0 + (uint32_t) lay.seqs[s].pools.size();
+ const uint32_t nv = (uint32_t) (std::upper_bound(pool_end.begin() + p0, pool_end.begin() + p1, p) - (pool_end.begin() + p0));
+ std::fill(row + p0, row + p0 + nv, keep);
+
+ // Finite visible pools occupy the first min(nv, n_top) ranked slots.
+ if (gm != nullptr) {
+ const uint32_t nvc = std::min(nv, n_top);
+ float * grow = gm + (size_t) i*n_sel;
+ std::fill(grow, grow + (size_t) nvc*kpool, 0.0f);
+ std::fill(grow + (size_t) nvc*kpool, grow + (size_t) n_top*kpool, -INFINITY);
+ }
+ }
+ };
+ if (pool_mask->type == GGML_TYPE_F16) {
+ fill_mask((ggml_fp16_t *) pool_mask->data);
+ } else {
+ fill_mask((float *) pool_mask->data);
+ }
+
+ int32_t * tidx = (int32_t *) tail_idxs->data;
+ for (uint32_t i = 0; i < n_tokens; ++i) {
+ const llama_seq_id s = ubatch->seq_id[i][0];
+ const llama_pos p = ubatch->pos[i];
+ const auto & sq = lay.seqs[s];
+
+ const uint32_t n_tail = (uint32_t) ((p - sq.pos_min + 1) % (llama_pos) kpool);
+
+ for (uint32_t k = 0; k < kpool - 1; ++k) {
+ int32_t cell = sentinel;
+ bool real = false;
+ if (k < n_tail) {
+ const llama_pos pt = p - (llama_pos) k;
+ auto it = std::lower_bound(sq.cells.begin(), sq.cells.end(), std::make_pair(pt, 0u));
+ if (it != sq.cells.end() && it->first == pt) {
+ cell = (int32_t) (gather ? gcell(sq, it->second) : (int64_t) it->second);
+ real = true;
+ }
+ }
+ tidx[(size_t) i*(kpool - 1) + k] = cell;
+
+ if (gm != nullptr && n_sel % kpool != 0) {
+ gm[(size_t) i*n_sel + (size_t) n_top*kpool + k] = real ? 0.0f : -INFINITY;
+ }
+ }
+ }
+}
diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h
index 3c3a788ac..66953cacf 100644
--- a/src/llama-memory-hybrid-idx.h
+++ b/src/llama-memory-hybrid-idx.h
@@ -2,6 +2,8 @@
#include "llama-memory-hybrid.h"
+#include <array>
+#include <limits>
#include <memory>
#include <vector>
@@ -41,7 +43,8 @@ public:
/* the indexer cache exists only if this is given */
const layer_filter_cb & filter_idx);
- ~llama_memory_hybrid_idx() = default;
+ // Defined out of line because kpool_layout is incomplete here.
+ ~llama_memory_hybrid_idx();
//
// llama_memory_i
@@ -90,6 +93,32 @@ public:
ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio,
bool blk_bias, bool causal_attn) const;
+ // The model's indexer pool size.
+ uint32_t get_kpool() const { return hparams_idx.indexer_kpool; }
+
+ // Which cells of a sequence make up which pool of kpool consecutive positions.
+ // It is kept here because it outlives the batch: pools are fixed by the positions relative to the
+ // sequence's first one, so a ubatch only ever appends to it. Sequence edits drop it, see mem_idx_stale.
+ struct kpool_layout;
+
+ const kpool_layout & kpool_layout_update();
+ const kpool_layout & kpool_layout_get() const;
+
+ // The pooled keys persist in the idx cache across batches. A sequence edit can regroup the pools
+ // from some position on, which stales every pooled key at or after it. POS_CLEAN means none.
+ using stale_pos_t = std::array<llama_pos, LLAMA_MAX_SEQ>;
+
+ static constexpr llama_pos POS_CLEAN = std::numeric_limits<llama_pos>::max();
+
+ static stale_pos_t stale_pos_clean() {
+ stale_pos_t res;
+ res.fill(POS_CLEAN);
+ return res;
+ }
+
+ const stale_pos_t & mem_idx_stale_get() const { return mem_idx_stale; }
+ void mem_idx_stale_clear() { mem_idx_stale.fill(POS_CLEAN); }
+
private:
// forget seq_id (all of it if seq_id < 0) in every cache at once, so a failed restore cannot leave the caches out of step
// seq_id < 0 drops the whole context, as the caches themselves do on a failed restore
@@ -100,10 +129,40 @@ private:
llama_hparams hparams_idx;
const std::unique_ptr<llama_kv_cache> mem_idx;
+
+ // unique_ptr because kpool_layout is incomplete here
+ std::unique_ptr<kpool_layout> kpool_lay;
+
+ // whether the current layout has cells shared between sequences (kpool_layout is incomplete here, so out of line)
+ bool kpool_layout_shared() const;
+
+ // seq_id < 0 stales every sequence, p0 < 0 stales the sequence from its first position
+ void mem_idx_stale_set(llama_seq_id seq_id, llama_pos p0);
+
+ // the position an edit at p0 stales the sequence from
+ llama_pos mem_idx_stale_pos(llama_seq_id seq_id, llama_pos p0) const;
+
+ stale_pos_t mem_idx_stale = stale_pos_clean();
};
class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context {
public:
+ class kpool_access {
+ public:
+ ggml_tensor * gather_key_gate(ggml_tensor * idxs) const;
+ ggml_tensor * scatter_pooled(ggml_tensor * values, ggml_tensor * idxs) const;
+ ggml_tensor * gather_pooled(ggml_tensor * idxs) const;
+
+ private:
+ friend class llama_memory_hybrid_idx_context;
+
+ kpool_access(ggml_context * ctx, ggml_tensor * k, int64_t n_embd);
+
+ ggml_context * ctx;
+ ggml_tensor * key_gate;
+ ggml_tensor * pooled;
+ };
+
using slot_info_vec_t = llama_kv_cache::slot_info_vec_t;
// used for errors
@@ -125,7 +184,7 @@ public:
slot_info_vec_t sinfos_idx,
std::vector<llama_ubatch> ubatches);
- ~llama_memory_hybrid_idx_context() = default;
+ ~llama_memory_hybrid_idx_context(); // Defined out of line because kpool_state is incomplete here.
//
// llama_memory_context_i
@@ -144,12 +203,21 @@ public:
// streams in the current slot info, the `ns` of get_k/get_v; 1 if unified
uint32_t get_n_stream() const;
+ // glm5-next, complete pools of kpool consecutive positions per sequence, scored as whole pools.
+ uint32_t get_n_kpool () const; // Padded pool count, where the last pool is always unused.
+ uint32_t get_n_kpool_new() const; // Exact count of pools completed by the current ubatch.
+ bool get_kpool_cache_safe() const;
+ kpool_access get_kpool_access(ggml_context * ctx, int32_t il, int64_t n_embd) const;
+ ggml_tensor * gather_mla_rows(ggml_context * ctx, ggml_tensor * idxs, int64_t n_rows, int64_t n_embd, int32_t il) const;
+ void set_input_kpool(ggml_tensor * pool_cells, ggml_tensor * pool_idxs, ggml_tensor * pool_mask, ggml_tensor * tail_idxs,
+ ggml_tensor * gather_mask, bool gather, ggml_tensor * new_pool_idxs, ggml_tensor * new_pool_rep,
+ const llama_ubatch * ubatch) const;
void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos,
ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio,
bool blk_bias, bool causal_attn) const;
private:
- const llama_memory_hybrid_idx * mem = nullptr;
+ llama_memory_hybrid_idx * mem = nullptr;
// streams per ubatch, read from the slot infos before ctx_idx takes them
// declared first, so it is initialised while sinfos_idx is still intact
@@ -160,4 +228,22 @@ private:
// mirrors the base class's ubatch cursor, which is private there
size_t i_cur = 0;
+
+ // Which pools of the layout this ubatch must re-pool. The layout itself belongs to the memory.
+ struct kpool_state;
+ kpool_state kpool_build_sizes() const;
+ void kpool_build_state(const llama_ubatch & ubatch);
+ const kpool_state & kpool_cur() const;
+
+ // unique_ptr because kpool_state is incomplete here.
+ std::unique_ptr<kpool_state> kpool_st;
+
+ // The ubatch kpool_st was built for, guards against reads before apply.
+ size_t i_kpool = SIZE_MAX;
+
+ // Whether this context tracks k-pool states.
+ bool kpool_track() const;
+
+ // Positions each sequence must re-pool from, cleared only after the first ubatch succeeds
+ llama_memory_hybrid_idx::stale_pos_t mem_idx_stale_batch = llama_memory_hybrid_idx::stale_pos_clean();
};
diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp
index 6f58dd150..eeecadac0 100644
--- a/src/llama-model-saver.cpp
+++ b/src/llama-model-saver.cpp
@@ -309,6 +309,8 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
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_KPOOL, hparams.indexer_kpool);
+ add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, hparams.indexer_kpool_select_tail);
add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);
add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true);
add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true);
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 151a3a2a8..6ca536f22 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -340,6 +340,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_kimi_linear(params);
case LLM_ARCH_KIMI_K3:
return new llama_model_kimi_k3(params);
+ case LLM_ARCH_GLM5_NEXT:
+ return new llama_model_glm5_next(params);
case LLM_ARCH_STEP35:
return new llama_model_step35(params);
case LLM_ARCH_SPARK2_5:
@@ -1013,6 +1015,7 @@ const char * llm_type_name(llm_type type) {
case LLM_TYPE_685B_A37B: return "685B.A37B";
case LLM_TYPE_744B_A40B: return "744B.A40B";
case LLM_TYPE_2_8T_A50B: return "2.8T.A50B";
+ case LLM_TYPE_320B_A18B: return "320B.A18B";
case LLM_TYPE_E2B: return "E2B";
case LLM_TYPE_E4B: return "E4B";
default: return "?B";
@@ -2437,6 +2440,50 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
nullptr);
}
} break;
+ case LLM_ARCH_GLM5_NEXT:
+ {
+ // KDA layers are recurrent, the DSA layers use a K-only MLA cache plus an indexer cache.
+ // tThe Nextn block is never attended by the trunk graph
+ llama_memory_hybrid_idx::layer_filter_cb filter_attn = [&](uint32_t il) {
+ return il < hparams.n_layer() && !hparams.is_recr(il);
+ };
+ llama_memory_hybrid_idx::layer_filter_cb filter_idx = [&](uint32_t il) {
+ return il < hparams.n_layer() && !hparams.is_recr(il) && hparams.is_indexer_full(il);
+ };
+ llama_memory_hybrid_idx::layer_filter_cb filter_recr = [&](uint32_t il) {
+ return il < hparams.n_layer() && hparams.is_recr(il);
+ };
+
+ // the draft head is a single DSA layer
+ if (params.ctx_type == LLAMA_CONTEXT_TYPE_MTP) {
+ if (hparams.n_layer_nextn == 0) {
+ throw std::runtime_error("GLM5-Next MTP requires the NextN block, convert without --no-mtp");
+ }
+ filter_attn = [&](uint32_t il) { return il >= hparams.n_layer(); };
+ filter_idx = [&](uint32_t il) { return il >= hparams.n_layer(); };
+ filter_recr = [&](uint32_t) { return false; };
+ }
+
+ res = new llama_memory_hybrid_idx(
+ /* model */ *this,
+ /* attn_type_k */ params.type_k,
+ /* attn_type_v */ params.type_v,
+ /* attn_v_trans */ !cparams.flash_attn,
+ /* attn_kv_size */ cparams.n_ctx_seq,
+ /* attn_n_pad */ 1,
+ /* attn_n_swa */ hparams.n_swa,
+ /* attn_swa_type */ hparams.swa_type,
+ /* recurrent_type_r */ GGML_TYPE_F32,
+ /* recurrent_type_s */ GGML_TYPE_F32,
+ /* recurrent_rs_size */ std::max((uint32_t) 1, cparams.n_seq_max),
+ /* n_seq_max */ cparams.n_seq_max,
+ /* n_rs_seq */ cparams.n_rs_seq,
+ /* offload */ cparams.offload_kqv,
+ /* unified */ cparams.kv_unified,
+ /* filter_attn */ std::move(filter_attn),
+ /* filter_recr */ std::move(filter_recr),
+ /* filter_idx */ std::move(filter_idx));
+ } break;
case LLM_ARCH_HY_V4:
{
if (hparams.indexer_top_k == 0) {
@@ -2991,6 +3038,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_NEMOTRON_H_MOE:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_KIMI_K3:
+ case LLM_ARCH_GLM5_NEXT:
return LLAMA_ROPE_TYPE_NONE;
// use what we call a normal RoPE, operating on pairs of consecutive head values
diff --git a/src/llama-model.h b/src/llama-model.h
index a0f9f1142..9c18ef045 100644
--- a/src/llama-model.h
+++ b/src/llama-model.h
@@ -154,6 +154,7 @@ enum llm_type {
LLM_TYPE_685B_A37B, // DeepSeek V3.2
LLM_TYPE_744B_A40B, // GLM-5
LLM_TYPE_2_8T_A50B, // Kimi-K3
+ LLM_TYPE_320B_A18B, // GLM-5.3-Flash
LLM_TYPE_E2B,
LLM_TYPE_E4B,
};
@@ -561,6 +562,10 @@ struct llama_layer {
struct ggml_tensor * indexer_attn_k = nullptr;
struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias
+ // glm5-next k-pool indexer
+ struct ggml_tensor * indexer_kpool_gate = nullptr;
+ struct ggml_tensor * indexer_kpool_ape = nullptr;
+
// MSA
struct ggml_tensor * index_q_proj = nullptr;
struct ggml_tensor * index_k_proj = nullptr;
diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp
index 34ff25db5..821b4798c 100644
--- a/src/llama-quant.cpp
+++ b/src/llama-quant.cpp
@@ -328,6 +328,24 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param
quantize &= name.find("indexer.k_proj.weight") == std::string::npos;
quantize &= name.find("indexer.q_proj.weight") == std::string::npos;
+ // glm5-next
+ if (arch == LLM_ARCH_GLM5_NEXT) {
+ quantize &= name.find("hc_") == std::string::npos;
+ quantize &= name.find("indexer.attn_q_b") == std::string::npos;
+ quantize &= name.find("indexer.attn_k") == std::string::npos;
+ quantize &= name.find("indexer.proj") == std::string::npos;
+ quantize &= name.find("indexer_compressor_gate") == std::string::npos;
+ quantize &= name.find("indexer_compressor_ape") == std::string::npos;
+ quantize &= name.find("ssm_f_a.weight") == std::string::npos;
+ quantize &= name.find("ssm_f_b.weight") == std::string::npos;
+ quantize &= name.find("ssm_g_a.weight") == std::string::npos;
+ quantize &= name.find("ssm_g_b.weight") == std::string::npos;
+ quantize &= name.find("ssm_beta.weight") == std::string::npos;
+ quantize &= name.find("attn_kv_a_mqa.weight") == std::string::npos;
+ quantize &= name.find("attn_k_b.weight") == std::string::npos;
+ quantize &= name.find("attn_v_b.weight") == std::string::npos;
+ }
+
// do not quantize RWKV's small yet 2D weights
quantize &= name.find("time_mix_first.weight") == std::string::npos;
quantize &= name.find("time_mix_w0.weight") == std::string::npos;
@@ -451,6 +469,22 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type
return std::make_pair(i_layer, n_layer);
};
+ // by default, for glm5-next, don't let these tensors be quantized below Q8_0
+ if (arch == LLM_ARCH_GLM5_NEXT && (
+ name.find("attn_q_a") != std::string::npos ||
+ name.find("attn_q_b") != std::string::npos ||
+ name.find("nextn.eh_proj") != std::string::npos))
+ {
+ switch (new_type) {
+ case GGML_TYPE_F32:
+ case GGML_TYPE_BF16:
+ case GGML_TYPE_F16:
+ break;
+ default:
+ return GGML_TYPE_Q8_0;
+ }
+ }
+
// for arches that share the same tensor between the token embeddings and the output, we quantize the token embeddings
// with the quantization of the output tensor
if (category == tensor_category::OUTPUT || (qs.has_tied_embeddings && category == tensor_category::TOKEN_EMBD)) {
diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp
index 9fdfe7372..93468b6b4 100644
--- a/src/llama-vocab.cpp
+++ b/src/llama-vocab.cpp
@@ -2283,10 +2283,15 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
pre_type = LLAMA_VOCAB_PRE_TYPE_PORO;
clean_spaces = false;
} else if (
- tokenizer_pre == "glm4" ||
tokenizer_pre == "chatglm-bpe") {
pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4;
special_bos_id = LLAMA_TOKEN_NULL;
+ } else if (
+ tokenizer_pre == "glm4" ||
+ tokenizer_pre == "glm5") {
+ pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4;
+ special_bos_id = LLAMA_TOKEN_NULL;
+ 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 6bf9d3444..a14388725 100644
--- a/src/models/deepseek4.cpp
+++ b/src/models/deepseek4.cpp
@@ -266,15 +266,14 @@ static dsv4_state_tensors dsv4_build_state_snapshot(
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_tensor * x) const {
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);
+ ggml_tensor * acc = ggml_view_2d(ctx0, x, x->ne[0], x->ne[2], x->nb[2], 0);
for (int64_t s = 1; s < hc; ++s) {
- acc = ggml_add(ctx, acc, ggml_view_2d(ctx, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1]));
+ acc = ggml_add(ctx0, acc, ggml_view_2d(ctx0, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1]));
}
- return ggml_scale(ctx, acc, 1.0f/hc);
+ return ggml_scale(ctx0, acc, 1.0f/hc);
}
static ggml_tensor * dsv4_hc_affine(
@@ -1238,7 +1237,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(inpL);
cb(res->t_layer_inp[il], "layer_inp", il);
ggml_build_forward_expand(gf, res->t_layer_inp[il]);
}
@@ -1327,7 +1326,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(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/glm5-next.cpp b/src/models/glm5-next.cpp
new file mode 100644
index 000000000..5b11c77c3
--- /dev/null
+++ b/src/models/glm5-next.cpp
@@ -0,0 +1,1005 @@
+#include "models.h"
+#include "llama-memory-hybrid-idx.h"
+
+// GLM5-Next (GLM-5.3-Flash): hybrid KDA (linear) + nope MLA with a k-pool DSA indexer,
+// mHC residual streams, DeepSeek-style MoE.
+
+void llama_model_glm5_next::load_arch_hparams(llama_model_loader & ml) {
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps, false);
+ ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
+ ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);
+ ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q);
+ ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);
+ ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv);
+ ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda);
+ ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false);
+
+ // the MLA cache holds the compressed latent
+ hparams.n_embd_head_v_full = hparams.n_lora_kv;
+
+ for (uint32_t i = 0; i < hparams.n_layer_all; ++i) {
+ hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0;
+ }
+
+ ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
+ ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
+ ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
+ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false);
+ if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) {
+ hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID;
+ }
+ ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false);
+ if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false)) {
+ hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp;
+ }
+
+ // DSA indexer with k-pool compression
+ 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);
+ ml.get_key(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, hparams.indexer_kpool_select_tail, false);
+ GGML_ASSERT(hparams.indexer_kpool > 1 && hparams.indexer_top_k % hparams.indexer_kpool == 0);
+ std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1);
+ ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false);
+
+ // mHC
+ ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult);
+ ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters);
+ ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps);
+ GGML_ASSERT(hparams.dsv4_hc_mult == 4 && "mHC with hc_mult != 4 is not supported");
+
+ switch (hparams.n_layer()) {
+ case 45: type = LLM_TYPE_320B_A18B; break; // GLM-5.3-Flash
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+}
+
+void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) {
+ LLAMA_LOAD_LOCALS;
+
+ const int64_t hc = hparams.dsv4_hc_mult;
+ const int64_t hc_mix_dim = (2 + hc)*hc;
+
+ // the NextN block is loaded but only used by the MTP graph.
+ // Separated trunk_only/mtp_only handling TODO with DECODER_MTP graph in the MTP follow up
+ int mtp_flags = 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}, 0);
+
+ for (int i = 0; i < n_layer_all; ++i) {
+ auto & layer = layers[i];
+
+ const int flags = (i >= n_layer) ? mtp_flags : 0;
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags);
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags);
+
+ if (i < n_layer) {
+ layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, 0);
+ layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0);
+ layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0);
+ layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, 0);
+ layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0);
+ layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0);
+ }
+
+ const int64_t head_dim = hparams.n_embd_head_kda;
+ const int64_t d_conv = hparams.ssm_d_conv;
+ const int64_t d_inner = head_dim * n_head;
+
+ if (hparams.is_recr(i)) {
+ auto conv = [&](llm_tensor tid) {
+ ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED);
+ return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0);
+ };
+ layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q);
+ layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K);
+ layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V);
+
+ create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0);
+
+ layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0);
+ layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0);
+ layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0);
+
+ layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, 0);
+ layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0);
+
+ layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", i), {n_embd, head_dim}, 0);
+ layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", i), {head_dim, d_inner}, 0);
+ layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0);
+ } else {
+ const int64_t q_lora_rank = hparams.n_lora_q;
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+ const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v = hparams.n_embd_head_v_mla();
+ const int64_t qk_rope_head_dim = hparams.n_rot();
+ const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim;
+
+ layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags);
+ layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags);
+
+ layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags);
+ layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, flags);
+
+ layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, flags);
+ layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, flags);
+ layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, flags);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, flags);
+
+ const int64_t n_indexer_head = hparams.indexer_n_head;
+ const int64_t n_embd_indexer = hparams.indexer_head_size;
+ const int64_t kpool = hparams.indexer_kpool;
+
+ const bool full = i >= n_layer || hparams.is_indexer_full(i);
+ const int iflags = flags | (full ? 0 : TENSOR_NOT_REQUIRED);
+
+ layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_embd_indexer}, iflags);
+ layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {n_embd_indexer}, iflags);
+ layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, n_indexer_head}, iflags);
+ layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, n_embd_indexer}, iflags);
+ layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_indexer_head * n_embd_indexer}, iflags);
+ layer.indexer_kpool_gate = create_tensor(tn(LLM_TENSOR_INDEXER_KPOOL_GATE, "weight", i), {n_embd, n_embd_indexer}, iflags);
+ layer.indexer_kpool_ape = create_tensor(tn(LLM_TENSOR_INDEXER_KPOOL_APE, "weight", i), {n_embd_indexer, kpool}, iflags);
+ }
+
+ if (i < (int) hparams.n_layer_dense_lead) {
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
+ } else {
+ const int64_t n_ff_exp = hparams.n_ff_exp(i);
+ const int64_t n_expert_shared = hparams.n_expert_shared;
+
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags);
+ layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags);
+
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags);
+
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, flags);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags);
+ }
+
+ if (i >= n_layer) {
+ layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, flags);
+ layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, flags);
+ layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, flags);
+ layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED);
+ layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, flags | TENSOR_NOT_REQUIRED);
+ }
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_glm5_next::build_arch_graph(const llm_graph_params & params) const {
+ if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) {
+ throw std::runtime_error("GLM5-Next NextN graph not implemented yet");
+ }
+ return std::make_unique<graph>(*this, params);
+}
+
+// Causal conv1d over one of Q/K/V
+static ggml_tensor * glm5_conv1d(ggml_cgraph * gf, ggml_context * ctx0,
+ ggml_tensor * conv_states_all, ggml_tensor * conv_state_all,
+ int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w,
+ int64_t d_conv, int64_t head_dim, int64_t n_head,
+ int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head,
+ int64_t mem_size, int64_t K_rs) {
+ const int64_t d_inner = head_dim * n_head;
+ const int64_t conv_state_size = (d_conv - 1) * d_inner;
+ const int64_t n_embd_r_total = 3 * conv_state_size;
+
+ ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs,
+ (d_conv - 1) * ggml_element_size(conv_state_all),
+ n_embd_r_total * ggml_element_size(conv_state_all),
+ qkv * conv_state_size * ggml_element_size(conv_state_all));
+
+ ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x);
+ ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs);
+ ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0);
+
+ // group s holds the conv window s tokens back.
+ // [TAG_RECURRENT_ROLLBACK_SPLITS]: the last K_rs tokens must share one ubatch.
+ for (int64_t s = 0; s < K_rs; ++s) {
+ const int64_t s_idx = std::max<int64_t>(0, n_seq_tokens - s);
+ ggml_tensor * conv_x_s = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs,
+ conv_x->nb[1], conv_x->nb[2], s_idx * conv_x->nb[0]);
+ ggml_build_forward_expand(gf,
+ ggml_cpy(ctx0, conv_x_s,
+ ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs,
+ (d_conv - 1) * ggml_element_size(conv_states_all),
+ n_embd_r_total * ggml_element_size(conv_states_all),
+ ((s * mem_size + kv_head) * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all))));
+ }
+
+ ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner);
+ ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight);
+ Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens);
+ Xcur = ggml_silu(ctx0, Xcur);
+
+ return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs);
+}
+
+
+// K-pool indexer inputs
+class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i {
+public:
+ llm_graph_input_kpool(const llama_memory_hybrid_idx_context * mctx, uint32_t kpool) : mctx(mctx), kpool(kpool) {}
+ virtual ~llm_graph_input_kpool() = default;
+
+ void set_input(const llama_ubatch * ubatch) override {
+ mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch);
+ mctx->set_input_kpool(pool_cells, pool_idxs, pool_mask, tail_idxs, gather_mask, gather, new_pool_idxs, new_pool_rep, ubatch);
+ }
+
+ bool can_reuse(const llm_graph_params & params) override {
+ mctx = static_cast<const llama_memory_hybrid_idx_context *>(params.mctx);
+
+ const auto * idx = mctx->get_idx();
+ if (idx == nullptr) {
+ return false;
+ }
+
+ bool res = true;
+
+ res &= k_idxs->ne[0] == params.ubatch.n_tokens;
+ res &= pool_cells->ne[0] == mctx->get_n_kpool();
+ res &= pool_mask->ne[1] == params.ubatch.n_tokens;
+ res &= tail_idxs->ne[1] == params.ubatch.n_tokens;
+ // The scatter mask shape follows n_kv.
+ res &= n_kv == idx->get_n_kv();
+ res &= n_new == std::max(mctx->get_n_kpool_new(), 1u);
+ res &= cache_safe == mctx->get_kpool_cache_safe();
+
+ return res;
+ }
+
+ ggml_tensor * k_idxs = nullptr; // I64 [n_tokens]
+ ggml_tensor * pool_cells = nullptr; // I32 [n_pool] cell caching each pool's pooled key
+ ggml_tensor * pool_idxs = nullptr; // I32 [kpool, n_pool] member cells per pool, n_kv sentinel for the padded pools
+ ggml_tensor * pool_mask = nullptr; // F32/F16 [n_pool, n_tokens]
+ ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens]
+ ggml_tensor * gather_mask = nullptr; // F32 [n_sel, 1, 1, n_tokens]
+ // n_new is never below 1, see build_inp_kpool
+ ggml_tensor * new_pool_idxs = nullptr; // I32 [kpool, n_new] members of the pools completed this ubatch
+ ggml_tensor * new_pool_rep = nullptr; // I64 [n_new] cell to write each new pooled key into
+
+ const llama_memory_hybrid_idx_context * mctx;
+ const uint32_t kpool;
+ uint32_t n_new = 0;
+ uint32_t n_sel = 0;
+ bool cache_safe = true;
+ bool gather = false;
+ uint32_t n_kv = 0;
+};
+
+llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::build_inp_kpool(const llama_memory_hybrid_idx_context * mctx_hyb) {
+ const auto * mctx_idx = mctx_hyb->get_idx();
+ GGML_ASSERT(mctx_idx != nullptr);
+
+ const uint32_t kpool = hparams.indexer_kpool;
+ const uint32_t n_pool = mctx_hyb->get_n_kpool();
+ const uint32_t n_kv = mctx_idx->get_n_kv();
+ // a ubatch that completes no pool still builds one dummy entry, so the graph does not
+ // change shape every kpool tokens
+ const uint32_t n_new = std::max(mctx_hyb->get_n_kpool_new(), 1u);
+ const bool cache_safe = mctx_hyb->get_kpool_cache_safe();
+
+ // the fused lightning indexer wants an f16 mask
+ const auto type_mask = cparams.fused_lid ? GGML_TYPE_F16 : GGML_TYPE_F32;
+
+ auto inp = std::make_unique<llm_graph_input_kpool>(mctx_hyb, kpool);
+
+ inp->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch);
+ inp->pool_cells = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pool);
+ inp->pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_pool);
+ inp->pool_mask = ggml_new_tensor_2d(ctx0, type_mask, n_pool, n_tokens);
+ inp->tail_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool - 1, n_tokens);
+ ggml_set_input(inp->pool_cells);
+ ggml_set_input(inp->pool_idxs);
+ ggml_set_input(inp->pool_mask);
+ ggml_set_input(inp->tail_idxs);
+
+ ggml_build_forward_expand(gf, inp->pool_cells);
+ ggml_build_forward_expand(gf, inp->pool_idxs);
+ ggml_build_forward_expand(gf, inp->pool_mask);
+ ggml_build_forward_expand(gf, inp->tail_idxs);
+
+ inp->n_kv = n_kv;
+
+ // Gather selected latents for small decode batches when n_kv exceeds n_sel.
+ {
+ constexpr int64_t max_ub = 16;
+
+ const int64_t n_top_pool = std::min<int64_t>(n_pool, hparams.indexer_top_k / kpool);
+ const int64_t n_sel = kpool*n_top_pool + (hparams.indexer_kpool_select_tail ? kpool - 1 : 0);
+ inp->n_sel = (uint32_t) n_sel;
+ inp->gather = (int64_t) n_tokens <= max_ub && (int64_t) n_kv > n_sel;
+
+ if (inp->gather) {
+ inp->gather_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens);
+ ggml_set_input(inp->gather_mask);
+ // Keep the mask allocated even when no op reads it, because set_input_kpool always fills it.
+ ggml_build_forward_expand(gf, inp->gather_mask);
+ }
+ }
+
+ inp->n_new = n_new;
+ inp->cache_safe = cache_safe;
+
+ inp->new_pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_new);
+ ggml_set_input(inp->new_pool_idxs);
+ if (cache_safe) {
+ inp->new_pool_rep = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new);
+ ggml_set_input(inp->new_pool_rep);
+ }
+
+ return (llm_graph_input_kpool *) res->add_input(std::move(inp));
+}
+
+static size_t dsv4_elem_offset(const ggml_tensor * t, int64_t i) {
+ return ggml_row_size(t->type, i);
+}
+
+static ggml_tensor * dsv4_view_1d(ggml_context * ctx, ggml_tensor * t, int64_t ne0, int64_t i0) {
+ return ggml_view_1d(ctx, t, ne0, dsv4_elem_offset(t, i0));
+}
+
+static ggml_tensor * dsv4_view_2d(
+ ggml_context * ctx,
+ ggml_tensor * t,
+ int64_t ne0,
+ int64_t ne1,
+ int64_t i0) {
+ return ggml_view_2d(ctx, t, ne0, ne1, t->nb[1], dsv4_elem_offset(t, i0));
+}
+
+ggml_tensor * llama_model_glm5_next::graph::build_hc_mean(ggml_tensor * x) const {
+ const int64_t hc = x->ne[1];
+
+ ggml_tensor * acc = ggml_view_2d(ctx0, x, x->ne[0], x->ne[2], x->nb[2], 0);
+ for (int64_t s = 1; s < hc; ++s) {
+ acc = ggml_add(ctx0, acc, ggml_view_2d(ctx0, x, x->ne[0], x->ne[2], x->nb[2], s*x->nb[1]));
+ }
+ return ggml_scale(ctx0, acc, 1.0f/hc);
+}
+
+static ggml_tensor * dsv4_hc_affine(
+ ggml_context * ctx,
+ ggml_tensor * x,
+ ggml_tensor * scale,
+ ggml_tensor * base) {
+ x = ggml_mul(ctx, x, scale);
+ x = ggml_add(ctx, x, base);
+ return x;
+}
+
+ggml_tensor * llama_model_glm5_next::graph::build_hc_pre(
+ ggml_tensor * x,
+ ggml_tensor * weights,
+ int il) const {
+ GGML_ASSERT(x->ne[0] == n_embd);
+ GGML_ASSERT(x->ne[1] == hparams.dsv4_hc_mult);
+
+ const int64_t hc = hparams.dsv4_hc_mult;
+ const int64_t nt = x->ne[2];
+
+ if (cparams.fused_dsv4_hc_pre && il >= 0) {
+ ggml_tensor * result = ggml_dsv4_hc_pre(ctx0, x, weights);
+ res->add_fused_node({LLM_FUSED_OP_DSV4_HC_PRE, result, il});
+ return result;
+ }
+
+ ggml_tensor * result = nullptr;
+ for (int64_t ih = 0; ih < hc; ++ih) {
+ ggml_tensor * xh = ggml_view_2d(ctx0, x, n_embd, nt, x->nb[2], ih*x->nb[1]);
+ ggml_tensor * wh = ggml_view_2d(ctx0, weights, 1, nt, weights->nb[1], ih*weights->nb[0]);
+ ggml_tensor * cur = ggml_mul(ctx0, xh, wh);
+ result = result ? ggml_add(ctx0, result, cur) : cur;
+ }
+
+ return result;
+}
+
+ggml_tensor * llama_model_glm5_next::graph::build_hc_sinkhorn(
+ ggml_tensor * comb,
+ int il) const {
+ GGML_UNUSED(il);
+
+ // comb is [dst_hc, src_hc, n_tokens]. Sinkhorn follows the reference:
+ // row softmax over dst, one column normalization, then repeated row/column normalization.
+ comb = ggml_soft_max(ctx0, comb);
+
+ ggml_tensor * eps = ggml_new_tensor_1d(ctx0, GGML_TYPE_F32, 1);
+ eps = ggml_fill(ctx0, eps, hparams.dsv4_hc_eps);
+
+ comb = ggml_add(ctx0, comb, eps);
+
+ auto norm_cols = [&]() {
+ ggml_tensor * comb_src_dst = ggml_cont(ctx0, ggml_permute(ctx0, comb, 1, 0, 2, 3));
+ ggml_tensor * col_sum = ggml_sum_rows(ctx0, comb_src_dst);
+ col_sum = ggml_add(ctx0, col_sum, eps);
+ col_sum = ggml_permute(ctx0, col_sum, 1, 0, 2, 3);
+ comb = ggml_div(ctx0, comb, col_sum);
+ };
+
+ auto norm_rows = [&]() {
+ ggml_tensor * row_sum = ggml_sum_rows(ctx0, comb);
+ row_sum = ggml_add(ctx0, row_sum, eps);
+ comb = ggml_div(ctx0, comb, row_sum);
+ };
+
+ norm_cols();
+ for (uint32_t i = 1; i < hparams.dsv4_hc_sinkhorn_iters; ++i) {
+ norm_rows();
+ norm_cols();
+ }
+
+ return comb;
+}
+
+ggml_tensor * llama_model_glm5_next::graph::build_hc_pre(
+ ggml_tensor * x,
+ ggml_tensor * hc_fn,
+ ggml_tensor * hc_scale,
+ ggml_tensor * hc_base,
+ ggml_tensor ** post,
+ ggml_tensor ** comb,
+ int il) const {
+ const int64_t hc = hparams.dsv4_hc_mult;
+ const int64_t hc_dim = hc*n_embd;
+ const int64_t hc_mix_dim = (2 + hc)*hc;
+ const int64_t nt = x->ne[2];
+
+ GGML_ASSERT(hc == 4);
+ GGML_ASSERT(hc_fn->ne[1] == hc_mix_dim);
+
+ ggml_tensor * flat = ggml_reshape_2d(ctx0, x, hc_dim, nt);
+ ggml_tensor * flat_norm = ggml_rms_norm(ctx0, flat, norm_rms_eps);
+ ggml_tensor * mixes = ggml_mul_mat(ctx0, hc_fn, flat_norm);
+ cb(mixes, "hc_mixes", il);
+
+ ggml_tensor * scale_pre = dsv4_view_1d(ctx0, hc_scale, 1, 0);
+ ggml_tensor * scale_post = dsv4_view_1d(ctx0, hc_scale, 1, 1);
+
+ ggml_tensor * base_pre = dsv4_view_1d(ctx0, hc_base, hc, 0);
+ ggml_tensor * base_post = dsv4_view_1d(ctx0, hc_base, hc, hc);
+
+ ggml_tensor * pre = dsv4_view_2d(ctx0, mixes, hc, nt, 0);
+ pre = dsv4_hc_affine(ctx0, pre, scale_pre, base_pre);
+ pre = ggml_sigmoid(ctx0, pre);
+ pre = ggml_scale_bias(ctx0, pre, 1.0f, hparams.dsv4_hc_eps);
+ cb(pre, "hc_pre", il);
+
+ *post = dsv4_view_2d(ctx0, mixes, hc, nt, hc);
+ *post = dsv4_hc_affine(ctx0, *post, scale_post, base_post);
+ *post = ggml_sigmoid(ctx0, *post);
+ *post = ggml_scale(ctx0, *post, 2.0f);
+ cb(*post, "hc_post", il);
+
+ if (cparams.fused_dsv4_hc_comb) {
+ *comb = ggml_dsv4_hc_comb(ctx0, mixes, hc_scale, hc_base, hparams.dsv4_hc_eps,
+ (int32_t) hparams.dsv4_hc_sinkhorn_iters);
+ res->add_fused_node({LLM_FUSED_OP_DSV4_HC_COMB, *comb, il});
+ } else {
+ ggml_tensor * scale_comb = dsv4_view_1d(ctx0, hc_scale, 1, 2);
+ ggml_tensor * base_comb = dsv4_view_1d(ctx0, hc_base, hc*hc, 2*hc);
+
+ *comb = dsv4_view_2d(ctx0, mixes, hc*hc, nt, 2*hc);
+ *comb = dsv4_hc_affine(ctx0, *comb, scale_comb, base_comb);
+ *comb = ggml_reshape_3d(ctx0, *comb, hc, hc, nt);
+ *comb = build_hc_sinkhorn(*comb, il);
+ }
+ cb(*comb, "hc_comb", il);
+
+ ggml_tensor * result = build_hc_pre(x, pre, il);
+ return result;
+}
+
+ggml_tensor * llama_model_glm5_next::graph::build_hc_post(
+ ggml_tensor * x,
+ ggml_tensor * residual,
+ ggml_tensor * post,
+ ggml_tensor * comb,
+ int il) const {
+ GGML_ASSERT(x->ne[0] == n_embd);
+ GGML_ASSERT(residual->ne[1] == hparams.dsv4_hc_mult);
+
+ if (cparams.fused_dsv4_hc_post) {
+ ggml_tensor * result = ggml_dsv4_hc_post(ctx0, x, residual, post, comb);
+ res->add_fused_node({LLM_FUSED_OP_DSV4_HC_POST, result, il});
+ return result;
+ }
+
+ const int64_t hc = hparams.dsv4_hc_mult;
+ const int64_t nt = x->ne[1];
+
+ ggml_tensor * out = nullptr;
+ for (int64_t dst = 0; dst < hc; ++dst) {
+ ggml_tensor * post_dst = ggml_view_2d(ctx0, post, 1, nt, post->nb[1], dst*post->nb[0]);
+ ggml_tensor * cur = ggml_mul(ctx0, x, post_dst);
+
+ for (int64_t src = 0; src < hc; ++src) {
+ ggml_tensor * res_src = ggml_view_2d(ctx0, residual, n_embd, nt, residual->nb[2], src*residual->nb[1]);
+ ggml_tensor * comb_src_dst = ggml_view_2d(ctx0, comb, 1, nt, comb->nb[2],
+ dst*comb->nb[0] + src*comb->nb[1]);
+ cur = ggml_add(ctx0, cur, ggml_mul(ctx0, res_src, comb_src_dst));
+ }
+
+ cur = ggml_reshape_3d(ctx0, cur, n_embd, 1, nt);
+ out = out ? ggml_concat(ctx0, out, cur, 1) : cur;
+ }
+
+ return out;
+}
+
+llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params) :
+ llm_build_delta_net_base(params), model(model) {
+
+ ggml_tensor * cur;
+
+ ggml_tensor * inp = build_inp_embd(model.tok_embd);
+ cb(inp, "inp_embd", -1);
+
+ // recurrent state + K-only MLA cache through the generic hybrid input, plus the indexer cache
+ const auto * mctx_hyb = static_cast<const llama_memory_hybrid_idx_context *>(mctx);
+
+ auto * inp_hyb = build_inp_mem_hybrid_k();
+ auto * inp_rs = inp_hyb->get_recr();
+ auto * inp_attn = inp_hyb->get_attn();
+ auto * inp_kpool = build_inp_kpool(mctx_hyb);
+
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ const int64_t n_head_kda = hparams.n_head();
+ const int64_t head_dim = hparams.n_embd_head_kda;
+ const int64_t d_conv = hparams.ssm_d_conv;
+ const int64_t d_inner = n_head_kda * head_dim;
+ const int64_t n_seqs = ubatch.n_seqs;
+ const int64_t n_seq_tokens = ubatch.n_seq_tokens;
+
+ GGML_ASSERT(n_seqs != 0);
+ GGML_ASSERT(ubatch.equal_seqs());
+ GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs);
+
+ const int64_t hc = hparams.dsv4_hc_mult;
+ 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);
+
+ ggml_tensor * prev_sel = nullptr;
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ ggml_tensor * residual = inpL;
+ ggml_tensor * post = nullptr;
+ ggml_tensor * comb = nullptr;
+
+ cur = build_hc_pre(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, il);
+ cb(cur, "hc_attn_pre", il);
+
+ cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+ ggml_build_forward_expand(gf, cur);
+
+ if (hparams.is_recr(il)) {
+ cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda,
+ d_inner, n_seq_tokens, n_seqs, il);
+ } else {
+ cur = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, &prev_sel, il);
+ }
+
+ inpL = build_hc_post(cur, residual, post, comb, il);
+ cb(inpL, "hc_attn_post", il);
+
+ residual = inpL;
+ cur = build_hc_pre(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, il);
+ cb(cur, "hc_ffn_pre", il);
+
+ cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ if ((uint32_t) il < hparams.n_layer_dense_lead) {
+ cur = 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);
+ cb(cur, "ffn_out", il);
+ } else {
+ ggml_tensor * moe_out = build_moe_ffn(cur,
+ layer.ffn_gate_inp,
+ layer.ffn_up_exps,
+ layer.ffn_gate_exps,
+ layer.ffn_down_exps,
+ layer.ffn_exp_probs_b,
+ n_expert, n_expert_used,
+ LLM_FFN_SILU, hparams.expert_weights_norm,
+ hparams.expert_weights_scale,
+ (llama_expert_gating_func_type) hparams.expert_gating_func,
+ il);
+ cb(moe_out, "ffn_moe_out", il);
+
+ ggml_tensor * ffn_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(ffn_shexp, "ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "ffn_out", il);
+ }
+
+ inpL = build_hc_post(cur, residual, post, comb, il);
+ inpL = build_cvec(inpL, il);
+ cb(inpL, "l_out", il);
+ }
+
+ // narrow to the output tokens, then collapse the streams
+ // Unmasked nextn embeddings need all rows.
+ const bool narrow_early = inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked);
+ if (narrow_early) {
+ ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens);
+ flat = ggml_get_rows(ctx0, flat, inp_out_ids);
+ inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs);
+ }
+
+ cur = build_hc_mean(inpL);
+ cb(cur, "hc_head", -1);
+
+ cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+
+ // the post-norm hidden state feeds the draft head
+ cb(cur, "h_nextn", -1);
+ res->t_h_nextn = cur;
+
+ if (inp_out_ids && !narrow_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);
+}
+
+// KDA layer, g_a/g_b output gate
+
+ggml_tensor * llama_model_glm5_next::graph::build_kda_layer(
+ ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs,
+ int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
+ int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) {
+
+ const auto * mctx_cur = inp_rs->mctx;
+ const auto kv_head = mctx_cur->get_head();
+
+ ggml_tensor * conv_states_all = mctx_cur->get_r_l(il);
+ ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs);
+
+ const int64_t mem_size = mctx_cur->get_size();
+ const int64_t K_rs = (int64_t) cparams.n_rs_seq + 1;
+
+ ggml_tensor * Qcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
+ ggml_tensor * Kcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
+ ggml_tensor * Vcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head, mem_size, K_rs);
+ cb(Qcur, "kda_q_conv", il);
+ cb(Kcur, "kda_k_conv", il);
+ cb(Vcur, "kda_v_conv", il);
+
+ // Decay gate, ssm_a holds -exp(A_log)
+ ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur);
+ ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a);
+ g1 = ggml_add(ctx0, g1, layer.ssm_dt_b);
+
+ ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1);
+
+ if (hparams.kda_gate_lower_bound > -INFINITY) {
+ g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);
+ g1 = ggml_mul(ctx0, g1, A);
+ g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f));
+ g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound);
+ } else {
+ g1 = ggml_softplus(ctx0, g1);
+ g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens);
+ g1 = ggml_mul(ctx0, g1, A);
+ }
+ cb(g1, "kda_g1", il);
+
+ g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs);
+
+ ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur);
+ beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs);
+ beta = ggml_sigmoid(ctx0, beta);
+ 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_kda, n_seqs);
+
+ // Match FLA l2 norm
+ constexpr float l2_eps = 1e-6f;
+ Qcur = build_gdn_l2_norm(ctx0, Qcur, l2_eps);
+ Kcur = build_gdn_l2_norm(ctx0, Kcur, l2_eps);
+
+ ggml_tensor * output = build_recurrent_attn(inp_rs, ssm_states_all, Qcur, Kcur, Vcur, g1, beta, state, il);
+ output = ggml_cont(ctx0, output);
+ cb(output, "kda_scan_out", il);
+
+ // output gate, then RMSNorm(o) * Sigmoid(g2)
+ ggml_tensor * g_a = ggml_mul_mat(ctx0, layer.ssm_g_a, cur);
+ ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g_b, g_a);
+ g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_tokens);
+
+ ggml_tensor * o = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_tokens);
+ ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il);
+ cb(g2, "kda_g2", il);
+ cb(normed, "kda_normed", il);
+ ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2));
+
+ gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens);
+ cur = ggml_mul_mat(ctx0, layer.wo, gated);
+ cb(cur, "kda_out", il);
+
+ return cur;
+}
+
+// Scores pools of kpool consecutive tokens, expands the selected pools and the incomplete tail into an additive mask
+
+ggml_tensor * llama_model_glm5_next::graph::build_kpool_select(
+ ggml_tensor * cur, ggml_tensor * qr, ggml_tensor * kq_mask, const llama_layer & layer,
+ const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il) {
+
+ const auto * mctx_lid = mctx_hyb->get_idx();
+
+ const int64_t n_indexer_head = hparams.indexer_n_head;
+ const int64_t n_embd_indexer = hparams.indexer_head_size;
+ const int64_t kpool = hparams.indexer_kpool;
+ const int64_t n_pool = inp_kpool->pool_cells->ne[0];
+ const int64_t n_new = inp_kpool->n_new;
+
+ ggml_tensor * iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr);
+ iq = ggml_reshape_3d(ctx0, iq, n_embd_indexer, n_indexer_head, n_tokens);
+ cb(iq, "indexer_q", il);
+
+ // Per-token key and pool gate scores, cached together
+ ggml_tensor * ik = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur);
+ ik = build_norm(ik, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il);
+ cb(ik, "indexer_k", il);
+
+ ggml_tensor * ig = ggml_mul_mat(ctx0, layer.indexer_kpool_gate, cur);
+ cb(ig, "indexer_gate", il);
+
+ // Cache rows store key | gate | pooled
+ ggml_tensor * pzero = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_indexer, n_tokens), 0.0f);
+ ggml_tensor * packed = ggml_concat(ctx0, ggml_concat(ctx0, ik, ig, 0), pzero, 0);
+ packed = ggml_reshape_3d(ctx0, packed, 3*n_embd_indexer, 1, n_tokens);
+ ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, packed, inp_kpool->k_idxs, il));
+
+ auto kpool_cache = mctx_hyb->get_kpool_access(ctx0, il, n_embd_indexer);
+ const int64_t n_kv = mctx_lid->get_n_kv();
+
+ // Pool the entries completed by this ubatch. The last one is a dummy when the ubatch completes none.
+ ggml_tensor * rows = kpool_cache.gather_key_gate(ggml_reshape_1d(ctx0, inp_kpool->new_pool_idxs, kpool*n_new));
+ rows = ggml_reshape_3d(ctx0, rows, 2*n_embd_indexer, kpool, n_new);
+
+ ggml_tensor * pk = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_new, rows->nb[1], rows->nb[2], 0);
+ ggml_tensor * pg = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_new, rows->nb[1], rows->nb[2], ggml_row_size(rows->type, n_embd_indexer));
+
+ ggml_tensor * logits = ggml_add(ctx0, pg, layer.indexer_kpool_ape);
+ logits = ggml_cont(ctx0, ggml_permute(ctx0, logits, 1, 0, 2, 3)); // [kpool, head_dim, n_new]
+ // fold the pool axis into rows, soft_max maps ne[2] to CUDA gridDim.y which caps at 65535
+ ggml_tensor * probs = ggml_soft_max(ctx0, ggml_reshape_2d(ctx0, logits, kpool, n_embd_indexer*n_new));
+ probs = ggml_reshape_3d(ctx0, probs, kpool, n_embd_indexer, n_new);
+
+ pk = ggml_cont(ctx0, ggml_permute(ctx0, pk, 1, 0, 2, 3));
+ ggml_tensor * pooled_new = ggml_sum_rows(ctx0, ggml_mul(ctx0, probs, pk)); // [1, head_dim, n_new]
+ pooled_new = ggml_reshape_2d(ctx0, pooled_new, n_embd_indexer, n_new);
+ cb(pooled_new, "indexer_pool_k_new", il);
+
+ if (inp_kpool->cache_safe) {
+ // Write before the pool gather.
+ ggml_build_forward_expand(gf, kpool_cache.scatter_pooled(pooled_new, inp_kpool->new_pool_rep));
+ }
+
+ ggml_tensor * pooled = nullptr;
+ if (inp_kpool->cache_safe) {
+ pooled = kpool_cache.gather_pooled(inp_kpool->pool_cells);
+ } else {
+ GGML_ASSERT(n_new <= n_pool);
+ ggml_tensor * pad = ggml_fill(ctx0,
+ ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_indexer, n_pool - n_new), 0.0f);
+ pooled = ggml_concat(ctx0, pooled_new, pad, 1);
+ }
+ pooled = ggml_reshape_3d(ctx0, pooled, n_embd_indexer, 1, n_pool);
+ cb(pooled, "indexer_pool_k", il);
+
+ ggml_tensor * sel_idx = nullptr;
+ {
+ ggml_tensor * weights = ggml_mul_mat(ctx0, layer.indexer_proj, cur);
+ weights = ggml_scale(ctx0, weights, 1.0f / sqrtf(float(n_embd_indexer * n_indexer_head)));
+ cb(weights, "indexer_weights", il);
+
+ ggml_tensor * score = nullptr;
+ if (cparams.fused_lid) {
+ score = ggml_lightning_indexer(ctx0, iq, pooled, weights, inp_kpool->pool_mask);
+ res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il});
+ } else {
+ ggml_tensor * q_p = ggml_permute(ctx0, iq, 0, 2, 1, 3); // [head_dim, n_tokens, n_head]
+ ggml_tensor * k_p = ggml_permute(ctx0, pooled, 0, 2, 1, 3); // [head_dim, n_pool, 1]
+
+ ggml_tensor * kq = ggml_mul_mat(ctx0, k_p, q_p); // [n_pool, n_tokens, n_head]
+ kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); // [n_head, n_tokens, n_pool]
+ score = ggml_relu(ctx0, kq);
+ score = ggml_mul(ctx0, score, weights);
+ score = ggml_sum_rows(ctx0, score); // [1, n_tokens, n_pool]
+ score = ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3)); // [n_pool, n_tokens, 1]
+ score = ggml_add(ctx0, score, inp_kpool->pool_mask);
+ }
+ cb(score, "indexer_score", il);
+
+ const int64_t n_top_pool = std::min<int64_t>(n_pool, hparams.indexer_top_k / kpool);
+ ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], UNORDERED
+
+ // The gather mask marks the first min(nv, n_top_pool) slots as the visible pools, so order the set by descending score.
+ ggml_tensor * sel_score = ggml_get_rows(ctx0,
+ ggml_reshape_3d(ctx0, score, 1, n_pool, n_tokens), top_k); // [1, n_top_pool, n_tokens]
+ ggml_tensor * sel_order = ggml_argsort(ctx0,
+ ggml_reshape_2d(ctx0, sel_score, n_top_pool, n_tokens), GGML_SORT_ORDER_DESC);
+ top_k = ggml_get_rows(ctx0,
+ ggml_reshape_3d(ctx0, ggml_cast(ctx0, top_k, GGML_TYPE_F32), 1, n_top_pool, n_tokens), sel_order);
+ top_k = ggml_cast(ctx0, ggml_cont(ctx0, ggml_reshape_2d(ctx0, top_k, n_top_pool, n_tokens)), GGML_TYPE_I32);
+ cb(top_k, "indexer_top_k", il);
+
+ sel_idx = ggml_get_rows(ctx0, inp_kpool->pool_idxs,
+ ggml_reshape_1d(ctx0, top_k, n_top_pool*n_tokens)); // [kpool, n_top_pool*n_tokens]
+ sel_idx = ggml_reshape_2d(ctx0, sel_idx, kpool*n_top_pool, n_tokens);
+
+ if (hparams.indexer_kpool_select_tail) {
+ // Append the incomplete tail with n_kv for missing cells.
+ sel_idx = ggml_concat(ctx0, sel_idx, inp_kpool->tail_idxs, 0);
+ }
+ }
+ const int64_t n_sel = sel_idx->ne[0];
+
+ // Gather returns selected cell indices and masks padding separately.
+ if (inp_kpool->gather) {
+ GGML_ASSERT(inp_kpool->gather_mask->ne[0] == n_sel && inp_kpool->gather_mask->ne[3] == n_tokens);
+ cb(sel_idx, "indexer_sel_idx", il);
+ return sel_idx;
+ }
+
+ // Tie scatter storage lifetime to this layer's selected indices.
+ ggml_tensor * seed = ggml_cast(ctx0, ggml_view_1d(ctx0, sel_idx, 1, 0), GGML_TYPE_F32);
+
+ ggml_tensor * mask_seed = kq_mask->type == GGML_TYPE_F32 ? seed : ggml_cast(ctx0, seed, kq_mask->type);
+ mask_seed = ggml_fill(ctx0, mask_seed, -INFINITY);
+ ggml_tensor * mask_all = ggml_repeat_4d(ctx0, mask_seed, 1, n_kv + 1, n_tokens, 1);
+ mask_all = ggml_reshape_3d(ctx0, mask_all, 1, n_kv + 1, n_tokens);
+
+ ggml_tensor * zero_seed = ggml_fill(ctx0, seed, 0.0f);
+ ggml_tensor * zeros = ggml_repeat_4d(ctx0, zero_seed, 1, n_sel, n_tokens, 1);
+ zeros = ggml_reshape_3d(ctx0, zeros, 1, n_sel, n_tokens);
+
+ ggml_tensor * sel = ggml_set_rows(ctx0, mask_all, zeros, ggml_reshape_3d(ctx0, sel_idx, n_sel, n_tokens, 1));
+ sel = ggml_view_2d(ctx0, sel, n_kv, n_tokens, sel->nb[2], 0);
+
+ // Fold causal visibility before shared-indexer reuse.
+ GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[1]*kq_mask->ne[2]*kq_mask->ne[3] == n_tokens);
+ sel = ggml_add(ctx0, sel, ggml_reshape_2d(ctx0, kq_mask, n_kv, n_tokens));
+ cb(sel, "indexer_sel", il);
+
+ return sel;
+}
+
+// Nope MLA layer with sparse attention over the indexer selection
+
+ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer(
+ ggml_tensor * cur, const llama_layer & layer,
+ const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_attn_k * inp_attn,
+ llm_graph_input_kpool * inp_kpool, ggml_tensor ** prev_sel, int il) {
+
+ const auto * mctx_mla = mctx_hyb->get_attn();
+
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+ const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - hparams.n_rot();
+ const float kq_scale = 1.0f / sqrtf((float) n_embd_head_k_mla);
+
+ GGML_ASSERT(hparams.n_rot() == 0 && "GLM5-Next MLA is nope-only");
+
+ 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, "q_resid", il);
+
+ ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr);
+ q = ggml_reshape_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens);
+
+ ggml_tensor * kv_cmpr = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
+ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
+ cb(kv_cmpr, "kv_cmpr", il);
+
+ // absorb wk_b so the cache holds only the latent
+ ggml_tensor * q_absorbed = ggml_permute(ctx0, q, 0, 2, 1, 3);
+ q_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_absorbed);
+ q_absorbed = ggml_permute(ctx0, q_absorbed, 0, 2, 1, 3);
+ cb(q_absorbed, "q_absorbed", il);
+
+ ggml_tensor * kq_mask = inp_attn->get_kq_mask();
+
+ ggml_tensor * sel = nullptr;
+ if (il >= (int) hparams.n_layer() || hparams.is_indexer_full(il)) { // the NextN block always has a full indexer
+ sel = build_kpool_select(cur, qr, kq_mask, layer, mctx_hyb, inp_kpool, il);
+ *prev_sel = sel;
+ } else {
+ GGML_ASSERT(*prev_sel != nullptr && "shared indexer layer must follow a full indexer layer");
+ sel = *prev_sel;
+ }
+
+ ggml_build_forward_expand(gf, q_absorbed);
+ ggml_build_forward_expand(gf, kv_cmpr);
+ ggml_build_forward_expand(gf, mctx_mla->cpy_k(ctx0, kv_cmpr, inp_attn->get_k_idxs(), il));
+
+ ggml_tensor * out = nullptr;
+ if (inp_kpool->gather) {
+ // Attend over gathered latents with the token dimension in ne[3].
+
+ ggml_build_forward_expand(gf, kq_mask);
+
+ ggml_tensor * sel_idx = sel; // I32 [n_sel, n_tokens]
+ const int64_t n_sel = sel_idx->ne[0];
+
+ ggml_tensor * k_g = mctx_hyb->gather_mla_rows(ctx0, sel_idx, n_sel*n_tokens, kv_lora_rank, il);
+ k_g = ggml_reshape_4d(ctx0, k_g, kv_lora_rank, n_sel, 1, n_tokens); // F32 [kv_lora_rank, n_sel, 1, n_tokens]
+ cb(k_g, "kv_gathered", il);
+
+ ggml_tensor * q_g = ggml_permute(ctx0, q_absorbed, 0, 2, 3, 1); // [kv_lora_rank, 1, n_head, n_tokens]
+
+ ggml_tensor * kq = ggml_mul_mat(ctx0, k_g, q_g); // [n_sel, 1, n_head, n_tokens]
+ ggml_prec_set_acc(kq, GGML_PREC_F32);
+ kq = ggml_soft_max_ext(ctx0, kq, inp_kpool->gather_mask, kq_scale, 0.0f);
+ cb(kq, "kq_soft_max_gathered", il);
+
+ ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, k_g)); // [n_sel, kv_lora_rank, 1, n_tokens]
+ ggml_tensor * kqv = ggml_mul_mat(ctx0, v_t, kq); // [kv_lora_rank, 1, n_head, n_tokens]
+ kqv = ggml_mul_mat(ctx0, layer.wv_b, kqv); // [n_embd_head_v, 1, n_head, n_tokens]
+ cb(kqv, "kqv_gathered", il);
+
+ out = ggml_cont(ctx0, ggml_permute(ctx0, kqv, 0, 2, 1, 3)); // [n_embd_head_v, n_head, 1, n_tokens]
+ out = ggml_reshape_2d(ctx0, out, kqv->ne[0]*n_head, n_tokens);
+ } else {
+ // The scatter selection already includes the causal mask.
+ ggml_tensor * mask = ggml_reshape_4d(ctx0, sel, kq_mask->ne[0], kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]);
+ cb(mask, "kq_mask_dsa", il);
+
+ ggml_tensor * k = mctx_mla->get_k(ctx0, il);
+ ggml_tensor * v = ggml_view_4d(ctx0, k, kv_lora_rank, k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
+
+ out = build_attn_mha(q_absorbed, k, v, nullptr, mask, nullptr, layer.wv_b, inp_kpool->n_sel, kq_scale, il);
+ }
+ cb(out, "kqv_out", il);
+
+ out = ggml_mul_mat(ctx0, layer.wo, out);
+ cb(out, "attn_out", il);
+
+ return out;
+}
diff --git a/src/models/models.h b/src/models/models.h
index 3f9c67c63..0e7d59a73 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -1200,6 +1200,16 @@ struct llama_model_deepseek4 : public llama_model_base {
graph(const llm_graph_params & params) : llm_graph_context(params) {}
graph(const llama_model & model, const llm_graph_params & params);
+ // collapse the hc streams with per-stream weights
+ ggml_tensor * build_hc_pre(
+ ggml_tensor * x,
+ ggml_tensor * weights,
+ int il) const;
+
+ // mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens]
+ ggml_tensor * build_hc_mean(ggml_tensor * x) const;
+
+ // returns the collapsed input and fills the post / comb weights
ggml_tensor * build_hc_pre(
ggml_tensor * x,
ggml_tensor * hc_fn,
@@ -1216,6 +1226,10 @@ struct llama_model_deepseek4 : public llama_model_base {
ggml_tensor * comb,
int il) const;
+ ggml_tensor * build_hc_sinkhorn(
+ ggml_tensor * comb,
+ int il) const;
+
ggml_tensor * build_hc_head(
ggml_tensor * x,
ggml_tensor * hc_fn,
@@ -1309,14 +1323,6 @@ struct llama_model_deepseek4 : public llama_model_base {
float kq_scale,
int il) const;
- ggml_tensor * build_hc_pre(
- ggml_tensor * x,
- ggml_tensor * weights,
- int il) const;
-
- ggml_tensor * build_hc_sinkhorn(
- ggml_tensor * comb,
- int il) const;
};
struct graph_mtp : public graph {
@@ -2596,6 +2602,68 @@ struct llama_model_kimi_k3 : public llama_model_base {
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
+struct llama_model_glm5_next : public llama_model_base {
+ llama_model_glm5_next(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;
+
+ // k-pool indexer inputs on top of the generic hybrid input
+ class llm_graph_input_kpool;
+
+ struct graph : public llm_build_delta_net_base {
+ graph(const llama_model & model, const llm_graph_params & params);
+
+ // collapse the hc streams with per-stream weights
+ ggml_tensor * build_hc_pre(
+ ggml_tensor * x,
+ ggml_tensor * weights,
+ int il) const;
+
+ // mean over the hyper-connection streams: [n_embd, hc, n_tokens] -> [n_embd, n_tokens]
+ ggml_tensor * build_hc_mean(ggml_tensor * x) const;
+
+ // returns the collapsed input and fills the post / comb weights
+ ggml_tensor * build_hc_pre(
+ ggml_tensor * x,
+ ggml_tensor * hc_fn,
+ ggml_tensor * hc_scale,
+ ggml_tensor * hc_base,
+ ggml_tensor ** post,
+ ggml_tensor ** comb,
+ int il) const;
+
+ ggml_tensor * build_hc_post(
+ ggml_tensor * x,
+ ggml_tensor * residual,
+ ggml_tensor * post,
+ ggml_tensor * comb,
+ int il) const;
+
+ ggml_tensor * build_hc_sinkhorn(
+ ggml_tensor * comb,
+ int il) const;
+
+ const llama_model & model;
+
+ llm_graph_input_kpool * build_inp_kpool(const llama_memory_hybrid_idx_context * mctx_hyb);
+
+ ggml_tensor * build_kda_layer(ggml_tensor * cur, const llama_layer & layer,
+ llm_graph_input_rs * inp_rs,
+ int64_t d_conv, int64_t head_dim, int64_t n_head_kda,
+ int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il);
+
+ ggml_tensor * build_kpool_select(ggml_tensor * cur, ggml_tensor * qr, ggml_tensor * kq_mask, const llama_layer & layer,
+ const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il);
+
+ ggml_tensor * build_dsa_layer(ggml_tensor * cur, const llama_layer & layer,
+ const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_attn_k * inp_attn,
+ llm_graph_input_kpool * inp_kpool, ggml_tensor ** prev_sel, int il);
+
+ };
+
+ std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
struct llama_model_kimi_linear : public llama_model_base {
llama_model_kimi_linear(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 5a8a196c1..cebda026b 100644
--- a/tests/test-llama-archs.cpp
+++ b/tests/test-llama-archs.cpp
@@ -152,6 +152,7 @@ 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_GLM5_NEXT
|| arch == LLM_ARCH_MISTRAL4
|| arch == LLM_ARCH_HY_V4) {
n_embd = 128;
@@ -203,14 +204,19 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
if (arch == LLM_ARCH_PLAMO2 || arch == LLM_ARCH_JAMBA || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE ||
arch == LLM_ARCH_GRANITE_HYBRID || arch == LLM_ARCH_LFM2 || arch == LLM_ARCH_LFM2MOE || arch == LLM_ARCH_KIMI_LINEAR ||
- arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3) {
+ arch == LLM_ARCH_BAILINGMOE3 || arch == LLM_ARCH_KIMI_K3 || arch == LLM_ARCH_GLM5_NEXT) {
GGML_ASSERT(n_layer >= 2);
std::vector<uint32_t> n_head_per_layer;
n_head_per_layer.reserve(n_layer);
for (uint32_t il = 0; il < n_layer; il++) {
n_head_per_layer.push_back(il == 1 ? 0 : n_head);
}
- ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head_per_layer);
+ // GLM5 next KDA heads come from the uniform head count, only head_count_kv is per layer.
+ if (arch == LLM_ARCH_GLM5_NEXT) {
+ ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
+ } else {
+ 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 {
ms.add_kv(LLM_KV_ATTENTION_HEAD_COUNT, n_head);
@@ -229,11 +235,13 @@ 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_HY_V4) {
- ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(576));
+ || arch == LLM_ARCH_GLM5_NEXT
+ || arch == LLM_ARCH_HY_V4
+ || arch == LLM_ARCH_MISTRAL4) {
+ // GLM5 next MLA is nope only, the cache row is the compressed latent alone.
+ ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, arch == LLM_ARCH_GLM5_NEXT ? uint32_t(512) : uint32_t(576));
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512));
- ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, uint32_t(64));
+ ms.add_kv(LLM_KV_ROPE_DIMENSION_COUNT, arch == LLM_ARCH_GLM5_NEXT ? uint32_t(0) : uint32_t(64));
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH_MLA, uint32_t(192));
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, uint32_t(128));
if (arch == LLM_ARCH_DOTS3NOTE) {
@@ -289,8 +297,10 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
// MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the
// indexer head count is independent of the main attention head count.
- if (arch == LLM_ARCH_QWEN4EXP) {
+ if (arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_GLM5_NEXT) {
ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4));
+ ms.add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, uint32_t(2));
+ ms.add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, 1.0e-6f);
ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8));
// without this the QSA layers fall back to dense and go uncovered
ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector<uint32_t>(n_layer, 4));
@@ -332,6 +342,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(131072));
ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4));
+ ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL, uint32_t(4));
+ ms.add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, true);
ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1));
// mrope sections count rope pairs; Ling 3.0 VL files carry [t, h, w] sections
// summing to n_rot / 2 (n_rot is 64 in this fixture)
@@ -573,6 +585,7 @@ static bool moe_mandatory(const llm_arch arch) {
case LLM_ARCH_MIMO2:
case LLM_ARCH_KIMI_LINEAR:
case LLM_ARCH_KIMI_K3:
+ case LLM_ARCH_GLM5_NEXT:
case LLM_ARCH_STEP35:
case LLM_ARCH_MISTRAL4:
case LLM_ARCH_MELLUM:
diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h
index d18c76bf8..9aec2a25b 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_SWIGLU_CLAMP "clip.vision.swiglu_clamp"
#define KEY_MM_PATCH_MERGE_TYPE "clip.vision.mm_patch_merge_type"
#define KEY_IMAGE_GRID_PINPOINTS "clip.vision.image_grid_pinpoints"
@@ -487,6 +488,7 @@ enum projector_type {
PROJECTOR_TYPE_DEEPSEEK4V,
PROJECTOR_TYPE_LFM2A,
PROJECTOR_TYPE_GLM4V,
+ PROJECTOR_TYPE_GLM5V,
PROJECTOR_TYPE_YOUTUVL,
PROJECTOR_TYPE_YASA2,
PROJECTOR_TYPE_KIMIK25,
@@ -553,6 +555,7 @@ static std::map<projector_type, std::string> PROJECTOR_TYPE_NAMES = {
{ PROJECTOR_TYPE_DEEPSEEK4V, "deepseek4v"},
{ PROJECTOR_TYPE_LFM2A, "lfm2a"},
{ PROJECTOR_TYPE_GLM4V, "glm4v"},
+ { PROJECTOR_TYPE_GLM5V, "glm5v"},
{ 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 f737ccc24..77248ca76 100644
--- a/tools/mtmd/clip-model.h
+++ b/tools/mtmd/clip-model.h
@@ -93,6 +93,8 @@ struct clip_hparams {
float eps = 1e-6;
float rope_theta = 0.0;
+
+ float swiglu_clamp = 0.0f;
int32_t n_expert_used = 0;
std::vector<int32_t> feature_layers;
int32_t attn_window_size = 0;
diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp
index 62d146e1d..572a7b987 100644
--- a/tools/mtmd/clip.cpp
+++ b/tools/mtmd/clip.cpp
@@ -648,7 +648,9 @@ ggml_tensor * clip_graph::build_ffn(
switch (type_op) {
case FFN_SILU:
if (gate) {
- cur = ggml_swiglu_split(ctx0, cur, tmp);
+ cur = hparams.swiglu_clamp > 0.0f
+ ? ggml_swiglu_clamp(ctx0, cur, tmp, hparams.swiglu_clamp)
+ : ggml_swiglu_split(ctx0, cur, tmp);
cb(cur, "ffn_swiglu", il);
} else {
cur = ggml_silu(ctx0, cur);
@@ -1086,6 +1088,7 @@ static std::unique_ptr<clip_graph> clip_get_graph_builder(clip_ctx * ctx, const
builder = std::make_unique<clip_graph_granite_speech>(ctx, img);
} break;
case PROJECTOR_TYPE_GLM4V:
+ case PROJECTOR_TYPE_GLM5V:
{
builder = std::make_unique<clip_graph_glm4v>(ctx, img);
} break;
@@ -1769,6 +1772,19 @@ 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_GLM5V:
+ {
+ // glm4v tower with clamped SwiGLU, ceil-aligned resize and its own token budget
+ hparams.rope_theta = 10000.0f;
+ hparams.n_merge = 2;
+ hparams.image_resize_algo = RESIZE_ALGO_BICUBIC;
+ get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false);
+ get_f32(KEY_SWIGLU_CLAMP, hparams.swiglu_clamp, true);
+ get_u32(KEY_IMAGE_MIN_PIXELS, hparams.image_min_pixels);
+ get_u32(KEY_IMAGE_MAX_PIXELS, hparams.image_max_pixels);
+ hparams.set_limit_image_tokens();
+ hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup
+ } break;
case PROJECTOR_TYPE_LLAMA4:
{
hparams.rope_theta = 10000.0f;
@@ -2595,6 +2611,7 @@ struct clip_model_loader {
}
} break;
case PROJECTOR_TYPE_GLM4V:
+ case PROJECTOR_TYPE_GLM5V:
{
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"));
@@ -4067,6 +4084,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_GLM5V:
case PROJECTOR_TYPE_PADDLEOCR:
case PROJECTOR_TYPE_HUNYUANVL:
case PROJECTOR_TYPE_YOUTUVL:
@@ -4094,6 +4112,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_GLM5V:
case PROJECTOR_TYPE_PADDLEOCR:
case PROJECTOR_TYPE_HUNYUANVL:
case PROJECTOR_TYPE_YOUTUVL:
@@ -4176,6 +4195,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_GLM5V:
case PROJECTOR_TYPE_YOUTUVL:
case PROJECTOR_TYPE_MUSE_GLIMMER:
{
@@ -4814,6 +4834,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) {
case PROJECTOR_TYPE_QWEN3VL:
case PROJECTOR_TYPE_LING3VL:
case PROJECTOR_TYPE_GLM4V:
+ case PROJECTOR_TYPE_GLM5V:
{
const int merge_ratio = hparams.n_merge;
const int pw = image_size_width / patch_size;
@@ -6030,6 +6051,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_GLM5V:
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 0e1d596b4..1cbcea942 100644
--- a/tools/mtmd/models/glm4v.cpp
+++ b/tools/mtmd/models/glm4v.cpp
@@ -41,8 +41,10 @@ ggml_cgraph * clip_graph_glm4v::build() {
inp = ggml_add(ctx0, inp, model.patch_bias);
cb(inp, "patch_bias", -1);
- // pos-conv norm
- inp = build_norm(inp, model.norm_embd_w, model.norm_embd_b, norm_t, eps, -1);
+ // pos-conv norm (absent in GLM-5.3-Flash)
+ if (model.norm_embd_w) {
+ 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/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp
index 9b2e98862..ea67ec3b8 100644
--- a/tools/mtmd/mtmd-image.cpp
+++ b/tools/mtmd/mtmd-image.cpp
@@ -792,6 +792,75 @@ mtmd_image_preproc_out mtmd_image_preprocessor_dyn_size::preprocess(const clip_i
return output;
}
+//
+// mtmd_image_preprocessor_glm5v
+//
+
+// The canvas is ceil-aligned to patch_size*n_merge and fitted to the token budget.
+// Only rescaled to meet the budget and sits top-left, with black padding on the right and bottom
+mtmd_image_preproc_out mtmd_image_preprocessor_glm5v::preprocess(const clip_image_u8 & img) const {
+ GGML_ASSERT(hparams.image_min_pixels > 0 && hparams.image_max_pixels > 0);
+
+ const int64_t factor = hparams.patch_size * hparams.n_merge;
+ const int64_t min_px = hparams.image_min_pixels; // single-frame pixel counts
+ const int64_t max_px = hparams.image_max_pixels;
+ const int64_t height = img.get_size().height;
+ const int64_t width = img.get_size().width;
+
+ auto align = [factor](int64_t v) { return (v + factor - 1) / factor * factor; };
+
+ // aligned canvas within the budget
+ int64_t canvas_h = align(height);
+ int64_t canvas_w = align(width);
+
+ if (canvas_h * canvas_w < min_px) {
+ const double scale = std::sqrt((double) min_px / (double) (height * width));
+ canvas_h = align(std::max<int64_t>(1, (int64_t) std::ceil(height * scale)));
+ canvas_w = align(std::max<int64_t>(1, (int64_t) std::ceil(width * scale)));
+ }
+
+ if (canvas_h * canvas_w > max_px) {
+ // largest content height whose aligned canvas fits the budget
+ int64_t lo = 1, hi = height;
+ int64_t best_h = factor, best_w = factor;
+ while (lo <= hi) {
+ const int64_t ch = (lo + hi) / 2;
+ const int64_t cw = std::max<int64_t>(1, width * ch / height);
+ const int64_t ah = align(ch);
+ const int64_t aw = align(cw);
+ if (ah * aw <= max_px) {
+ best_h = ah;
+ best_w = aw;
+ lo = ch + 1;
+ } else {
+ hi = ch - 1;
+ }
+ }
+ canvas_h = best_h;
+ canvas_w = best_w;
+ }
+
+ // Scaled to fit the canvas, and never upscaled, unless below the min budget
+ double scale = std::min((double) canvas_h / height, (double) canvas_w / width);
+ if (height * width >= min_px) {
+ scale = std::min(1.0, scale);
+ }
+ const int content_h = (int) std::max<int64_t>(1, std::min<int64_t>(canvas_h, (int64_t) std::floor(height * scale)));
+ const int content_w = (int) std::max<int64_t>(1, std::min<int64_t>(canvas_w, (int64_t) std::floor(width * scale)));
+
+ clip_image_u8 content;
+ img_tool::resize(img, content, clip_image_size{content_w, content_h}, hparams.image_resize_algo, PAD_NONE);
+
+ clip_image_u8 canvas;
+ canvas.set_size(clip_image_size{(int) canvas_w, (int) canvas_h}, img.is_placeholder());
+ img_tool::fill(canvas, {0, 0, 0});
+ img_tool::composite(canvas, content, 0, 0);
+
+ mtmd_image_preproc_out output;
+ output.append(hparams, canvas, true);
+ return output;
+}
+
//
// mtmd_image_preprocessor_longest_edge
//
diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h
index e2cf69872..4fa6207d0 100644
--- a/tools/mtmd/mtmd-image.h
+++ b/tools/mtmd/mtmd-image.h
@@ -123,6 +123,12 @@ struct mtmd_image_preprocessor_dyn_size : mtmd_image_preprocessor {
mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override;
};
+// GLM 5.3 flash, similar to dyn_size, but each edge is aligned up to patch_size*n_merge, and max budget is met by a search over the height with the width scaled proportionally
+struct mtmd_image_preprocessor_glm5v : mtmd_image_preprocessor {
+ mtmd_image_preprocessor_glm5v(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {}
+ mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override;
+};
+
// similar to mtmd_image_preprocessor_dyn_size, but resize the image to have longest edge equal to hparams.image_longest_edge, while preserving aspect ratio
struct mtmd_image_preprocessor_longest_edge : mtmd_image_preprocessor {
mtmd_image_preprocessor_longest_edge(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {}
diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp
index 2efa4d65d..2adf34208 100644
--- a/tools/mtmd/mtmd.cpp
+++ b/tools/mtmd/mtmd.cpp
@@ -868,6 +868,13 @@ struct mtmd_context {
img_end = "<|end_of_image|>";
image_preproc = std::make_unique<mtmd_image_preprocessor_dyn_size>(ctx_v);
} break;
+ case PROJECTOR_TYPE_GLM5V:
+ {
+ // <|begin_of_image|> ... (image embeddings) ... <|end_of_image|>
+ img_beg = "<|begin_of_image|>";
+ img_end = "<|end_of_image|>";
+ image_preproc = std::make_unique<mtmd_image_preprocessor_glm5v>(ctx_v);
+ } break;
case PROJECTOR_TYPE_PADDLEOCR:
{
// <|IMAGE_START|> ... (image embeddings) ... <|IMAGE_END|>