diff --git a/conversion/__init__.py b/conversion/__init__.py index ba73192efa1..425ce92262b 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -104,6 +104,7 @@ "Glm4MoeLiteForCausalLM": "glm", "Glm4vForConditionalGeneration": "glm", "Glm4vMoeForConditionalGeneration": "glm", + "Glm5NextForConditionalGeneration": "glm", "GlmForCausalLM": "chatglm", "GlmMoeDsaForCausalLM": "glm", "GlmOcrForConditionalGeneration": "glm", @@ -298,6 +299,7 @@ "Gemma4UnifiedForConditionalGeneration": "gemma", "Glm4vForConditionalGeneration": "qwen3vl", "Glm4vMoeForConditionalGeneration": "qwen3vl", + "Glm5NextForConditionalGeneration": "qwen3vl", "Glm5vForConditionalGeneration": "kimivl", "GlmOcrForConditionalGeneration": "qwen3vl", "GlmasrModel": "ultravox", diff --git a/conversion/glm.py b/conversion/glm.py index 7544f850cb2..5b68f9ac6a8 100644 --- a/conversion/glm.py +++ b/conversion/glm.py @@ -402,3 +402,198 @@ def set_vocab(self): special_vocab._set_special_token("unk", tokenizer.get_added_vocab()[""]) # 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)) + self.gguf_writer.add_indexer_index_share_mtp(hp.get("index_share_for_mtp_iteration", False)) + if (indexer_types := hp.get("indexer_types")) is not None: + self.gguf_writer.add_indexer_types([t == "full" for t in indexer_types]) + + # mHC + assert hp.get("mhc", True) + self.gguf_writer.add_hyper_connection_count(hp["hc_mult"]) + self.gguf_writer.add_hyper_connection_sinkhorn_iterations(hp["hc_sinkhorn_iters"]) + self.gguf_writer.add_hyper_connection_epsilon(hp["hc_eps"]) + + # MoE + self.gguf_writer.add_leading_dense_block_count(hp["first_k_dense_replace"]) + self.gguf_writer.add_expert_feed_forward_length(hp["moe_intermediate_size"]) + self.gguf_writer.add_expert_shared_count(hp["n_shared_experts"]) + self.gguf_writer.add_expert_weights_scale(hp["routed_scaling_factor"]) + self.gguf_writer.add_expert_weights_norm(hp["norm_topk_prob"]) + if (limit := hp.get("swiglu_limit")) is not None: + self.gguf_writer.add_swiglu_clamp_exp([limit] * self.block_count) + self.gguf_writer.add_swiglu_clamp_shexp([limit] * self.block_count) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + if name == "lm_head.weight" and self.hparams.get("tie_word_embeddings", False): + return + + # routed experts + if ".mlp.experts." in name: + n_experts = self.hparams["n_routed_experts"] + assert bid is not None + if self._experts is None: + self._experts = [{} for _ in range(self.block_count)] + self._experts[bid][name] = data_torch + if len(self._experts[bid]) < n_experts * 3: + return + for w_name in ("down_proj", "gate_proj", "up_proj"): + datas: list[Tensor] = [] + for xid in range(n_experts): + ename = f"model.layers.{bid}.mlp.experts.{xid}.{w_name}.weight" + datas.append(self._experts[bid].pop(ename)) + merged = f"model.layers.{bid}.mlp.experts.{w_name}.weight" + yield from super().modify_tensors(torch.stack(datas, dim=0), merged, bid) + return + + # MLA absorption + if name.endswith("kv_b_proj.weight"): + n_head = self.hparams["num_attention_heads"] + v_head_dim = self.hparams["v_head_dim"] + qk_nope_head_dim = self.hparams["qk_nope_head_dim"] + assert data_torch.shape[0] == n_head * (v_head_dim + qk_nope_head_dim) + kv_b = data_torch.view(n_head, v_head_dim + qk_nope_head_dim, data_torch.shape[-1]) + k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1) + yield from super().modify_tensors(k_b.transpose(1, 2), name.replace("kv_b_proj", "k_b_proj"), bid) + yield from super().modify_tensors(v_b, name.replace("kv_b_proj", "v_b_proj"), bid) + return + + # KDA conv1d + if name.endswith((".q_conv1d.weight", ".k_conv1d.weight", ".v_conv1d.weight")): + if data_torch.ndim == 3: + d_inner, _, d_conv = data_torch.shape + elif data_torch.ndim == 2: + d_inner, d_conv = data_torch.shape + else: + raise ValueError(f"unexpected conv1d rank {data_torch.ndim} for {name}") + data_torch = data_torch.reshape(1, d_inner, 1, d_conv) + + if name.endswith(".A_log"): + n_head = self.hparams["num_attention_heads"] + data_torch = -torch.exp(data_torch.float().flatten()[:n_head]) + + if name.endswith(".dt_bias"): + name = name.rpartition(".dt_bias")[0] + ".dt_proj.bias" + + if re.search(r"\.(hc_(?:attn|ffn)_(?:fn|base|scale)|index_kpool_compress_(?:ape|gate))$", name): + yield self.map_tensor_name(name) + ".weight", data_torch + return + + yield from super().modify_tensors(data_torch, name, bid) + + def tensor_force_quant(self, name: str, new_name: str, bid: int | None, n_dims: int) -> gguf.GGMLQuantizationType | bool: + # keep the small mHC / gating parameters exact + exact_keys = ("hc_attn_", "hc_ffn_", "indexer_compressor_", "ssm_a", "ssm_dt", "exp_probs_b") + if new_name.startswith(("blk.", "output_hc")) and any(k in new_name for k in exact_keys): + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def prepare_metadata(self, vocab_only: bool): + from_dir = self.fname_out.is_dir() + super().prepare_metadata(vocab_only=vocab_only) + if not self.mtp_only or not from_dir: + return + output_type: str = self.ftype.name.partition("_")[2] + fname_default: str = gguf.naming_convention( + self.metadata.name, self.metadata.basename, self.metadata.finetune, + self.metadata.version, size_label=None, output_type=output_type, model_type=None) + self.fname_out = self.fname_out.parent / f"mtp-{fname_default}.gguf" + + def prepare_tensors(self): + super().prepare_tensors() + if self._experts is not None: + leftover = [k for d in self._experts for k in d.keys()] + if leftover: + raise ValueError(f"Unprocessed experts: {leftover}") diff --git a/conversion/qwen3vl.py b/conversion/qwen3vl.py index 4fec708c9ff..11ce68515b2 100644 --- a/conversion/qwen3vl.py +++ b/conversion/qwen3vl.py @@ -228,10 +228,12 @@ class Qwen3ASRMmprojModel(Qwen3OmniMmprojModel): @ModelBase.register("Glm4vForConditionalGeneration", "Glm4vMoeForConditionalGeneration", "GlmOcrForConditionalGeneration") @ModelBase.example("zai-org/GLM-4.1V-9B-Thinking", "zai-org/GLM-4.5V") class Glm4VVisionModel(Qwen3VLVisionModel): + projector_type = gguf.VisionProjectorType.GLM4V + def set_gguf_parameters(self): MmprojModel.set_gguf_parameters(self) # skip Qwen3VLVisionModel parameters assert self.hparams_vision is not None - self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.GLM4V) + self.gguf_writer.add_clip_projector_type(self.projector_type) hidden_act = str(self.hparams_vision.get("hidden_act", "")).lower() if hidden_act == "gelu": @@ -249,6 +251,32 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter yield from super().modify_tensors(data_torch, name, bid) +@ModelBase.register("Glm5NextForConditionalGeneration") +@ModelBase.example("zai-org/GLM-5.3-Flash") +class Glm5NextVisionModel(Glm4VVisionModel): + # GLM-5.3-Flash vision tower. glm4v layout with per-head qk-norm, no post-conv norm and no learned position embeddings. + # Images are placed on a ceil aligned canvas with padding. + + projector_type = gguf.VisionProjectorType.GLM5V + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_vision is not None + self.gguf_writer.add_vision_spatial_merge_size(int(self.hparams_vision.get("spatial_merge_size", 2))) + if (limit := self.hparams_vision.get("swiglu_limit")) is not None: + self.gguf_writer.add_vision_swiglu_clamp(float(limit)) + + # image token budget from the processor, stored as single-frame pixel counts + pc = self.preprocessor_config + patch = int(pc.get("patch_size", 14)) + merge = int(pc.get("merge_size", 2)) + pixels_per_token = (patch * merge) ** 2 + if (min_tok := pc.get("min_image_tokens")) is not None: + self.gguf_writer.add_vision_min_pixels(int(min_tok) * pixels_per_token) + if (max_tok := pc.get("max_image_tokens")) is not None: + self.gguf_writer.add_vision_max_pixels(int(max_tok) * pixels_per_token) + + @ModelBase.register("Qwen3VLForConditionalGeneration") @ModelBase.example("Qwen/Qwen3-VL-4B-Instruct") class Qwen3VLTextModel(Qwen3Model): diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index b85f62a3114..7909cb8fd00 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -225,6 +225,9 @@ class Indexer: BLOCK_SIZE = "{arch}.attention.indexer.block_size" # MSA LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" # MSA TYPES = "{arch}.attention.indexer.types" + KPOOL = "{arch}.attention.indexer.kpool" # GLM5-Next + INDEX_SHARE_MTP = "{arch}.attention.indexer.index_share_mtp" # GLM5-Next + KPOOL_SELECT_TAIL = "{arch}.attention.indexer.kpool_select_tail" # GLM5-Next class HyperConnection: COUNT = "{arch}.hyper_connection.count" @@ -382,6 +385,7 @@ class ClipVision: 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" @@ -559,6 +563,7 @@ class MODEL_ARCH(IntEnum): GLM4 = auto() GLM4_MOE = auto() GLM_DSA = auto() + GLM5_NEXT = auto() BITNET = auto() T5 = auto() T5ENCODER = auto() @@ -889,6 +894,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() @@ -1311,6 +1318,7 @@ class MODEL_TENSOR(IntEnum): 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", @@ -1640,6 +1648,8 @@ class MODEL_TENSOR(IntEnum): 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", @@ -4003,6 +4013,70 @@ class MODEL_TENSOR(IntEnum): 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, @@ -5661,6 +5735,7 @@ class VisionProjectorType: LFM2A = "lfm2a" # audio MUSIC_FLAMINGO = "musicflamingo" # audio GLM4V = "glm4v" + GLM5V = "glm5v" YOUTUVL = "youtuvl" NEMOTRON_V2_VL = "nemotron_v2_vl" QWEN3TTS_SPKENC = "qwen3tts_spkenc" # audio: ECAPA-TDNN speaker encoder diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 689c2fca111..33eb3a9927c 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -824,6 +824,15 @@ def add_indexer_types(self, value: Sequence[bool]) -> None: 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_indexer_index_share_mtp(self, value: bool) -> None: + self.add_bool(Keys.Attention.Indexer.INDEX_SHARE_MTP.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) @@ -1384,6 +1393,9 @@ def add_vision_image_mean(self, values: Sequence[float]) -> None: def add_vision_image_std(self, values: Sequence[float]) -> None: self.add_array(Keys.ClipVision.IMAGE_STD, values) + def add_vision_swiglu_clamp(self, value: float) -> None: + self.add_float32(Keys.ClipVision.SWIGLU_CLAMP, value) + def add_vision_spatial_merge_size(self, value: int) -> None: self.add_uint32(Keys.ClipVision.SPATIAL_MERGE_SIZE, value) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index d644d502eae..462a30ee474 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1317,6 +1317,38 @@ 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.HC_ATTN_FN: ( + "model.layers.{bid}.hc_attn_fn", # glm5-next + ), + + MODEL_TENSOR.HC_ATTN_BASE: ( + "model.layers.{bid}.hc_attn_base", # glm5-next + ), + + MODEL_TENSOR.HC_ATTN_SCALE: ( + "model.layers.{bid}.hc_attn_scale", # glm5-next + ), + + MODEL_TENSOR.HC_FFN_FN: ( + "model.layers.{bid}.hc_ffn_fn", # glm5-next + ), + + MODEL_TENSOR.HC_FFN_BASE: ( + "model.layers.{bid}.hc_ffn_base", # glm5-next + ), + + MODEL_TENSOR.HC_FFN_SCALE: ( + "model.layers.{bid}.hc_ffn_scale", # glm5-next + ), + MODEL_TENSOR.INDEXER_Q_PROJ: ( "model.layers.{bid}.self_attn.index_q_proj", # MSA ), diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 446de4ae25b..b450d1089d7 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -149,6 +149,7 @@ static const std::map 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" }, @@ -284,6 +285,9 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" }, { LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" }, { LLM_KV_ATTENTION_INDEXER_TYPES, "%s.attention.indexer.types" }, + { LLM_KV_ATTENTION_INDEXER_KPOOL, "%s.attention.indexer.kpool" }, + { LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, "%s.attention.indexer.kpool_select_tail" }, + { LLM_KV_ATTENTION_INDEXER_INDEX_SHARE_MTP, "%s.attention.indexer.index_share_mtp" }, { 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" }, @@ -679,6 +683,8 @@ static const std::map 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" }, @@ -952,6 +958,8 @@ static const std::map 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}}, @@ -1075,6 +1083,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: @@ -1144,6 +1153,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_BAILINGMOE3: case LLM_ARCH_KIMI_K3: + case LLM_ARCH_GLM5_NEXT: case LLM_ARCH_QWEN3TTS: case LLM_ARCH_QWEN4EXP: // TODO: fix test-llama-archs return false; diff --git a/src/llama-arch.h b/src/llama-arch.h index 0c0b994836f..1cbc147c57a 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -150,6 +150,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, @@ -289,6 +290,9 @@ 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_INDEXER_INDEX_SHARE_MTP, LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, LLM_KV_ATTENTION_OUTPUT_LORA_RANK, LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, @@ -678,6 +682,8 @@ enum llm_tensor { LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, LLM_TENSOR_INDEXER_COMPRESSOR_APE, LLM_TENSOR_INDEXER_COMPRESSOR_NORM, + LLM_TENSOR_INDEXER_KPOOL_GATE, // glm5-next: k-pool gate scores + LLM_TENSOR_INDEXER_KPOOL_APE, // glm5-next: k-pool position bias LLM_TENSOR_FFN_GATE_TID2EID, LLM_TENSOR_NEXTN_PROJ_PRE, LLM_TENSOR_NEXTN_PROJ_POST, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 3cc27717ece..b5b889c420a 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2304,7 +2304,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(n_tokens * 160, 64u * model.n_tensors()); } else if (model.arch == LLM_ARCH_QWEN3NEXT || diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 274a6264336..d9209b5a53c 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1776,7 +1776,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); @@ -2170,7 +2170,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_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, up, limit); } else { up = ggml_clamp(ctx0, up, -limit, limit); diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 2f238a1744c..fd74dd4612b 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -268,6 +268,9 @@ 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; + bool indexer_index_share_mtp = false; // MTP iterations reuse one indexer selection // MSA uint32_t indexer_block_size = 0; uint32_t indexer_local_blocks = 0; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index f22054c3d6a..0505b09dbb5 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -1247,6 +1247,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; @@ -2748,6 +2754,10 @@ ggml_tensor * llama_kv_cache_context::get_k(ggml_context * ctx, int32_t il) cons return kv->get_k(ctx, il, n_kv, sinfos[i_cur]); } +ggml_tensor * llama_kv_cache_context::get_k_storage(int32_t il) const { + return kv->get_k_storage(il); +} + ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) const { return kv->get_v(ctx, il, n_kv, sinfos[i_cur]); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index c4d8699def1..b6eb47e3a0f 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -168,6 +168,9 @@ class llama_kv_cache : public llama_memory_i { 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 @@ -398,6 +401,9 @@ class llama_kv_cache_context : public llama_memory_context_i { ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il) const; + // The full K storage tensor of the layer, spanning all streams. + ggml_tensor * get_k_storage(int32_t il) const; + // store k_cur and v_cur in the cache based on the provided head location // note: the heads in k_cur and v_cur should be laid out contiguously in memory // - k_cur [n_embd_head_k, n_head_k, n_tokens] diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp index 93b468784a3..8c9bbf1815e 100644 --- a/src/llama-memory-hybrid-idx.cpp +++ b/src/llama-memory-hybrid-idx.cpp @@ -1,5 +1,9 @@ #include "llama-memory-hybrid-idx.h" +#include +#include +#include + #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 @@ -152,6 +157,7 @@ bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_po if (mem_idx) { mem_idx->seq_rm(seq_id, p0, p1); + mem_idx_stale = true; } return get_mem_attn()->seq_rm(seq_id, p0, p1); @@ -162,6 +168,7 @@ 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); + mem_idx_stale = true; } } @@ -170,6 +177,7 @@ void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { if (mem_idx) { mem_idx->seq_keep(seq_id); + mem_idx_stale = true; } } @@ -178,6 +186,7 @@ void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_p if (mem_idx) { mem_idx->seq_add(seq_id, p0, p1, shift); + mem_idx_stale = true; } } @@ -186,6 +195,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 = true; } } @@ -600,6 +610,32 @@ static std::vector llama_memory_hybrid_idx_ns(const llama_kv_cache::sl return res; } +// The kpool layout of one ubatch. +struct llama_memory_hybrid_idx_context::kpool_state { + struct seq { + llama_pos pos_min = 0; + uint32_t strm = 0; // Stream holding this sequence's cells + std::vector> cells; // Position and stream local cell pairs, sorted by position. + std::vector pools; + std::vector is_new; + }; + + std::vector seqs; + + 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(64u, GGML_PAD(n_pool + 1, 64u)); +} + +} + llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_status status) : llama_memory_hybrid_context(status) {} @@ -611,7 +647,12 @@ llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hy ns_ubatch(mem->get_mem_idx() == nullptr ? std::vector() : std::vector{ 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()) { + kpool_st = std::make_unique(kpool_build_layout()); + i_kpool = 0; + } +} llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( llama_memory_hybrid_idx * mem, @@ -633,9 +674,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 require a full re-pool. + mem_idx_stale_batch = mem->mem_idx_is_stale(); +} + +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_idx_stale_batch && mem != nullptr) { + mem->mem_idx_stale_clear(); + } + if (ctx_idx) { ctx_idx->next(); } @@ -652,9 +703,20 @@ bool llama_memory_hybrid_idx_context::apply() { res = res & ctx_idx->apply(); } + // Fix the pool layout of this ubatch. + if (res && kpool_track()) { + kpool_st = std::make_unique(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(ctx_idx.get()); } @@ -677,3 +739,373 @@ void llama_memory_hybrid_idx_context::set_input_qsa( mem->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); } + +// k-pool DSA indexer (glm5-next) + +// Layout 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_layout() const { + GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); + + const uint32_t kpool = mem->get_kpool(); + + kpool_state st; + st.seqs.resize(LLAMA_MAX_SEQ); + + const auto * kv = mem->get_mem_idx(); + const uint32_t n_stream_kv = kv->get_n_stream(); + + if (n_stream_kv == 1) { + const auto & cells = kv->get_cells(0); + + // Scan only active sequences + std::vector active; + for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + if (cells.seq_pos_min(s) >= 0) { + active.push_back(s); + } + } + + for (uint32_t i = 0; i < cells.size(); ++i) { + if (cells.is_empty(i)) { + continue; + } + const llama_pos p = cells.pos_get(i); + uint32_t n_seq_cell = 0; + for (const llama_seq_id s : active) { + if (cells.seq_has(i, s)) { + st.seqs[s].cells.emplace_back(p, i); + ++n_seq_cell; + } + } + if (n_seq_cell > 1) { + st.cache_safe = false; + } + } + } else { + // When kv is non unified, one stream per sequence, so streams never share cells. Cell indices stay stream-local. + for (llama_seq_id s = 0; s < (llama_seq_id) n_stream_kv; ++s) { + const auto & cells = kv->get_cells(s); + if (cells.seq_pos_min(s) < 0) { + continue; + } + auto & sq = st.seqs[s]; + sq.strm = kv->get_stream(s); + for (uint32_t i = 0; i < cells.size(); ++i) { + if (!cells.is_empty(i) && cells.seq_has(i, s)) { + sq.cells.emplace_back(cells.pos_get(i), i); + } + } + } + } + + for (auto & sq : st.seqs) { + if (sq.cells.empty()) { + continue; + } + if (!std::is_sorted(sq.cells.begin(), sq.cells.end())) { + std::sort(sq.cells.begin(), sq.cells.end()); + } + + sq.pos_min = sq.cells.front().first; + + // Pools start at the first valid token + for (size_t j = 0; 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; + } + } + + st.n_pool_real += (uint32_t) sq.pools.size(); + } + + for (auto & sq : st.seqs) { + sq.is_new.assign(sq.pools.size(), 0); + } + + return st; +} + +// Layout of the cells as of this ubatch plus which pools it completes or rewrites. +// Pool cache lifecycle: +// 1. cpy_k writes each token's key | gate into its idx cache row, pooled slot are zeroed. +// 2. This scan derives the current grouping from mem_idx and marks the pools the ubatch touches or completes as is_new, during decode that's one pool every kpool tokens, zero elsewise. +// 3. The graph pools only the is_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. +// Any seq_* edit regroups the pools, so it sets mem_idx_stale and the first ubatch of the next batch re-pools everything 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 is derived from the current grouping every ubatch. +llama_memory_hybrid_idx_context::kpool_state llama_memory_hybrid_idx_context::kpool_build_state( + const llama_ubatch & ubatch) const { + kpool_state st = kpool_build_layout(); + + const uint32_t kpool = mem->get_kpool(); + + // Pools touched by this ubatch are re-pooled, shared cells cannot cache sequence relative pools. + const bool all_new = !st.cache_safe || (mem_idx_stale_batch && i_cur == 0); + + std::vector> upos(LLAMA_MAX_SEQ); + if (!all_new) { + for (uint32_t i = 0; i < ubatch.n_tokens; ++i) { + for (int32_t k = 0; k < ubatch.n_seq_id[i]; ++k) { + upos[ubatch.seq_id[i][k]].push_back(ubatch.pos[i]); + } + } + for (auto & v : upos) { + if (!std::is_sorted(v.begin(), v.end())) { + std::sort(v.begin(), v.end()); + } + } + } + + for (llama_seq_id s = 0; s < LLAMA_MAX_SEQ; ++s) { + auto & sq = st.seqs[s]; + + sq.is_new.assign(sq.pools.size(), all_new ? 1 : 0); + if (all_new) { + st.n_new += (uint32_t) sq.pools.size(); + continue; + } + + const auto & up = upos[s]; + if (up.empty()) { + continue; + } + + for (size_t pi = 0; pi < sq.pools.size(); ++pi) { + const llama_pos p0 = sq.cells[sq.pools[pi]].first; + + auto it = std::lower_bound(up.begin(), up.end(), p0); + if (it != up.end() && *it < p0 + (llama_pos) kpool) { + sq.is_new[pi] = 1; + st.n_new++; + } + } + } + + return st; +} + +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 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; + + GGML_ASSERT(n_pool == kpool_pad(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((n_new == 0) == (new_pool_idxs == nullptr)); + GGML_ASSERT(st.cache_safe || new_pool_rep == nullptr); + GGML_ASSERT(!st.cache_safe || (n_new == 0) == (new_pool_rep == nullptr)); + + if (n_new > 0) { + 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); + 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); + } + } + + 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 kpool_state::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 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 = st.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 seq_pool_start(LLAMA_MAX_SEQ, 0); + std::vector 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 = n_new > 0 ? (int32_t *) new_pool_idxs->data : nullptr; + 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 = st.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 (sq.is_new[pi]) { + 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); + + 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; + const T keep = llama_cast(0.0f); + const T drop = llama_cast(-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) st.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 = st.seqs[s]; + + const uint32_t n_tail = (uint32_t) ((p - sq.pos_min + 1) % (llama_pos) kpool); + + for (uint32_t k = 0; k < kpool - 1; ++k) { + int32_t cell = sentinel; + bool real = false; + if (k < n_tail) { + const llama_pos pt = p - (llama_pos) k; + auto it = std::lower_bound(sq.cells.begin(), sq.cells.end(), std::make_pair(pt, 0u)); + if (it != sq.cells.end() && it->first == pt) { + cell = (int32_t) (gather ? gcell(sq, it->second) : (int64_t) it->second); + real = true; + } + } + tidx[(size_t) i*(kpool - 1) + k] = cell; + + if (gm != nullptr && n_sel % kpool != 0) { + gm[(size_t) i*n_sel + (size_t) n_top*kpool + k] = real ? 0.0f : -INFINITY; + } + } + } +} diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h index 705189e7eb5..27069d8bd46 100644 --- a/src/llama-memory-hybrid-idx.h +++ b/src/llama-memory-hybrid-idx.h @@ -87,6 +87,14 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, bool blk_bias) const; + // The model's indexer pool size. + uint32_t get_kpool() const { return hparams_idx.indexer_kpool; } + + // The pooled keys persist in the idx cache across batches. + // Sequence edits shift the pool grid and stale the cached values. + bool mem_idx_is_stale() const { return mem_idx_stale; } + void mem_idx_stale_clear () { mem_idx_stale = false; } + 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 @@ -97,6 +105,8 @@ class llama_memory_hybrid_idx : public llama_memory_hybrid { llama_hparams hparams_idx; const std::unique_ptr mem_idx; + + bool mem_idx_stale = false; }; class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { @@ -122,7 +132,7 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { slot_info_vec_t sinfos_idx, std::vector 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 @@ -141,12 +151,19 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // 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; + 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) 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 @@ -157,4 +174,22 @@ class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { // mirrors the base class's ubatch cursor, which is private there size_t i_cur = 0; + + // K-pool layouts + struct kpool_state; + kpool_state kpool_build_layout() const; + kpool_state kpool_build_state(const llama_ubatch & ubatch) const; + const kpool_state & kpool_cur() const; + + // unique_ptr because kpool_state is incomplete here. + std::unique_ptr 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; + + // Clear a pending full re-pool only after the first ubatch succeeds + bool mem_idx_stale_batch = false; }; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 919e90ecccd..bba0799bbfc 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -297,6 +297,9 @@ 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_INDEX_SHARE_MTP, hparams.indexer_index_share_mtp); add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks); add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true); add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 6344f2d8aee..300d63eb828 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -334,6 +334,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); default: @@ -965,6 +967,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"; @@ -2325,6 +2328,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_DOTS3NOTE: { GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE); @@ -2837,6 +2884,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_NEMOTRON_H_MOE: case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_KIMI_K3: + case LLM_ARCH_GLM5_NEXT: return LLAMA_ROPE_TYPE_NONE; // use what we call a normal RoPE, operating on pairs of consecutive head values diff --git a/src/llama-model.h b/src/llama-model.h index 4c4a30e018b..1f038bccbf1 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -150,6 +150,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, }; @@ -557,6 +558,10 @@ struct llama_layer { struct ggml_tensor * indexer_attn_k = nullptr; struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias + // glm5-next k-pool indexer + struct ggml_tensor * indexer_kpool_gate = nullptr; + struct ggml_tensor * indexer_kpool_ape = nullptr; + // MSA struct ggml_tensor * index_q_proj = nullptr; struct ggml_tensor * index_k_proj = nullptr; diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 34ff25db57e..aa92c361b05 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -328,6 +328,25 @@ 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("hc_") == 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 +470,22 @@ static ggml_type llama_tensor_get_type_impl(quantize_state_impl & qs, ggml_type return std::make_pair(i_layer, n_layer); }; + // by default, for glm5-next, don't let these tensors be quantized below Q8_0 + if (arch == LLM_ARCH_GLM5_NEXT && ( + name.find("attn_q_a") != std::string::npos || + name.find("attn_q_b") != std::string::npos || + name.find("nextn.eh_proj") != std::string::npos)) + { + switch (new_type) { + case GGML_TYPE_F32: + case GGML_TYPE_BF16: + case GGML_TYPE_F16: + break; + default: + return GGML_TYPE_Q8_0; + } + } + // for arches that share the same tensor between the token embeddings and the output, we quantize the token embeddings // with the quantization of the output tensor if (category == tensor_category::OUTPUT || (qs.has_tied_embeddings && category == tensor_category::TOKEN_EMBD)) { diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index ff926ceecd1..ebd46aa7082 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -2253,10 +2253,15 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) { pre_type = LLAMA_VOCAB_PRE_TYPE_PORO; clean_spaces = false; } else if ( - tokenizer_pre == "glm4" || tokenizer_pre == "chatglm-bpe") { pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4; special_bos_id = LLAMA_TOKEN_NULL; + } else if ( + tokenizer_pre == "glm4" || + tokenizer_pre == "glm5") { + pre_type = LLAMA_VOCAB_PRE_TYPE_CHATGLM4; + special_bos_id = LLAMA_TOKEN_NULL; + ignore_merges = true; } else if ( tokenizer_pre == "viking") { pre_type = LLAMA_VOCAB_PRE_TYPE_VIKING; diff --git a/src/models/deepseek4.cpp b/src/models/deepseek4.cpp index 5bdf14b4860..66729b2dd1f 100644 --- a/src/models/deepseek4.cpp +++ b/src/models/deepseek4.cpp @@ -266,15 +266,15 @@ 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) { +template +ggml_tensor * llama_model_deepseek4::graph_base::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( @@ -287,7 +287,8 @@ static ggml_tensor * dsv4_hc_affine( return x; } -ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_pre( ggml_tensor * x, ggml_tensor * weights, int il) const { @@ -314,7 +315,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( return result; } -ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn( +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_sinkhorn( ggml_tensor * comb, int il) const { GGML_UNUSED(il); @@ -351,7 +353,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_sinkhorn( return comb; } -ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_pre( ggml_tensor * x, ggml_tensor * hc_fn, ggml_tensor * hc_scale, @@ -409,7 +412,8 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_pre( return result; } -ggml_tensor * llama_model_deepseek4::graph::build_hc_post( +template +ggml_tensor * llama_model_deepseek4::graph_base::build_hc_post( ggml_tensor * x, ggml_tensor * residual, ggml_tensor * post, @@ -446,6 +450,10 @@ ggml_tensor * llama_model_deepseek4::graph::build_hc_post( return out; } +// instantiate the mHC helpers for deepseek4 (and dflash) and glm5-next +template struct llama_model_deepseek4::graph_base; +template struct llama_model_deepseek4::graph_base; + ggml_tensor * llama_model_deepseek4::graph::build_hc_head( ggml_tensor * x, ggml_tensor * hc_fn, @@ -1221,7 +1229,7 @@ ggml_tensor * llama_model_deepseek4::graph::build_attention_impl( } llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_params & params) : - llm_graph_context(params) { + graph_base<>(params) { ggml_tensor * cur; ggml_tensor * inp = build_inp_embd(model.tok_embd); @@ -1238,7 +1246,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 +1335,7 @@ llama_model_deepseek4::graph::graph(const llama_model & model, const llm_graph_p } if ((size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]) { - res->t_layer_inp[n_layer] = dsv4_hc_mean(ctx0, inpL); + res->t_layer_inp[n_layer] = build_hc_mean(inpL); cb(res->t_layer_inp[n_layer], "layer_inp", n_layer); ggml_build_forward_expand(gf, res->t_layer_inp[n_layer]); } diff --git a/src/models/glm5-next.cpp b/src/models/glm5-next.cpp new file mode 100644 index 00000000000..e3780032860 --- /dev/null +++ b/src/models/glm5-next.cpp @@ -0,0 +1,820 @@ +#include "models.h" +#include "llama-memory-hybrid-idx.h" + +// GLM5-Next (GLM-5.3-Flash): hybrid KDA (linear) + nope MLA with a k-pool DSA indexer, +// mHC residual streams, DeepSeek-style MoE. + +void llama_model_glm5_next::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps, false); + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl); + ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q); + ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv); + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_KDA_HEAD_DIM, hparams.n_embd_head_kda); + ml.get_key(LLM_KV_KDA_GATE_LOWER_BOUND, hparams.kda_gate_lower_bound, false); + + // the MLA cache holds the compressed latent + hparams.n_embd_head_v_full = hparams.n_lora_kv; + + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; + } + + ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); + ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false); + ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func, false); + if (hparams.expert_gating_func == LLAMA_EXPERT_GATING_FUNC_TYPE_NONE) { + hparams.expert_gating_func = LLAMA_EXPERT_GATING_FUNC_TYPE_SIGMOID; + } + ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_EXP, hparams.swiglu_clamp_exp, hparams.n_layer_all, false); + if (!ml.get_key_or_arr(LLM_KV_SWIGLU_CLAMP_SHEXP, hparams.swiglu_clamp_shexp, hparams.n_layer_all, false)) { + hparams.swiglu_clamp_shexp = hparams.swiglu_clamp_exp; + } + + // DSA indexer with k-pool compression + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KPOOL, hparams.indexer_kpool); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, hparams.indexer_kpool_select_tail, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_INDEX_SHARE_MTP, hparams.indexer_index_share_mtp, false); + GGML_ASSERT(hparams.indexer_kpool > 1 && hparams.indexer_top_k % hparams.indexer_kpool == 0); + std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1); + ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false); + + // mHC + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + ml.get_key(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); + ml.get_key(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); + GGML_ASSERT(hparams.dsv4_hc_mult == 4 && "mHC with hc_mult != 4 is not supported"); + + switch (hparams.n_layer()) { + case 45: type = LLM_TYPE_320B_A18B; break; // GLM-5.3-Flash + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_glm5_next::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_mix_dim = (2 + hc)*hc; + + // the NextN block is loaded but only used by the MTP graph. + // Separated trunk_only/mtp_only handling TODO with DECODER_MTP graph in the MTP follow up + int mtp_flags = 0; + if (!ml.load_mtp) { + mtp_flags |= TENSOR_SKIP; + } + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0); + + for (int i = 0; i < n_layer_all; ++i) { + auto & layer = layers[i]; + + const int flags = (i >= n_layer) ? mtp_flags : 0; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, flags); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, flags); + + if (i < n_layer) { + layer.hc_attn_fn = create_tensor(tn(LLM_TENSOR_HC_ATTN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, 0); + layer.hc_attn_base = create_tensor(tn(LLM_TENSOR_HC_ATTN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_attn_scale = create_tensor(tn(LLM_TENSOR_HC_ATTN_SCALE, "weight", i), {3}, 0); + layer.hc_ffn_fn = create_tensor(tn(LLM_TENSOR_HC_FFN_FN, "weight", i), {hc*n_embd, hc_mix_dim}, 0); + layer.hc_ffn_base = create_tensor(tn(LLM_TENSOR_HC_FFN_BASE, "weight", i), {hc_mix_dim}, 0); + layer.hc_ffn_scale = create_tensor(tn(LLM_TENSOR_HC_FFN_SCALE, "weight", i), {3}, 0); + } + + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = head_dim * n_head; + + if (hparams.is_recr(i)) { + auto conv = [&](llm_tensor tid) { + ggml_tensor * t = create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner, 1}, TENSOR_NOT_REQUIRED); + return t ? t : create_tensor(tn(tid, "weight", i), {d_conv, 1, d_inner}, 0); + }; + layer.ssm_q_conv = conv(LLM_TENSOR_SSM_CONV1D_Q); + layer.ssm_k_conv = conv(LLM_TENSOR_SSM_CONV1D_K); + layer.ssm_v_conv = conv(LLM_TENSOR_SSM_CONV1D_V); + + create_tensor_qkv(layer, i, n_embd, d_inner, d_inner, d_inner, 0); + + layer.ssm_f_a = create_tensor(tn(LLM_TENSOR_SSM_F_A, "weight", i), {n_embd, head_dim}, 0); + layer.ssm_f_b = create_tensor(tn(LLM_TENSOR_SSM_F_B, "weight", i), {head_dim, d_inner}, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", i), {n_embd, n_head}, 0); + + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, i), {n_head}, 0); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {d_inner}, 0); + + layer.ssm_g_a = create_tensor(tn(LLM_TENSOR_SSM_G_A, "weight", i), {n_embd, head_dim}, 0); + layer.ssm_g_b = create_tensor(tn(LLM_TENSOR_SSM_G_B, "weight", i), {head_dim, d_inner}, 0); + layer.ssm_o_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {head_dim}, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {d_inner, n_embd}, 0); + } else { + const int64_t q_lora_rank = hparams.n_lora_q; + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_k = hparams.n_embd_head_k_mla(); + const int64_t n_embd_head_v = hparams.n_embd_head_v_mla(); + const int64_t qk_rope_head_dim = hparams.n_rot(); + const int64_t qk_nope_head_dim = n_embd_head_k - qk_rope_head_dim; + + layer.attn_q_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_A_NORM, "weight", i), {q_lora_rank}, flags); + layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, flags); + + layer.wq_a = create_tensor(tn(LLM_TENSOR_ATTN_Q_A, "weight", i), {n_embd, q_lora_rank}, flags); + layer.wq_b = create_tensor(tn(LLM_TENSOR_ATTN_Q_B, "weight", i), {q_lora_rank, n_head * n_embd_head_k}, flags); + + layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + qk_rope_head_dim}, flags); + layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {qk_nope_head_dim, kv_lora_rank, n_head}, flags); + layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v, n_head}, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v, n_embd}, flags); + + const int64_t n_indexer_head = hparams.indexer_n_head; + const int64_t n_embd_indexer = hparams.indexer_head_size; + const int64_t kpool = hparams.indexer_kpool; + + const bool full = i >= n_layer || hparams.is_indexer_full(i); + const int iflags = flags | (full ? 0 : TENSOR_NOT_REQUIRED); + + layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_embd_indexer}, iflags); + layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {n_embd_indexer}, iflags); + layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, n_indexer_head}, iflags); + layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, n_embd_indexer}, iflags); + layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {q_lora_rank, n_indexer_head * n_embd_indexer}, iflags); + layer.indexer_kpool_gate = create_tensor(tn(LLM_TENSOR_INDEXER_KPOOL_GATE, "weight", i), {n_embd, n_embd_indexer}, iflags); + layer.indexer_kpool_ape = create_tensor(tn(LLM_TENSOR_INDEXER_KPOOL_APE, "weight", i), {n_embd_indexer, kpool}, iflags); + } + + if (i < (int) hparams.n_layer_dense_lead) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), {n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + } else { + const int64_t n_ff_exp = hparams.n_ff_exp(i); + const int64_t n_expert_shared = hparams.n_expert_shared; + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, flags); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, flags); + + layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, flags); + layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, flags); + + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, flags); + } + + if (i >= n_layer) { + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", i), {2 * n_embd, n_embd}, flags); + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", i), {n_embd}, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", i), {n_embd}, flags); + layer.nextn.embed_tokens = create_tensor(tn(LLM_TENSOR_NEXTN_EMBED_TOKENS, "weight", i), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_head = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, "weight", i), {n_embd, n_vocab}, flags | TENSOR_NOT_REQUIRED); + layer.nextn.shared_head_norm = create_tensor(tn(LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, "weight", i), {n_embd}, flags | TENSOR_NOT_REQUIRED); + } + } +} + +std::unique_ptr llama_model_glm5_next::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + throw std::runtime_error("GLM5-Next NextN graph not implemented yet"); + } + return std::make_unique(*this, params); +} + +// Causal conv1d over one of Q/K/V +static ggml_tensor * glm5_conv1d(ggml_cgraph * gf, ggml_context * ctx0, + ggml_tensor * conv_states_all, ggml_tensor * conv_state_all, + int64_t qkv, ggml_tensor * x, ggml_tensor * proj_w, ggml_tensor * conv_w, + int64_t d_conv, int64_t head_dim, int64_t n_head, + int64_t n_seq_tokens, int64_t n_seqs, int64_t n_tokens, int64_t kv_head) { + const int64_t d_inner = head_dim * n_head; + const int64_t conv_state_size = (d_conv - 1) * d_inner; + const int64_t n_embd_r_total = 3 * conv_state_size; + + ggml_tensor * conv_state_x = ggml_view_3d(ctx0, conv_state_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_state_all), + n_embd_r_total * ggml_element_size(conv_state_all), + qkv * conv_state_size * ggml_element_size(conv_state_all)); + + ggml_tensor * x_proj = ggml_mul_mat(ctx0, proj_w, x); + ggml_tensor * x_3d = ggml_reshape_3d(ctx0, x_proj, d_inner, n_seq_tokens, n_seqs); + ggml_tensor * conv_x = ggml_concat(ctx0, conv_state_x, ggml_transpose(ctx0, x_3d), 0); + + ggml_tensor * last_conv_x = ggml_view_3d(ctx0, conv_x, d_conv - 1, d_inner, n_seqs, + conv_x->nb[1], conv_x->nb[2], n_seq_tokens * conv_x->nb[0]); + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, last_conv_x, + ggml_view_3d(ctx0, conv_states_all, d_conv - 1, d_inner, n_seqs, + (d_conv - 1) * ggml_element_size(conv_states_all), + n_embd_r_total * ggml_element_size(conv_states_all), + (kv_head * n_embd_r_total + qkv * conv_state_size) * ggml_element_size(conv_states_all)))); + + ggml_tensor * conv_weight = ggml_reshape_2d(ctx0, conv_w, d_conv, d_inner); + ggml_tensor * Xcur = ggml_ssm_conv(ctx0, conv_x, conv_weight); + Xcur = ggml_reshape_2d(ctx0, Xcur, d_inner, n_tokens); + Xcur = ggml_silu(ctx0, Xcur); + + return ggml_reshape_4d(ctx0, Xcur, head_dim, n_head, n_seq_tokens, n_seqs); +} + + +// K-pool indexer inputs +class llama_model_glm5_next::llm_graph_input_kpool : public llm_graph_input_i { +public: + llm_graph_input_kpool(const llama_memory_hybrid_idx_context * mctx, uint32_t kpool) : mctx(mctx), kpool(kpool) {} + virtual ~llm_graph_input_kpool() = default; + + void set_input(const llama_ubatch * ubatch) override { + mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); + mctx->set_input_kpool(pool_cells, pool_idxs, pool_mask, tail_idxs, gather_mask, gather, new_pool_idxs, new_pool_rep, ubatch); + } + + bool can_reuse(const llm_graph_params & params) override { + mctx = static_cast(params.mctx); + + const auto * idx = mctx->get_idx(); + if (idx == nullptr) { + return false; + } + + bool res = true; + + res &= k_idxs->ne[0] == params.ubatch.n_tokens; + res &= pool_cells->ne[0] == mctx->get_n_kpool(); + res &= pool_mask->ne[1] == params.ubatch.n_tokens; + res &= tail_idxs->ne[1] == params.ubatch.n_tokens; + // The scatter mask shape follows n_kv. + res &= n_kv == idx->get_n_kv(); + // The new pool path is sized exactly + res &= n_new == mctx->get_n_kpool_new(); + res &= cache_safe == mctx->get_kpool_cache_safe(); + + return res; + } + + ggml_tensor * k_idxs = nullptr; // I64 [n_tokens] + ggml_tensor * pool_cells = nullptr; // I32 [n_pool] cell caching each pool's pooled key + ggml_tensor * pool_idxs = nullptr; // I32 [kpool, n_pool] member cells per pool, n_kv sentinel for the padded pools + ggml_tensor * pool_mask = nullptr; // F32/F16 [n_pool, n_tokens] + ggml_tensor * tail_idxs = nullptr; // I32 [kpool - 1, n_tokens] + ggml_tensor * gather_mask = nullptr; // F32 [n_sel, 1, 1, n_tokens] + ggml_tensor * new_pool_idxs = nullptr; // I32 [kpool, n_new] members of the pools completed this ubatch + ggml_tensor * new_pool_rep = nullptr; // I64 [n_new] cell to write each new pooled key into + + const llama_memory_hybrid_idx_context * mctx; + const uint32_t kpool; + uint32_t n_new = 0; + uint32_t n_sel = 0; + bool cache_safe = true; + bool gather = false; + uint32_t n_kv = 0; +}; + +llama_model_glm5_next::llm_graph_input_kpool * llama_model_glm5_next::graph::build_inp_kpool(const llama_memory_hybrid_idx_context * mctx_hyb) { + const auto * mctx_idx = mctx_hyb->get_idx(); + GGML_ASSERT(mctx_idx != nullptr); + + const uint32_t kpool = hparams.indexer_kpool; + const uint32_t n_pool = mctx_hyb->get_n_kpool(); + const uint32_t n_kv = mctx_idx->get_n_kv(); + const uint32_t n_new = mctx_hyb->get_n_kpool_new(); + const bool cache_safe = mctx_hyb->get_kpool_cache_safe(); + + // the fused lightning indexer wants an f16 mask + const auto type_mask = cparams.fused_lid ? GGML_TYPE_F16 : GGML_TYPE_F32; + + auto inp = std::make_unique(mctx_hyb, kpool); + + inp->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + inp->pool_cells = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pool); + inp->pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_pool); + inp->pool_mask = ggml_new_tensor_2d(ctx0, type_mask, n_pool, n_tokens); + inp->tail_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool - 1, n_tokens); + ggml_set_input(inp->pool_cells); + ggml_set_input(inp->pool_idxs); + ggml_set_input(inp->pool_mask); + ggml_set_input(inp->tail_idxs); + + ggml_build_forward_expand(gf, inp->pool_cells); + ggml_build_forward_expand(gf, inp->pool_idxs); + ggml_build_forward_expand(gf, inp->pool_mask); + ggml_build_forward_expand(gf, inp->tail_idxs); + + inp->n_kv = n_kv; + + // Gather selected latents for small decode batches when n_kv exceeds n_sel. + { + constexpr int64_t max_ub = 16; + + const int64_t n_top_pool = std::min(n_pool, hparams.indexer_top_k / kpool); + const int64_t n_sel = kpool*n_top_pool + (hparams.indexer_kpool_select_tail ? kpool - 1 : 0); + inp->n_sel = (uint32_t) n_sel; + inp->gather = (int64_t) n_tokens <= max_ub && (int64_t) n_kv > n_sel; + + if (inp->gather) { + inp->gather_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_sel, 1, 1, n_tokens); + ggml_set_input(inp->gather_mask); + // Keep the mask allocated even when no op reads it, because set_input_kpool always fills it. + ggml_build_forward_expand(gf, inp->gather_mask); + } + } + + inp->n_new = n_new; + inp->cache_safe = cache_safe; + if (n_new > 0) { + inp->new_pool_idxs = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, kpool, n_new); + ggml_set_input(inp->new_pool_idxs); + if (cache_safe) { + inp->new_pool_rep = ggml_new_tensor_1d(ctx0, GGML_TYPE_I64, n_new); + ggml_set_input(inp->new_pool_rep); + } + } + + return (llm_graph_input_kpool *) res->add_input(std::move(inp)); +} + +llama_model_glm5_next::graph::graph(const llama_model & model, const llm_graph_params & params) : + llama_model_deepseek4::graph_base(params), model(model) { + + ggml_tensor * cur; + + ggml_tensor * inp = build_inp_embd(model.tok_embd); + cb(inp, "inp_embd", -1); + + // recurrent state + K-only MLA cache through the generic hybrid input, plus the indexer cache + const auto * mctx_hyb = static_cast(mctx); + + auto * inp_hyb = build_inp_mem_hybrid_k(); + auto * inp_rs = inp_hyb->get_recr(); + auto * inp_attn = inp_hyb->get_attn(); + auto * inp_kpool = build_inp_kpool(mctx_hyb); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + const int64_t n_head_kda = hparams.n_head(); + const int64_t head_dim = hparams.n_embd_head_kda; + const int64_t d_conv = hparams.ssm_d_conv; + const int64_t d_inner = n_head_kda * head_dim; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + + const int64_t hc = hparams.dsv4_hc_mult; + ggml_tensor * inpL = ggml_reshape_3d(ctx0, inp, n_embd, 1, n_tokens); + inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1); + cb(inpL, "hc_init", -1); + + ggml_tensor * prev_sel = nullptr; + + for (int il = 0; il < n_layer; ++il) { + const auto & layer = model.layers[il]; + + ggml_tensor * residual = inpL; + ggml_tensor * post = nullptr; + ggml_tensor * comb = nullptr; + + cur = build_hc_pre(inpL, layer.hc_attn_fn, layer.hc_attn_scale, layer.hc_attn_base, &post, &comb, il); + cb(cur, "hc_attn_pre", il); + + cur = build_norm(cur, layer.attn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recr(il)) { + cur = build_kda_layer(cur, layer, inp_rs, d_conv, head_dim, n_head_kda, + d_inner, n_seq_tokens, n_seqs, il); + } else { + cur = build_dsa_layer(cur, layer, mctx_hyb, inp_attn, inp_kpool, &prev_sel, il); + } + + inpL = build_hc_post(cur, residual, post, comb, il); + cb(inpL, "hc_attn_post", il); + + residual = inpL; + cur = build_hc_pre(inpL, layer.hc_ffn_fn, layer.hc_ffn_scale, layer.hc_ffn_base, &post, &comb, il); + cb(cur, "hc_ffn_pre", il); + + cur = build_norm(cur, layer.ffn_norm, nullptr, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + cur = build_ffn(cur, + layer.ffn_up, nullptr, nullptr, + layer.ffn_gate, nullptr, nullptr, + layer.ffn_down, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + } else { + ggml_tensor * moe_out = build_moe_ffn(cur, + layer.ffn_gate_inp, + layer.ffn_up_exps, + layer.ffn_gate_exps, + layer.ffn_down_exps, + layer.ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SILU, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * ffn_shexp = build_ffn(cur, + layer.ffn_up_shexp, nullptr, nullptr, + layer.ffn_gate_shexp, nullptr, nullptr, + layer.ffn_down_shexp, nullptr, nullptr, + nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + } + + inpL = build_hc_post(cur, residual, post, comb, il); + inpL = build_cvec(inpL, il); + cb(inpL, "l_out", il); + } + + // narrow to the output tokens, then collapse the streams + // Unmasked nextn embeddings need all rows. + const bool narrow_early = inp_out_ids && (!cparams.embeddings_nextn || cparams.embeddings_nextn_masked); + if (narrow_early) { + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, n_embd*hc, n_tokens); + flat = ggml_get_rows(ctx0, flat, inp_out_ids); + inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs); + } + + cur = build_hc_mean(inpL); + cb(cur, "hc_head", -1); + + cur = build_norm(cur, model.output_norm, nullptr, LLM_NORM_RMS, -1); + + // the post-norm hidden state feeds the draft head + cb(cur, "h_nextn", -1); + res->t_h_nextn = cur; + + if (inp_out_ids && !narrow_early) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = ggml_mul_mat(ctx0, model.output, cur); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +// KDA layer, g_a/g_b output gate + +ggml_tensor * llama_model_glm5_next::graph::build_kda_layer( + ggml_tensor * cur, const llama_layer & layer, llm_graph_input_rs * inp_rs, + int64_t d_conv, int64_t head_dim, int64_t n_head_kda, + int64_t d_inner, int64_t n_seq_tokens, int64_t n_seqs, int il) { + + const auto * mctx_cur = inp_rs->mctx; + const auto kv_head = mctx_cur->get_head(); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * conv_state_all = build_rs(inp_rs, conv_states_all, hparams.n_embd_r(), n_seqs); + + ggml_tensor * Qcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 0, cur, layer.wq, layer.ssm_q_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Kcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 1, cur, layer.wk, layer.ssm_k_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + ggml_tensor * Vcur = glm5_conv1d(gf, ctx0, conv_states_all, conv_state_all, 2, cur, layer.wv, layer.ssm_v_conv, d_conv, head_dim, n_head_kda, n_seq_tokens, n_seqs, n_tokens, kv_head); + cb(Qcur, "kda_q_conv", il); + cb(Kcur, "kda_k_conv", il); + cb(Vcur, "kda_v_conv", il); + + // Decay gate, ssm_a holds -exp(A_log) + ggml_tensor * f_a = ggml_mul_mat(ctx0, layer.ssm_f_a, cur); + ggml_tensor * g1 = ggml_mul_mat(ctx0, layer.ssm_f_b, f_a); + g1 = ggml_add(ctx0, g1, layer.ssm_dt_b); + + ggml_tensor * A = ggml_reshape_3d(ctx0, layer.ssm_a, 1, n_head_kda, 1); + + if (hparams.kda_gate_lower_bound > -INFINITY) { + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); + g1 = ggml_sigmoid(ctx0, ggml_scale(ctx0, g1, -1.0f)); + g1 = ggml_scale(ctx0, g1, hparams.kda_gate_lower_bound); + } else { + g1 = ggml_softplus(ctx0, g1); + g1 = ggml_reshape_3d(ctx0, g1, head_dim, n_head_kda, n_tokens); + g1 = ggml_mul(ctx0, g1, A); + } + cb(g1, "kda_g1", il); + + g1 = ggml_reshape_4d(ctx0, g1, head_dim, n_head_kda, n_seq_tokens, n_seqs); + + ggml_tensor * beta = ggml_mul_mat(ctx0, layer.ssm_beta, cur); + beta = ggml_reshape_4d(ctx0, beta, 1, n_head_kda, n_seq_tokens, n_seqs); + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "kda_beta", il); + + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + ggml_tensor * state = build_rs(inp_rs, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_dim, head_dim, n_head_kda, n_seqs); + + // Match FLA l2 norm + constexpr float l2_eps = 1e-6f; + const float l2_scale = 1.0f / std::sqrt((float) head_dim); + Qcur = ggml_scale(ctx0, ggml_rms_norm(ctx0, Qcur, l2_eps / (float) head_dim), l2_scale); + Kcur = ggml_scale(ctx0, ggml_rms_norm(ctx0, Kcur, l2_eps / (float) head_dim), l2_scale); + + auto attn_out = build_delta_net(Qcur, Kcur, Vcur, g1, beta, state, il); + + ggml_tensor * output = ggml_cont(ctx0, attn_out.first); + ggml_tensor * new_state = attn_out.second; + cb(output, "kda_scan_out", il); + + ggml_build_forward_expand(gf, + ggml_cpy(ctx0, new_state, + ggml_view_1d(ctx0, ssm_states_all, hparams.n_embd_s() * n_seqs, + kv_head * hparams.n_embd_s() * ggml_element_size(ssm_states_all)))); + + // output gate, then RMSNorm(o) * Sigmoid(g2) + ggml_tensor * g_a = ggml_mul_mat(ctx0, layer.ssm_g_a, cur); + ggml_tensor * g2 = ggml_mul_mat(ctx0, layer.ssm_g_b, g_a); + g2 = ggml_reshape_3d(ctx0, g2, head_dim, n_head_kda, n_tokens); + + ggml_tensor * o = ggml_reshape_3d(ctx0, output, head_dim, n_head_kda, n_tokens); + ggml_tensor * normed = build_norm(o, layer.ssm_o_norm, nullptr, LLM_NORM_RMS, il); + cb(g2, "kda_g2", il); + cb(normed, "kda_normed", il); + ggml_tensor * gated = ggml_mul(ctx0, normed, ggml_sigmoid(ctx0, g2)); + + gated = ggml_cont_2d(ctx0, gated, d_inner, n_tokens); + cur = ggml_mul_mat(ctx0, layer.wo, gated); + cb(cur, "kda_out", il); + + return cur; +} + +// Scores pools of kpool consecutive tokens, expands the selected pools and the incomplete tail into an additive mask + +ggml_tensor * llama_model_glm5_next::graph::build_kpool_select( + ggml_tensor * cur, ggml_tensor * qr, ggml_tensor * kq_mask, const llama_layer & layer, + const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_kpool * inp_kpool, int il) { + + const auto * mctx_lid = mctx_hyb->get_idx(); + + const int64_t n_indexer_head = hparams.indexer_n_head; + const int64_t n_embd_indexer = hparams.indexer_head_size; + const int64_t kpool = hparams.indexer_kpool; + const int64_t n_pool = inp_kpool->pool_cells->ne[0]; + const int64_t n_new = inp_kpool->n_new; + + ggml_tensor * iq = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, qr); + iq = ggml_reshape_3d(ctx0, iq, n_embd_indexer, n_indexer_head, n_tokens); + cb(iq, "indexer_q", il); + + // Per-token key and pool gate scores, cached together + ggml_tensor * ik = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur); + ik = build_norm(ik, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il); + cb(ik, "indexer_k", il); + + ggml_tensor * ig = ggml_mul_mat(ctx0, layer.indexer_kpool_gate, cur); + cb(ig, "indexer_gate", il); + + // Cache rows store key | gate | pooled + ggml_tensor * pzero = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_indexer, n_tokens), 0.0f); + ggml_tensor * packed = ggml_concat(ctx0, ggml_concat(ctx0, ik, ig, 0), pzero, 0); + packed = ggml_reshape_3d(ctx0, packed, 3*n_embd_indexer, 1, n_tokens); + ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, packed, inp_kpool->k_idxs, il)); + + ggml_tensor * k_store = mctx_lid->get_k_storage(il); // [3*n_embd_indexer, kv_size, n_stream] + GGML_ASSERT(k_store->ne[0] == 3*n_embd_indexer); + const int64_t n_cells = k_store->ne[1]*k_store->ne[2]; + const int64_t n_kv = mctx_lid->get_n_kv(); + + ggml_tensor * kg_all = ggml_view_2d(ctx0, k_store, 2*n_embd_indexer, n_cells, k_store->nb[1], 0); + // View into the persistent pooled slots of the idx cache. Guarded by mem_idx_stale. + ggml_tensor * pooled_all = ggml_view_2d(ctx0, k_store, n_embd_indexer, n_cells, k_store->nb[1], + ggml_row_size(k_store->type, 2*n_embd_indexer)); + + ggml_tensor * pooled_new = nullptr; + // Pool only entries completed by this ubatch. + if (n_new > 0) { + ggml_tensor * rows = ggml_get_rows(ctx0, kg_all, ggml_reshape_1d(ctx0, inp_kpool->new_pool_idxs, kpool*n_new)); + rows = ggml_reshape_3d(ctx0, rows, 2*n_embd_indexer, kpool, n_new); + + ggml_tensor * pk = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_new, rows->nb[1], rows->nb[2], 0); + ggml_tensor * pg = ggml_view_3d(ctx0, rows, n_embd_indexer, kpool, n_new, rows->nb[1], rows->nb[2], ggml_row_size(rows->type, n_embd_indexer)); + + ggml_tensor * logits = ggml_add(ctx0, pg, layer.indexer_kpool_ape); + logits = ggml_cont(ctx0, ggml_permute(ctx0, logits, 1, 0, 2, 3)); // [kpool, head_dim, n_new] + ggml_tensor * probs = ggml_soft_max(ctx0, logits); + + pk = ggml_cont(ctx0, ggml_permute(ctx0, pk, 1, 0, 2, 3)); + pooled_new = ggml_sum_rows(ctx0, ggml_mul(ctx0, probs, pk)); // [1, head_dim, n_new] + pooled_new = ggml_reshape_2d(ctx0, pooled_new, n_embd_indexer, n_new); + cb(pooled_new, "indexer_pool_k_new", il); + + if (inp_kpool->cache_safe) { + // Write before the pool gather. + ggml_build_forward_expand(gf, ggml_set_rows(ctx0, pooled_all, pooled_new, inp_kpool->new_pool_rep)); + } + } + + ggml_tensor * pooled = nullptr; + if (inp_kpool->cache_safe) { + pooled = ggml_get_rows(ctx0, pooled_all, inp_kpool->pool_cells); + } else { + GGML_ASSERT(n_new <= n_pool); + ggml_tensor * pad = ggml_fill(ctx0, + ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_indexer, n_pool - n_new), 0.0f); + pooled = n_new > 0 ? ggml_concat(ctx0, pooled_new, pad, 1) : pad; + } + pooled = ggml_reshape_3d(ctx0, pooled, n_embd_indexer, 1, n_pool); + cb(pooled, "indexer_pool_k", il); + + ggml_tensor * sel_idx = nullptr; + { + ggml_tensor * weights = ggml_mul_mat(ctx0, layer.indexer_proj, cur); + weights = ggml_scale(ctx0, weights, 1.0f / sqrtf(float(n_embd_indexer * n_indexer_head))); + cb(weights, "indexer_weights", il); + + ggml_tensor * score = nullptr; + if (cparams.fused_lid) { + score = ggml_lightning_indexer(ctx0, iq, pooled, weights, inp_kpool->pool_mask); + res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, score, il}); + } else { + ggml_tensor * q_p = ggml_permute(ctx0, iq, 0, 2, 1, 3); // [head_dim, n_tokens, n_head] + ggml_tensor * k_p = ggml_permute(ctx0, pooled, 0, 2, 1, 3); // [head_dim, n_pool, 1] + + ggml_tensor * kq = ggml_mul_mat(ctx0, k_p, q_p); // [n_pool, n_tokens, n_head] + kq = ggml_cont(ctx0, ggml_permute(ctx0, kq, 2, 1, 0, 3)); // [n_head, n_tokens, n_pool] + score = ggml_relu(ctx0, kq); + score = ggml_mul(ctx0, score, weights); + score = ggml_sum_rows(ctx0, score); // [1, n_tokens, n_pool] + score = ggml_cont(ctx0, ggml_permute(ctx0, score, 2, 1, 0, 3)); // [n_pool, n_tokens, 1] + score = ggml_add(ctx0, score, inp_kpool->pool_mask); + } + cb(score, "indexer_score", il); + + const int64_t n_top_pool = std::min(n_pool, hparams.indexer_top_k / kpool); + ggml_tensor * top_k = ggml_top_k(ctx0, score, n_top_pool); // [n_top_pool, n_tokens], UNORDERED + + // The gather mask marks the first min(nv, n_top_pool) slots as the visible pools, so order the set by descending score. + ggml_tensor * sel_score = ggml_get_rows(ctx0, + ggml_reshape_3d(ctx0, score, 1, n_pool, n_tokens), top_k); // [1, n_top_pool, n_tokens] + ggml_tensor * sel_order = ggml_argsort(ctx0, + ggml_reshape_2d(ctx0, sel_score, n_top_pool, n_tokens), GGML_SORT_ORDER_DESC); + top_k = ggml_get_rows(ctx0, + ggml_reshape_3d(ctx0, ggml_cast(ctx0, top_k, GGML_TYPE_F32), 1, n_top_pool, n_tokens), sel_order); + top_k = ggml_cast(ctx0, ggml_cont(ctx0, ggml_reshape_2d(ctx0, top_k, n_top_pool, n_tokens)), GGML_TYPE_I32); + cb(top_k, "indexer_top_k", il); + + sel_idx = ggml_get_rows(ctx0, inp_kpool->pool_idxs, + ggml_reshape_1d(ctx0, top_k, n_top_pool*n_tokens)); // [kpool, n_top_pool*n_tokens] + sel_idx = ggml_reshape_2d(ctx0, sel_idx, kpool*n_top_pool, n_tokens); + + if (hparams.indexer_kpool_select_tail) { + // Append the incomplete tail with n_kv for missing cells. + sel_idx = ggml_concat(ctx0, sel_idx, inp_kpool->tail_idxs, 0); + } + } + const int64_t n_sel = sel_idx->ne[0]; + + // Gather returns selected cell indices and masks padding separately. + if (inp_kpool->gather) { + GGML_ASSERT(inp_kpool->gather_mask->ne[0] == n_sel && inp_kpool->gather_mask->ne[3] == n_tokens); + cb(sel_idx, "indexer_sel_idx", il); + return sel_idx; + } + + // Tie scatter storage lifetime to this layer's selected indices. + ggml_tensor * seed = ggml_cast(ctx0, ggml_view_1d(ctx0, sel_idx, 1, 0), GGML_TYPE_F32); + + ggml_tensor * mask_seed = kq_mask->type == GGML_TYPE_F32 ? seed : ggml_cast(ctx0, seed, kq_mask->type); + mask_seed = ggml_fill(ctx0, mask_seed, -INFINITY); + ggml_tensor * mask_all = ggml_repeat_4d(ctx0, mask_seed, 1, n_kv + 1, n_tokens, 1); + mask_all = ggml_reshape_3d(ctx0, mask_all, 1, n_kv + 1, n_tokens); + + ggml_tensor * zero_seed = ggml_fill(ctx0, seed, 0.0f); + ggml_tensor * zeros = ggml_repeat_4d(ctx0, zero_seed, 1, n_sel, n_tokens, 1); + zeros = ggml_reshape_3d(ctx0, zeros, 1, n_sel, n_tokens); + + ggml_tensor * sel = ggml_set_rows(ctx0, mask_all, zeros, ggml_reshape_3d(ctx0, sel_idx, n_sel, n_tokens, 1)); + sel = ggml_view_2d(ctx0, sel, n_kv, n_tokens, sel->nb[2], 0); + + // Fold causal visibility before shared-indexer reuse. + GGML_ASSERT(kq_mask->ne[0] == n_kv && kq_mask->ne[1]*kq_mask->ne[2]*kq_mask->ne[3] == n_tokens); + sel = ggml_add(ctx0, sel, ggml_reshape_2d(ctx0, kq_mask, n_kv, n_tokens)); + cb(sel, "indexer_sel", il); + + return sel; +} + +// Nope MLA layer with sparse attention over the indexer selection + +ggml_tensor * llama_model_glm5_next::graph::build_dsa_layer( + ggml_tensor * cur, const llama_layer & layer, + const llama_memory_hybrid_idx_context * mctx_hyb, llm_graph_input_attn_k * inp_attn, + llm_graph_input_kpool * inp_kpool, ggml_tensor ** prev_sel, int il) { + + const auto * mctx_mla = mctx_hyb->get_attn(); + + const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla(); + const int64_t kv_lora_rank = hparams.n_lora_kv; + const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - hparams.n_rot(); + const float kq_scale = 1.0f / sqrtf((float) n_embd_head_k_mla); + + GGML_ASSERT(hparams.n_rot() == 0 && "GLM5-Next MLA is nope-only"); + + ggml_tensor * qr = ggml_mul_mat(ctx0, layer.wq_a, cur); + qr = build_norm(qr, layer.attn_q_a_norm, nullptr, LLM_NORM_RMS, il); + cb(qr, "q_resid", il); + + ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq_b, qr); + q = ggml_reshape_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens); + + ggml_tensor * kv_cmpr = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur); + kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il); + kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens); + cb(kv_cmpr, "kv_cmpr", il); + + // absorb wk_b so the cache holds only the latent + ggml_tensor * q_absorbed = ggml_permute(ctx0, q, 0, 2, 1, 3); + q_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_absorbed); + q_absorbed = ggml_permute(ctx0, q_absorbed, 0, 2, 1, 3); + cb(q_absorbed, "q_absorbed", il); + + ggml_tensor * kq_mask = inp_attn->get_kq_mask(); + + ggml_tensor * sel = nullptr; + if (il >= (int) hparams.n_layer() || hparams.is_indexer_full(il)) { // the NextN block always has a full indexer + sel = build_kpool_select(cur, qr, kq_mask, layer, mctx_hyb, inp_kpool, il); + *prev_sel = sel; + } else { + GGML_ASSERT(*prev_sel != nullptr && "shared indexer layer must follow a full indexer layer"); + sel = *prev_sel; + } + + ggml_build_forward_expand(gf, q_absorbed); + ggml_build_forward_expand(gf, kv_cmpr); + ggml_build_forward_expand(gf, mctx_mla->cpy_k(ctx0, kv_cmpr, inp_attn->get_k_idxs(), il)); + + ggml_tensor * out = nullptr; + if (inp_kpool->gather) { + // Attend over gathered latents with the token dimension in ne[3]. + + ggml_build_forward_expand(gf, kq_mask); + + ggml_tensor * sel_idx = sel; // I32 [n_sel, n_tokens] + const int64_t n_sel = sel_idx->ne[0]; + + ggml_tensor * k = mctx_mla->get_k_storage(il); // [kv_lora_rank, kv_size, n_stream] + GGML_ASSERT(k->ne[0] == kv_lora_rank && "GLM5-Next MLA cache holds a single latent head"); + + ggml_tensor * rows = ggml_view_2d(ctx0, k, k->ne[0], k->ne[1]*k->ne[2], k->nb[1], 0); + ggml_tensor * k_g = ggml_get_rows(ctx0, rows, ggml_reshape_1d(ctx0, sel_idx, n_sel*n_tokens)); + k_g = ggml_reshape_4d(ctx0, k_g, k->ne[0], n_sel, 1, n_tokens); // F32 [kv_lora_rank, n_sel, 1, n_tokens] + cb(k_g, "kv_gathered", il); + + ggml_tensor * q_g = ggml_permute(ctx0, q_absorbed, 0, 2, 3, 1); // [kv_lora_rank, 1, n_head, n_tokens] + + ggml_tensor * kq = ggml_mul_mat(ctx0, k_g, q_g); // [n_sel, 1, n_head, n_tokens] + ggml_mul_mat_set_prec(kq, GGML_PREC_F32); + kq = ggml_soft_max_ext(ctx0, kq, inp_kpool->gather_mask, kq_scale, 0.0f); + cb(kq, "kq_soft_max_gathered", il); + + ggml_tensor * v_t = ggml_cont(ctx0, ggml_transpose(ctx0, k_g)); // [n_sel, kv_lora_rank, 1, n_tokens] + ggml_tensor * kqv = ggml_mul_mat(ctx0, v_t, kq); // [kv_lora_rank, 1, n_head, n_tokens] + kqv = ggml_mul_mat(ctx0, layer.wv_b, kqv); // [n_embd_head_v, 1, n_head, n_tokens] + cb(kqv, "kqv_gathered", il); + + out = ggml_cont(ctx0, ggml_permute(ctx0, kqv, 0, 2, 1, 3)); // [n_embd_head_v, n_head, 1, n_tokens] + out = ggml_reshape_2d(ctx0, out, kqv->ne[0]*n_head, n_tokens); + } else { + // The scatter selection already includes the causal mask. + ggml_tensor * mask = ggml_reshape_4d(ctx0, sel, kq_mask->ne[0], kq_mask->ne[1], kq_mask->ne[2], kq_mask->ne[3]); + cb(mask, "kq_mask_dsa", il); + + ggml_tensor * k = mctx_mla->get_k(ctx0, il); + ggml_tensor * v = ggml_view_4d(ctx0, k, kv_lora_rank, k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0); + + out = build_attn_mha(q_absorbed, k, v, nullptr, mask, nullptr, layer.wv_b, inp_kpool->n_sel, kq_scale, il); + } + cb(out, "kqv_out", il); + + out = ggml_mul_mat(ctx0, layer.wo, out); + cb(out, "attn_out", il); + + return out; +} diff --git a/src/models/models.h b/src/models/models.h index 9b87a40d5af..bcbc1cc670f 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1176,10 +1176,30 @@ struct llama_model_deepseek4 : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; - struct graph : public llm_graph_context { - graph(const llm_graph_params & params) : llm_graph_context(params) {} - graph(const llama_model & model, const llm_graph_params & params); + // manifold-constrained hyper-connections (mHC), shared by deepseek4 and derived model graphs like glm5-next. + template + struct graph_base : public Base { + graph_base(const llm_graph_params & params) : Base(params) {} + + // members of the dependent base used by the mHC helpers + using Base::ctx0; + using Base::res; + using Base::hparams; + using Base::cparams; + using Base::n_embd; + using Base::norm_rms_eps; + using Base::cb; + + // 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, @@ -1196,6 +1216,15 @@ 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; + }; + + struct graph : public graph_base<> { + graph(const llm_graph_params & params) : graph_base<>(params) {} + graph(const llama_model & model, const llm_graph_params & params); + ggml_tensor * build_hc_head( ggml_tensor * x, ggml_tensor * hc_fn, @@ -1289,14 +1318,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 { @@ -2492,6 +2513,39 @@ struct llama_model_kimi_k3 : public llama_model_base { std::unique_ptr 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; + + // mHC helpers from deepseek4, stacked on the delta net helpers + struct graph : public llama_model_deepseek4::graph_base { + graph(const llama_model & model, const llm_graph_params & params); + + 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 build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_kimi_linear : public llama_model_base { llama_model_kimi_linear(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index b2ea245ab84..2304242b8f5 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -118,6 +118,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) { n_embd = 128; n_head = 1; @@ -165,14 +166,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 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); @@ -191,10 +197,12 @@ 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) { - ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, uint32_t(576)); + // 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) { @@ -249,8 +257,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(n_layer, 4)); @@ -288,6 +298,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(8)); 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)); ms.add_kv(LLM_KV_ROPE_DIMENSION_SECTIONS, std::vector({n_embd_head/4, n_embd_head/4, n_embd_head/4, n_embd_head/4})); @@ -481,6 +493,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_MIMO2: case LLM_ARCH_KIMI_LINEAR: case LLM_ARCH_KIMI_K3: + case LLM_ARCH_GLM5_NEXT: case LLM_ARCH_STEP35: case LLM_ARCH_MISTRAL4: case LLM_ARCH_MELLUM: diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 72148a4d9a9..11eec2da1e7 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -63,6 +63,7 @@ #define KEY_PROJ_SAMPLE_WINDOW_SIDE "clip.vision.projector.window_side" #define KEY_PROJ_SPATIAL_OFFSETS "clip.vision.projector.spatial_offsets" #define KEY_SPATIAL_MERGE_SIZE "clip.vision.spatial_merge_size" +#define KEY_SWIGLU_CLAMP "clip.vision.swiglu_clamp" #define KEY_MM_PATCH_MERGE_TYPE "clip.vision.mm_patch_merge_type" #define KEY_IMAGE_GRID_PINPOINTS "clip.vision.image_grid_pinpoints" @@ -486,6 +487,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, @@ -551,6 +553,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_DEEPSEEK4V, "deepseek4v"}, { PROJECTOR_TYPE_LFM2A, "lfm2a"}, { PROJECTOR_TYPE_GLM4V, "glm4v"}, + { PROJECTOR_TYPE_GLM5V, "glm5v"}, { PROJECTOR_TYPE_YOUTUVL, "youtuvl"}, { PROJECTOR_TYPE_YASA2, "yasa2"}, { PROJECTOR_TYPE_KIMIK25, "kimik25"}, diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index f737ccc2452..2ba9bd1fe04 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -93,6 +93,13 @@ struct clip_hparams { float eps = 1e-6; float rope_theta = 0.0; + + std::pair swiglu_clamp_gate = {0.0f, 0.0f}; + std::pair swiglu_clamp_up = {0.0f, 0.0f}; + + bool has_swiglu_clamp() const { + return swiglu_clamp_gate.second > 0.0f || swiglu_clamp_up.second > 0.0f; + } int32_t n_expert_used = 0; std::vector feature_layers; int32_t attn_window_size = 0; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index f2e48753464..0c1256d00ca 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -648,6 +648,11 @@ ggml_tensor * clip_graph::build_ffn( switch (type_op) { case FFN_SILU: if (gate) { + if (hparams.has_swiglu_clamp()) { + cur = ggml_clamp(ctx0, cur, hparams.swiglu_clamp_gate.first, hparams.swiglu_clamp_gate.second); + tmp = ggml_clamp(ctx0, tmp, hparams.swiglu_clamp_up.first, hparams.swiglu_clamp_up.second); + cb(cur, "ffn_gate_clamped", il); + } cur = ggml_swiglu_split(ctx0, cur, tmp); cb(cur, "ffn_swiglu", il); } else { @@ -1082,6 +1087,7 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const builder = std::make_unique(ctx, img); } break; case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: { builder = std::make_unique(ctx, img); } break; @@ -1755,6 +1761,23 @@ struct clip_model_loader { hparams.set_limit_image_tokens(8, 4096); hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup } break; + case PROJECTOR_TYPE_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); + float swiglu_clamp = 0.0f; + get_f32(KEY_SWIGLU_CLAMP, swiglu_clamp, true); + if (swiglu_clamp > 0.0f) { + hparams.swiglu_clamp_gate = { -INFINITY, swiglu_clamp }; + hparams.swiglu_clamp_up = { -swiglu_clamp, swiglu_clamp }; + } + get_u32(KEY_IMAGE_MIN_PIXELS, hparams.image_min_pixels); + get_u32(KEY_IMAGE_MAX_PIXELS, hparams.image_max_pixels); + hparams.set_warmup_n_tokens(46*46); // avoid OOM on warmup + } break; case PROJECTOR_TYPE_LLAMA4: { hparams.rope_theta = 10000.0f; @@ -2572,6 +2595,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")); @@ -4042,6 +4066,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: @@ -4068,6 +4093,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: @@ -4149,6 +4175,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: { @@ -4781,6 +4808,7 @@ bool clip_encode(struct clip_ctx * ctx, struct clip_encode_params * params) { case PROJECTOR_TYPE_QWEN2VL: case PROJECTOR_TYPE_QWEN3VL: case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: { const int merge_ratio = hparams.n_merge; const int pw = image_size_width / patch_size; @@ -5995,6 +6023,7 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { case PROJECTOR_TYPE_GRANITE4_VISION: return ctx->model.qf_proj_blocks.size() * ctx->model.hparams.projection_dim; case PROJECTOR_TYPE_GLM4V: + case PROJECTOR_TYPE_GLM5V: return ctx->model.mm_ffn_down_w->ne[1]; case PROJECTOR_TYPE_MIMO_AUDIO: return ctx->model.mm_2_w->ne[1]; diff --git a/tools/mtmd/models/glm4v.cpp b/tools/mtmd/models/glm4v.cpp index 0e1d596b41b..1cbcea942a3 100644 --- a/tools/mtmd/models/glm4v.cpp +++ b/tools/mtmd/models/glm4v.cpp @@ -41,8 +41,10 @@ ggml_cgraph * clip_graph_glm4v::build() { inp = ggml_add(ctx0, inp, model.patch_bias); cb(inp, "patch_bias", -1); - // pos-conv norm - inp = build_norm(inp, model.norm_embd_w, model.norm_embd_b, norm_t, eps, -1); + // pos-conv norm (absent in GLM-5.3-Flash) + if (model.norm_embd_w) { + inp = build_norm(inp, model.norm_embd_w, model.norm_embd_b, norm_t, eps, -1); + } ggml_tensor * learned_pos_embd = nullptr; // Note: GLM-OCR does not have learned position embeddings diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index c11d35c87d7..a1121550932 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -786,6 +786,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(1, (int64_t) std::ceil(height * scale))); + canvas_w = align(std::max(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(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(1, std::min(canvas_h, (int64_t) std::floor(height * scale))); + const int content_w = (int) std::max(1, std::min(canvas_w, (int64_t) std::floor(width * scale))); + + clip_image_u8 content; + img_tool::resize(img, content, clip_image_size{content_w, content_h}, hparams.image_resize_algo, PAD_NONE); + + clip_image_u8 canvas; + canvas.set_size(clip_image_size{(int) canvas_w, (int) canvas_h}, img.is_placeholder()); + img_tool::fill(canvas, {0, 0, 0}); + img_tool::composite(canvas, content, 0, 0); + + mtmd_image_preproc_out output; + output.append(hparams, canvas, true); + return output; +} + // // mtmd_image_preprocessor_longest_edge // diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index e2cf6987232..4fa6207d03b 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -123,6 +123,12 @@ struct mtmd_image_preprocessor_dyn_size : mtmd_image_preprocessor { mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override; }; +// GLM 5.3 flash, similar to dyn_size, but each edge is aligned up to patch_size*n_merge, and max budget is met by a search over the height with the width scaled proportionally +struct mtmd_image_preprocessor_glm5v : mtmd_image_preprocessor { + mtmd_image_preprocessor_glm5v(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} + mtmd_image_preproc_out preprocess(const clip_image_u8 & img) const override; +}; + // similar to mtmd_image_preprocessor_dyn_size, but resize the image to have longest edge equal to hparams.image_longest_edge, while preserving aspect ratio struct mtmd_image_preprocessor_longest_edge : mtmd_image_preprocessor { mtmd_image_preprocessor_longest_edge(const clip_ctx * ctx) : mtmd_image_preprocessor(ctx) {} diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index e7f5f114ec0..05ab1b40a48 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -867,6 +867,13 @@ struct mtmd_context { img_end = "<|end_of_image|>"; image_preproc = std::make_unique(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(ctx_v); + } break; case PROJECTOR_TYPE_PADDLEOCR: { // <|IMAGE_START|> ... (image embeddings) ... <|IMAGE_END|>