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>
This commit is contained in:
32 files changed
+2377
-39
No files matched your search
@@ -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",
|
||||
|
||||
@@ -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}")
|
||||
+29
-1
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
@@ -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: (
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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) {
|
||||
|
||||
+6
-3
@@ -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);
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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;
|
||||
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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();
|
||||
};
|
||||
@@ -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);
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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;
|
||||
|
||||
@@ -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)) {
|
||||
|
||||
+6
-1
@@ -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;
|
||||
|
||||
@@ -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]);
|
||||
}
|
||||
|
||||
File diff suppressed because it is too large.
Load diff
+76
-8
@@ -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;
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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"},
|
||||
|
||||
@@ -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;
|
||||
|
||||
+23
-1
@@ -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];
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
//
|
||||
|
||||
@@ -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) {}
|
||||
|
||||
@@ -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|>
|
||||
|
||||
Reference in new issue
Block a user