From dfa0c0fee2b704fd2ac228d365d40502c3006c40 Mon Sep 17 00:00:00 2001 From: JJJYmmm <1650675829@qq.com> Date: Wed, 26 Aug 2026 21:06:05 +0800 Subject: [PATCH] model : add Qwen4-Exp --- conversion/__init__.py | 3 + conversion/qwen.py | 14 +- conversion/qwen3vl.py | 4 +- conversion/qwen4exp.py | 199 +++++ gguf-py/gguf/constants.py | 101 +++ gguf-py/gguf/gguf_writer.py | 27 + gguf-py/gguf/tensor_mapping.py | 76 +- src/llama-arch.cpp | 45 + src/llama-arch.h | 26 + src/llama-hparams.cpp | 8 + src/llama-hparams.h | 32 + src/llama-memory-hybrid-iswa.cpp | 1 + src/llama-memory-hybrid.cpp | 239 ++++- src/llama-memory-hybrid.h | 45 +- src/llama-memory-recurrent.cpp | 3 +- src/llama-memory-recurrent.h | 1 + src/llama-model-loader.cpp | 11 +- src/llama-model.cpp | 41 +- src/llama-model.h | 22 + src/llama-quant.cpp | 4 + src/models/delta-net-base.cpp | 12 +- src/models/models.h | 100 +++ src/models/qwen4exp.cpp | 1407 ++++++++++++++++++++++++++++++ 23 files changed, 2398 insertions(+), 23 deletions(-) create mode 100644 conversion/qwen4exp.py create mode 100644 src/models/qwen4exp.cpp diff --git a/conversion/__init__.py b/conversion/__init__.py index 8de97e95969a..be3e49196c59 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -226,6 +226,8 @@ "Qwen3Model": "qwen", "Qwen3MoeForCausalLM": "qwen", "Qwen3NextForCausalLM": "qwen", + "Qwen4ExpForCausalLM": "qwen4exp", + "Qwen4ExpForConditionalGeneration": "qwen4exp", "Qwen3OmniMoeForConditionalGeneration": "qwen3vl", "PocketTTSModel": "pockettts", "Qwen3TTSForConditionalGeneration": "qwen3tts", @@ -332,6 +334,7 @@ "Qwen3VLMoeForConditionalGeneration": "qwen3vl", "Qwen3_5ForConditionalGeneration": "qwen3vl", "Qwen3_5MoeForConditionalGeneration": "qwen3vl", + "Qwen4ExpForConditionalGeneration": "qwen3vl", "RADIOModel": "nemotron", "Sarashina2VisionForCausalLM": "sarashina2", "SmolVLMForConditionalGeneration": "smolvlm", diff --git a/conversion/qwen.py b/conversion/qwen.py index cdba8a63e9c9..4e4ac79ecdd4 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -292,16 +292,24 @@ class _QwenMtpMixin: mtp_only: bool _original_block_count: int | None = None opt_num_mtp_layers: int = 0 + # any mtp.* tensor at all, so "no MTP block" can be told from "MTP block not counted" + saw_mtp_tensor: bool = False def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.block_count = self.hparams["num_hidden_layers"] if not self.no_mtp: n_mtp = self.hparams.get("mtp_num_hidden_layers", 0) - # Qwen-3-Next doesn't include `mtp_num_hidden_layers` in config. + # Qwen-3-Next doesn't include `mtp_num_hidden_layers` in config, so fall back to the + # count of mtp.* tensors. That count is also 0 when the checkpoint has no MTP block, + # which is not an error - only an explicit --mtp export needs one if n_mtp == 0: - assert self.opt_num_mtp_layers != 0 n_mtp = self.opt_num_mtp_layers + if n_mtp == 0 and self.saw_mtp_tensor: + raise ValueError("the checkpoint carries mtp.* tensors but none of them name a " + "layer, so the MTP block count is unknown") + if n_mtp == 0 and self.mtp_only: + raise ValueError("--mtp: this checkpoint has no MTP block to export") self.block_count += n_mtp self.tensor_map = gguf.get_tensor_name_map(self.model_arch, self.block_count) @@ -310,6 +318,7 @@ def index_tensors(self, remote_hf_model_id: str | None = None) -> dict[str, Call key = next((k for k in ["n_layers", "num_hidden_layers", "n_layer", "num_layers"] if k in hparams), None) type(self)._original_block_count = hparams.get(key) type(self).opt_num_mtp_layers = 0 + type(self).saw_mtp_tensor = False return super().index_tensors(remote_hf_model_id=remote_hf_model_id) # ty: ignore[unresolved-attribute] @classmethod @@ -322,6 +331,7 @@ def filter_tensors(cls, item): if name.startswith("model.mtp."): name = name.replace("model.", "", 1) if name.startswith("mtp."): + cls.saw_mtp_tensor = True if cls.no_mtp: return None remapper = { diff --git a/conversion/qwen3vl.py b/conversion/qwen3vl.py index 4fec708c9ff3..a177d5e16bc5 100644 --- a/conversion/qwen3vl.py +++ b/conversion/qwen3vl.py @@ -13,8 +13,8 @@ from .qwenvl import Qwen25AudioModel -@ModelBase.register("Qwen3VLForConditionalGeneration", "Qwen3VLMoeForConditionalGeneration", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration") -@ModelBase.example("Qwen/Qwen3-VL-4B-Instruct", "Qwen/Qwen3-VL-30B-A3B-Instruct", "Qwen/Qwen3.5-9B", "Qwen/Qwen3.5-35B-A3B") +@ModelBase.register("Qwen3VLForConditionalGeneration", "Qwen3VLMoeForConditionalGeneration", "Qwen3_5ForConditionalGeneration", "Qwen3_5MoeForConditionalGeneration", "Qwen4ExpForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3-VL-4B-Instruct", "Qwen/Qwen3-VL-30B-A3B-Instruct", "Qwen/Qwen3.5-9B", "Qwen/Qwen3.5-35B-A3B", "Qwen/Qwen4-Exp") class Qwen3VLVisionModel(MmprojModel): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py new file mode 100644 index 000000000000..d411f0375581 --- /dev/null +++ b/conversion/qwen4exp.py @@ -0,0 +1,199 @@ +from __future__ import annotations + +from typing import Callable, Iterable, TYPE_CHECKING + +import torch + +if TYPE_CHECKING: + from torch import Tensor + +from .base import LazyTorchTensor, ModelBase, gguf, logger + +from .qwen import _LinearAttentionVReorderBase, _Qwen35MRopeMixin + + +@ModelBase.register("Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM") +@ModelBase.example("Qwen/Qwen4-Exp") +class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): + model_arch = gguf.MODEL_ARCH.QWEN4EXP + + # n-gram embedding tables are stored as many row shards, joined into one tensor here. + # Keyed by (layer, shard) so two PLE layers cannot mix their shards + _ngram_shards: dict[tuple[int, int], Tensor] | None = None + + def set_gguf_parameters(self): + super().set_gguf_parameters() + + self.gguf_writer.add_hyper_connection_count(self.hparams["hc_count"]) + + # the MTP block consumes the target's hyper-connection streams, not the collapsed hidden + # state, so both files declare the wider row: the target for the nextn read-back and the + # draft for common/speculative.cpp. Same as DeepSeek-V4, see conversion/deepseek.py. + self.gguf_writer.add_embedding_length_out( + self.hparams["hc_count"]*self.hparams["hidden_size"]) + self.gguf_writer.add_hyper_connection_lowrank(self.hparams["hc_lowrank"]) + + # the sigmoid output gate of the linear attention layers is hardcoded in the graph + gate = self.hparams.get("output_gate_type") or self.hparams.get("hidden_act") + if gate != "sigmoid": + raise ValueError(f"unsupported output_gate_type {gate!r} (only 'sigmoid' is supported)") + + # the MTP block has no PLE layer (the reference clears ple_layer_ids for it) + if not self.mtp_only and self.hparams.get("ple_layer_ids"): + if len(self.hparams["ple_layer_ids"]) != 1: + raise ValueError("only a single PLE layer is supported") + self.gguf_writer.add_ple_embedding_length(self.hparams.get("ple_embed_dim") or self.hparams["hidden_size"]) + self.gguf_writer.add_ple_conv_kernel(self.hparams["ple_conv_kernel_size"]) + + # the n-gram hash constants are needed on the host, keep them as metadata. They are read + # here and not in modify_tensors because that already casts int64 buffers to float32. + self.gguf_writer.add_ple_ngram_multipliers(self._read_u64(".ple_embedding.layer_multipliers")) + self.gguf_writer.add_ple_ngram_vocab_sizes(self._read_u64(".ple_embedding.ngram_heads_vocab_sizes")) + self.gguf_writer.add_ple_ngram_offsets(self._read_u64(".ple_embedding.ngram_heads_offsets")) + + # export the layer types instead of letting llama.cpp infer them from + # full_attention_interval, so an irregular layout cannot misalign silently + if layer_types := self.hparams.get("layer_types"): + n_layer = self.hparams["num_hidden_layers"] + if len(layer_types) != n_layer: + raise ValueError(f"layer_types has {len(layer_types)} entries, " + f"expected num_hidden_layers ({n_layer})") + + recurrent = [] + for t in layer_types: + if t == "linear_attention": + recurrent.append(True) + elif t.endswith("sparse_attention"): + recurrent.append(False) + else: + raise ValueError(f"unsupported qwen4exp layer type {t!r}") + + # llama.cpp reads this with get_key_or_arr(.., n_layer_all), which wants exactly + # block_count entries. The MTP block is sparse attention, so it is not recurrent + recurrent += [False] * (self.block_count - n_layer) + self.gguf_writer.add_recurrent_layers(recurrent) + + if self.hparams.get("indexer_n_heads") is not None: + self.gguf_writer.add_indexer_head_count(self.hparams["indexer_n_heads"]) + self.gguf_writer.add_indexer_key_length(self.hparams["indexer_head_dim"]) + self.gguf_writer.add_indexer_top_k (self.hparams["indexer_budget"]) + self.gguf_writer.add_indexer_block_size(self.hparams["indexer_compress_ratio"]) + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + yield from super().generate_extra_tensors() + + # the reference adds the two MTP input projections, and A*e + B*h == [A|B]*concat(e, h), + # so they join into the single eh_proj the tensor map already knows. + # ref: conversion/deepseek.py, which joins the DeepSeek-V4 pair the same way + e_name = "mtp.fc_embedding.weight" + h_name = "mtp.fc_hidden.weight" + + have_e = e_name in self.model_tensors + have_h = h_name in self.model_tensors + if not have_e and not have_h: + return + if not have_e or not have_h: + raise KeyError(f"unpaired MTP input projection: need both {e_name} and {h_name}") + + e = LazyTorchTensor.to_eager(self.model_tensors[e_name]()) + h = LazyTorchTensor.to_eager(self.model_tensors[h_name]()) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, + self.hparams["num_hidden_layers"]), + torch.cat([e, h], dim=1).contiguous()) + + del self.model_tensors[e_name] + del self.model_tensors[h_name] + + def tensor_force_quant(self, name, new_name, bid, n_dims): + # the graph slices the PLE conv weight and multiplies it, which needs F32 + if ".ple_conv1d.weight" in new_name: + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name = item[0] + + # the vision tower goes to the mmproj file + if name.startswith("model.visual."): + return None + + # the MTP block brings its own hyper-connection mixer, which takes the model-level slot of + # an MTP-only file. Rename it here, before _QwenMtpMixin drops it as a non-MTP tensor. + if name.startswith("mtp.hyper_connection_mixer."): + return None if cls.no_mtp else (name.replace("mtp.", "model.", 1), item[1]) + + return super().filter_tensors(item) + + def _read_u64(self, suffix: str) -> list[int]: + name = next((n for n in self.model_tensors if n.endswith(suffix)), None) + if name is None: + raise ValueError(f"missing tensor *{suffix}, needed for the PLE n-gram hash") + return [int(v) for v in LazyTorchTensor.to_eager(self.model_tensors[name]()).tolist()] + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # hash constants, already exported as metadata by set_gguf_parameters + if ".ple_embedding." in name and not name.endswith(".weight"): + return + + if ".ple_embedding.ngram_embedding.shard_" in name: + shard = int(name.rpartition(".ple_embedding.ngram_embedding.shard_")[2].partition(".")[0]) + n_shard = self.hparams["split_ngram_parts"] + + assert bid is not None + + if self._ngram_shards is None: + self._ngram_shards = {} + self._ngram_shards[(bid, shard)] = data_torch + + shards = {s: t for (b, s), t in self._ngram_shards.items() if b == bid} + if len(shards) < n_shard: + return + + # the shards are joined in shard order, so the numbering has to be 0-based and dense + if sorted(shards) != list(range(n_shard)): + raise ValueError(f"layer {bid}: expected n-gram shards 0..{n_shard - 1}, " + f"got {sorted(shards)}") + + logger.info(f"joining {n_shard} n-gram embedding shards") + rows = sum(shards[i].shape[0] for i in range(n_shard)) + + # llama.cpp indexes the joined table as offs[h] + 0..vocab[h]-1, so a table that is + # too short has to fail here and not hours later as a shape mismatch + offs = self._read_u64(".ple_embedding.ngram_heads_offsets") + vocab = self._read_u64(".ple_embedding.ngram_heads_vocab_sizes") + need = max(o + v for o, v in zip(offs, vocab)) + if rows < need: + raise ValueError(f"joined n-gram table has {rows} rows, " + f"the hash offsets need {need}") + + joined = torch.empty((rows, data_torch.shape[1]), dtype=data_torch.dtype) + row = 0 + for i in range(n_shard): + part = LazyTorchTensor.to_eager(self._ngram_shards.pop((bid, i))) + joined[row:row + part.shape[0]] = part + row += part.shape[0] + del part + + yield (self.format_tensor_name(gguf.MODEL_TENSOR.PLE_NGRAM_EMBD, bid), joined) + return + + # the QSA indexer packs its query and key projections together + if name.endswith(".indexer.index_qk_proj.weight"): + n_q = self.hparams["indexer_n_heads"]*self.hparams["indexer_head_dim"] + assert bid is not None + yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid), data_torch[:n_q]) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid), data_torch[n_q:]) + return + + # PLE norms are also 1-centered, the base class only catches names ending in "norm.weight" + if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight")): + data_torch = data_torch + 1 + + yield from super().modify_tensors(data_torch, name, bid) + + def prepare_tensors(self): + super().prepare_tensors() + + if self._ngram_shards: + raise ValueError(f"unprocessed n-gram embedding shards: {sorted(self._ngram_shards)}") diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index 8f6f55519c33..10681c33b33c 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -212,6 +212,7 @@ class Attention: SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" TEMPERATURE_SCALE = "{arch}.attention.temperature_scale" ROPE_PATTERN = "{arch}.attention.rope_pattern" + RECURRENT_LAYERS = "{arch}.attention.recurrent_layers" class Indexer: HEAD_COUNT = "{arch}.attention.indexer.head_count" @@ -223,9 +224,17 @@ class Indexer: class HyperConnection: COUNT = "{arch}.hyper_connection.count" + LOWRANK = "{arch}.hyper_connection.lowrank" SINKHORN_ITERATIONS = "{arch}.hyper_connection.sinkhorn_iterations" EPSILON = "{arch}.hyper_connection.epsilon" + class Ple: + EMBD_LENGTH = "{arch}.ple.embedding_length" + CONV_KERNEL = "{arch}.ple.conv_kernel" + NGRAM_MULTIPLIERS = "{arch}.ple.ngram_multipliers" + NGRAM_VOCAB_SIZES = "{arch}.ple.ngram_vocab_sizes" + NGRAM_OFFSETS = "{arch}.ple.ngram_offsets" + class Rope: DIMENSION_COUNT = "{arch}.rope.dimension_count" DIMENSION_COUNT_SWA = "{arch}.rope.dimension_count_swa" @@ -494,6 +503,7 @@ class MODEL_ARCH(IntEnum): QWEN3VLMOE = auto() QWEN35 = auto() QWEN35MOE = auto() + QWEN4EXP = auto() PHI2 = auto() PHI3 = auto() PHIMOE = auto() @@ -780,6 +790,24 @@ class MODEL_TENSOR(IntEnum): HC_FFN_FN = auto() HC_FFN_BASE = auto() HC_FFN_SCALE = auto() + HC_ATTN_NORM = auto() # qwen4exp + HC_ATTN_DOWN = auto() # qwen4exp + HC_ATTN_UP = auto() # qwen4exp + HC_ATTN_INJECT = auto() # qwen4exp + HC_FFN_NORM = auto() # qwen4exp + HC_FFN_DOWN = auto() # qwen4exp + HC_FFN_UP = auto() # qwen4exp + HC_FFN_INJECT = auto() # qwen4exp + HC_MIX_NORM = auto() # qwen4exp + HC_MIX_DOWN = auto() # qwen4exp + HC_MIX_UP = auto() # qwen4exp + PLE_NGRAM_EMBD = auto() # qwen4exp + PLE_KEY = auto() # qwen4exp + PLE_VALUE = auto() # qwen4exp + PLE_KEY_NORM = auto() # qwen4exp + PLE_QUERY_NORM = auto() # qwen4exp + PLE_CONV_NORM = auto() # qwen4exp + PLE_CONV1D = auto() # qwen4exp ATTN_COMPRESSOR_WKV = auto() ATTN_COMPRESSOR_WGATE = auto() ATTN_COMPRESSOR_APE = auto() @@ -1217,6 +1245,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.QWEN3VLMOE: "qwen3vlmoe", MODEL_ARCH.QWEN35: "qwen35", MODEL_ARCH.QWEN35MOE: "qwen35moe", + MODEL_ARCH.QWEN4EXP: "qwen4exp", MODEL_ARCH.PHI2: "phi2", MODEL_ARCH.PHI3: "phi3", MODEL_ARCH.PHIMOE: "phimoe", @@ -1502,6 +1531,24 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.HC_FFN_FN: "blk.{bid}.hc_ffn_fn", MODEL_TENSOR.HC_FFN_BASE: "blk.{bid}.hc_ffn_base", MODEL_TENSOR.HC_FFN_SCALE: "blk.{bid}.hc_ffn_scale", + MODEL_TENSOR.HC_ATTN_NORM: "blk.{bid}.hc_attn_norm", + MODEL_TENSOR.HC_ATTN_DOWN: "blk.{bid}.hc_attn_down", + MODEL_TENSOR.HC_ATTN_UP: "blk.{bid}.hc_attn_up", + MODEL_TENSOR.HC_ATTN_INJECT: "blk.{bid}.hc_attn_inject", + MODEL_TENSOR.HC_FFN_NORM: "blk.{bid}.hc_ffn_norm", + MODEL_TENSOR.HC_FFN_DOWN: "blk.{bid}.hc_ffn_down", + MODEL_TENSOR.HC_FFN_UP: "blk.{bid}.hc_ffn_up", + MODEL_TENSOR.HC_FFN_INJECT: "blk.{bid}.hc_ffn_inject", + MODEL_TENSOR.HC_MIX_NORM: "output_hc_norm", + MODEL_TENSOR.HC_MIX_DOWN: "output_hc_down", + MODEL_TENSOR.HC_MIX_UP: "output_hc_up", + MODEL_TENSOR.PLE_NGRAM_EMBD: "blk.{bid}.ple_ngram_embd", + MODEL_TENSOR.PLE_KEY: "blk.{bid}.ple_key", + MODEL_TENSOR.PLE_VALUE: "blk.{bid}.ple_value", + MODEL_TENSOR.PLE_KEY_NORM: "blk.{bid}.ple_key_norm", + MODEL_TENSOR.PLE_QUERY_NORM: "blk.{bid}.ple_query_norm", + MODEL_TENSOR.PLE_CONV_NORM: "blk.{bid}.ple_conv_norm", + MODEL_TENSOR.PLE_CONV1D: "blk.{bid}.ple_conv1d", MODEL_TENSOR.ATTN_COMPRESSOR_WKV: "blk.{bid}.attn_compressor_kv", MODEL_TENSOR.ATTN_COMPRESSOR_WGATE: "blk.{bid}.attn_compressor_gate", MODEL_TENSOR.ATTN_COMPRESSOR_APE: "blk.{bid}.attn_compressor_ape", @@ -2795,6 +2842,60 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.QWEN4EXP: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.INDEXER_Q_PROJ, + MODEL_TENSOR.INDEXER_K_PROJ, + MODEL_TENSOR.INDEXER_Q_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_INP_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_ALPHA, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.HC_ATTN_NORM, + MODEL_TENSOR.HC_ATTN_DOWN, + MODEL_TENSOR.HC_ATTN_UP, + MODEL_TENSOR.HC_ATTN_INJECT, + MODEL_TENSOR.HC_FFN_NORM, + MODEL_TENSOR.HC_FFN_DOWN, + MODEL_TENSOR.HC_FFN_UP, + MODEL_TENSOR.HC_FFN_INJECT, + MODEL_TENSOR.HC_MIX_NORM, + MODEL_TENSOR.HC_MIX_DOWN, + MODEL_TENSOR.HC_MIX_UP, + MODEL_TENSOR.PLE_NGRAM_EMBD, + MODEL_TENSOR.PLE_KEY, + MODEL_TENSOR.PLE_VALUE, + MODEL_TENSOR.PLE_KEY_NORM, + MODEL_TENSOR.PLE_QUERY_NORM, + MODEL_TENSOR.PLE_CONV_NORM, + MODEL_TENSOR.PLE_CONV1D, + # the MTP block reuses the normal per-layer tensors at index n_layer, plus these + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_EH_PROJ, + ], MODEL_ARCH.PLAMO: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index d8a96a27bdd5..7f5a96ac0754 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -836,6 +836,11 @@ def add_sliding_window_pattern(self, value: int | Sequence[bool]) -> None: def add_rope_pattern(self, value: Sequence[bool]) -> None: self.add_array(Keys.Attention.ROPE_PATTERN.format(arch=self.arch), value) + # true means the layer is recurrent (linear attention), false means full attention. + # llama.cpp falls back to full_attention_interval when this is absent + def add_recurrent_layers(self, value: Sequence[bool]) -> None: + self.add_array(Keys.Attention.RECURRENT_LAYERS.format(arch=self.arch), value) + def add_dense_features_dims(self, dense:str, in_f:int, out_f:int) -> None: self.add_uint32(Keys.LLM.DENSE_FEAT_IN_SIZE.format(arch=self.arch, dense=dense), in_f) self.add_uint32(Keys.LLM.DENSE_FEAT_OUT_SIZE.format(arch=self.arch, dense=dense), out_f) @@ -1023,6 +1028,28 @@ def add_attention_compress_rope_freq_base(self, value: float) -> None: def add_hyper_connection_count(self, count: int) -> None: self.add_uint32(Keys.HyperConnection.COUNT.format(arch=self.arch), count) + def add_hyper_connection_lowrank(self, lowrank: int) -> None: + self.add_uint32(Keys.HyperConnection.LOWRANK.format(arch=self.arch), lowrank) + + def add_ple_embedding_length(self, length: int) -> None: + self.add_uint32(Keys.Ple.EMBD_LENGTH.format(arch=self.arch), length) + + def add_ple_conv_kernel(self, size: int) -> None: + self.add_uint32(Keys.Ple.CONV_KERNEL.format(arch=self.arch), size) + + def _add_i64_array(self, key: str, values: Sequence[int]) -> None: + # signed: the values come from a torch int64 buffer and the hash is done in int64 + self.add_key_value(key, list(values), GGUFValueType.ARRAY, GGUFValueType.INT64) + + def add_ple_ngram_multipliers(self, values: Sequence[int]) -> None: + self._add_i64_array(Keys.Ple.NGRAM_MULTIPLIERS.format(arch=self.arch), values) + + def add_ple_ngram_vocab_sizes(self, values: Sequence[int]) -> None: + self._add_i64_array(Keys.Ple.NGRAM_VOCAB_SIZES.format(arch=self.arch), values) + + def add_ple_ngram_offsets(self, values: Sequence[int]) -> None: + self._add_i64_array(Keys.Ple.NGRAM_OFFSETS.format(arch=self.arch), values) + def add_hyper_connection_sinkhorn_iterations(self, count: int) -> None: self.add_uint32(Keys.HyperConnection.SINKHORN_ITERATIONS.format(arch=self.arch), count) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index ef580518e97d..eb5cca9116a9 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -93,6 +93,18 @@ class TensorNameMap: MODEL_TENSOR.DENSE_3_OUT: ( "dense_3_out", # embeddinggemma ), + MODEL_TENSOR.HC_MIX_NORM: ( + "model.hyper_connection_mixer.hc_norm", # qwen4exp + ), + + MODEL_TENSOR.HC_MIX_DOWN: ( + "model.hyper_connection_mixer.input_mix_weight_down", # qwen4exp + ), + + MODEL_TENSOR.HC_MIX_UP: ( + "model.hyper_connection_mixer.input_mix_weight_up", # qwen4exp + ), + # Output norm MODEL_TENSOR.OUTPUT_NORM: ( "gpt_neox.final_layer_norm", # gptneox @@ -1301,8 +1313,9 @@ class TensorNameMap: ), MODEL_TENSOR.INDEXER_K_NORM: ( - "model.layers.{bid}.self_attn.indexer.k_norm", # DSA - "model.layers.{bid}.self_attn.index_k_norm", # MSA + "model.layers.{bid}.self_attn.indexer.k_norm", # DSA + "model.layers.{bid}.self_attn.index_k_norm", # MSA + "model.layers.{bid}.self_attn.indexer.k_layernorm", # qwen4exp ), MODEL_TENSOR.INDEXER_PROJ: ( @@ -1326,7 +1339,64 @@ class TensorNameMap: ), MODEL_TENSOR.INDEXER_Q_NORM: ( - "model.layers.{bid}.self_attn.index_q_norm", # MSA + "model.layers.{bid}.self_attn.index_q_norm", # MSA + "model.layers.{bid}.self_attn.indexer.q_layernorm", # qwen4exp + ), + + MODEL_TENSOR.HC_ATTN_NORM: ( + "model.layers.{bid}.attn_hyper_connection.hc_norm", # qwen4exp + ), + + MODEL_TENSOR.HC_ATTN_DOWN: ( + "model.layers.{bid}.attn_hyper_connection.input_mix_weight_down", # qwen4exp + ), + + MODEL_TENSOR.HC_ATTN_UP: ( + "model.layers.{bid}.attn_hyper_connection.input_mix_weight_up", # qwen4exp + ), + + MODEL_TENSOR.HC_ATTN_INJECT: ( + "model.layers.{bid}.attn_hyper_connection.block_inject_weight", # qwen4exp + ), + + MODEL_TENSOR.HC_FFN_NORM: ( + "model.layers.{bid}.mlp_hyper_connection.hc_norm", # qwen4exp + ), + + MODEL_TENSOR.HC_FFN_DOWN: ( + "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_down", # qwen4exp + ), + + MODEL_TENSOR.HC_FFN_UP: ( + "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_up", # qwen4exp + ), + + MODEL_TENSOR.HC_FFN_INJECT: ( + "model.layers.{bid}.mlp_hyper_connection.block_inject_weight", # qwen4exp + ), + + MODEL_TENSOR.PLE_KEY: ( + "model.layers.{bid}.ple.key_proj", # qwen4exp + ), + + MODEL_TENSOR.PLE_VALUE: ( + "model.layers.{bid}.ple.value_proj", # qwen4exp + ), + + MODEL_TENSOR.PLE_KEY_NORM: ( + "model.layers.{bid}.ple.norm_key", # qwen4exp + ), + + MODEL_TENSOR.PLE_QUERY_NORM: ( + "model.layers.{bid}.ple.norm_query", # qwen4exp + ), + + MODEL_TENSOR.PLE_CONV_NORM: ( + "model.layers.{bid}.ple.norm_conv", # qwen4exp + ), + + MODEL_TENSOR.PLE_CONV1D: ( + "model.layers.{bid}.ple.conv1d", # qwen4exp ), ############################################################################ diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 025f9fb54dec..07bf881661c7 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -40,6 +40,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_QWEN3VLMOE, "qwen3vlmoe" }, { LLM_ARCH_QWEN35, "qwen35" }, { LLM_ARCH_QWEN35MOE, "qwen35moe" }, + { LLM_ARCH_QWEN4EXP, "qwen4exp" }, { LLM_ARCH_PHI2, "phi2" }, { LLM_ARCH_PHI3, "phi3" }, { LLM_ARCH_PHIMOE, "phimoe" }, @@ -291,9 +292,16 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_RECURRENT_LAYERS, "%s.attention.recurrent_layers" }, { LLM_KV_HYPER_CONNECTION_COUNT, "%s.hyper_connection.count" }, + { LLM_KV_HYPER_CONNECTION_LOWRANK, "%s.hyper_connection.lowrank" }, { LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, "%s.hyper_connection.sinkhorn_iterations" }, { LLM_KV_HYPER_CONNECTION_EPSILON, "%s.hyper_connection.epsilon" }, + { LLM_KV_PLE_EMBD_LENGTH, "%s.ple.embedding_length" }, + { LLM_KV_PLE_CONV_KERNEL, "%s.ple.conv_kernel" }, + { LLM_KV_PLE_NGRAM_MULTIPLIERS, "%s.ple.ngram_multipliers" }, + { LLM_KV_PLE_NGRAM_VOCAB_SIZES, "%s.ple.ngram_vocab_sizes" }, + { LLM_KV_PLE_NGRAM_OFFSETS, "%s.ple.ngram_offsets" }, + { LLM_KV_HASH_LAYER_COUNT, "%s.hash_layer_count" }, { LLM_KV_ROPE_DIMENSION_COUNT, "%s.rope.dimension_count" }, @@ -506,6 +514,24 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_HC_FFN_FN, "blk.%d.hc_ffn_fn" }, { LLM_TENSOR_HC_FFN_BASE, "blk.%d.hc_ffn_base" }, { LLM_TENSOR_HC_FFN_SCALE, "blk.%d.hc_ffn_scale" }, + { LLM_TENSOR_HC_ATTN_NORM, "blk.%d.hc_attn_norm" }, + { LLM_TENSOR_HC_ATTN_DOWN, "blk.%d.hc_attn_down" }, + { LLM_TENSOR_HC_ATTN_UP, "blk.%d.hc_attn_up" }, + { LLM_TENSOR_HC_ATTN_INJECT, "blk.%d.hc_attn_inject" }, + { LLM_TENSOR_HC_FFN_NORM, "blk.%d.hc_ffn_norm" }, + { LLM_TENSOR_HC_FFN_DOWN, "blk.%d.hc_ffn_down" }, + { LLM_TENSOR_HC_FFN_UP, "blk.%d.hc_ffn_up" }, + { LLM_TENSOR_HC_FFN_INJECT, "blk.%d.hc_ffn_inject" }, + { LLM_TENSOR_HC_MIX_NORM, "output_hc_norm" }, + { LLM_TENSOR_HC_MIX_DOWN, "output_hc_down" }, + { LLM_TENSOR_HC_MIX_UP, "output_hc_up" }, + { LLM_TENSOR_PLE_NGRAM_EMBD, "blk.%d.ple_ngram_embd" }, + { LLM_TENSOR_PLE_KEY, "blk.%d.ple_key" }, + { LLM_TENSOR_PLE_VALUE, "blk.%d.ple_value" }, + { LLM_TENSOR_PLE_KEY_NORM, "blk.%d.ple_key_norm" }, + { LLM_TENSOR_PLE_QUERY_NORM, "blk.%d.ple_query_norm" }, + { LLM_TENSOR_PLE_CONV_NORM, "blk.%d.ple_conv_norm" }, + { LLM_TENSOR_PLE_CONV1D, "blk.%d.ple_conv1d" }, { LLM_TENSOR_ATTN_COMPRESSOR_WKV, "blk.%d.attn_compressor_kv" }, { LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "blk.%d.attn_compressor_gate" }, { LLM_TENSOR_ATTN_COMPRESSOR_APE, "blk.%d.attn_compressor_ape" }, @@ -710,6 +736,24 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_HC_FFN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_HC_FFN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, {LLM_TENSOR_HC_FFN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_ATTN_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_ATTN_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_ATTN_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_ATTN_INJECT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_FFN_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_INJECT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_MIX_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_HC_MIX_DOWN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_MIX_UP, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_PLE_NGRAM_EMBD, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_PLE_KEY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_PLE_VALUE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_PLE_KEY_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_QUERY_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_CONV_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_CONV1D, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_ATTN_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, @@ -995,6 +1039,7 @@ bool llm_arch_is_recurrent(const llm_arch & arch) { bool llm_arch_is_hybrid(const llm_arch & arch) { switch (arch) { + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_JAMBA: case LLM_ARCH_FALCON_H1: case LLM_ARCH_PLAMO2: diff --git a/src/llama-arch.h b/src/llama-arch.h index 7159e23bf7ab..5e81f6e9c568 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -45,6 +45,7 @@ enum llm_arch { LLM_ARCH_QWEN3VLMOE, LLM_ARCH_QWEN35, LLM_ARCH_QWEN35MOE, + LLM_ARCH_QWEN4EXP, LLM_ARCH_PHI2, LLM_ARCH_PHI3, LLM_ARCH_PHIMOE, @@ -296,9 +297,16 @@ enum llm_kv { LLM_KV_ATTENTION_RECURRENT_LAYERS, LLM_KV_HYPER_CONNECTION_COUNT, + LLM_KV_HYPER_CONNECTION_LOWRANK, LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, LLM_KV_HYPER_CONNECTION_EPSILON, + LLM_KV_PLE_EMBD_LENGTH, + LLM_KV_PLE_CONV_KERNEL, + LLM_KV_PLE_NGRAM_MULTIPLIERS, + LLM_KV_PLE_NGRAM_VOCAB_SIZES, + LLM_KV_PLE_NGRAM_OFFSETS, + LLM_KV_HASH_LAYER_COUNT, LLM_KV_ROPE_DIMENSION_COUNT, @@ -571,6 +579,24 @@ enum llm_tensor { LLM_TENSOR_HC_FFN_FN, LLM_TENSOR_HC_FFN_BASE, LLM_TENSOR_HC_FFN_SCALE, + LLM_TENSOR_HC_ATTN_NORM, + LLM_TENSOR_HC_ATTN_DOWN, + LLM_TENSOR_HC_ATTN_UP, + LLM_TENSOR_HC_ATTN_INJECT, + LLM_TENSOR_HC_FFN_NORM, + LLM_TENSOR_HC_FFN_DOWN, + LLM_TENSOR_HC_FFN_UP, + LLM_TENSOR_HC_FFN_INJECT, + LLM_TENSOR_HC_MIX_NORM, + LLM_TENSOR_HC_MIX_DOWN, + LLM_TENSOR_HC_MIX_UP, + LLM_TENSOR_PLE_NGRAM_EMBD, + LLM_TENSOR_PLE_KEY, + LLM_TENSOR_PLE_VALUE, + LLM_TENSOR_PLE_KEY_NORM, + LLM_TENSOR_PLE_QUERY_NORM, + LLM_TENSOR_PLE_CONV_NORM, + LLM_TENSOR_PLE_CONV1D, LLM_TENSOR_ATTN_COMPRESSOR_WKV, LLM_TENSOR_ATTN_COMPRESSOR_WGATE, LLM_TENSOR_ATTN_COMPRESSOR_APE, diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index cbe31134ff45..36b087c9b1ac 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -181,6 +181,10 @@ uint32_t llama_hparams::n_embd_v_gqa_max() const { } uint32_t llama_hparams::n_embd_r() const { + if (n_embd_r_override != 0) { + return n_embd_r_override; + } + if (wkv_head_size != 0) { // for RWKV models return token_shift_count * n_embd; @@ -205,6 +209,10 @@ uint32_t llama_hparams::n_embd_r() const { } uint32_t llama_hparams::n_embd_s() const { + if (n_embd_s_override != 0) { + return n_embd_s_override; + } + if (wkv_head_size != 0) { // corresponds to RWKV's wkv_states size return n_embd * wkv_head_size; diff --git a/src/llama-hparams.h b/src/llama-hparams.h index c3c14292c32d..fc550d2b04f4 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -7,6 +7,9 @@ #include // bump if necessary +#define LLAMA_MAX_NGRAM 8 +#define LLAMA_MAX_NGRAM_HEADS 64 + #define LLAMA_MAX_LAYERS 512 #define LLAMA_MAX_EXPERTS 1024 // Kimi K3 @@ -257,6 +260,12 @@ struct llama_hparams { uint32_t indexer_local_blocks = 0; // Indexer is "full" (1) or "shared" (0) + // QSA (Qwen4-Exp): the indexer selects whole blocks of indexer_block_size keys, up to + // indexer_top_k keys in total. All four values are needed to run it. + bool has_qsa() const { + return indexer_n_head > 0 && indexer_head_size > 0 && indexer_block_size > 0 && indexer_top_k > 0; + } + // Shared indexers reuse top-k from previous full layer std::array is_indexer_full_impl; @@ -287,6 +296,29 @@ struct llama_hparams { // gemma4 per-layer embedding uint32_t n_embd_per_layer = 0; + // Qwen4-Exp hyper-connections and per-layer embedding (PLE) + uint32_t hc_count = 0; + uint32_t hc_lowrank = 0; + uint32_t n_embd_ple = 0; + uint32_t ple_conv_kernel = 0; + uint32_t ple_ngram_size = 0; // 0 = no PLE + uint32_t ple_ngram_heads = 0; + // n_embd_r()/n_embd_s() are derived from the ssm sizes. A model that keeps a second recurrent + // state hands llama_memory_recurrent an hparams copy with these set to that state's geometry. + // Keep them 0 in the model's own hparams, or the primary state changes shape too + uint32_t n_embd_r_override = 0; + uint32_t n_embd_s_override = 0; + + // geometry of that second state, which the hybrid copies into the overrides above + uint32_t n_embd_r_2nd = 0; + uint32_t n_embd_s_2nd = 0; + + // the single layer that owns the second state, -1 when there is none + int32_t il_2nd = -1; + std::array ple_ngram_mult; // hash multiplier per n-gram position + std::array ple_ngram_vocab; // vocab size of each hashed n-gram head + std::array ple_ngram_offs; // row offset of each hashed n-gram head + // needed by encoder-decoder models (e.g. T5, FLAN-T5) // ref: https://github.com/ggml-org/llama.cpp/pull/8141 llama_token dec_start_token_id = LLAMA_TOKEN_NULL; diff --git a/src/llama-memory-hybrid-iswa.cpp b/src/llama-memory-hybrid-iswa.cpp index 06f7fd5428c4..949e681e0882 100644 --- a/src/llama-memory-hybrid-iswa.cpp +++ b/src/llama-memory-hybrid-iswa.cpp @@ -52,6 +52,7 @@ llama_memory_hybrid_iswa::llama_memory_hybrid_iswa( )), mem_recr(new llama_memory_recurrent( model, + model.hparams, type_r, type_s, offload, diff --git a/src/llama-memory-hybrid.cpp b/src/llama-memory-hybrid.cpp index 42c7381a9e6f..24ce93df7d34 100644 --- a/src/llama-memory-hybrid.cpp +++ b/src/llama-memory-hybrid.cpp @@ -4,6 +4,8 @@ #include "llama-model.h" #include "llama-context.h" +#include + // // llama_memory_hybrid // @@ -29,8 +31,12 @@ llama_memory_hybrid::llama_memory_hybrid( bool unified, /* layer filters */ const layer_filter_cb & filter_attn, - const layer_filter_cb & filter_recr) : + const layer_filter_cb & filter_recr, + const layer_filter_cb & filter_lid, + const layer_filter_cb & filter_ple) : hparams(model.hparams), + hparams_lid(model.hparams), + hparams_ple(model.hparams), mem_attn(new llama_kv_cache( model, model.hparams, @@ -53,6 +59,7 @@ llama_memory_hybrid::llama_memory_hybrid( )), mem_recr(new llama_memory_recurrent( model, + model.hparams, type_r, type_s, offload, @@ -62,7 +69,31 @@ llama_memory_hybrid::llama_memory_hybrid( filter_recr == nullptr ? [&](int32_t il) { return hparams.is_recr(il); } : filter_recr - )) {} + )) { + if (filter_lid) { + // the indexer stores one key per token, indexer_head_size wide + std::fill(hparams_lid.n_head_kv_arr.begin(), hparams_lid.n_head_kv_arr.end(), 1); + hparams_lid.n_embd_head_k_full = model.hparams.indexer_head_size; + // the indexer has no values, but llama_kv_cache always allocates V - keep it narrow + hparams_lid.n_embd_head_v_full = model.hparams.indexer_head_size; + + mem_lid = std::make_unique( + model, hparams_lid, type_k, type_v, v_trans, offload, unified, + kv_size, n_seq_max, n_pad, n_swa, swa_type, nullptr, filter_lid, nullptr, nullptr); + } + + if (filter_ple) { + // its geometry is its own, so override what n_embd_r()/n_embd_s() would derive from the + // ssm sizes. Same trick as hparams_lid above + GGML_ASSERT(hparams.n_embd_r_2nd != 0 && hparams.n_embd_s_2nd != 0 && + "a second recurrent state needs n_embd_r_2nd/n_embd_s_2nd"); + hparams_ple.n_embd_r_override = hparams.n_embd_r_2nd; + hparams_ple.n_embd_s_override = hparams.n_embd_s_2nd; + + mem_ple = std::make_unique( + model, hparams_ple, type_r, type_s, offload, rs_size, n_seq_max, n_rs_seq, filter_ple); + } +} llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { do { @@ -108,6 +139,11 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } + if (mem_ple && !mem_ple->prepare(ubatches)) { + LLAMA_LOG_ERROR("%s: failed to prepare the second recurrent ubatches\n", __func__); + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + // prepare the attention cache auto heads_attn = mem_attn->prepare(ubatches); if (heads_attn.empty()) { @@ -115,8 +151,19 @@ llama_memory_context_ptr llama_memory_hybrid::init_batch(llama_batch_allocr & ba return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); } + // the indexer cache mirrors the attention cache, so it uses the same ubatch splits + llama_kv_cache::slot_info_vec_t heads_lid; + if (mem_lid) { + heads_lid = mem_lid->prepare(ubatches); + if (heads_lid.empty()) { + LLAMA_LOG_ERROR("%s: failed to prepare indexer ubatches\n", __func__); + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + GGML_ASSERT(heads_lid.size() == heads_attn.size()); + } + return std::make_unique( - this, std::move(heads_attn), std::move(ubatches)); + this, std::move(heads_attn), std::move(heads_lid), std::move(ubatches)); } while(false); return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); @@ -138,6 +185,12 @@ bool llama_memory_hybrid::get_can_shift() const { void llama_memory_hybrid::clear(bool data) { mem_attn->clear(data); mem_recr->clear(data); + if (mem_lid) { + mem_lid->clear(data); + } + if (mem_ple) { + mem_ple->clear(data); + } } bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { @@ -146,27 +199,57 @@ bool llama_memory_hybrid::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1 if (!mem_recr->seq_rm(seq_id, p0, p1)) { return false; } + if (mem_lid) { + mem_lid->seq_rm(seq_id, p0, p1); + } + if (mem_ple) { + mem_ple->seq_rm(seq_id, p0, p1); + } return mem_attn->seq_rm(seq_id, p0, p1); } void llama_memory_hybrid::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { mem_attn->seq_cp(seq_id_src, seq_id_dst, p0, p1); mem_recr->seq_cp(seq_id_src, seq_id_dst, p0, p1); + if (mem_lid) { + mem_lid->seq_cp(seq_id_src, seq_id_dst, p0, p1); + } + if (mem_ple) { + mem_ple->seq_cp(seq_id_src, seq_id_dst, p0, p1); + } } void llama_memory_hybrid::seq_keep(llama_seq_id seq_id) { mem_attn->seq_keep(seq_id); mem_recr->seq_keep(seq_id); + if (mem_lid) { + mem_lid->seq_keep(seq_id); + } + if (mem_ple) { + mem_ple->seq_keep(seq_id); + } } void llama_memory_hybrid::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { mem_attn->seq_add(seq_id, p0, p1, shift); mem_recr->seq_add(seq_id, p0, p1, shift); + if (mem_lid) { + mem_lid->seq_add(seq_id, p0, p1, shift); + } + if (mem_ple) { + mem_ple->seq_add(seq_id, p0, p1, shift); + } } void llama_memory_hybrid::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { mem_attn->seq_div(seq_id, p0, p1, d); mem_recr->seq_div(seq_id, p0, p1, d); + if (mem_lid) { + mem_lid->seq_div(seq_id, p0, p1, d); + } + if (mem_ple) { + mem_ple->seq_div(seq_id, p0, p1, d); + } } llama_pos llama_memory_hybrid::seq_pos_min(llama_seq_id seq_id) const { @@ -184,12 +267,28 @@ std::map llama_memory_hybrid::memory_breakdo for (const auto & buft_size : mem_recr->memory_breakdown()) { mb[buft_size.first] += buft_size.second; } + if (mem_ple) { + for (const auto & buft_size : mem_ple->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + } + if (mem_lid) { + for (const auto & buft_size : mem_lid->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + } return mb; } void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_write(io, seq_id, flags); + if (mem_lid) { + mem_lid->state_write(io, seq_id, flags); + } + } + if (mem_ple) { + mem_ple->state_write(io, seq_id, flags); } mem_recr->state_write(io, seq_id, flags); } @@ -197,6 +296,12 @@ void llama_memory_hybrid::state_write(llama_io_write_i & io, llama_seq_id seq_id void llama_memory_hybrid::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { mem_attn->state_read(io, seq_id, flags); + if (mem_lid) { + mem_lid->state_read(io, seq_id, flags); + } + } + if (mem_ple) { + mem_ple->state_read(io, seq_id, flags); } mem_recr->state_read(io, seq_id, flags); } @@ -209,32 +314,55 @@ llama_memory_recurrent * llama_memory_hybrid::get_mem_recr() const { return mem_recr.get(); } +llama_kv_cache * llama_memory_hybrid::get_mem_lid() const { + return mem_lid.get(); +} + +llama_memory_recurrent * llama_memory_hybrid::get_mem_ple() const { + return mem_ple.get(); +} + llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_status status) : status(status) {} llama_memory_hybrid_context::llama_memory_hybrid_context(llama_memory_hybrid * mem) : + mem(mem), ctx_attn(mem->get_mem_attn()->init_full()), ctx_recr(mem->get_mem_recr()->init_full()), - status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { + ctx_lid(mem->get_mem_lid() ? mem->get_mem_lid()->init_full() : nullptr), + ctx_ple(mem->get_mem_ple() ? mem->get_mem_ple()->init_full() : nullptr), + status(llama_memory_status_combine(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status()), + ctx_lid ? ctx_lid->get_status() : LLAMA_MEMORY_STATUS_SUCCESS)) { } llama_memory_hybrid_context::llama_memory_hybrid_context( llama_memory_hybrid * mem, llama_context * lctx, bool optimize) : + mem(mem), ctx_attn(mem->get_mem_attn()->init_update(lctx, optimize)), ctx_recr(mem->get_mem_recr()->init_update(lctx, optimize)), - status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { + ctx_lid(mem->get_mem_lid() ? mem->get_mem_lid()->init_update(lctx, optimize) : nullptr), + ctx_ple(mem->get_mem_ple() ? mem->get_mem_ple()->init_update(lctx, optimize) : nullptr), + status(llama_memory_status_combine(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status()), + ctx_lid ? ctx_lid->get_status() : LLAMA_MEMORY_STATUS_SUCCESS)) { } llama_memory_hybrid_context::llama_memory_hybrid_context( llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, + slot_info_vec_t sinfos_lid, std::vector ubatches) : + mem(mem), ubatches(std::move(ubatches)), // note: here we copy the ubatches. not sure if this is ideal ctx_attn(new llama_kv_cache_context(mem->get_mem_attn(), std::move(sinfos_attn), this->ubatches)), ctx_recr(new llama_memory_recurrent_context(mem->get_mem_recr(), this->ubatches)), - status(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status())) { + ctx_lid(mem->get_mem_lid() ? + std::make_unique(mem->get_mem_lid(), std::move(sinfos_lid), this->ubatches) : nullptr), + ctx_ple(mem->get_mem_ple() ? + std::make_unique(mem->get_mem_ple(), this->ubatches) : nullptr), + status(llama_memory_status_combine(llama_memory_status_combine(ctx_attn->get_status(), ctx_recr->get_status()), + ctx_lid ? ctx_lid->get_status() : LLAMA_MEMORY_STATUS_SUCCESS)) { } bool llama_memory_hybrid_context::next() { @@ -242,6 +370,12 @@ bool llama_memory_hybrid_context::next() { ctx_attn->next(); ctx_recr->next(); + if (ctx_lid) { + ctx_lid->next(); + } + if (ctx_ple) { + ctx_ple->next(); + } if (++i_next >= ubatches.size()) { return false; @@ -257,6 +391,12 @@ bool llama_memory_hybrid_context::apply() { res = res & ctx_attn->apply(); res = res & ctx_recr->apply(); + if (ctx_lid) { + res = res & ctx_lid->apply(); + } + if (ctx_ple) { + res = res & ctx_ple->apply(); + } return res; } @@ -274,6 +414,93 @@ const llama_kv_cache_context * llama_memory_hybrid_context::get_attn() const { return static_cast(ctx_attn.get()); } +const llama_kv_cache_context * llama_memory_hybrid_context::get_lid() const { + return static_cast(ctx_lid.get()); +} + +const llama_memory_recurrent_context * llama_memory_hybrid_context::get_ple() const { + return static_cast(ctx_ple.get()); +} + +uint32_t llama_memory_hybrid_context::get_n_pos(uint32_t n_pad, uint32_t n_seq_max) const { + llama_pos pos_max = -1; + + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq_max; ++seq_id) { + pos_max = std::max(pos_max, mem->get_mem_attn()->seq_pos_max(seq_id)); + } + + return std::max(n_pad, GGML_PAD((uint32_t) (pos_max + 1), n_pad)); +} + +void llama_memory_hybrid_context::set_input_pos_cell(ggml_tensor * dst, const llama_ubatch * ubatch, bool lid, int64_t n_rows) const { + GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); + GGML_ASSERT(dst->type == GGML_TYPE_I32 || dst->type == GGML_TYPE_F32); + + const llama_kv_cache * kv = lid ? mem->get_mem_lid() : mem->get_mem_attn(); + GGML_ASSERT(kv); + + const int64_t n_pos = dst->ne[0]; + const int64_t ns = dst->ne[1]; + + GGML_ASSERT(ubatch->n_tokens % ns == 0); + const int64_t n_tps = ubatch->n_tokens/ns; + + for (int64_t s = 0; s < ns; ++s) { + const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0]; + + const auto & cells = kv->get_cells(seq_id); + + std::vector map(n_pos, 0); + + for (uint32_t j = 0; j < cells.size(); ++j) { + if (cells.is_empty(j) || !cells.seq_has(j, seq_id)) { + continue; + } + + const llama_pos p0 = cells.pos_get(j); + + if (p0 >= 0 && p0 < n_pos && (int64_t) j < n_rows) { + map[p0] = (int32_t) j; + } + } + + if (dst->type == GGML_TYPE_I32) { + std::copy(map.begin(), map.end(), (int32_t *) dst->data + s*n_pos); + } else { + float * data = (float *) dst->data + s*n_pos; + for (int64_t p = 0; p < n_pos; ++p) { + data[p] = (float) map[p]; + } + } + } +} + +void llama_memory_hybrid_context::set_input_cell_blk(ggml_tensor * dst, const llama_ubatch * ubatch, int32_t div) const { + GGML_ASSERT(ggml_backend_buffer_is_host(dst->buffer)); + GGML_ASSERT(dst->type == GGML_TYPE_I32); + + const int64_t n_kv = dst->ne[0]; + const int64_t ns = dst->ne[1]; + + GGML_ASSERT(ubatch->n_tokens % ns == 0); + const int64_t n_tps = ubatch->n_tokens/ns; + + int32_t * data = (int32_t *) dst->data; + + for (int64_t s = 0; s < ns; ++s) { + const llama_seq_id seq_id = ubatch->seq_id[s*n_tps][0]; + + const auto & cells = mem->get_mem_attn()->get_cells(seq_id); + + for (int64_t j = 0; j < n_kv; ++j) { + data[s*n_kv + j] = + cells.is_empty(j) || !cells.seq_has(j, seq_id) + ? 0 + : (int32_t) (cells.pos_get(j)/div); + } + } +} + const llama_memory_recurrent_context * llama_memory_hybrid_context::get_recr() const { return static_cast(ctx_recr.get()); } diff --git a/src/llama-memory-hybrid.h b/src/llama-memory-hybrid.h index 484eafb74991..1e6b22d9408b 100644 --- a/src/llama-memory-hybrid.h +++ b/src/llama-memory-hybrid.h @@ -39,7 +39,11 @@ class llama_memory_hybrid : public llama_memory_i { bool unified, /* layer filters */ const layer_filter_cb & filter_attn = nullptr, - const layer_filter_cb & filter_recr = nullptr); + const layer_filter_cb & filter_recr = nullptr, + /* optional per-token indexer keys (QSA), attention layers only */ + const layer_filter_cb & filter_lid = nullptr, + /* optional second recurrent state (qwen4exp PLE), its own layers */ + const layer_filter_cb & filter_ple = nullptr); ~llama_memory_hybrid() = default; @@ -83,11 +87,21 @@ class llama_memory_hybrid : public llama_memory_i { llama_kv_cache * get_mem_attn() const; llama_memory_recurrent * get_mem_recr() const; + // null unless the model asked for an indexer cache + llama_kv_cache * get_mem_lid() const; + llama_memory_recurrent * get_mem_ple() const; + private: const llama_hparams & hparams; + // the indexer cache stores one key per token, so it needs its own head layout + llama_hparams hparams_lid; + llama_hparams hparams_ple; + const std::unique_ptr mem_attn; const std::unique_ptr mem_recr; + std::unique_ptr mem_lid; + std::unique_ptr mem_ple; }; class llama_memory_hybrid_context : public llama_memory_context_i { @@ -110,6 +124,7 @@ class llama_memory_hybrid_context : public llama_memory_context_i { llama_memory_hybrid_context( llama_memory_hybrid * mem, slot_info_vec_t sinfos_attn, + slot_info_vec_t sinfos_lid, std::vector ubatches); ~llama_memory_hybrid_context() = default; @@ -127,7 +142,33 @@ class llama_memory_hybrid_context : public llama_memory_context_i { const llama_kv_cache_context * get_attn() const; const llama_memory_recurrent_context * get_recr() const; + // null unless the model asked for an indexer cache + const llama_kv_cache_context * get_lid() const; + + // null unless the model asked for a second recurrent state + const llama_memory_recurrent_context * get_ple() const; + + // QSA selects blocks of keys over token positions, but the caches store them by cell, and the + // two only agree for an append-only single-sequence cache. These maps let the graph translate. + // ref: src/llama-kv-cache-msa.cpp + + // length of the position axis: past the largest position in the cache, padded so that the + // graph shape stays constant across batches and can be reused + uint32_t get_n_pos(uint32_t n_pad, uint32_t n_seq_max) const; + + // I32/F32 [n_pos, n_stream] pos -> cell, of the indexer cache when lid, else of the attention + // cache. Positions with no cell, and cells past the first n_rows (all the graph views), get 0, + // which aliases cell 0. That is only harmless while a sequence occupies a contiguous range of + // positions, which is what an append-only cache gives. A hole in the middle of the range (a + // mid-range seq_rm, a ranged seq_cp) makes a consumer read cell 0 for the missing positions. + void set_input_pos_cell(ggml_tensor * dst, const llama_ubatch * ubatch, bool lid, int64_t n_rows) const; + + // I32 [n_kv, n_stream] attention cell -> pos/div. Empty and other-sequence cells get 0. + void set_input_cell_blk(ggml_tensor * dst, const llama_ubatch * ubatch, int32_t div) const; + private: + llama_memory_hybrid * mem = nullptr; + // the index of the next ubatch to process size_t i_next = 0; @@ -135,6 +176,8 @@ class llama_memory_hybrid_context : public llama_memory_context_i { const llama_memory_context_ptr ctx_attn; const llama_memory_context_ptr ctx_recr; + const llama_memory_context_ptr ctx_lid; + const llama_memory_context_ptr ctx_ple; const llama_memory_status status; }; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index ef82eb976ca7..9b093e3c6c11 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -19,13 +19,14 @@ llama_memory_recurrent::llama_memory_recurrent( const llama_model & model, + const llama_hparams & hparams, ggml_type type_r, ggml_type type_s, bool offload, uint32_t mem_size, uint32_t n_seq_max, uint32_t n_rs_seq, - const layer_filter_cb & filter) : hparams(model.hparams), n_seq_max(n_seq_max) { + const layer_filter_cb & filter) : hparams(hparams), n_seq_max(n_seq_max) { const int32_t n_layer = hparams.n_layer(); head = 0; diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f5e..013fe3de3518 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -18,6 +18,7 @@ class llama_memory_recurrent : public llama_memory_i { public: llama_memory_recurrent( const llama_model & model, + const llama_hparams & hparams, ggml_type type_r, ggml_type type_s, bool offload, diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 9b22cb05f29b..9b23a75aeba4 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -321,10 +321,13 @@ namespace GGUFMeta { case GGUF_TYPE_UINT32: case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || (std::is_same::value); break; + case GGUF_TYPE_UINT64: + case GGUF_TYPE_INT64: type_ok = (std::is_same::value) || + (std::is_same::value); break; case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: - throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); + throw std::runtime_error(format("%s is not a string/float32/uint64/int64/uint32/int32 array", key.c_str())); } if (!type_ok) { throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); @@ -367,10 +370,13 @@ namespace GGUFMeta { case GGUF_TYPE_UINT32: case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || (std::is_same::value); break; + case GGUF_TYPE_UINT64: + case GGUF_TYPE_INT64: type_ok = (std::is_same::value) || + (std::is_same::value); break; case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: - throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); + throw std::runtime_error(format("%s is not a string/float32/uint64/int64/uint32/int32 array", key.c_str())); } if (!type_ok) { throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); @@ -409,6 +415,7 @@ namespace GGUFMeta { template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template diff --git a/src/llama-model.cpp b/src/llama-model.cpp index de0d3c1a6f6f..78366fa4bf84 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -319,6 +319,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_qwen35(params); case LLM_ARCH_QWEN35MOE: return new llama_model_qwen35moe(params); + case LLM_ARCH_QWEN4EXP: + return new llama_model_qwen4exp(params); case LLM_ARCH_MISTRAL3: return new llama_model_mistral3(params); case LLM_ARCH_EAGLE3: @@ -538,7 +540,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str }; auto get_split_segments = [&](int axis, uint32_t il) -> std::vector> { - if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) { + if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || + ud->model->arch == LLM_ARCH_QWEN4EXP) { const int64_t head_k_dim = hparams.ssm_d_state; const int64_t head_v_dim = hparams.ssm_d_state; const int64_t n_k_heads = hparams.ssm_n_group; @@ -651,7 +654,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) { GGML_ASSERT(segments.size() == 1); // some models have Q gate tensors, for those cases the granularity needs to be doubled: - if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) { + if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || + ud->model->arch == LLM_ARCH_QWEN4EXP) { return {std::lcm(2*n_embd_q, blck_size_perf)}; } return {granularity_q}; @@ -1897,6 +1901,7 @@ void llama_model::print_info() const { arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || + arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) { LLAMA_LOG_INFO("%s: ssm_d_conv = %u\n", __func__, hparams.ssm_d_conv); @@ -2327,7 +2332,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, const bool mtp_on_hybrid_qwen = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || - arch == LLM_ARCH_BAILINGMOE3); + arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_BAILINGMOE3); const bool mtp_on_hybrid_nemotron = params.ctx_type == LLAMA_CONTEXT_TYPE_MTP && arch == LLM_ARCH_NEMOTRON_H_MOE; @@ -2335,6 +2340,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, if (llm_arch_is_recurrent(arch)) { res = new llama_memory_recurrent( *this, + hparams, GGML_TYPE_F32, GGML_TYPE_F32, cparams.offload_kqv, @@ -2347,6 +2353,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; + llama_memory_hybrid::layer_filter_cb filter_lid = nullptr; + llama_memory_hybrid::layer_filter_cb filter_ple = nullptr; if (arch == LLM_ARCH_FALCON_H1) { filter_attn = [&](uint32_t) { return true; }; filter_recr = [&](uint32_t) { return true; }; @@ -2357,13 +2365,33 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_recr = [&](uint32_t il) { return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; - } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01) { + } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || + arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_MINIMAX_01) { filter_attn = [&](uint32_t il) { return il < hparams.n_layer() && !hparams.is_recr(il); }; filter_recr = [&](uint32_t il) { return il < hparams.n_layer() && hparams.is_recr(il); }; + if (arch == LLM_ARCH_QWEN4EXP && hparams.n_embd_r_2nd > 0) { + // PLE keeps its conv and n-gram history out of the delta-net row + filter_ple = [&](uint32_t il) { return (int32_t) il == hparams.il_2nd; }; + } + + if (hparams.has_qsa()) { + // QSA selects blocks over token positions and maps them to cells per stream, + // so every sequence needs its own cell space. A unified cache shares one + // between them, so skip the indexer cache and run those layers dense. + // ref: src/models/minimax-m3.cpp + if (cparams.n_seq_max == 1 || !cparams.kv_unified) { + // the QSA indexer keeps one key per token on the attention layers + filter_lid = filter_attn; + } else { + LLAMA_LOG_WARN("%s: unified KV cache with n_seq_max > 1; QSA needs per-sequence " + "streams -> running DENSE attention (output may be degraded). " + "Drop --kv-unified to enable QSA.\n", __func__); + } + } } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { @@ -2404,7 +2432,9 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* offload */ cparams.offload_kqv, /* unified */ cparams.kv_unified, /* filter_attn */ std::move(filter_attn), - /* filter_recr */ std::move(filter_recr)); + /* filter_recr */ std::move(filter_recr), + /* filter_lid */ std::move(filter_lid), + /* filter_ple */ std::move(filter_ple)); } } else { llama_kv_cache::layer_filter_cb filter = nullptr; @@ -2807,6 +2837,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_QWEN3VLMOE: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_QWEN3TTS: return LLAMA_ROPE_TYPE_IMROPE; diff --git a/src/llama-model.h b/src/llama-model.h index 44bd9675754b..07134252a907 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -500,6 +500,23 @@ struct llama_layer { struct ggml_tensor * hc_ffn_fn = nullptr; struct ggml_tensor * hc_ffn_base = nullptr; struct ggml_tensor * hc_ffn_scale = nullptr; + // Qwen4-Exp hyper-connections and per-layer embedding (PLE) + struct ggml_tensor * hc_attn_norm = nullptr; + struct ggml_tensor * hc_attn_down = nullptr; + struct ggml_tensor * hc_attn_up = nullptr; + struct ggml_tensor * hc_attn_inject = nullptr; + struct ggml_tensor * hc_ffn_norm = nullptr; + struct ggml_tensor * hc_ffn_down = nullptr; + struct ggml_tensor * hc_ffn_up = nullptr; + struct ggml_tensor * hc_ffn_inject = nullptr; + struct ggml_tensor * ple_ngram_embd = nullptr; + struct ggml_tensor * ple_key = nullptr; + struct ggml_tensor * ple_value = nullptr; + struct ggml_tensor * ple_key_norm = nullptr; + struct ggml_tensor * ple_query_norm = nullptr; + struct ggml_tensor * ple_conv_norm = nullptr; + struct ggml_tensor * ple_conv1d = nullptr; + struct ggml_tensor * attn_comp_wkv = nullptr; struct ggml_tensor * attn_comp_wgate = nullptr; struct ggml_tensor * attn_comp_ape = nullptr; @@ -621,6 +638,11 @@ struct llama_model { struct ggml_tensor * hc_head_base = nullptr; struct ggml_tensor * hc_head_scale = nullptr; + // Qwen4-Exp final hyper-connection mixer + struct ggml_tensor * hc_mix_norm = nullptr; + struct ggml_tensor * hc_mix_down = nullptr; + struct ggml_tensor * hc_mix_up = nullptr; + // classifier struct ggml_tensor * cls = nullptr; struct ggml_tensor * cls_b = nullptr; diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 20252815d5c3..744a55930423 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -324,6 +324,10 @@ static bool tensor_allows_quantization(const llama_model_quantize_params * param // do not quantize Mamba/Kimi's small conv1d weights // NOTE: can't use LLM_TN here because the layer number is not known quantize &= name.find("ssm_conv1d") == std::string::npos; + quantize &= name.find("ple_conv1d") == std::string::npos; + // the hyper-connection projections are tiny (hc_lowrank wide) and their names would otherwise + // be picked up by the ffn_up/ffn_down substring match in tensor_get_category() + quantize &= name.find("hc_ffn_") == std::string::npos; quantize &= name.find("shortconv.conv.weight") == std::string::npos; // do not quantize MiniMax's indexer projection weights, they are tiny diff --git a/src/models/delta-net-base.cpp b/src/models/delta-net-base.cpp index ad6612647736..e103e9c99478 100644 --- a/src/models/delta-net-base.cpp +++ b/src/models/delta-net-base.cpp @@ -460,7 +460,17 @@ ggml_tensor * llm_build_delta_net_base::build_conv_state( const int64_t n_seqs = ubatch.n_seqs; - ggml_tensor * conv_states = build_rs(inp, conv_states_all, hparams.n_embd_r(), n_seqs); + // the row width comes from the target: this helper is also used for a model's second + // recurrent state, whose rows are narrower than n_embd_r() + const int64_t n_embd_row = conv_states_all->ne[0]; + + // ... and the row has to be exactly the conv state, since that is what gets reshaped and what + // the cell stride below assumes. A wider row (state packed next to something else) needs its + // own slice, not this helper + GGML_ASSERT(n_embd_row == (conv_kernel_size - 1) * conv_channels && + "build_conv_state: the recurrent row must hold only the conv state"); + + ggml_tensor * conv_states = build_rs(inp, conv_states_all, n_embd_row, n_seqs); cb(conv_states, "conv_states", il); conv_states = ggml_reshape_3d(ctx0, conv_states, conv_kernel_size - 1, conv_channels, n_seqs); diff --git a/src/models/models.h b/src/models/models.h index 157b05dc050f..2954a472a66f 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2314,6 +2314,106 @@ struct llama_model_qwen35moe : public llama_model_base { }; +class llm_graph_input_qsa; + +struct llama_model_qwen4exp : public llama_model_base { + llama_model_qwen4exp(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; + + struct graph : public llm_build_delta_net_base { + // the MTP block reuses the builders below, so it chains to this ctor and builds its own graph + graph(const llama_model & model, const llm_graph_params & params, bool /*mtp*/) + : llm_build_delta_net_base(params), model(model) {} + graph(const llama_model & model, const llm_graph_params & params); + protected: + // hyper-connection input mixer: returns the block input and the per-stream injection weights + ggml_tensor * build_hc_pre( + ggml_tensor * x, + ggml_tensor * norm, + ggml_tensor * down, + ggml_tensor * up, + ggml_tensor * inject, + ggml_tensor ** inject_out, + int il); + + // scatter the block output back into the hyper-connection streams + ggml_tensor * build_hc_post( + ggml_tensor * residual, + ggml_tensor * y, + ggml_tensor * inject, + int il); + + // RMS norm over each hyper-connection stream separately + ggml_tensor * build_hc_norm( + ggml_tensor * x, + ggml_tensor * w, + int il); + + ggml_tensor * build_ple( + ggml_tensor * x, + ggml_tensor * ngram_ids, + ggml_tensor * tokens, + llm_graph_input_rs * inp, + int il); + + // the host-side QSA maps, or nullptr when this context has no indexer cache + llm_graph_input_qsa * build_qsa_inp(); + + // relu(q . k_blk) summed over the indexer heads, for one token against every block + ggml_tensor * build_qsa_scores( + llm_graph_input_qsa * inp_qsa, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il); + + // QSA block scores -> the kq mask of the cells the attention may read + ggml_tensor * build_qsa_mask( + llm_graph_input_qsa * inp_qsa, + ggml_tensor * sc, + ggml_tensor * kq_mask, + int il); + + // QSA block scores -> the selected keys and values, gathered into one compact buffer + void build_qsa_gather( + llm_graph_input_qsa * inp_qsa, + ggml_tensor * sc, + ggml_tensor * k_all, + ggml_tensor * v_all, + ggml_tensor * kq_mask, + ggml_tensor ** k_out, + ggml_tensor ** v_out, + ggml_tensor ** m_out); + + ggml_tensor * build_layer_attn( + llm_graph_input_attn_kv * inp_attn, + llm_graph_input_qsa * inp_qsa, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il); + + ggml_tensor * build_layer_attn_linear( + llm_graph_input_rs * inp, + ggml_tensor * cur, + int il); + + ggml_tensor * build_layer_ffn( + ggml_tensor * cur, + int il); + + const llama_model & model; + }; + + struct graph_mtp : public graph { + graph_mtp(const llama_model & model, const llm_graph_params & params); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_mistral3 : public llama_model_base { llama_model_mistral3(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp new file mode 100644 index 000000000000..f906af2116ff --- /dev/null +++ b/src/models/qwen4exp.cpp @@ -0,0 +1,1407 @@ +#include "models.h" +#include "llama-impl.h" +#include "llama-memory-hybrid.h" +#include "llama-memory-recurrent.h" + +#include +#include +#include + +// Qwen4-Exp: Qwen3.5-MoE trunk (gated delta-net + gated full attention + MoE) on top of +// hyper-connections, plus a hashed n-gram per-layer embedding (PLE) on selected linear layers. +// +// The full-attention layers use QSA: an indexer scores blocks of `indexer_block_size` keys and +// the attention reads only the top `indexer_top_k`/`indexer_block_size` blocks plus the trailing +// incomplete block. Blocks are anchored to cache cells - see llm_graph_input_qsa. + +void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true); + + // linear attention (gated delta net) + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner); + ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state); + ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); + ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); + + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.hc_count); + ml.get_key(LLM_KV_HYPER_CONNECTION_LOWRANK, hparams.hc_lowrank); + GGML_ASSERT(hparams.hc_count > 1); + + ml.get_key(LLM_KV_NEXTN_PREDICT_LAYERS, hparams.n_layer_nextn, false); + + // the MTP block reads the hyper-connection streams, not the collapsed hidden state, so the + // nextn read-back and common/speculative.cpp move hc_count*n_embd per token + hparams.n_embd_out_impl = hparams.hc_count*hparams.n_embd; + + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size, false); + + // per-layer embedding (PLE) + hparams.n_embd_ple = hparams.n_embd; + ml.get_key(LLM_KV_PLE_EMBD_LENGTH, hparams.n_embd_ple, false); + ml.get_key(LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel, false); + // the n-gram hash constants are needed on the host, they come in as metadata + auto load_u64 = [&ml](llm_kv kid, auto & dst) -> uint32_t { + std::vector tmp; + if (!ml.get_arr(kid, tmp, false)) { + return 0; + } + GGML_ASSERT(tmp.size() <= dst.size()); + std::copy(tmp.begin(), tmp.end(), dst.begin()); + return tmp.size(); + }; + + hparams.ple_ngram_size = load_u64(LLM_KV_PLE_NGRAM_MULTIPLIERS, hparams.ple_ngram_mult); + hparams.ple_ngram_heads = load_u64(LLM_KV_PLE_NGRAM_VOCAB_SIZES, hparams.ple_ngram_vocab); + GGML_ASSERT(hparams.ple_ngram_heads == load_u64(LLM_KV_PLE_NGRAM_OFFSETS, hparams.ple_ngram_offs)); + + if (hparams.ple_ngram_size > 0) { + // the hash spreads the heads over the 2..n_ngram grams and gives each head one slice of + // the PLE row. Without these the head loop divides by zero or truncates + GGML_ASSERT(hparams.ple_ngram_size > 1 && "PLE needs at least 2-grams"); + GGML_ASSERT(hparams.ple_ngram_heads > 0 && "PLE needs at least one n-gram head"); + GGML_ASSERT(hparams.ple_ngram_heads % (hparams.ple_ngram_size - 1) == 0 && + "PLE n-gram heads must be a multiple of ple_ngram_size - 1"); + GGML_ASSERT(hparams.n_embd_ple % hparams.ple_ngram_heads == 0 && + "PLE embedding length must be a multiple of the n-gram heads"); + } + + if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) { + uint32_t full_attn_interval = 4; + ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = (i + 1) % full_attn_interval != 0; + } + } + + // the MTP block is a sparse-attention layer, but the interval above covers n_layer_all and + // would call it recurrent. Nothing reads it there today, so say the truth before it does + for (uint32_t i = hparams.n_layer(); i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = false; + } + + switch (hparams.n_layer()) { + case 48: type = LLM_TYPE_122B_A10B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + const int64_t hc = hparams.hc_count; + const int64_t hc_dim = hc * n_embd; + const int64_t n_ngram = hparams.ple_ngram_size; + const int64_t n_head_ngram = hparams.ple_ngram_heads; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + + // the hyper-connection mixer replaces the usual output norm + hc_mix_norm = create_tensor(tn(LLM_TENSOR_HC_MIX_NORM, "weight"), { hc_dim }, 0); + hc_mix_down = create_tensor(tn(LLM_TENSOR_HC_MIX_DOWN, "weight"), { hc_dim, hparams.hc_lowrank }, 0); + hc_mix_up = create_tensor(tn(LLM_TENSOR_HC_MIX_UP, "weight"), { hparams.hc_lowrank, hc_dim }, 0); + + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + if (output == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); + } + + uint32_t n_embd_r_ple = 0; // PLE conv taps -> the second recurrent state's r + uint32_t n_embd_s_ple = 0; // n-gram history -> its s + + // an MTP-only file carries just the MTP block, see conversion --mtp + const bool mtp_only = hparams.n_layer_nextn > 0 && ml.get_weight("blk.0.hc_attn_norm.weight") == nullptr; + + for (int il = 0; il < (mtp_only ? 0 : n_layer); ++il) { + auto & layer = layers[il]; + + const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; + const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff; + + const int64_t head_k_dim = hparams.ssm_d_state; + const int64_t head_v_dim = hparams.ssm_d_state; + const int64_t n_k_heads = hparams.ssm_n_group; + const int64_t n_v_heads = hparams.ssm_dt_rank; + const int64_t key_dim = head_k_dim * n_k_heads; + const int64_t value_dim = head_v_dim * n_v_heads; + const int64_t conv_dim = key_dim * 2 + value_dim; + + layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { hc_dim }, 0); + layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hparams.hc_lowrank }, 0); + layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hparams.hc_lowrank, hc_dim }, 0); + layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, 0); + layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { hc_dim }, 0); + layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hparams.hc_lowrank }, 0); + layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hparams.hc_lowrank, hc_dim }, 0); + layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, 0); + + if (!hparams.is_recr(il)) { + create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0); + + // QSA indexer, skipped when the file has no indexer at all + const int64_t n_embd_idx = hparams.indexer_head_size; + const int idx_flags = hparams.has_qsa() ? 0 : (TENSOR_NOT_REQUIRED | TENSOR_SKIP); + layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, n_embd_idx * hparams.indexer_n_head }, idx_flags); + layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, n_embd_idx }, idx_flags); + layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { n_embd_idx }, idx_flags); + layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { n_embd_idx }, idx_flags); + } else { + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, conv_dim }, 0); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, 0); + layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, 0); + layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { n_v_heads }, 0); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { n_v_heads }, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, 0); + layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, 0); + layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, 0); + layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, 0); + } + + if (n_ngram > 0) { + const int64_t n_embd_ngram = hparams.n_embd_ple / n_head_ngram; + + // the hash reads row offs[h] + 0..vocab[h]-1 per head, so the table has to be at + // least this tall + int64_t n_row_ngram = 0; + for (int64_t h = 0; h < n_head_ngram; ++h) { + n_row_ngram = std::max(n_row_ngram, + (int64_t) (hparams.ple_ngram_offs[h] + hparams.ple_ngram_vocab[h])); + } + + // the checkpoint pads the table and does not export the divisor, so read the real + // height off the file instead of assuming one + const auto * w_ngram = ml.get_weight(tn(LLM_TENSOR_PLE_NGRAM_EMBD, "weight", il).str().c_str()); + if (w_ngram) { + GGML_ASSERT(w_ngram->tensor->ne[1] >= n_row_ngram && + "PLE n-gram table is shorter than the hash offsets need"); + n_row_ngram = w_ngram->tensor->ne[1]; + } + + layer.ple_ngram_embd = create_tensor(tn(LLM_TENSOR_PLE_NGRAM_EMBD, "weight", il), { n_embd_ngram, n_row_ngram }, TENSOR_NOT_REQUIRED); + + if (layer.ple_ngram_embd) { + layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { hparams.n_embd_ple, hc_dim }, 0); + layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { hparams.n_embd_ple, n_embd }, 0); + layer.ple_key_norm = create_tensor(tn(LLM_TENSOR_PLE_KEY_NORM, "weight", il), { hc_dim }, 0); + layer.ple_query_norm = create_tensor(tn(LLM_TENSOR_PLE_QUERY_NORM, "weight", il), { hc_dim }, 0); + layer.ple_conv_norm = create_tensor(tn(LLM_TENSOR_PLE_CONV_NORM, "weight", il), { hc_dim }, 0); + layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0); + + GGML_ASSERT(hparams.is_recr(il) && "PLE is only supported on linear attention layers"); + n_embd_r_ple = (hparams.ple_conv_kernel - 1) * n_ngram * hc_dim; + n_embd_s_ple = n_ngram - 1; + + GGML_ASSERT(hparams.il_2nd < 0 && "only one PLE layer is supported"); + hparams.il_2nd = il; + } + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0); + + layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, 0); + } + + // the MTP block is one sparse-attention layer with its own input fusion. It has no PLE and no + // linear attention, so it is built here instead of in the loop above. + for (int il = n_layer; il < n_layer + (int) hparams.n_layer_nextn; ++il) { + auto & layer = layers[il]; + + const int flags = ml.load_mtp ? 0 : TENSOR_SKIP; + + const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; + const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff; + + layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { hc_dim }, flags); + layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hparams.hc_lowrank }, flags); + layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hparams.hc_lowrank, hc_dim }, flags); + layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, flags); + layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { hc_dim }, flags); + layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hparams.hc_lowrank }, flags); + layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hparams.hc_lowrank, hc_dim }, flags); + layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, flags); + + create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, flags); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, flags); + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, flags); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, flags); + + // the MTP attention runs dense, so the indexer tensors are loaded but unused + const int64_t n_embd_idx = hparams.indexer_head_size; + const int idx_flags = flags | TENSOR_NOT_REQUIRED | TENSOR_SKIP; + layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, n_embd_idx * hparams.indexer_n_head }, idx_flags); + layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, n_embd_idx }, idx_flags); + layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { n_embd_idx }, idx_flags); + layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { n_embd_idx }, idx_flags); + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, flags); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, flags); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, flags); + + layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, flags); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, flags); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, flags); + + // hnorm is over the whole hyper-connection row, not per stream + layer.nextn.enorm = create_tensor(tn(LLM_TENSOR_NEXTN_ENORM, "weight", il), { n_embd }, flags); + layer.nextn.hnorm = create_tensor(tn(LLM_TENSOR_NEXTN_HNORM, "weight", il), { hc_dim }, flags); + layer.nextn.eh_proj = create_tensor(tn(LLM_TENSOR_NEXTN_EH_PROJ, "weight", il), { 2 * n_embd, n_embd }, flags); + } + + // PLE keeps a state of its own, so the delta-net row stays exactly the delta-net conv state + hparams.n_embd_r_2nd = n_embd_r_ple; + hparams.n_embd_s_2nd = n_embd_s_ple; +} + + +// A second recurrent state needs its own graph input. The plain llm_graph_input_rs cannot be +// used as a top-level input here: its can_reuse() casts params.mctx straight to a recurrent +// context, which is wrong when the context is hybrid - on a reused graph it would rebind to +// garbage. Rebind from the hybrid's second state instead. ref: llm_graph_input_mem_hybrid +class llm_graph_input_rs_2nd : public llm_graph_input_rs { +public: + using llm_graph_input_rs::llm_graph_input_rs; + virtual ~llm_graph_input_rs_2nd() = default; + + bool can_reuse(const llm_graph_params & params) override { + const auto * mctx_cur = static_cast(params.mctx)->get_ple(); + + mctx = mctx_cur; + + bool res = true; + + res &= s_copy->ne[0] == mctx_cur->get_n_rs(); + res &= s_copy_main->ne[0] == params.ubatch.n_seqs; + res &= s_copy_extra->ne[0] == mctx_cur->get_n_rs() - params.ubatch.n_seqs; + + res &= head == mctx_cur->get_head(); + res &= rs_z == mctx_cur->get_rs_z(); + + return res; + } +}; + +// QSA: the indexer scores blocks of `indexer_block_size` consecutive *positions*, so the host says +// which blocks each query may see and where every position lives in the caches. Selecting in +// position space keeps this correct whatever cells find_slot handed out. +class llm_graph_input_qsa : public llm_graph_input_i { +public: + llm_graph_input_qsa(const llama_memory_hybrid_context * mctx, int64_t ratio, uint32_t pad) : + mctx(mctx), ratio(ratio), pad(pad) {} + virtual ~llm_graph_input_qsa() = default; + + void set_input(const llama_ubatch * ubatch) override; + + bool can_reuse(const llm_graph_params & params) override { + mctx = static_cast(params.mctx); + + const int64_t ns = params.cparams.kv_unified ? 1 : params.ubatch.n_seqs_unq; + + n_kv_lid = mctx->get_lid()->get_n_kv(); + + if (pos_cell_f && (pos_cell_f->ne[0] != (int64_t) mctx->get_n_pos(pad, params.cparams.n_seq_max) || + pos_cell_f->ne[1] != ns)) { + return false; + } + + return pos_cell->ne[0] == (int64_t) mctx->get_n_pos(pad, params.cparams.n_seq_max) && + pos_cell->ne[1] == ns && + (!cell_blk || cell_blk->ne[0] == (int64_t) mctx->get_attn()->get_n_kv()) && + blk_mask->ne[1] == (int64_t) params.ubatch.n_tokens/ns && + blk_mask->ne[3] == ns; + } + + ggml_tensor * k_idxs = nullptr; // I64 [n_tokens] destination cells of this ubatch + ggml_tensor * pos_cell = nullptr; // I32 [n_pos, ns] position -> indexer cell + ggml_tensor * cell_blk = nullptr; // I32 [n_kv, ns] attention cell -> position block + ggml_tensor * blk_mask = nullptr; // F16 [n_blk, n_tps, 1, ns] 0 / -inf, block complete and visible + ggml_tensor * blk_cur = nullptr; // I32 [1, n_tps, 1, ns] the query's own block, always read + + // only for the decode gather, see build_qsa_gather() + ggml_tensor * pos_cell_f = nullptr; // F32 [n_pos, ns] position -> attention cell, as values + ggml_tensor * pos_vis = nullptr; // F32 [n_pos, ns] 0 / -inf, the query may read this position + + const llama_memory_hybrid_context * mctx; + const int64_t ratio; + const uint32_t pad; + + // rows of the indexer cache the graph views, see set_input_pos_cell() + int64_t n_kv_lid = 0; +}; + +void llm_graph_input_qsa::set_input(const llama_ubatch * ubatch) { + mctx->get_lid()->set_input_k_idxs(k_idxs, ubatch); + mctx->set_input_pos_cell(pos_cell, ubatch, true, n_kv_lid); + if (cell_blk) { + mctx->set_input_cell_blk(cell_blk, ubatch, ratio); + } + + const int64_t n_blk = blk_mask->ne[0]; + const int64_t n_tokens = ubatch->n_tokens; + + std::vector mask(n_blk*n_tokens, ggml_fp32_to_fp16(-INFINITY)); + std::vector cur (n_tokens, 0); + + for (int64_t i = 0; i < n_tokens; ++i) { + const int64_t n_vis = ubatch->pos[i] + 1; + const int64_t n_full = n_vis/ratio; + + // only complete blocks compete in the top-k, as in the reference + for (int64_t b = 0; b < std::min(n_full, n_blk); ++b) { + mask[i*n_blk + b] = ggml_fp32_to_fp16(0.0f); + } + + // the trailing incomplete block is always read on top of the top-k. Opening all of it is + // fine - the cells past the query are closed again by the causal mask. When the last block + // is complete there is no tail, so point at the last block instead: no cell maps to it, so + // opening it does nothing. + cur[i] = (int32_t) (n_vis%ratio != 0 ? std::min(n_full, n_blk - 1) : n_blk - 1); + } + + ggml_backend_tensor_set(blk_mask, mask.data(), 0, ggml_nbytes(blk_mask)); + ggml_backend_tensor_set(blk_cur, cur.data(), 0, ggml_nbytes(blk_cur)); + + if (pos_cell_f) { + // the gather reads whole blocks, so the positions past the query have to be closed here. + // Unmapped positions default to cell 0, which holds real but unrelated keys - they are only + // ever past the query, so the same test covers both + mctx->set_input_pos_cell(pos_cell_f, ubatch, false, mctx->get_attn()->get_n_kv()); + + const int64_t n_pos = pos_vis->ne[0]; + const int64_t ns = pos_vis->ne[1]; + + std::vector vis(n_pos*ns, -INFINITY); + + for (int64_t s = 0; s < ns; ++s) { + const llama_pos p1 = ubatch->pos[s*(n_tokens/ns) + n_tokens/ns - 1]; + + for (int64_t p = 0; p <= p1 && p < n_pos; ++p) { + vis[s*n_pos + p] = 0.0f; + } + } + + ggml_backend_tensor_set(pos_vis, vis.data(), 0, ggml_nbytes(pos_vis)); + } + +} + +// Computes the n-gram embedding rows on the host: the hash needs 64-bit integer math, and the +// previous tokens of a sequence come from the tail of its recurrent state row. +class llm_graph_input_ple : public llm_graph_input_i { +public: + llm_graph_input_ple(const llama_hparams & hparams, const llama_memory_recurrent_context * mctx, + llama_token tok_eos, int il) : + hparams(hparams), mctx(mctx), tok_eos(tok_eos), il(il) {} + virtual ~llm_graph_input_ple() = default; + + void set_input(const llama_ubatch * ubatch) override; + + bool can_reuse(const llm_graph_params & params) override { + mctx = static_cast(params.mctx)->get_ple(); + + return ngram_ids->ne[0] == (int64_t) params.ubatch.n_tokens*hparams.ple_ngram_heads; + } + + ggml_tensor * ngram_ids = nullptr; // I32 [n_head_ngram*n_tokens] rows of the n-gram table + ggml_tensor * tokens = nullptr; // F32 [1, n_tokens] token ids kept as the next history + + const llama_hparams & hparams; + const llama_memory_recurrent_context * mctx; + const llama_token tok_eos; + const int il; +}; + +void llm_graph_input_ple::set_input(const llama_ubatch * ubatch) { + // media chunks arrive as embeddings without token ids, so eos is hashed there. The reference + // hashes the media placeholder token instead - gemma4 has the same gap in build_inp_per_layer. + // TODO: needs the placeholder id to reach the ubatch, see llm_graph_context::build_inp_embd + auto token_at = [ubatch, this](int64_t i) { + return ubatch->token ? ubatch->token[i] : tok_eos; + }; + + const int64_t n_ngram = hparams.ple_ngram_size; + const int64_t n_head_ngram = hparams.ple_ngram_heads; + const int64_t n_per_ngram = n_head_ngram / (n_ngram - 1); + const int64_t n_seqs = ubatch->n_seqs; + const int64_t n_seq_tokens = ubatch->n_seq_tokens; + + ggml_tensor * s_l = mctx->get_s_l(il); + + std::vector ids(n_head_ngram*ubatch->n_tokens); + std::vector toks(ubatch->n_tokens); + + // history of the (n_ngram - 1) tokens before this ubatch, eos where a sequence starts + std::vector hist(n_ngram - 1, tok_eos); + std::vector hist_f(n_ngram - 1); + + std::vector shifted(n_ngram); + + for (int64_t s = 0; s < n_seqs; ++s) { + std::fill(hist.begin(), hist.end(), tok_eos); + + // how many tokens the state row already holds. This reads pos as the history length, + // which context shift breaks: it renumbers pos and leaves the state alone. Context shift + // is off by default and mmproj forces it off, so it cannot happen today + const int64_t n_prev = ubatch->pos[s*n_seq_tokens]; + + if (n_prev > 0) { + const size_t row = mctx->s_copy(s); + ggml_backend_tensor_get(s_l, hist_f.data(), + row*hist_f.size()*sizeof(float), hist_f.size()*sizeof(float)); + // the state row starts as zeros, so only the last n_prev entries hold tokens + for (int64_t k = 0; k < std::min(n_ngram - 1, n_prev); ++k) { + hist[k] = (llama_token) hist_f[n_ngram - 2 - k]; + } + } + + for (int64_t t = 0; t < n_seq_tokens; ++t) { + const int64_t i = s*n_seq_tokens + t; + + // tokens at t, t-1, t-2 ...; a sequence restarts after an eos, so anything at or + // before the closest previous eos reads as eos + bool cut = false; + for (int64_t k = 0; k < n_ngram; ++k) { + const int64_t src = t - k; + const llama_token tok = cut ? tok_eos + : src >= 0 ? token_at(s*n_seq_tokens + src) + : hist[-src - 1]; + shifted[k] = tok; + cut = cut || (k > 0 && tok == tok_eos); + } + + toks[i] = (float) token_at(i); + + for (int64_t n = 2; n <= n_ngram; ++n) { + // unsigned: the products overflow by design, and the remainder has to stay + // non-negative to be a row index + uint64_t mixed = (uint64_t) shifted[0]*hparams.ple_ngram_mult[0]; + for (int64_t k = 1; k < n; ++k) { + mixed ^= (uint64_t) shifted[k]*hparams.ple_ngram_mult[k]; + } + + const int64_t h0 = (n - 2)*n_per_ngram; + for (int64_t h = h0; h < h0 + n_per_ngram; ++h) { + ids[i*n_head_ngram + h] = + (int32_t) (mixed % hparams.ple_ngram_vocab[h] + hparams.ple_ngram_offs[h]); + } + } + } + } + + ggml_backend_tensor_set(ngram_ids, ids.data(), 0, ids.size() *ggml_element_size(ngram_ids)); + ggml_backend_tensor_set(tokens, toks.data(), 0, toks.size()*ggml_element_size(tokens)); +} + +// The MTP block predicts the next-next token from the target's hyper-connection streams and the +// embedding of the token the target just produced. One sparse-attention layer, no PLE. +// ref: sglang qwen4_exp_mtp.py, and src/models/deepseek4.cpp for the same shape of MTP block +llama_model_qwen4exp::graph_mtp::graph_mtp(const llama_model & model, const llm_graph_params & params) + : graph(model, params, true) { + // one block also keeps t_h_nextn below correct: the draft context reads it masked, and only + // a chained head would put non-output tokens in a draft batch. See common/speculative.cpp + GGML_ASSERT(hparams.n_layer_nextn == 1 && "QWEN4EXP MTP supports a single MTP block"); + GGML_ASSERT(ubatch.token && "QWEN4EXP MTP requires token input"); + + const int64_t hc = hparams.hc_count; + const int64_t n_embd = hparams.n_embd; + GGML_ASSERT(hparams.n_embd_out() == (uint32_t) (n_embd*hc) && "QWEN4EXP MTP hidden width mismatch"); + + const int il = hparams.n_layer(); + const auto & layer = model.layers[il]; + + GGML_ASSERT(layer.nextn.eh_proj && "MTP block missing nextn.eh_proj"); + GGML_ASSERT(layer.nextn.enorm && "MTP block missing nextn.enorm"); + GGML_ASSERT(layer.nextn.hnorm && "MTP block missing nextn.hnorm"); + + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + auto inp = std::make_unique(hparams.n_embd_out()); + + inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); + ggml_set_input(inp->tokens); + + inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens); + ggml_set_input(inp->embd); + + inp->h = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, hparams.n_embd_out(), n_tokens); + ggml_set_input(inp->h); + ggml_set_name(inp->h, "mtp_h_input"); + + ggml_tensor * tok_embd = ggml_get_rows(ctx0, model.tok_embd, inp->tokens); + cb(tok_embd, "mtp_tok_embd", il); + + ggml_tensor * h = inp->h; + + res->add_input(std::move(inp)); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + auto * inp_attn = build_attn_inp_kv(); + + // the reference norms the whole hyper-connection row and only then splits it into streams, + // so hnorm is hc_count*n_embd wide - unlike deepseek4, which norms each stream + ggml_tensor * h_norm = build_norm(h, layer.nextn.hnorm, nullptr, LLM_NORM_RMS, il); + h_norm = ggml_reshape_3d(ctx0, h_norm, n_embd, hc, n_tokens); + cb(h_norm, "mtp_hnorm", il); + + // the embedding term is shared by every stream + ggml_tensor * e_norm = build_norm(tok_embd, layer.nextn.enorm, nullptr, LLM_NORM_RMS, il); + e_norm = ggml_repeat_4d(ctx0, ggml_reshape_3d(ctx0, e_norm, n_embd, 1, n_tokens), n_embd, hc, n_tokens, 1); + cb(e_norm, "mtp_enorm", il); + + // fc_embedding(e) + fc_hidden(h) is one projection of the concatenation, see conversion + ggml_tensor * inpL = build_lora_mm(layer.nextn.eh_proj, + ggml_concat(ctx0, e_norm, h_norm, 0), layer.nextn.eh_proj_s); + cb(inpL, "mtp_eh_proj", il); + + ggml_tensor * inject = nullptr; + ggml_tensor * cur = build_hc_pre(inpL, + layer.hc_attn_norm, layer.hc_attn_down, layer.hc_attn_up, layer.hc_attn_inject, &inject, il); + cb(cur, "mtp_hc_attn_pre", il); + + // the draft runs the attention dense: a QSA indexer would need its own cache, and one layer + // spends its time reading weights, not attending - see the measurements in the notes + cur = build_layer_attn(inp_attn, nullptr, cur, inp_pos, sections, il); + + inpL = build_hc_post(inpL, cur, inject, il); + cb(inpL, "mtp_hc_attn_post", il); + + cur = build_hc_pre(inpL, + layer.hc_ffn_norm, layer.hc_ffn_down, layer.hc_ffn_up, layer.hc_ffn_inject, &inject, il); + cb(cur, "mtp_hc_ffn_pre", il); + + cur = build_layer_ffn(cur, il); + cb(cur, "mtp_ffn_out", il); + + inpL = build_hc_post(inpL, cur, inject, il); + cb(inpL, "mtp_l_out", il); + + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, hc*n_embd, n_tokens); + + // chained heads read the streams back, as the trunk hands them over + res->t_h_nextn = flat; + + if (inp_out_ids) { + flat = ggml_get_rows(ctx0, flat, inp_out_ids); + inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs); + } + + cur = build_hc_pre(inpL, model.hc_mix_norm, model.hc_mix_down, model.hc_mix_up, nullptr, nullptr, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +std::unique_ptr llama_model_qwen4exp::build_arch_graph(const llm_graph_params & params) const { + if (params.gtype == LLM_GRAPH_TYPE_DECODER_MTP) { + return std::make_unique(*this, params); + } + + return std::make_unique(*this, params); +} + +llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_params & params) : + graph(model, params, true) { + GGML_ASSERT(hparams.n_embd_head_v() == hparams.n_embd_head_k()); + + const int64_t hc = hparams.hc_count; + + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + ggml_tensor * inp = build_inp_embd(model.tok_embd); + + auto * inp_mem = build_inp_mem_hybrid(); + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // every hyper-connection stream starts as a copy of the token embedding + 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); + + llm_graph_input_qsa * inp_qsa = build_qsa_inp(); + + const auto * mctx_hyb = static_cast(mctx); + + llm_graph_input_ple * inp_ple = nullptr; + for (int il = 0; il < n_layer; ++il) { + if (!model.layers[il].ple_ngram_embd) { + continue; + } + + // the hash constants are derived per PLE layer, only one such layer is supported + GGML_ASSERT(!inp_ple && "only one PLE layer is supported"); + + const int64_t n_head_ngram = hparams.ple_ngram_heads; + + // eos pads the history and cuts it at sequence breaks, so it has to exist + GGML_ASSERT(model.vocab.token_eos() != LLAMA_TOKEN_NULL && + "PLE needs eos to pad the n-gram history"); + + auto ple = std::make_unique(model.hparams, mctx_hyb->get_ple(), + model.vocab.token_eos(), il); + + ple->ngram_ids = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_head_ngram*n_tokens); + ple->tokens = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, 1, n_tokens); + ggml_set_input(ple->ngram_ids); + ggml_set_input(ple->tokens); + + inp_ple = (llm_graph_input_ple *) res->add_input(std::move(ple)); + } + + // PLE has a recurrent state of its own, so the delta-net row below stays exactly the + // delta-net conv state and both can use the shared build_conv_state + llm_graph_input_rs * inp_ple_rs = nullptr; + if (inp_ple) { + const auto * mctx_ple = mctx_hyb->get_ple(); + + auto rs = std::make_unique(mctx_ple); + + const int64_t n_rs = mctx_ple->get_n_rs(); + + rs->s_copy = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_rs); + ggml_set_input(rs->s_copy); + + rs->s_copy_main = ggml_view_1d(ctx0, rs->s_copy, ubatch.n_seqs, 0); + rs->s_copy_extra = ggml_view_1d(ctx0, rs->s_copy, n_rs - ubatch.n_seqs, + ubatch.n_seqs*rs->s_copy->nb[0]); + + rs->head = mctx_ple->get_head(); + rs->rs_z = mctx_ple->get_rs_z(); + + inp_ple_rs = (llm_graph_input_rs *) res->add_input(std::move(rs)); + } + + for (int il = 0; il < n_layer; ++il) { + if (model.layers[il].ple_ngram_embd) { + inpL = ggml_add(ctx0, inpL, + build_ple(inpL, inp_ple->ngram_ids, inp_ple->tokens, inp_ple_rs, il)); + cb(inpL, "ple_out", il); + } + + ggml_tensor * inject = nullptr; + ggml_tensor * cur = build_hc_pre(inpL, + model.layers[il].hc_attn_norm, + model.layers[il].hc_attn_down, + model.layers[il].hc_attn_up, + model.layers[il].hc_attn_inject, &inject, il); + cb(cur, "hc_attn_pre", il); + + if (hparams.is_recr(il)) { + cur = build_layer_attn_linear(inp_mem->get_recr(), cur, il); + } else { + cur = build_layer_attn(inp_mem->get_attn(), inp_qsa, cur, inp_pos, sections, il); + } + + inpL = build_hc_post(inpL, cur, inject, il); + cb(inpL, "hc_attn_post", il); + + cur = build_hc_pre(inpL, + model.layers[il].hc_ffn_norm, + model.layers[il].hc_ffn_down, + model.layers[il].hc_ffn_up, + model.layers[il].hc_ffn_inject, &inject, il); + cb(cur, "hc_ffn_pre", il); + + cur = build_layer_ffn(cur, il); + cb(cur, "ffn_out", il); + + inpL = build_hc_post(inpL, cur, inject, il); + inpL = build_cvec(inpL, il); + cb(inpL, "l_out", il); + } + + ggml_tensor * flat = ggml_reshape_2d(ctx0, inpL, hc*n_embd, n_tokens); + + // the MTP block consumes the streams, so hand them over before the mixer collapses them. The + // read-back covers every token, so it has to see the rows before the output masking below. + res->t_h_nextn = flat; + + if (inp_out_ids) { + flat = ggml_get_rows(ctx0, flat, inp_out_ids); + inpL = ggml_reshape_3d(ctx0, flat, n_embd, hc, n_outputs); + } + + // the mixer without an injection gate is also the final norm + ggml_tensor * cur = build_hc_pre(inpL, model.hc_mix_norm, model.hc_mix_down, model.hc_mix_up, nullptr, nullptr, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +ggml_tensor * llama_model_qwen4exp::graph::build_hc_norm( + ggml_tensor * x, + ggml_tensor * w, + int il) { + ggml_tensor * cur = ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps); + cur = ggml_mul(ctx0, cur, ggml_reshape_2d(ctx0, w, n_embd, hparams.hc_count)); + cb(cur, "hc_norm", il); + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_hc_pre( + ggml_tensor * x, + ggml_tensor * norm, + ggml_tensor * down, + ggml_tensor * up, + ggml_tensor * inject, + ggml_tensor ** inject_out, + int il) { + const int64_t hc = hparams.hc_count; + const int64_t nt = x->ne[2]; + + ggml_tensor * xn = build_hc_norm(x, norm, il); + ggml_tensor * flat = ggml_reshape_2d(ctx0, xn, hc*n_embd, nt); + + ggml_tensor * mix = build_lora_mm(down, flat); + mix = ggml_silu(ctx0, ggml_scale(ctx0, mix, 1.0f/hc)); + mix = ggml_sigmoid(ctx0, build_lora_mm(up, mix)); + mix = ggml_reshape_3d(ctx0, mix, n_embd, hc, nt); + cb(mix, "hc_mix", il); + + // mean over the streams + ggml_tensor * cur = ggml_mul(ctx0, mix, xn); + cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 0, 2, 3)); + cur = ggml_mean(ctx0, cur); + cur = ggml_reshape_2d(ctx0, cur, n_embd, nt); + + if (inject_out) { + ggml_tensor * w = build_lora_mm(inject, flat); + w = ggml_scale(ctx0, ggml_sigmoid(ctx0, ggml_scale(ctx0, w, 1.0f/hc)), 2.0f); + cb(w, "hc_inject", il); + *inject_out = w; + } + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_hc_post( + ggml_tensor * residual, + ggml_tensor * y, + ggml_tensor * inject, + int il) { + const int64_t hc = hparams.hc_count; + const int64_t nt = y->ne[1]; + + ggml_tensor * cur = ggml_reshape_3d(ctx0, y, n_embd, 1, nt); + cur = ggml_repeat_4d(ctx0, cur, n_embd, hc, nt, 1); + cur = ggml_mul(ctx0, cur, ggml_reshape_3d(ctx0, inject, 1, hc, nt)); + cb(cur, "hc_injected", il); + + return ggml_add(ctx0, residual, cur); +} + +ggml_tensor * llama_model_qwen4exp::graph::build_ple( + ggml_tensor * x, + ggml_tensor * ngram_ids, + ggml_tensor * tokens, + llm_graph_input_rs * inp, + int il) { + const auto & layer = model.layers[il]; + + const int64_t hc = hparams.hc_count; + const int64_t hc_dim = hc*n_embd; + const int64_t n_ngram = hparams.ple_ngram_size; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + ggml_tensor * emb = ggml_get_rows(ctx0, layer.ple_ngram_embd, ngram_ids); + emb = ggml_reshape_2d(ctx0, emb, hparams.n_embd_ple, n_tokens); + cb(emb, "ple_embd", il); + + ggml_tensor * key = build_lora_mm(layer.ple_key, emb); + key = build_hc_norm(ggml_reshape_3d(ctx0, key, n_embd, hc, n_tokens), layer.ple_key_norm, il); + + ggml_tensor * val = build_lora_mm(layer.ple_value, emb); + ggml_tensor * query = build_hc_norm(x, layer.ple_query_norm, il); + + ggml_tensor * gate = ggml_sum_rows(ctx0, ggml_mul(ctx0, key, query)); + gate = ggml_scale(ctx0, gate, 1.0f/sqrtf((float) n_embd)); + + // signed square root keeps the gate in a narrow range + gate = ggml_mul(ctx0, + ggml_sqrt(ctx0, ggml_clamp(ctx0, ggml_abs(ctx0, gate), 1e-6f, INFINITY)), + ggml_sgn(ctx0, gate)); + gate = ggml_sigmoid(ctx0, gate); + cb(gate, "ple_gate", il); + + ggml_tensor * cur = ggml_reshape_3d(ctx0, val, n_embd, 1, n_tokens); + cur = ggml_repeat_4d(ctx0, cur, n_embd, hc, n_tokens, 1); + cur = ggml_mul(ctx0, cur, gate); + cb(cur, "ple_value", il); + + // dilated depthwise convolution over the normalized values + const int64_t n_kernel = layer.ple_conv1d->ne[0]; + const int64_t n_taps = (n_kernel - 1)*n_ngram; + const int64_t n_gdn = (hparams.ssm_d_conv - 1)*(hparams.ssm_d_inner + 2*hparams.ssm_n_group*hparams.ssm_d_state); + + GGML_UNUSED(n_gdn); + ggml_tensor * conv_in = build_conv_state(inp, inp->mctx->get_r_l(il), + ggml_reshape_3d(ctx0, build_hc_norm(cur, layer.ple_conv_norm, il), hc_dim, n_seq_tokens, n_seqs), + n_taps + 1, hc_dim, il); + + const size_t es_in = ggml_element_size(conv_in); + const size_t es_w = ggml_element_size(layer.ple_conv1d); + + ggml_tensor * conv = nullptr; + for (int64_t k = 0; k < n_kernel; ++k) { + ggml_tensor * w = ggml_view_2d(ctx0, layer.ple_conv1d, 1, hc_dim, layer.ple_conv1d->nb[1], k*es_w); + ggml_tensor * v = ggml_cont(ctx0, ggml_view_3d(ctx0, conv_in, n_seq_tokens, hc_dim, n_seqs, + conv_in->nb[1], conv_in->nb[2], k*n_ngram*es_in)); + v = ggml_mul(ctx0, v, w); + + conv = conv ? ggml_add(ctx0, conv, v) : v; + } + conv = ggml_silu(ctx0, conv); + conv = ggml_cont(ctx0, ggml_transpose(ctx0, conv)); + cb(conv, "ple_conv", il); + + // remember the tokens of this ubatch as the n-gram history of the next one. These are token + // ids, not a conv state, so build_conv_state does not apply - it reshapes the whole row + { + ggml_tensor * s_all = inp->mctx->get_s_l(il); + + // gather the source rows directly: build_rs would also clear a row and propagate the + // extra states, which is what a delta-net state needs but not what this is - the history + // is moved by hand below, and the clearing would wipe it on every decode step + ggml_tensor * s_row = ggml_get_rows(ctx0, + ggml_reshape_2d(ctx0, s_all, n_ngram - 1, s_all->ne[1]), inp->s_copy_main); + + ggml_tensor * hist = ggml_reshape_3d(ctx0, s_row, n_ngram - 1, 1, n_seqs); + ggml_tensor * both = ggml_concat(ctx0, hist, + ggml_transpose(ctx0, ggml_reshape_3d(ctx0, tokens, 1, n_seq_tokens, n_seqs)), 0); + + const size_t es = ggml_element_size(s_all); + + ggml_tensor * tail = ggml_view_3d(ctx0, both, n_ngram - 1, 1, n_seqs, + both->nb[1], both->nb[2], (both->ne[0] - (n_ngram - 1))*ggml_element_size(both)); + + ggml_tensor * dst = ggml_view_3d(ctx0, s_all, n_ngram - 1, 1, n_seqs, + (n_ngram - 1)*es, (n_ngram - 1)*es, + (inp->mctx->get_head()*(n_ngram - 1))*es); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, tail, dst)); + } + + return ggml_add(ctx0, cur, ggml_reshape_3d(ctx0, conv, n_embd, hc, n_tokens)); +} + +// No indexer cache means QSA cannot run for this context (see llama_model::create_memory), and +// the sparse layers fall back to dense. +llm_graph_input_qsa * llama_model_qwen4exp::graph::build_qsa_inp() { + const auto * mctx_hyb = static_cast(mctx); + + if (!hparams.has_qsa() || !mctx_hyb->get_lid()) { + return nullptr; + } + + const int64_t ratio = hparams.indexer_block_size; + // round 256 up to a whole number of blocks; GGML_PAD would need ratio to be a power of two + const uint32_t pad = ((256u + (uint32_t) ratio - 1)/(uint32_t) ratio)*(uint32_t) ratio; + const int64_t n_pos = mctx_hyb->get_n_pos(pad, cparams.n_seq_max); + const int64_t ns = cparams.kv_unified ? 1 : ubatch.n_seqs_unq; + const int64_t n_tps = n_tokens/ns; + + auto qsa = std::make_unique(mctx_hyb, ratio, pad); + + qsa->n_kv_lid = mctx_hyb->get_lid()->get_n_kv(); + + // one query per stream is what lets the attention read a gathered buffer: the keys of a flash + // attention call are shared by all of its queries, and each query picks its own blocks. Longer + // ubatches - a prefill, or a speculative verification batch - take the mask path below + const bool gather = n_tps == 1 && cparams.flash_attn; + + qsa->k_idxs = mctx_hyb->get_lid()->build_input_k_idxs(ctx0, ubatch); + qsa->pos_cell = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_pos, ns); + qsa->blk_mask = ggml_new_tensor_4d(ctx0, GGML_TYPE_F16, n_pos/ratio, n_tps, 1, ns); + qsa->blk_cur = ggml_new_tensor_4d(ctx0, GGML_TYPE_I32, 1, n_tps, 1, ns); + + ggml_set_input(qsa->pos_cell); + ggml_set_input(qsa->blk_mask); + ggml_set_input(qsa->blk_cur); + + // only the mask path walks the cells, and an input that nothing reads gets no buffer + if (!gather) { + qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, mctx_hyb->get_attn()->get_n_kv(), ns); + ggml_set_input(qsa->cell_blk); + } + + // The gather indexes V by cell, so it needs V stored untransposed, which is what flash + // attention asks for anyway (llama_kv_cache is built with v_trans = !flash_attn) + if (gather) { + qsa->pos_cell_f = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, ns); + qsa->pos_vis = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_pos, ns); + + ggml_set_input(qsa->pos_cell_f); + ggml_set_input(qsa->pos_vis); + } + + return (llm_graph_input_qsa *) res->add_input(std::move(qsa)); +} + +// Block scores, shared by the mask and the gather path: relu(q . k_blk) summed over the indexer +// heads, with the incomplete and invisible blocks already closed by blk_mask. +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_scores( + llm_graph_input_qsa * inp_qsa, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il) { + const auto & layer = model.layers[il]; + + const int64_t n_idx = hparams.indexer_head_size; + const int64_t n_ih = hparams.indexer_n_head; + const int64_t ratio = hparams.indexer_block_size; + const int64_t n_blk = inp_qsa->blk_mask->ne[0]; + const int64_t n_tps = inp_qsa->blk_mask->ne[1]; + const int64_t ns = inp_qsa->blk_mask->ne[3]; + + // indexer queries: normed, then the same partial IMRoPE as the attention queries + ggml_tensor * q = build_lora_mm(layer.index_q_proj, cur); + q = ggml_reshape_3d(ctx0, q, n_idx, n_ih, n_tokens); + q = build_norm(q, layer.index_q_norm, nullptr, LLM_NORM_RMS, il); + q = ggml_rope_multi(ctx0, q, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(q, "qsa_q", il); + + // indexer keys are cached raw: the norm and the RoPE come after the block pooling + ggml_tensor * k = build_lora_mm(layer.index_k_proj, cur); + k = ggml_reshape_3d(ctx0, k, n_idx, 1, n_tokens); + ggml_build_forward_expand(gf, inp_qsa->mctx->get_lid()->cpy_k(ctx0, k, inp_qsa->k_idxs, il)); + + // read the cached keys by position, so a block always covers the same `ratio` positions + ggml_tensor * kc = inp_qsa->mctx->get_lid()->get_k(ctx0, il); + kc = ggml_view_3d(ctx0, kc, n_idx, kc->ne[2], ns, kc->nb[2], kc->nb[3], 0); + + ggml_tensor * kb = ggml_get_rows(ctx0, kc, inp_qsa->pos_cell); + kb = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, kb, n_idx, ratio, n_blk*ns), 1, 0, 2, 3)); + kb = ggml_reshape_3d(ctx0, ggml_mean(ctx0, kb), n_idx, 1, n_blk*ns); + kb = build_norm(kb, layer.index_k_norm, nullptr, LLM_NORM_RMS, il); + + // a block covers positions [b*ratio, b*ratio + ratio), so it is anchored to b*ratio + ggml_tensor * bpos = ggml_repeat_4d(ctx0, + ggml_arange(ctx0, 0.0f, (float) n_blk*ratio, (float) ratio), + n_blk*ns*hparams.n_pos_per_embd(), 1, 1, 1); + bpos = ggml_cast(ctx0, bpos, GGML_TYPE_I32); + kb = ggml_rope_multi(ctx0, kb, bpos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(kb, "qsa_k_blk", il); + + // score = sum_h relu(q_h . k_blk)/sqrt(n_idx), the indexer weights are prescaled + ggml_tensor * w = ggml_scale(ctx0, + ggml_repeat_4d(ctx0, ggml_arange(ctx0, 1.0f, 2.0f, 1.0f), n_ih, n_tps, 1, ns), + 1.0f/sqrtf((float) n_idx)); + ggml_tensor * sc = ggml_lightning_indexer(ctx0, + ggml_reshape_4d(ctx0, q, n_idx, n_ih, n_tps, ns), + ggml_reshape_4d(ctx0, kb, n_idx, 1, n_blk, ns), + w, inp_qsa->blk_mask); + cb(sc, "qsa_score", il); + + return sc; +} + +// The mask path: turn the block scores into the kq mask of the cells the attention may read. Every +// cell is walked whatever the mask says - see build_qsa_gather() for the decode path. +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_mask( + llm_graph_input_qsa * inp_qsa, + ggml_tensor * sc, + ggml_tensor * kq_mask, + int il) { + const int64_t ratio = hparams.indexer_block_size; + const int64_t n_blk = inp_qsa->blk_mask->ne[0]; + const int64_t n_tps = inp_qsa->blk_mask->ne[1]; + const int64_t ns = inp_qsa->blk_mask->ne[3]; + const int64_t n_kv = inp_qsa->cell_blk->ne[0]; + const int64_t n_top = std::min(n_blk, (int64_t) hparams.indexer_top_k/ratio); + + // fully masked block axis, then open the top-k blocks and the query's own block + ggml_tensor * bm = ggml_fill(ctx0, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_blk, n_tps*ns), -INFINITY); + + ggml_tensor * idx = ggml_top_k(ctx0, sc, n_top); + bm = ggml_set_rows(ctx0, bm, + ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, n_top, n_tps*ns), 0.0f), + ggml_reshape_2d(ctx0, idx, n_top, n_tps*ns)); + bm = ggml_set_rows(ctx0, bm, + ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, n_tps*ns), 0.0f), + ggml_reshape_2d(ctx0, inp_qsa->blk_cur, 1, n_tps*ns)); + + // block -> cell: gather the block mask with the block each cell belongs to + bm = ggml_cont(ctx0, ggml_permute(ctx0, ggml_reshape_3d(ctx0, bm, n_blk, n_tps, ns), 1, 0, 2, 3)); + + ggml_tensor * mask = ggml_get_rows(ctx0, bm, inp_qsa->cell_blk); + mask = ggml_cont(ctx0, ggml_permute(ctx0, mask, 1, 0, 2, 3)); + mask = ggml_reshape_4d(ctx0, mask, n_kv, n_tps, 1, ns); + + // the causal mask also closes empty cells and the cells of the other sequences + mask = ggml_add(ctx0, mask, + kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32)); + if (kq_mask->type != GGML_TYPE_F32) { + mask = ggml_cast(ctx0, mask, kq_mask->type); + } + cb(mask, "qsa_kq_mask", il); + + return mask; +} + +// The decode path: gather the selected keys and values into one compact buffer and attend over +// that, instead of masking the whole cache. The CUDA flash attention only trims the tail of the KV +// range (flash_attn_mask_to_KV_max scans back to the first tile that is not all -inf), so a mask +// with holes in it costs exactly as much as dense attention. One query per stream is what makes +// this possible: the keys a flash attention call reads are shared by all of its queries, and each +// query picks its own blocks. +// ref: src/models/minimax-m3.cpp, the msa_decode branch +void llama_model_qwen4exp::graph::build_qsa_gather( + llm_graph_input_qsa * inp_qsa, + ggml_tensor * sc, + ggml_tensor * k_all, + ggml_tensor * v_all, + ggml_tensor * kq_mask, + ggml_tensor ** k_out, + ggml_tensor ** v_out, + ggml_tensor ** m_out) { + const int64_t ratio = hparams.indexer_block_size; + const int64_t n_blk = inp_qsa->blk_mask->ne[0]; + const int64_t ns = inp_qsa->blk_mask->ne[3]; + const int64_t n_pos = inp_qsa->pos_cell_f->ne[0]; + const int64_t n_top = std::min(n_blk, (int64_t) hparams.indexer_top_k/ratio); + + const int64_t d_head = k_all->ne[0]; + const int64_t n_hkv = k_all->ne[1]; + const int64_t n_kv = k_all->ne[2]; + + // The query's own block has to be read on top of the top-k. Do not append it: below the indexer + // budget the top-k reaches into the masked blocks and can pick that same block, and a key + // gathered twice takes two shares of the softmax. Score it above everything instead, so it is + // always rank 0 of a top-k of n_top + 1 - the same visible set as the mask path, no duplicates. + // ref: the bias input of src/models/minimax-m3.cpp + // n_top can already be every block, and top_k cannot ask for more rows than there are. Then + // the own block is still rank 0 and the rest is all of them, which is what the mask path opens + const int64_t n_sel = std::min(n_blk, n_top + 1); + const int64_t n_g = ratio*n_sel; + + ggml_tensor * scb = ggml_cont(ctx0, ggml_permute(ctx0, + ggml_reshape_3d(ctx0, sc, n_blk, 1, ns), 1, 0, 2, 3)); + scb = ggml_set_rows(ctx0, scb, + ggml_fill(ctx0, ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, 1, 1, ns), 1e30f), + ggml_reshape_2d(ctx0, inp_qsa->blk_cur, 1, ns)); + scb = ggml_cont(ctx0, ggml_permute(ctx0, scb, 1, 0, 2, 3)); + + ggml_tensor * sel = ggml_reshape_3d(ctx0, + ggml_cast(ctx0, ggml_top_k(ctx0, ggml_reshape_3d(ctx0, scb, n_blk, 1, ns), n_sel), + GGML_TYPE_F32), n_sel, 1, ns); + + // block -> the `ratio` positions it covers: tj[t, b] = sel[b]*ratio + t + ggml_tensor * tj = ggml_add(ctx0, + ggml_repeat_4d(ctx0, ggml_scale(ctx0, ggml_reshape_4d(ctx0, sel, 1, n_sel, 1, ns), (float) ratio), + ratio, n_sel, 1, ns), + ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.0f, (float) ratio, 1.0f), ratio, 1, 1)); + ggml_tensor * tokj = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tj, n_g, ns), GGML_TYPE_I32); + + // position -> cell. The map has to be read as values here, so it is the F32 copy + ggml_tensor * cs = ggml_get_rows(ctx0, + ggml_reshape_3d(ctx0, inp_qsa->pos_cell_f, 1, n_pos, ns), tokj); + + // cell -> the n_hkv rows of that cell: tr[h, l] = cs[l]*n_hkv + h + ggml_tensor * tr = ggml_add(ctx0, + ggml_repeat_4d(ctx0, ggml_scale(ctx0, ggml_reshape_4d(ctx0, cs, 1, n_g, ns, 1), (float) n_hkv), + n_hkv, n_g, ns, 1), + ggml_reshape_2d(ctx0, ggml_arange(ctx0, 0.0f, (float) n_hkv, 1.0f), n_hkv, 1)); + ggml_tensor * tokr = ggml_cast(ctx0, ggml_reshape_2d(ctx0, tr, n_hkv*n_g, ns), GGML_TYPE_I32); + + ggml_tensor * k3 = ggml_view_3d(ctx0, k_all, d_head, n_hkv*n_kv, ns, k_all->nb[1], k_all->nb[3], 0); + ggml_tensor * v3 = ggml_view_3d(ctx0, v_all, d_head, n_hkv*n_kv, ns, v_all->nb[1], v_all->nb[3], 0); + + ggml_tensor * kg = ggml_get_rows(ctx0, k3, tokr); + ggml_tensor * vg = ggml_get_rows(ctx0, v3, tokr); + + // the rows came out with the heads fastest, so this is already the layout the cache hands to + // build_attn_mha - one gathered position where a cell used to be + ggml_tensor * kfa = ggml_reshape_4d(ctx0, kg, d_head, n_hkv, n_g, ns); + ggml_tensor * vfa = ggml_reshape_4d(ctx0, vg, d_head, n_hkv, n_g, ns); + + // a quantized cache is gathered as it is stored, and flash attention wants F16 + if (ggml_is_quantized(kfa->type)) { kfa = ggml_cast(ctx0, kfa, GGML_TYPE_F16); } + if (ggml_is_quantized(vfa->type)) { vfa = ggml_cast(ctx0, vfa, GGML_TYPE_F16); } + + // two terms. pos_vis closes the positions the query may not read, which the gather can pick up + // because it reads whole blocks - the trailing one runs past the query and those positions have + // no cell, so they would otherwise land on cell 0. The causal mask then closes the empty cells + // and the cells of the other sequences, which position space knows nothing about + ggml_tensor * mg = ggml_get_rows(ctx0, + ggml_reshape_3d(ctx0, inp_qsa->pos_vis, 1, n_pos, ns), tokj); + + ggml_tensor * km = ggml_view_3d(ctx0, kq_mask, kq_mask->ne[0], 1, ns, + kq_mask->nb[1], kq_mask->nb[3], 0); + km = ggml_cont(ctx0, ggml_permute(ctx0, km, 1, 0, 2, 3)); + if (km->type != GGML_TYPE_F32) { + km = ggml_cast(ctx0, km, GGML_TYPE_F32); + } + + ggml_tensor * tokc = ggml_cast(ctx0, ggml_reshape_2d(ctx0, cs, n_g, ns), GGML_TYPE_I32); + + mg = ggml_add(ctx0, mg, ggml_get_rows(ctx0, km, tokc)); + + *k_out = kfa; + *v_out = vfa; + *m_out = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mg, n_g, 1, 1, ns), GGML_TYPE_F16); +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( + llm_graph_input_attn_kv * inp, + llm_graph_input_qsa * inp_qsa, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + + // a single Q projection produces the query and the output gate + ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); + cb(Qcur_full, "Qcur_full", il); + + ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, 0); + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + + ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, + ggml_element_size(Qcur_full) * n_embd_head); + gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); + cb(gate, "gate", il); + + ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "Kcur_normed", il); + + ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + Qcur = ggml_rope_multi(ctx0, Qcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + Kcur = ggml_rope_multi(ctx0, Kcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + if (inp_qsa == nullptr) { + cur = build_attn(inp, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + } else { + // same as build_attn() above, except that the mask only lets the selected cells through + // ref: llm_graph_context::build_attn() for the k_dsa input + const auto * mctx = inp->mctx; + + // a quantized KV cache rotates K/V before storing them, see llm_graph_context::build_attn + if (inp->self_k_rot) { + Qcur = llama_mul_mat_hadamard(ctx0, Qcur, inp->self_k_rot); + Kcur = llama_mul_mat_hadamard(ctx0, Kcur, inp->self_k_rot); + } + if (inp->self_v_rot) { + Vcur = llama_mul_mat_hadamard(ctx0, Vcur, inp->self_v_rot); + } + + ggml_build_forward_expand(gf, mctx->cpy_k(ctx0, Kcur, inp->get_k_idxs(), il)); + ggml_build_forward_expand(gf, mctx->cpy_v(ctx0, Vcur, inp->get_v_idxs(), il)); + + ggml_tensor * sc = build_qsa_scores(inp_qsa, cur, inp_pos, sections, il); + + if (inp_qsa->pos_cell_f) { + ggml_tensor * kg = nullptr; + ggml_tensor * vg = nullptr; + ggml_tensor * mg = nullptr; + + build_qsa_gather(inp_qsa, sc, mctx->get_k(ctx0, il), mctx->get_v(ctx0, il), + inp->get_kq_mask(), &kg, &vg, &mg); + + cur = build_attn_mha(Qcur, kg, vg, nullptr, mg, nullptr, nullptr, kq_scale, il); + } else { + ggml_tensor * mask = build_qsa_mask(inp_qsa, sc, inp->get_kq_mask(), il); + + cur = build_attn_mha(Qcur, mctx->get_k(ctx0, il), mctx->get_v(ctx0, il), + nullptr, mask, nullptr, nullptr, kq_scale, il); + } + + if (inp->self_v_rot) { + cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); + } + } + cb(cur, "attn_pregate", il); + + cur = ggml_mul(ctx0, cur, ggml_sigmoid(ctx0, gate)); + cb(cur, "attn_gated", il); + + cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); + cb(cur, "attn_output", il); + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( + llm_graph_input_rs * inp, + ggml_tensor * cur, + int il) { + const auto * mctx_cur = inp->mctx; + + const int64_t d_inner = hparams.ssm_d_inner; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t head_k_dim = hparams.ssm_d_state; + const int64_t num_k_heads = hparams.ssm_n_group; + const int64_t num_v_heads = hparams.ssm_dt_rank; + const int64_t head_v_dim = d_inner / num_v_heads; + 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); + + ggml_tensor * qkv_mixed = build_lora_mm(model.layers[il].wqkv, cur, model.layers[il].wqkv_s); + qkv_mixed = ggml_reshape_3d(ctx0, qkv_mixed, qkv_mixed->ne[0], n_seq_tokens, n_seqs); + cb(qkv_mixed, "linear_attn_qkv_mixed", il); + + ggml_tensor * z = build_lora_mm(model.layers[il].wqkv_gate, cur, model.layers[il].wqkv_gate_s); + cb(z, "z", il); + + ggml_tensor * beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s); + beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs); + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "beta_sigmoid", il); + + ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s); + alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs); + + ggml_tensor * gate = ggml_mul(ctx0, ggml_softplus(ctx0, ggml_add(ctx0, alpha, model.layers[il].ssm_dt)), + model.layers[il].ssm_a); + gate = ggml_reshape_4d(ctx0, gate, 1, num_v_heads, n_seq_tokens, n_seqs); + cb(gate, "gate", il); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + + ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d; + const int64_t conv_kernel_size = conv_kernel->ne[0]; + const int64_t conv_channels = d_inner + 2*num_k_heads*head_k_dim; + + ggml_tensor * conv_input = build_conv_state(inp, conv_states_all, qkv_mixed, + conv_kernel_size, conv_channels, il); + + ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); + + ggml_tensor * conv_qkv_mix = ggml_silu(ctx0, ggml_ssm_conv(ctx0, conv_input, conv_kernel)); + cb(conv_qkv_mix, "conv_output_silu", il); + + const int64_t qkv_dim = head_k_dim*num_k_heads*2 + head_v_dim*num_v_heads; + const int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, qkv_dim); + + ggml_tensor * q_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), nb1_qkv, nb1_qkv*n_seq_tokens, 0); + + ggml_tensor * k_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), nb1_qkv, nb1_qkv*n_seq_tokens, + ggml_row_size(conv_qkv_mix->type, head_k_dim*num_k_heads)); + + ggml_tensor * v_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_v_dim, num_v_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_v_dim), nb1_qkv, nb1_qkv*n_seq_tokens, + ggml_row_size(conv_qkv_mix->type, 2*head_k_dim*num_k_heads)); + + q_conv = ggml_l2_norm(ctx0, q_conv, hparams.f_norm_rms_eps); + k_conv = ggml_l2_norm(ctx0, k_conv, hparams.f_norm_rms_eps); + + if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) { + GGML_ASSERT(num_v_heads % num_k_heads == 0); + q_conv = ggml_repeat_4d(ctx0, q_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); + k_conv = ggml_repeat_4d(ctx0, k_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); + } + + ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il); + + // gated output norm, Qwen4-Exp uses a sigmoid gate here + ggml_tensor * z_4d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); + output = ggml_mul(ctx0, build_norm(output, model.layers[il].ssm_norm, nullptr, LLM_NORM_RMS, il), + ggml_sigmoid(ctx0, z_4d)); + + output = ggml_reshape_3d(ctx0, output, head_v_dim*num_v_heads, n_seq_tokens, n_seqs); + + cur = build_lora_mm(model.layers[il].ssm_out, output, model.layers[il].ssm_out_s); + cb(cur, "linear_attn_out", il); + + return ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens*n_seqs); +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, const int il) { + ggml_tensor * moe_out = + build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, true, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, il, + nullptr, model.layers[il].ffn_gate_up_exps, + model.layers[il].ffn_up_exps_s, + model.layers[il].ffn_gate_exps_s, + model.layers[il].ffn_down_exps_s); + cb(moe_out, "ffn_moe_out", il); + + ggml_tensor * ffn_shexp = + build_ffn(cur, + model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s, + model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s, + model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + ggml_tensor * shared_gate = ggml_sigmoid(ctx0, build_lora_mm(model.layers[il].ffn_gate_inp_shexp, cur)); + cb(shared_gate, "shared_expert_gate_sigmoid", il); + + return ggml_add(ctx0, moe_out, ggml_mul(ctx0, ffn_shexp, shared_gate)); +}