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