Commit 78feec062 for llama.cpp
commit 78feec062d50ec9e8940244cac918ac4480b60cb
Author: Xuan-Son Nguyen <son@huggingface.co>
Date: Thu Oct 29 15:43:33 2026 +0100
model: mistral le chonk (#31376)
* llama: extend DSA and sdpark to new model
* fix ci
diff --git a/conversion/__init__.py b/conversion/__init__.py
index d4e9f0441..3976e03ad 100644
--- a/conversion/__init__.py
+++ b/conversion/__init__.py
@@ -213,6 +213,7 @@ TEXT_MODEL_MAP: dict[str, str] = {
"Ministral3Model": "mistral3",
"Mistral3ForConditionalGeneration": "mistral3",
"Mistral3Model": "mistral3",
+ "MistralLarge4ForConditionalGeneration": "mistral4",
"VoxtralRealtimeForConditionalGeneration": "mistral3",
"MistralForCausalLM": "llama",
"MixtralForCausalLM": "llama",
diff --git a/conversion/base.py b/conversion/base.py
index d8420eb32..066be61b2 100644
--- a/conversion/base.py
+++ b/conversion/base.py
@@ -350,6 +350,12 @@ class ModelBase:
def dequant_simple(weight: Tensor, scale: Tensor, block_size: Sequence[int] | None = None) -> Tensor:
scale = scale.float()
+ # ModelOpt block scales are pre-shaped for broadcasting, e.g. [out/128, 1, in/128, 1]
+ if block_size is None and scale.ndim == 2 * weight.ndim and all(n == 1 for n in scale.shape[1::2]):
+ assert all(w % n == 0 for w, n in zip(weight.shape, scale.shape[0::2]))
+ block_size = [w // n for w, n in zip(weight.shape, scale.shape[0::2])]
+ scale = scale.reshape(tuple(scale.shape[0::2]))
+
if block_size is not None:
dim_offset = scale.ndim - len(block_size)
for i, size in enumerate(block_size):
@@ -1035,7 +1041,7 @@ class ModelBase:
expert_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
expert_input_scales: dict[tuple[int, str], list[tuple[int, float]]] = {}
expert_shapes: dict[tuple[int, str], list[int]] = {}
- n_experts = self.find_hparam(["num_local_experts", "num_experts"], optional=True) or 0
+ n_experts = self.find_hparam(["num_local_experts", "num_experts", "n_routed_experts"], optional=True) or 0
consumed: list[str] = []
for name in self.model_tensors.keys():
@@ -2672,7 +2678,13 @@ class TextModel(ModelBase):
)
if vocab.tokenizer_type == MistralTokenizerType.tekken:
- self.gguf_writer.add_tokenizer_pre("tekken")
+ with open(vocab.fname_tokenizer, encoding="utf-8") as f:
+ pattern = json.load(f)["config"]["pattern"]
+ # newer tekken.json split numbers into right-aligned groups of 3 digits instead of single digits
+ if "\\p{N}{1,3}(?=(?:\\p{N}{3})*(?!\\p{N}))" in pattern:
+ self.gguf_writer.add_tokenizer_pre("tekken-num3")
+ else:
+ self.gguf_writer.add_tokenizer_pre("tekken")
self.gguf_writer.add_token_merges(
vocab.extract_vocab_merges_from_model()
)
diff --git a/conversion/mistral4.py b/conversion/mistral4.py
new file mode 100644
index 000000000..dea6b3ecb
--- /dev/null
+++ b/conversion/mistral4.py
@@ -0,0 +1,217 @@
+from __future__ import annotations
+
+from typing import Callable, Iterable, TYPE_CHECKING
+
+import torch
+
+if TYPE_CHECKING:
+ from torch import Tensor
+
+from .base import ModelBase, gguf
+
+from .mistral import MistralModel, MistralMoeModel
+
+
+@ModelBase.register("MistralLarge4ForConditionalGeneration")
+class MistralLarge4Model(MistralMoeModel):
+ model_arch = gguf.MODEL_ARCH.MISTRAL_LARGE4
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ n_layer = self.hparams["num_hidden_layers"]
+ dsa = self.hparams["dsa"]
+
+ def expand(pattern: str) -> list[str]:
+ values = pattern.split(",")
+ return [values[i % len(values)] for i in range(n_layer)]
+
+ self.layer_has_rope = [v == "rope" for v in expand(self.hparams.get("nope") or "rope")]
+ self.layer_indexer_full = [v == "full" for v in expand(dsa.get("sharing") or "full")]
+
+ kv_sharing = self.hparams.get("kv_sharing") or {}
+ if "source_layers_str" in kv_sharing:
+ self.layer_kv_source = [int(v) for v in kv_sharing["source_layers_str"].split(",")]
+ else:
+ assert not kv_sharing, f"unsupported kv_sharing config: {kv_sharing}"
+ self.layer_kv_source = list(range(n_layer))
+ assert len(self.layer_kv_source) == n_layer
+
+ # note: the rope dim of the indexer must match the main attention
+ assert dsa.get("rope_head_dim", -1) in (-1, self.hparams["qk_rope_head_dim"])
+ assert dsa.get("match_nope", False), "indexer must follow the NoPE pattern of the main attention"
+ assert dsa.get("ssmax_scaling", False), "indexer without SSMax is not supported"
+ assert self.hparams.get("q_lora_rank") is None
+ assert not self.hparams["yarn"].get("apply_scale", True), "only YaRN without magnitude scaling is supported"
+
+ def dequant_model(self):
+ quant_config = self.hparams.get("quantization")
+ if quant_config is not None:
+ assert quant_config["qformat_weight"] == "fp8_e4m3"
+ assert quant_config.get("qscheme_weight") == "BLOCK"
+ self.hparams["quantization_config"] = {
+ "quant_method": "fp8",
+ "weight_block_size": [quant_config["block_size_out"], quant_config["block_size_in"]],
+ }
+ return super().dequant_model()
+
+ def set_gguf_parameters(self):
+ # skip the MistralMoeModel overrides, they apply the deepseek2 YaRN scaling
+ super(MistralMoeModel, self).set_gguf_parameters()
+ MistralModel.set_mistral_config(self.gguf_writer, self.hparams)
+
+ dsa = self.hparams["dsa"]
+ self.gguf_writer.add_indexer_head_count(dsa["n_heads"])
+ self.gguf_writer.add_indexer_key_length(dsa["head_dim"])
+ self.gguf_writer.add_indexer_top_k(dsa["topk"])
+ self.gguf_writer.add_indexer_types(self.layer_indexer_full)
+ self.gguf_writer.add_layer_norm_eps(dsa.get("norm_eps", 1e-6))
+
+ self.gguf_writer.add_rope_pattern(self.layer_has_rope)
+ self.gguf_writer.add_kv_source_map(self.layer_kv_source)
+ self.gguf_writer.add_expert_gating_func(gguf.ExpertGatingFuncType.SOFTMAX)
+ self.gguf_writer.add_final_logit_softcapping(self.hparams["softmax_tanh_gating"])
+
+ @classmethod
+ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+ name, gen = item
+
+ # exported separately with --dspark
+ if name.startswith("mtp_"):
+ return None
+
+ # experts need HF names for the merging logic in DeepseekV2Model
+ # note: bare "layers.N.experts" comes from the NVFP4 quantized_layers config
+ if ".experts." in name or name.endswith(".experts"):
+ name = "model." + name.replace(".experts", ".mlp.experts")
+ name = name.replace(".w1.", ".gate_proj.").replace(".w2.", ".down_proj.").replace(".w3.", ".up_proj.")
+
+ # skip the MistralMoeModel renaming, the qscale names are needed for block FP8 dequant
+ return super(MistralMoeModel, cls).filter_tensors((name, gen))
+
+ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+ if name == "lerope.alpha":
+ # LeRoPE scales inv_freq by exp(alpha), ggml divides theta by the freq factors
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.ROPE_FREQS), torch.exp(-data_torch.float()))
+ return
+
+ if bid is not None and name.endswith((".ssmax_scaling.s", ".ssmax_scaling.b")):
+ tensor = gguf.MODEL_TENSOR.INDEXER_SSMAX if ".indexer." in name else gguf.MODEL_TENSOR.ATTN_SSMAX
+ suffix = ".weight" if name.endswith(".s") else ".bias"
+ yield (self.format_tensor_name(tensor, bid, suffix), data_torch.float())
+ return
+
+ if bid is not None and ".indexer.q_input_norm." in name:
+ suffix = ".weight" if name.endswith(".weight") else ".bias"
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_NORM, bid, suffix), data_torch)
+ return
+
+ if bid is not None and ".indexer." in name:
+ name = "model." + name.replace(".attention.", ".self_attn.")
+
+ # NoPE layers have a larger non-rotary head, so wkv_b is split per layer
+ if bid is not None and name.endswith(".wkv_b.weight"):
+ n_head = self.hparams["num_attention_heads"]
+ v_head_dim = self.hparams["v_head_dim"]
+ qk_nope_head_dim = self.hparams["qk_nope_head_dim"]
+ if not self.layer_has_rope[bid]:
+ qk_nope_head_dim += self.hparams["qk_rope_head_dim"]
+
+ assert data_torch.shape[0] == n_head * (qk_nope_head_dim + v_head_dim)
+
+ kv_b = data_torch.view(n_head, qk_nope_head_dim + v_head_dim, data_torch.shape[-1])
+ k_b, v_b = torch.split(kv_b, [qk_nope_head_dim, v_head_dim], dim=1)
+ k_b = k_b.transpose(1, 2)
+
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_K_B, bid), k_b)
+ yield (self.format_tensor_name(gguf.MODEL_TENSOR.ATTN_V_B, bid), v_b)
+ return
+
+ yield from super().modify_tensors(data_torch, name, bid)
+
+
+class MistralLarge4DSparkModel(MistralLarge4Model):
+ """Mistral Large 4 DSpark draft (mtp_dspark.*) as a DFLASH GGUF, export with --dspark"""
+ model_arch = gguf.MODEL_ARCH.DFLASH
+
+ _DSPARK_ROOT: dict[str, gguf.MODEL_TENSOR] = {
+ "backbone.target_projection.weight": gguf.MODEL_TENSOR.FC,
+ "backbone.target_norm.weight": gguf.MODEL_TENSOR.ENC_OUTPUT_NORM,
+ "backbone.final_norm.weight": gguf.MODEL_TENSOR.OUTPUT_NORM,
+ "backbone.mask_embedding": gguf.MODEL_TENSOR.DSPARK_MASK_EMBD,
+ "autoregressive_head.embedding.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W1,
+ "autoregressive_head.output.weight": gguf.MODEL_TENSOR.DSPARK_MARKOV_W2,
+ "autoregressive_head.joint_proj.weight": gguf.MODEL_TENSOR.DSPARK_JOINT_PROJ,
+ "autoregressive_head.joint_proj.bias": gguf.MODEL_TENSOR.DSPARK_JOINT_PROJ,
+ }
+
+ def __init__(self, *args, **kwargs):
+ super().__init__(*args, **kwargs)
+ mtp = self.hparams.get("mtp") or {}
+ dspark = mtp.get("dspark")
+ if dspark is None:
+ raise ValueError("Mistral Large 4 DSpark export needs the mtp.dspark config")
+ assert dspark.get("head_type") == "rnn", f"unsupported DSpark head type: {dspark.get('head_type')}"
+ assert dspark.get("draft_input") == "position_specific", f"unsupported DSpark draft input: {dspark.get('draft_input')}"
+
+ windows = {int(v) for v in (mtp.get("ragged_attention") or "").split(",") if v.strip()}
+ assert len(windows) == 1, f"DSpark draft layers must share one sliding window, got {mtp.get('ragged_attention')}"
+
+ self._n_stage = int(mtp["n_layers"])
+ self._depth = int(mtp["depth"])
+ self._window = windows.pop()
+ self._target_layers = [int(v) for v in mtp["target_activation_layers"].split(",")]
+
+ self.layer_has_rope = [True] * self._n_stage
+
+ @classmethod
+ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None:
+ name, gen = item
+ if not name.startswith("mtp_dspark."):
+ return None
+ name = name.removeprefix("mtp_dspark.")
+ if name.startswith("backbone.layers."):
+ return name.removeprefix("backbone."), gen
+ return f"dspark.{name}", gen
+
+ def set_vocab(self):
+ super().set_vocab()
+ # block uses mask_embd, mask token id is not used
+ self.gguf_writer.add_mask_token_id(0)
+
+ def set_gguf_parameters(self):
+ hp = self.hparams
+ w = self.gguf_writer
+ w.add_block_count(self._n_stage)
+ w.add_context_length(hp["max_position_embeddings"])
+ w.add_embedding_length(hp["hidden_size"])
+ w.add_feed_forward_length(hp["intermediate_size"])
+ w.add_head_count(hp["num_attention_heads"])
+ w.add_head_count_kv(1)
+ w.add_key_length(hp["kv_lora_rank"] + hp["qk_rope_head_dim"])
+ w.add_value_length(hp["kv_lora_rank"])
+ w.add_key_length_mla(hp["qk_nope_head_dim"] + hp["qk_rope_head_dim"])
+ w.add_value_length_mla(hp["v_head_dim"])
+ w.add_kv_lora_rank(hp["kv_lora_rank"])
+ w.add_rope_dimension_count(hp["qk_rope_head_dim"])
+ w.add_rope_freq_base(float(hp["rope_theta"]))
+ w.add_layer_norm_rms_eps(hp["rms_norm_eps"])
+ MistralModel.set_mistral_config(w, hp)
+ w.add_sliding_window(self._window)
+ w.add_final_logit_softcapping(hp["softmax_tanh_gating"])
+ w.add_block_size(self._depth)
+ # draft reads layer output = input of next layer
+ w.add_target_layers([layer + 1 for layer in self._target_layers])
+ w.add_sample_from_anchor(True)
+ w.add_has_confidence_head(False)
+ w.add_file_type(self.ftype)
+
+ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]:
+ if name.startswith("dspark."):
+ key = name.removeprefix("dspark.")
+ if key not in self._DSPARK_ROOT:
+ raise ValueError(f"Unexpected DSpark tensor {name!r}")
+ suffix = ".bias" if key.endswith(".bias") else ".weight"
+ yield (self.format_tensor_name(self._DSPARK_ROOT[key], None, suffix), data_torch)
+ return
+
+ yield from super().modify_tensors(data_torch, name, bid)
diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py
index 5e263fd16..57ba49dbc 100755
--- a/convert_hf_to_gguf.py
+++ b/convert_hf_to_gguf.py
@@ -127,7 +127,7 @@ def parse_args() -> argparse.Namespace:
)
parser.add_argument(
"--dspark", action="store_true",
- help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.",
+ help="Export only the DeepSeek-V4 or Mistral Large 4 DSpark draft tensors as a separate GGUF.",
)
parser.add_argument(
"--mistral-format", action="store_true",
@@ -255,6 +255,9 @@ def main() -> None:
assert hparams.get("vision_encoder") is not None, "This model does not support multimodal"
from conversion.pixtral import PixtralModel
model_class = PixtralModel
+ elif hparams.get("moe") is not None and hparams.get("dsa") is not None:
+ from conversion.mistral4 import MistralLarge4Model
+ model_class = MistralLarge4Model
elif hparams.get("moe") is not None:
from conversion.mistral import MistralMoeModel
model_class = MistralMoeModel
@@ -267,10 +270,16 @@ def main() -> None:
sys.exit(1)
if args.dspark:
- if is_mistral_format or model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"):
- logger.error("--dspark is only supported for DeepseekV4ForCausalLM and DeepseekV41ForCausalLM")
+ if is_mistral_format or model_architecture == "MistralLarge4ForConditionalGeneration":
+ if model_class.__name__ != "MistralLarge4Model":
+ logger.error("--dspark is only supported for Mistral Large 4 among Mistral-format models")
+ sys.exit(1)
+ from conversion.mistral4 import MistralLarge4DSparkModel
+ model_class = MistralLarge4DSparkModel
+ elif model_architecture not in ("DeepseekV4ForCausalLM", "DeepseekV41ForCausalLM"):
+ logger.error("--dspark is only supported for DeepseekV4ForCausalLM, DeepseekV41ForCausalLM and Mistral Large 4")
sys.exit(1)
- if model_architecture == "DeepseekV41ForCausalLM":
+ elif model_architecture == "DeepseekV41ForCausalLM":
from conversion.deepseek41 import DeepseekV41DSparkModel
model_class = DeepseekV41DSparkModel
else:
diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py
index 1ca2415cd..7884b9841 100644
--- a/gguf-py/gguf/constants.py
+++ b/gguf-py/gguf/constants.py
@@ -222,6 +222,7 @@ class Keys:
VALUE_LENGTH_MLA_SWA = "{arch}.attention.value_length_mla_swa"
KV_LORA_RANK_SWA = "{arch}.attention.kv_lora_rank_swa"
SHARED_KV_LAYERS = "{arch}.attention.shared_kv_layers"
+ KV_SOURCE_MAP = "{arch}.attention.kv_source_map"
SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern"
RECURRENT_LAYERS = "{arch}.attention.recurrent_layers"
TEMPERATURE_SCALE = "{arch}.attention.temperature_scale"
@@ -722,6 +723,7 @@ class MODEL_ARCH(IntEnum):
EAGLE3 = auto()
DFLASH = auto()
MISTRAL4 = auto()
+ MISTRAL_LARGE4 = auto()
PADDLEOCR = auto()
MIMO2 = auto()
STEP35 = auto()
@@ -788,6 +790,7 @@ class MODEL_TENSOR(IntEnum):
ATTN_POST_NORM = auto()
ATTN_ROT_EMBD = auto()
ATTN_SINKS = auto()
+ ATTN_SSMAX = auto()
ATTN_GATE = auto()
ATTN_RESID_ALPHA = auto()
ATTN_RESID_BETA = auto()
@@ -1039,6 +1042,7 @@ class MODEL_TENSOR(IntEnum):
INDEXER_Q_PROJ = auto()
INDEXER_K_PROJ = auto()
INDEXER_Q_NORM = auto()
+ INDEXER_SSMAX = auto()
INDEXER_COMPRESSOR_WKV = auto()
INDEXER_COMPRESSOR_WGATE = auto()
INDEXER_COMPRESSOR_APE = auto()
@@ -1368,6 +1372,8 @@ class MODEL_TENSOR(IntEnum):
DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed
DSPARK_MARKOV_W2 = auto() # markov head: bias projection
DSPARK_CONF_PROJ = auto() # confidence head
+ DSPARK_JOINT_PROJ = auto() # rnn markov head
+ DSPARK_MASK_EMBD = auto() # per-slot mask embeddings
ENGRAM_KV = auto()
ENGRAM_GATE = auto()
DFLASH_ATTN_CONV_BASE = auto()
@@ -1576,6 +1582,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = {
MODEL_ARCH.EAGLE3: "eagle3",
MODEL_ARCH.DFLASH: "dflash",
MODEL_ARCH.MISTRAL4: "mistral4",
+ MODEL_ARCH.MISTRAL_LARGE4: "mistral-large4",
MODEL_ARCH.PADDLEOCR: "paddleocr",
MODEL_ARCH.MIMO2: "mimo2",
MODEL_ARCH.STEP35: "step35",
@@ -1638,6 +1645,7 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.ATTN_OUT: "blk.{bid}.attn_output",
MODEL_TENSOR.ATTN_ROT_EMBD: "blk.{bid}.attn_rot_embd",
MODEL_TENSOR.ATTN_SINKS: "blk.{bid}.attn_sinks",
+ MODEL_TENSOR.ATTN_SSMAX: "blk.{bid}.attn_ssmax",
MODEL_TENSOR.ATTN_GATE: "blk.{bid}.attn_gate",
MODEL_TENSOR.ATTN_RESID_ALPHA: "blk.{bid}.attn_resid_alpha",
MODEL_TENSOR.ATTN_RESID_BETA: "blk.{bid}.attn_resid_beta",
@@ -1891,6 +1899,7 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.INDEXER_Q_PROJ: "blk.{bid}.indexer.q_proj",
MODEL_TENSOR.INDEXER_K_PROJ: "blk.{bid}.indexer.k_proj",
MODEL_TENSOR.INDEXER_Q_NORM: "blk.{bid}.indexer.q_norm",
+ MODEL_TENSOR.INDEXER_SSMAX: "blk.{bid}.indexer.ssmax",
MODEL_TENSOR.INDEXER_COMPRESSOR_WKV: "blk.{bid}.indexer_compressor_kv",
MODEL_TENSOR.INDEXER_COMPRESSOR_WGATE: "blk.{bid}.indexer_compressor_gate",
MODEL_TENSOR.INDEXER_COMPRESSOR_APE: "blk.{bid}.indexer_compressor_ape",
@@ -2244,6 +2253,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = {
MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1",
MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2",
MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj",
+ MODEL_TENSOR.DSPARK_JOINT_PROJ: "joint_proj",
+ MODEL_TENSOR.DSPARK_MASK_EMBD: "mask_embd",
MODEL_TENSOR.ENGRAM_KV: "blk.{bid}.engram_kv",
MODEL_TENSOR.ENGRAM_GATE: "blk.{bid}.engram_gate",
MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base",
@@ -5872,6 +5883,10 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.ATTN_KV_NORM,
MODEL_TENSOR.ATTN_OUT_A,
MODEL_TENSOR.ATTN_OUT_B,
+ MODEL_TENSOR.ATTN_KV_A_MQA,
+ MODEL_TENSOR.ATTN_KV_A_NORM,
+ MODEL_TENSOR.ATTN_K_B,
+ MODEL_TENSOR.ATTN_V_B,
MODEL_TENSOR.HC_ATTN_FN,
MODEL_TENSOR.HC_ATTN_BASE,
MODEL_TENSOR.HC_ATTN_SCALE,
@@ -5900,6 +5915,8 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.DSPARK_MARKOV_W1,
MODEL_TENSOR.DSPARK_MARKOV_W2,
MODEL_TENSOR.DSPARK_CONF_PROJ,
+ MODEL_TENSOR.DSPARK_JOINT_PROJ,
+ MODEL_TENSOR.DSPARK_MASK_EMBD,
MODEL_TENSOR.DFLASH_ATTN_CONV_BASE,
MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ,
MODEL_TENSOR.DFLASH_FFN_CONV_BASE,
@@ -5952,6 +5969,38 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
MODEL_TENSOR.FFN_UP_SHEXP,
MODEL_TENSOR.FFN_EXP_PROBS_B,
],
+ MODEL_ARCH.MISTRAL_LARGE4: [
+ MODEL_TENSOR.TOKEN_EMBD,
+ MODEL_TENSOR.OUTPUT_NORM,
+ MODEL_TENSOR.OUTPUT,
+ MODEL_TENSOR.ROPE_FREQS,
+ MODEL_TENSOR.ATTN_NORM,
+ MODEL_TENSOR.ATTN_Q,
+ MODEL_TENSOR.ATTN_KV_A_MQA,
+ MODEL_TENSOR.ATTN_KV_B,
+ MODEL_TENSOR.ATTN_K_B,
+ MODEL_TENSOR.ATTN_V_B,
+ MODEL_TENSOR.ATTN_KV_A_NORM,
+ MODEL_TENSOR.ATTN_OUT,
+ MODEL_TENSOR.ATTN_SSMAX,
+ MODEL_TENSOR.FFN_GATE_INP,
+ MODEL_TENSOR.FFN_NORM,
+ MODEL_TENSOR.FFN_GATE,
+ MODEL_TENSOR.FFN_DOWN,
+ MODEL_TENSOR.FFN_UP,
+ MODEL_TENSOR.FFN_GATE_EXP,
+ MODEL_TENSOR.FFN_DOWN_EXP,
+ MODEL_TENSOR.FFN_UP_EXP,
+ MODEL_TENSOR.FFN_GATE_SHEXP,
+ MODEL_TENSOR.FFN_DOWN_SHEXP,
+ MODEL_TENSOR.FFN_UP_SHEXP,
+ MODEL_TENSOR.INDEXER_Q_NORM,
+ MODEL_TENSOR.INDEXER_K_NORM,
+ MODEL_TENSOR.INDEXER_PROJ,
+ MODEL_TENSOR.INDEXER_ATTN_K,
+ MODEL_TENSOR.INDEXER_ATTN_Q_B,
+ MODEL_TENSOR.INDEXER_SSMAX,
+ ],
MODEL_ARCH.MIMO2: [
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 b3dabc0bd..a0800e24c 100644
--- a/gguf-py/gguf/gguf_writer.py
+++ b/gguf-py/gguf/gguf_writer.py
@@ -846,6 +846,9 @@ class GGUFWriter:
def add_shared_kv_layers(self, value: int) -> None:
self.add_uint32(Keys.Attention.SHARED_KV_LAYERS.format(arch=self.arch), value)
+ def add_kv_source_map(self, value: Sequence[int]) -> None:
+ self.add_array(Keys.Attention.KV_SOURCE_MAP.format(arch=self.arch), value)
+
# if input is array, true means SWA and false means full_attention for each layer
def add_sliding_window_pattern(self, value: int | Sequence[bool]) -> None:
key = Keys.Attention.SLIDING_WINDOW_PATTERN.format(arch=self.arch)
diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp
index 3a19646eb..8a0d6b124 100644
--- a/src/llama-arch.cpp
+++ b/src/llama-arch.cpp
@@ -155,6 +155,7 @@ static const std::map<llm_arch, const char *> LLM_ARCH_NAMES = {
{ LLM_ARCH_EAGLE3, "eagle3" },
{ LLM_ARCH_DFLASH, "dflash" },
{ LLM_ARCH_MISTRAL4, "mistral4" },
+ { LLM_ARCH_MISTRAL_LARGE4, "mistral-large4" },
{ LLM_ARCH_PADDLEOCR, "paddleocr" },
{ LLM_ARCH_MIMO2, "mimo2" },
{ LLM_ARCH_STEP35, "step35" },
@@ -330,6 +331,7 @@ static const std::map<llm_kv, const char *> LLM_KV_NAMES = {
{ LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" },
{ LLM_KV_ATTENTION_COMPRESS_RATIOS, "%s.attention.compress_ratios" },
{ LLM_KV_ATTENTION_SHARED_KV_LAYERS, "%s.attention.shared_kv_layers" },
+ { LLM_KV_ATTENTION_KV_SOURCE_MAP, "%s.attention.kv_source_map" },
{ LLM_KV_ATTENTION_RECURRENT_LAYERS, "%s.attention.recurrent_layers" },
{ LLM_KV_HYPER_CONNECTION_COUNT, "%s.hyper_connection.count" },
@@ -779,6 +781,7 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_POS_NET_ATTN_V, "posnet.%d.attn_v" },
{ LLM_TENSOR_POS_NET_ATTN_OUT, "posnet.%d.attn_output" },
{ LLM_TENSOR_ATTN_SINKS, "blk.%d.attn_sinks" },
+ { LLM_TENSOR_ATTN_SSMAX, "blk.%d.attn_ssmax" },
{ LLM_TENSOR_SHORTCONV_CONV, "blk.%d.shortconv.conv" },
{ LLM_TENSOR_SHORTCONV_INPROJ, "blk.%d.shortconv.in_proj" },
{ LLM_TENSOR_SHORTCONV_OUTPROJ, "blk.%d.shortconv.out_proj" },
@@ -797,6 +800,7 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_INDEXER_Q_PROJ, "blk.%d.indexer.q_proj" },
{ LLM_TENSOR_INDEXER_K_PROJ, "blk.%d.indexer.k_proj" },
{ LLM_TENSOR_INDEXER_Q_NORM, "blk.%d.indexer.q_norm" },
+ { LLM_TENSOR_INDEXER_SSMAX, "blk.%d.indexer.ssmax" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_WKV, "blk.%d.indexer_compressor_kv" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, "blk.%d.indexer_compressor_gate" },
{ LLM_TENSOR_INDEXER_COMPRESSOR_APE, "blk.%d.indexer_compressor_ape" },
@@ -811,6 +815,8 @@ static const std::map<llm_tensor, const char *> LLM_TENSOR_NAMES = {
{ LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" },
{ LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" },
{ LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" },
+ { LLM_TENSOR_DSPARK_JOINT_PROJ, "joint_proj" },
+ { LLM_TENSOR_DSPARK_MASK_EMBD, "mask_embd" },
{ LLM_TENSOR_ENGRAM_KV, "blk.%d.engram_kv" },
{ LLM_TENSOR_ENGRAM_GATE, "blk.%d.engram_gate" },
{ LLM_TENSOR_DFLASH_ATTN_CONV_BASE, "blk.%d.attn_conv_base" },
@@ -943,6 +949,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_ATTN_K_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_ATTN_V_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_ATTN_SINKS, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SCALE}},
+ {LLM_TENSOR_ATTN_SSMAX, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_DEC_ATTN_Q, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_DEC_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_DEC_ATTN_V, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
@@ -1124,6 +1131,7 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_INDEXER_Q_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_K_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_Q_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
+ {LLM_TENSOR_INDEXER_SSMAX, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_INDEXER_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_INDEXER_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}},
@@ -1162,6 +1170,8 @@ static const std::map<llm_tensor, llm_tensor_info> LLM_TENSOR_INFOS = {
{LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}},
{LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
{LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_DSPARK_JOINT_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}},
+ {LLM_TENSOR_DSPARK_MASK_EMBD, {LLM_TENSOR_LAYER_INPUT, GGML_OP_GET_ROWS}},
{LLM_TENSOR_ENGRAM_KV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}},
{LLM_TENSOR_ENGRAM_GATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
{LLM_TENSOR_DFLASH_ATTN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}},
@@ -1370,6 +1380,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) {
case LLM_ARCH_MINIMAX_M2:
case LLM_ARCH_MINIMAX_M3:
case LLM_ARCH_MISTRAL4:
+ case LLM_ARCH_MISTRAL_LARGE4:
case LLM_ARCH_MELLUM: // [TAG_MOE_AMAX_SCALE]
case LLM_ARCH_FLINT: // [TAG_MOE_AMAX_SCALE]
case LLM_ARCH_LAGUNA: // [TAG_MOE_AMAX_SCALE]
diff --git a/src/llama-arch.h b/src/llama-arch.h
index 554156813..072d3bcde 100644
--- a/src/llama-arch.h
+++ b/src/llama-arch.h
@@ -153,6 +153,7 @@ enum llm_arch {
LLM_ARCH_PANGU_EMBED,
LLM_ARCH_MISTRAL3,
LLM_ARCH_MISTRAL4,
+ LLM_ARCH_MISTRAL_LARGE4,
LLM_ARCH_PADDLEOCR,
LLM_ARCH_MIMO2,
LLM_ARCH_STEP35,
@@ -334,6 +335,7 @@ enum llm_kv {
LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE,
LLM_KV_ATTENTION_COMPRESS_RATIOS,
LLM_KV_ATTENTION_SHARED_KV_LAYERS,
+ LLM_KV_ATTENTION_KV_SOURCE_MAP,
LLM_KV_ATTENTION_RECURRENT_LAYERS,
LLM_KV_HYPER_CONNECTION_COUNT,
@@ -528,6 +530,7 @@ enum llm_tensor {
LLM_TENSOR_ATTN_POST_NORM,
LLM_TENSOR_ATTN_ROT_EMBD,
LLM_TENSOR_ATTN_SINKS,
+ LLM_TENSOR_ATTN_SSMAX,
LLM_TENSOR_ATTN_GATE,
LLM_TENSOR_ATTN_RESID_ALPHA,
LLM_TENSOR_ATTN_RESID_BETA,
@@ -790,6 +793,7 @@ enum llm_tensor {
LLM_TENSOR_INDEXER_Q_PROJ,
LLM_TENSOR_INDEXER_K_PROJ,
LLM_TENSOR_INDEXER_Q_NORM,
+ LLM_TENSOR_INDEXER_SSMAX,
LLM_TENSOR_INDEXER_COMPRESSOR_WKV,
LLM_TENSOR_INDEXER_COMPRESSOR_WGATE,
LLM_TENSOR_INDEXER_COMPRESSOR_APE,
@@ -816,6 +820,8 @@ enum llm_tensor {
LLM_TENSOR_DSPARK_MARKOV_W1,
LLM_TENSOR_DSPARK_MARKOV_W2,
LLM_TENSOR_DSPARK_CONF_PROJ,
+ LLM_TENSOR_DSPARK_JOINT_PROJ,
+ LLM_TENSOR_DSPARK_MASK_EMBD,
LLM_TENSOR_ENGRAM_KV,
LLM_TENSOR_ENGRAM_GATE,
LLM_TENSOR_DFLASH_ATTN_CONV_BASE,
diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp
index cffd083d3..45bed2a42 100644
--- a/src/llama-graph.cpp
+++ b/src/llama-graph.cpp
@@ -3429,13 +3429,20 @@ ggml_tensor * llm_graph_context::build_attn(
// by doing so, the number of splits in the graph is reduced
// expand k later to enable rope fusion which directly writes into k-v cache
ggml_build_forward_expand(gf, q_cur);
- ggml_build_forward_expand(gf, v_cur);
- ggml_build_forward_expand(gf, k_cur);
+
+ if (v_cur) {
+ ggml_build_forward_expand(gf, v_cur);
+ }
+
+ if (k_cur) {
+ ggml_build_forward_expand(gf, k_cur);
+ }
const auto * mctx_cur = inp->mctx->get_mla();
// store to KV cache
- {
+ // note: k_cur is null for layers that reuse the KV cache of another layer
+ if (k_cur) {
const auto & k_idxs = inp->get_k_idxs_mla();
ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il));
@@ -3481,7 +3488,10 @@ ggml_tensor * llm_graph_context::build_attn(
ggml_tensor * q = q_cur;
ggml_tensor * k = mctx_cur->get_k(ctx0, il);
- ggml_tensor * v = ggml_view_4d(ctx0, k, v_cur->ne[0], k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
+
+ // MLA: V is a view of the first n_embd_v elements of K
+ const int64_t n_embd_v = v_cur ? v_cur->ne[0] : v_mla->ne[0];
+ ggml_tensor * v = ggml_view_4d(ctx0, k, n_embd_v, k->ne[1], k->ne[2], k->ne[3], k->nb[1], k->nb[2], k->nb[3], 0);
ggml_tensor * cur = build_attn_mha(q, k, v, kq_b, kq_mask_top_k, sinks, v_mla, top_k->ne[0], kq_scale, il);
cb(cur, "kqv_out", il);
diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp
index b83f45ba9..9b4244567 100644
--- a/src/llama-hparams.cpp
+++ b/src/llama-hparams.cpp
@@ -319,6 +319,10 @@ uint32_t llama_hparams::n_embd_head_v_mla() const {
}
bool llama_hparams::has_kv(uint32_t il) const {
+ if (kv_source(il) != il) {
+ return false;
+ }
+
if (n_layer_kv_from_start >= 0) {
if (il < (uint32_t) n_layer_kv_from_start) {
return true;
@@ -331,6 +335,14 @@ bool llama_hparams::has_kv(uint32_t il) const {
return true;
}
+uint32_t llama_hparams::kv_source(uint32_t il) const {
+ if (il < n_layer_all) {
+ return kv_source_impl[il];
+ }
+
+ GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all);
+}
+
bool llama_hparams::has_rope(uint32_t il) const {
// the router layer stores adapter routing signal, not positional info,
// so it must not be RoPE-shifted
diff --git a/src/llama-hparams.h b/src/llama-hparams.h
index 2ac7bdeae..4d2cb8931 100644
--- a/src/llama-hparams.h
+++ b/src/llama-hparams.h
@@ -85,6 +85,9 @@ struct llama_hparams {
// TODO: this needs to be reworked
int32_t n_layer_kv_from_start = -1; // if non-negative, the first n_layer_kv_from_start layers have KV cache
+ // layer il reads the KV cache of layer kv_source_impl[il] (own cache if equal to il)
+ std::array<uint32_t, LLAMA_MAX_LAYERS> kv_source_impl;
+
// different head size for full_attention and SWA layers
uint32_t n_embd_head_k_full; // dimension of keys (d_k). d_q is assumed to be the same, but there are n_head q heads, and only n_head_kv k-v heads
uint32_t n_embd_head_v_full; // dimension of values (d_v) aka n_embd_head
@@ -515,6 +518,8 @@ struct llama_hparams {
bool has_kv(uint32_t il) const;
+ uint32_t kv_source(uint32_t il) const;
+
bool has_rope(uint32_t il) const;
// number of effective layers (excludes nextn layers)
diff --git a/src/llama-kv-cache-dsa.cpp b/src/llama-kv-cache-dsa.cpp
index 2602ae3c3..405eee36b 100644
--- a/src/llama-kv-cache-dsa.cpp
+++ b/src/llama-kv-cache-dsa.cpp
@@ -6,6 +6,7 @@
#include <algorithm>
#include <cassert>
+#include <numeric>
//
// llama_kv_cache_dsa
@@ -45,12 +46,15 @@ llama_kv_cache_dsa::llama_kv_cache_dsa(
hparams_lid.n_embd_head_k_full = model.hparams.indexer_head_size;
hparams_lid.rope_type = LLAMA_ROPE_TYPE_NEOX;
+ // KV sharing applies to the MLA cache only, each indexer owns its keys
+ std::iota(hparams_lid.kv_source_impl.begin(), hparams_lid.kv_source_impl.end(), 0);
+
LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size);
kv_lid = std::make_unique<llama_kv_cache>(
model, no_alloc, hparams_lid, type_k, type_v,
v_trans, offload, unified, kv_size, n_seq_max,
- n_swa, swa_type, nullptr, filter_lid, reuse, nullptr);
+ n_swa, swa_type, nullptr, filter_lid, nullptr, nullptr);
}
void llama_kv_cache_dsa::clear(bool data) {
diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp
index 0e808a0be..00dbf7c94 100644
--- a/src/llama-kv-cache.cpp
+++ b/src/llama-kv-cache.cpp
@@ -315,7 +315,8 @@ llama_kv_cache::llama_kv_cache(
// always create Hadamard rotation tensors for DeepSeek lightning indexers
if ((model.arch == LLM_ARCH_DEEPSEEK32 || model.arch == LLM_ARCH_DEEPSEEK4 ||
- model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_DOTS3NOTE) &&
+ model.arch == LLM_ARCH_GLM_DSA || model.arch == LLM_ARCH_DOTS3NOTE ||
+ model.arch == LLM_ARCH_MISTRAL_LARGE4) &&
hparams.n_embd_head_k_full == hparams.indexer_head_size) {
attn_rot_k = true;
}
diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp
index 93836ed96..5c2d7a822 100644
--- a/src/llama-model-saver.cpp
+++ b/src/llama-model-saver.cpp
@@ -307,6 +307,12 @@ void llama_model_saver::add_kv_from_model() {
add_kv(LLM_KV_ATTENTION_INDEXER_KPOOL_SELECT_TAIL, hparams.indexer_kpool_select_tail);
add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks);
add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, true);
+ for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
+ if (hparams.kv_source(il) != il) {
+ add_kv(LLM_KV_ATTENTION_KV_SOURCE_MAP, hparams.kv_source_impl, true);
+ break;
+ }
+ }
add_kv(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, true);
add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, hparams.dsv4_o_group_count);
add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, hparams.dsv4_o_lora_rank);
diff --git a/src/llama-model.cpp b/src/llama-model.cpp
index 8966bc946..5fa3dfad9 100644
--- a/src/llama-model.cpp
+++ b/src/llama-model.cpp
@@ -218,6 +218,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params
return new llama_model_glm_dsa(params);
case LLM_ARCH_MISTRAL4:
return new llama_model_mistral4(params);
+ case LLM_ARCH_MISTRAL_LARGE4:
+ return new llama_model_mistral_large4(params);
case LLM_ARCH_CHATGLM:
return new llama_model_chatglm(params);
case LLM_ARCH_GLM4:
@@ -1467,6 +1469,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) {
std::fill(hparams.is_swa_impl.begin(), hparams.is_swa_impl.end(), 0);
std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), llm_arch_is_recurrent(ml.get_arch()) ? 1 : 0);
std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 0);
+ std::iota(hparams.kv_source_impl.begin(), hparams.kv_source_impl.end(), 0);
std::fill(hparams.xielu_alpha_n.begin(), hparams.xielu_alpha_n.end(), 0.0f);
std::fill(hparams.xielu_alpha_p.begin(), hparams.xielu_alpha_p.end(), 0.0f);
@@ -2629,7 +2632,7 @@ void llama_model::print_info() const {
if (arch == LLM_ARCH_DEEPSEEK2 || arch == LLM_ARCH_DEEPSEEK2OCR ||
arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA ||
arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_MISTRAL4 ||
- arch == LLM_ARCH_HY_V4) {
+ arch == LLM_ARCH_MISTRAL_LARGE4 || arch == LLM_ARCH_HY_V4) {
LLAMA_LOG_INFO("%s: n_layer_dense_lead = %d\n", __func__, hparams.n_layer_dense_lead);
LLAMA_LOG_INFO("%s: n_lora_q = %d\n", __func__, hparams.n_lora_q);
LLAMA_LOG_INFO("%s: n_lora_kv = %d\n", __func__, hparams.n_lora_kv);
@@ -2908,6 +2911,32 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
nullptr);
}
} break;
+ case LLM_ARCH_MISTRAL_LARGE4:
+ {
+ // only "full" layers own an indexer, so the shared layers need no indexer cache
+ llama_kv_cache::layer_filter_cb filter_lid = [&](uint32_t il) { return hparams.is_indexer_full(il); };
+
+ // layers without own KV read the MLA cache of their source layer
+ llama_kv_cache::layer_reuse_cb reuse = [&](int32_t il) -> int32_t {
+ return hparams.has_kv(il) ? -1 : (int32_t) hparams.kv_source(il);
+ };
+
+ res = new llama_kv_cache_dsa(
+ *this,
+ params.no_alloc,
+ params.type_k,
+ params.type_v,
+ !cparams.flash_attn,
+ cparams.offload_kqv,
+ cparams.kv_unified,
+ cparams.n_ctx_seq,
+ cparams.n_seq_max,
+ hparams.n_swa,
+ hparams.swa_type,
+ nullptr,
+ filter_lid,
+ reuse);
+ } break;
case LLM_ARCH_GLM5_NEXT:
{
// KDA layers are recurrent, the DSA layers use a K-only MLA cache plus an indexer cache.
@@ -3091,8 +3120,8 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params,
} break;
case LLM_ARCH_DFLASH:
{
- // DSV4 DSpark stages store a single MLA-style K per position (window = the draft ring)
- if (hparams.dsv4_hc_mult > 0) {
+ // DSV4 and ML4 DSpark stages store a single MLA-style K per position (window = the draft ring)
+ if (hparams.dsv4_hc_mult > 0 || hparams.n_lora_kv > 0) {
GGML_ASSERT(hparams.swa_type != LLAMA_SWA_TYPE_NONE);
res = new llama_kv_cache_iswa(
@@ -3571,6 +3600,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
case LLM_ARCH_MISTRAL3:
case LLM_ARCH_EAGLE3:
case LLM_ARCH_MISTRAL4:
+ case LLM_ARCH_MISTRAL_LARGE4:
case LLM_ARCH_LLAMA_EMBED:
case LLM_ARCH_MAINCODER:
case LLM_ARCH_GLM_DSA:
@@ -3675,8 +3705,8 @@ llama_rope_type llama_model_rope_type(const llama_model * model) {
if (const auto & s = model->hparams.rope_sections; s[0] || s[1] || s[2] || s[3]) {
return LLAMA_ROPE_TYPE_MROPE;
}
- // DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX
- return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX;
+ // DSV4 and ML4 DSpark drafters use normal RoPE; legacy DFlash backbones are NeoX
+ return model->hparams.dsv4_hc_mult > 0 || model->hparams.n_lora_kv > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX;
case LLM_ARCH_SENSENOVA_U1:
case LLM_ARCH_QWEN2VL:
diff --git a/src/llama-model.h b/src/llama-model.h
index 3af8b054f..da4964c98 100644
--- a/src/llama-model.h
+++ b/src/llama-model.h
@@ -600,6 +600,14 @@ struct llama_layer {
struct ggml_tensor * indexer_proj = nullptr;
struct ggml_tensor * indexer_attn_k = nullptr;
struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias
+ struct ggml_tensor * indexer_q_norm = nullptr;
+ struct ggml_tensor * indexer_q_norm_b = nullptr;
+
+ // SSMax: per-head query scale s * log(pos + 1) + b
+ struct ggml_tensor * attn_ssmax_s = nullptr;
+ struct ggml_tensor * attn_ssmax_b = nullptr;
+ struct ggml_tensor * indexer_ssmax_s = nullptr;
+ struct ggml_tensor * indexer_ssmax_b = nullptr;
// glm5-next k-pool indexer
struct ggml_tensor * indexer_kpool_gate = nullptr;
@@ -813,6 +821,9 @@ struct llama_model {
struct ggml_tensor * dspark_markov_w2_s = nullptr;
struct ggml_tensor * dspark_conf_proj = nullptr;
struct ggml_tensor * dspark_conf_proj_b = nullptr;
+ struct ggml_tensor * dspark_joint_proj = nullptr; // rnn markov head
+ struct ggml_tensor * dspark_joint_proj_b = nullptr;
+ struct ggml_tensor * dspark_mask_embd = nullptr; // per-slot mask embeddings
struct ggml_tensor * dflash_selector_prev = nullptr;
struct ggml_tensor * dflash_selector_next = nullptr;
diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp
index 65b569762..229a051e0 100644
--- a/src/llama-vocab.cpp
+++ b/src/llama-vocab.cpp
@@ -430,6 +430,13 @@ struct llm_tokenizer_bpe : llm_tokenizer {
"[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))*((?=[\\p{L}\\p{M}])([^A-Z]))+|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))+((?=[\\p{L}\\p{M}])([^A-Z]))*|\\p{N}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
};
break;
+ case LLAMA_VOCAB_PRE_TYPE_TEKKEN_NUM3:
+ // same as tekken, but numbers are split into right-aligned groups of 3 digits
+ // "...|\\p{N}{1,3}(?=(?:\\p{N}{3})*(?!\\p{N}))|\\p{N}{1,2}|..."
+ regex_exprs = {
+ "[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))*((?=[\\p{L}\\p{M}])([^A-Z]))+|[^\\r\\n\\p{L}\\p{N}]?((?=[\\p{L}\\p{M}])([^a-z]))+((?=[\\p{L}\\p{M}])([^A-Z]))*|\\p{N}{1,3}(?=(?:\\p{N}{3})*(?!\\p{N}))|\\p{N}{1,2}| ?[^\\s\\p{L}\\p{N}]+[\\r\\n/]*|\\s*[\\r\\n]+|\\s+(?!\\S)|\\s+",
+ };
+ break;
case LLAMA_VOCAB_PRE_TYPE_CHAMELEON:
// Note: in theory, the special token (sentinel and image token) regex_exprs below
// are unnecessary, as they are split in `tokenizer_st_partition` anyway.
@@ -2619,6 +2626,12 @@ void llama_vocab::impl::load(llama_model_loader & ml, const LLM_KV & kv) {
clean_spaces = false;
ignore_merges = true;
add_bos = true;
+ } else if (
+ tokenizer_pre == "tekken-num3") {
+ pre_type = LLAMA_VOCAB_PRE_TYPE_TEKKEN_NUM3;
+ clean_spaces = false;
+ ignore_merges = true;
+ add_bos = true;
} else if (
tokenizer_pre == "smollm") {
pre_type = LLAMA_VOCAB_PRE_TYPE_SMOLLM;
diff --git a/src/llama-vocab.h b/src/llama-vocab.h
index 14a90fbdb..1ae5b65f6 100644
--- a/src/llama-vocab.h
+++ b/src/llama-vocab.h
@@ -74,6 +74,7 @@ enum llama_vocab_pre_type {
LLAMA_VOCAB_PRE_TYPE_DOMYN_SMALL = 63,
LLAMA_VOCAB_PRE_TYPE_NEMOTRON_3_5_EMBD = 64,
LLAMA_VOCAB_PRE_TYPE_SOPRANO = 65,
+ LLAMA_VOCAB_PRE_TYPE_TEKKEN_NUM3 = 66,
};
// HF Metaspace pre-tokenizer prepend_scheme
diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp
index e3c2f1f29..194f17ec4 100644
--- a/src/models/dflash.cpp
+++ b/src/models/dflash.cpp
@@ -112,6 +112,34 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) {
return;
}
+ // Mistral Large 4 DSpark: dense MLA stages, one sliding window, plain YaRN
+ if (ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv, false)) {
+ ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
+ ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);
+ ml.get_key(LLM_KV_ROPE_SCALING_YARN_LOG_MUL, hparams.rope_yarn_log_mul, false);
+ ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa);
+
+ if (hparams.n_swa == 0) {
+ throw std::runtime_error("DSpark Mistral Large 4 draft expects a sliding window");
+ }
+ if (hparams.n_embd_head_k_full != hparams.n_lora_kv + hparams.n_rot()) {
+ throw std::runtime_error("DSpark Mistral Large 4 draft expects key_length == kv_lora_rank + rope dim");
+ }
+
+ // same as DSV4.1 DSpark: block uses the window of the anchor
+ hparams.n_swa += 1;
+ hparams.swa_type = LLAMA_SWA_TYPE_STANDARD;
+ hparams.set_swa_pattern(0);
+ for (uint32_t il = 0; il < hparams.n_layer_all; ++il) {
+ hparams.is_swa_impl[il] = true;
+ }
+ hparams.rope_freq_base_train_swa = hparams.rope_freq_base_train;
+ hparams.rope_freq_scale_train_swa = hparams.rope_freq_scale_train;
+
+ type = LLM_TYPE_UNKNOWN;
+ return;
+ }
+
// optional interleaved sliding-window attention with per-layer pattern array.
// DFlash has a single rope, so the SWA rope == main rope.
if (ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa, false) && hparams.n_swa > 0) {
@@ -152,6 +180,9 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED);
dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED);
+ dspark_joint_proj = create_tensor(tn(LLM_TENSOR_DSPARK_JOINT_PROJ, "weight"), { n_embd + 2*dspark_markov_rank, 3*dspark_markov_rank }, TENSOR_NOT_REQUIRED);
+ dspark_joint_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_JOINT_PROJ, "bias"), { 3*dspark_markov_rank }, dspark_joint_proj ? 0 : TENSOR_NOT_REQUIRED);
+
LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank);
}
@@ -252,6 +283,33 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) {
// a draft with its own embeddings + head references no target tensors and can run on devices the target does not use (e.g. -devd with a tensor-split target)
output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab_draft }, TENSOR_NOT_REQUIRED);
+ if (hparams.n_lora_kv > 0) {
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
+ const int64_t n_embd_head_nope = n_embd_head_k_mla - n_rot;
+
+ dspark_mask_embd = create_tensor(tn(LLM_TENSOR_DSPARK_MASK_EMBD, "weight"), { n_embd, hparams.dflash_block_size - 1 }, 0);
+
+ for (int i = 0; i < n_layer; ++i) {
+ auto & layer = layers[i];
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), { n_embd }, 0);
+ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), { n_embd, n_head * n_embd_head_k_mla }, 0);
+ layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), { n_embd, kv_lora_rank + n_rot }, 0);
+ layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), { kv_lora_rank }, 0);
+ layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), { n_embd_head_nope, kv_lora_rank, n_head }, 0);
+ layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), { kv_lora_rank, n_embd_head_v_mla, n_head }, 0);
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_head * n_embd_head_v_mla, n_embd }, 0);
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 0);
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), { n_embd, n_ff }, 0);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd }, 0);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), { n_embd, n_ff }, 0);
+ }
+ return;
+ }
+
if (hparams.dsv4_hc_mult > 0) {
const int64_t q_lora_rank = hparams.n_lora_q;
const int64_t n_ff_exp = hparams.n_ff_exp();
@@ -377,7 +435,7 @@ llama_model_dflash::graph<true>::graph(const llama_model & model, const llm_grap
}
// DSpark (DFlash + Markov & Confidence head): Markov bias on the draft logits, chained per block position
-static void build_dspark_markov_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) {
+static void build_dspark_markov_head(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens, ggml_tensor * hidden = nullptr) {
ggml_context * ctx0 = g.ctx0;
auto & res = g.res;
@@ -388,6 +446,10 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
// confidence head is optional
const bool has_conf = model.dspark_conf_proj != nullptr;
+ // rnn head: gated state carried across block positions
+ const bool is_rnn = model.dspark_joint_proj != nullptr;
+ GGML_ASSERT(!is_rnn || hidden != nullptr);
+
ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens]
const int64_t n_vocab = base->ne[0];
const int64_t n_tok = base->ne[1];
@@ -428,11 +490,37 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
}
}
+ ggml_tensor * state = nullptr; // [R, n_blocks]
+
// TODO: the in-graph chain is greedy (argmax); sampling params affect only the final
// token pick, not the Markov conditioning path
for (int64_t i = i_draft_beg; i < block_drafts; ++i) {
- ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks]
- ggml_tensor * bias = g.build_lora_mm(w2, w1_prev, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks]
+ ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks]
+ ggml_tensor * markov = w1_prev;
+
+ if (is_rnn) {
+ const int64_t rank = w1_prev->ne[0];
+ if (state == nullptr) {
+ state = ggml_fill(ctx0, ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, rank, n_blocks), 0.0f);
+ }
+
+ ggml_tensor * hidden_i = ggml_view_2d(ctx0, hidden, hidden->ne[0], n_blocks,
+ (size_t) block_drafts * hidden->nb[1], i*hidden->nb[1]);
+ ggml_tensor * joint = ggml_concat(ctx0, ggml_concat(ctx0, state, w1_prev, 0), ggml_cont(ctx0, hidden_i), 0);
+ joint = ggml_add(ctx0, ggml_mul_mat(ctx0, model.dspark_joint_proj, joint), model.dspark_joint_proj_b); // [3R, n_blocks]
+
+ const auto chunk = [&](int64_t k) {
+ return ggml_cont(ctx0, ggml_view_2d(ctx0, joint, rank, n_blocks, joint->nb[1], k*rank*ggml_element_size(joint)));
+ };
+ ggml_tensor * gate = ggml_sigmoid(ctx0, chunk(0));
+ ggml_tensor * cand = ggml_tanh(ctx0, chunk(1));
+
+ // state = gate * state + (1 - gate) * cand
+ state = ggml_add(ctx0, cand, ggml_mul(ctx0, gate, ggml_sub(ctx0, state, cand)));
+ markov = ggml_tanh(ctx0, chunk(2));
+ }
+
+ ggml_tensor * bias = g.build_lora_mm(w2, markov, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks]
if (model.d2t) {
// reduced draft vocab: scatter the bias to the target rows (base is -inf on the others)
const int64_t n_draft_vocab = bias->ne[0];
@@ -455,7 +543,7 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model &
// conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks]
ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks,
(size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]);
- ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0);
+ ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), markov, 0);
ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat);
if (model.dspark_conf_proj_b) {
conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b);
@@ -1291,6 +1379,23 @@ llama_model_dflash::graph_dsv4::graph_dsv4(const llama_model & model, const llm_
// DeepSeek-V4.1 DSpark decoder (mtp.0..2), dual-mode by batch type like graph_dsv4:
// * embd batch -> writes every stage's kv_norm(wkv(main_proj(target features))) into its window ring
// * token batch -> [anchor, noise...] block through the stages, single-pass mHC, no per-head q norm
+// window of a block = window of its first query, see n_swa
+static void build_dspark_block_mask(llm_graph_context & g, llm_graph_input_attn_k_iswa * inp_attn) {
+ ggml_context * ctx0 = g.ctx0;
+
+ ggml_tensor * m = inp_attn->self_kq_mask_swa_cnv;
+ const int64_t n_tokens = g.n_tokens;
+ const int64_t n_blocks = g.ubatch.n_seqs_unq;
+ const int64_t bd = n_blocks > 0 ? n_tokens / n_blocks : 0;
+ if (m && bd > 1 && n_tokens % n_blocks == 0 && m->ne[2] == 1 && m->ne[1] % bd == 0 && ggml_is_contiguous(m)) {
+ ggml_tensor * m4 = ggml_reshape_4d(ctx0, m, m->ne[0], bd, m->ne[1]/bd, m->ne[3]);
+ ggml_tensor * r0 = ggml_view_4d(ctx0, m4, m4->ne[0], 1, m4->ne[2], m4->ne[3], m4->nb[1], m4->nb[2], m4->nb[3], 0);
+ m4 = ggml_repeat(ctx0, r0, m4);
+ inp_attn->self_kq_mask_swa_cnv = ggml_reshape_4d(ctx0, m4, m->ne[0], m->ne[1], 1, m->ne[3]);
+ g.cb(inp_attn->self_kq_mask_swa_cnv, "dspark_block_mask", -1);
+ }
+}
+
llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const llm_graph_params & params) :
llama_model_deepseek41::graph(params) {
const int64_t n_embd_inp = hparams.n_embd_inp_enc();
@@ -1382,19 +1487,7 @@ llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const ll
inpL = ggml_repeat_4d(ctx0, inpL, n_embd, hc, n_tokens, 1);
cb(inpL, "hc_init", -1);
- // window of a block = window of its first query, see n_swa
- {
- ggml_tensor * m = inp_attn->self_kq_mask_swa_cnv;
- const int64_t n_blocks = ubatch.n_seqs_unq;
- const int64_t bd = n_blocks > 0 ? n_tokens / n_blocks : 0;
- if (m && bd > 1 && n_tokens % n_blocks == 0 && m->ne[2] == 1 && m->ne[1] % bd == 0 && ggml_is_contiguous(m)) {
- ggml_tensor * m4 = ggml_reshape_4d(ctx0, m, m->ne[0], bd, m->ne[1]/bd, m->ne[3]);
- ggml_tensor * r0 = ggml_view_4d(ctx0, m4, m4->ne[0], 1, m4->ne[2], m4->ne[3], m4->nb[1], m4->nb[2], m4->nb[3], 0);
- m4 = ggml_repeat(ctx0, r0, m4);
- inp_attn->self_kq_mask_swa_cnv = ggml_reshape_4d(ctx0, m4, m->ne[0], m->ne[1], 1, m->ne[3]);
- cb(inp_attn->self_kq_mask_swa_cnv, "dspark_block_mask", -1);
- }
- }
+ build_dspark_block_mask(*this, inp_attn);
ggml_tensor * pre_prev = nullptr;
@@ -1529,12 +1622,230 @@ llama_model_dflash::graph_dsv41::graph_dsv41(const llama_model & model, const ll
}
}
+// Mistral Large 4 DSpark, dual-mode like graph_dsv41:
+// - embd batch: write latent K of target_norm(fc(features)) to the cache, no attn norm
+// - token batch: decode [anchor, mask_1, ...] block, non-causal
+llama_model_dflash::graph_mistral_large4::graph_mistral_large4(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) {
+ const int64_t n_embd_inp = hparams.n_embd_inp_enc();
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
+ const int64_t n_embd_head_nope = n_embd_head_k_mla - n_rot;
+
+ const float kq_scale = 1.0f/sqrtf(float(n_embd_head_k_mla));
+
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ llm_graph_input_attn_k_iswa * inp_attn = build_attn_inp_k_iswa();
+
+ // no LeRoPE in the draft
+ const auto rope = [&](ggml_tensor * t) {
+ return ggml_rope_ext(ctx0, t, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ };
+
+ const auto build_kv = [&](ggml_tensor * x, int il) -> std::pair<ggml_tensor *, ggml_tensor *> {
+ const auto & layer = model.layers[il];
+
+ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, x);
+ cb(kv_cmpr_pe, "kv_cmpr_pe", il);
+
+ ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot), 0);
+ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
+ cb(kv_cmpr, "kv_cmpr", il);
+
+ ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_rot, 1, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
+ k_pe = rope(k_pe);
+ cb(k_pe, "k_pe", il);
+
+ return { ggml_concat(ctx0, kv_cmpr, k_pe, 0), kv_cmpr };
+ };
+
+ ASSERT_EMBD_OR_TOKEN(ubatch);
+ if (ubatch.embd) {
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd_inp);
+
+ inp->embd = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, n_embd_inp, n_tokens);
+ ggml_set_input(inp->embd);
+
+ ggml_tensor * inp_target = inp->embd;
+ cb(inp_target, "inp_target_features", -1);
+
+ res->add_input(std::move(inp));
+
+ ggml_tensor * inp_g = build_lora_mm(model.fc, inp_target, model.fc_s);
+ inp_g = build_norm(inp_g, model.output_norm_enc, nullptr, LLM_NORM_RMS, -1);
+ cb(inp_g, "main_x", -1);
+
+ for (int il = 0; il < n_layer; ++il) {
+ ggml_tensor * kv = build_kv(inp_g, il).first;
+ cb(kv, "kv_injected", il);
+
+ if (inp_attn->self_k_rot_swa) {
+ kv = llama_mul_mat_hadamard(ctx0, kv, inp_attn->self_k_rot_swa);
+ }
+ ggml_build_forward_expand(gf, inp_attn->mctx->get_swa()->cpy_k(ctx0, kv, inp_attn->get_k_idxs_swa(), il));
+ }
+
+ res->t_embd = inp_g;
+
+ ggml_build_forward_expand(gf, inp_g);
+ return;
+ }
+
+ auto * tok_embd = model.tok_embd;
+ if (tok_embd == nullptr) {
+ GGML_ASSERT(cparams.ctx_other != nullptr);
+ const auto * model_other = llama_get_model(cparams.ctx_other);
+ GGML_ASSERT(model_other->tok_embd != nullptr && "DSpark decoder requires the target model's token embeddings");
+ tok_embd = model_other->tok_embd;
+ }
+
+ auto inp = std::make_unique<llm_graph_input_embd>(n_embd);
+
+ inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens);
+ ggml_set_input(inp->tokens);
+
+ ggml_tensor * inp_tokens = inp->tokens;
+
+ res->add_input(std::move(inp));
+
+ const int64_t n_blocks = ubatch.n_seqs_unq;
+ const int64_t block_drafts = n_blocks > 0 ? n_tokens / n_blocks : 0;
+
+ ggml_tensor * inpL;
+ if (block_drafts > 0 && n_tokens % n_blocks == 0 && block_drafts <= (int64_t) hparams.dflash_block_size) {
+ // slot k > 0 uses mask row k - 1
+ ggml_tensor * anchors = ggml_view_2d(ctx0, inp_tokens, 1, n_blocks, (size_t) block_drafts * inp_tokens->nb[0], 0);
+ anchors = ggml_get_rows(ctx0, tok_embd, ggml_cont_1d(ctx0, anchors, n_blocks));
+ inpL = ggml_reshape_3d(ctx0, anchors, n_embd, 1, n_blocks);
+ if (block_drafts > 1) {
+ ggml_tensor * masks = ggml_view_3d(ctx0, model.dspark_mask_embd, n_embd, block_drafts - 1, 1,
+ model.dspark_mask_embd->nb[1], model.dspark_mask_embd->nb[2], 0);
+ masks = ggml_repeat_4d(ctx0, ggml_cast(ctx0, masks, GGML_TYPE_F32), n_embd, block_drafts - 1, n_blocks, 1);
+ inpL = ggml_concat(ctx0, inpL, masks, 1);
+ }
+ inpL = ggml_reshape_2d(ctx0, inpL, n_embd, n_tokens);
+ } else {
+ // not a draft block (e.g. graph reserve)
+ inpL = ggml_get_rows(ctx0, tok_embd, inp_tokens);
+ }
+ cb(inpL, "inp_noise_embd", -1);
+
+ build_dspark_block_mask(*this, inp_attn);
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ ggml_tensor * cur = build_norm(inpL, layer.attn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ {
+ ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq, cur);
+ cb(q, "q", il);
+
+ // {n_embd_head_nope, n_tokens, n_head}
+ ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_nope, n_head, n_tokens,
+ ggml_row_size(q->type, n_embd_head_k_mla),
+ ggml_row_size(q->type, n_embd_head_k_mla) * n_head, 0);
+ q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
+
+ // {kv_lora_rank, n_head, n_tokens}
+ ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
+ q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
+ cb(q_nope_absorbed, "q_nope_absorbed", il);
+
+ ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_rot, n_head, n_tokens,
+ ggml_row_size(q->type, n_embd_head_k_mla),
+ ggml_row_size(q->type, n_embd_head_k_mla) * n_head,
+ ggml_row_size(q->type, n_embd_head_nope));
+ q_pe = rope(q_pe);
+ cb(q_pe, "q_pe", il);
+
+ ggml_tensor * Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
+ cb(Qcur, "Qcur", il);
+
+ auto [Kcur, Vcur] = build_kv(cur, il);
+ cb(Kcur, "Kcur", il);
+
+ // apply wv_b after attn, so the K rotation is already undone
+ ggml_tensor * out = build_attn(inp_attn,
+ nullptr, nullptr, nullptr,
+ Qcur, Kcur, Vcur,
+ nullptr, nullptr, nullptr,
+ kq_scale, il);
+ cb(out, "attn_latent", il);
+
+ // {kv_lora_rank, n_tokens, n_head} -> {n_embd_head_v_mla, n_tokens, n_head}
+ out = ggml_reshape_3d(ctx0, out, kv_lora_rank, n_head, n_tokens);
+ out = ggml_mul_mat(ctx0, layer.wv_b, ggml_permute(ctx0, out, 0, 2, 1, 3));
+ out = ggml_cont_2d(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3), n_embd_head_v_mla*n_head, n_tokens);
+
+ cur = build_lora_mm(layer.wo, out, layer.wo_s);
+ cb(cur, "attn_out", il);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL);
+ cb(ffn_inp, "ffn_inp", il);
+
+ cur = build_norm(ffn_inp, layer.ffn_norm, nullptr, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ cur = build_ffn(cur,
+ layer.ffn_up, nullptr, nullptr,
+ layer.ffn_gate, nullptr, nullptr,
+ layer.ffn_down, nullptr, nullptr,
+ nullptr, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+
+ inpL = ggml_add(ctx0, cur, ffn_inp);
+ cb(inpL, "l_out", il);
+ }
+
+ ggml_tensor * cur = build_norm(inpL, model.output_norm, nullptr, LLM_NORM_RMS, -1);
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ auto * output = model.output;
+ auto * output_s = model.output_s;
+ if (output == nullptr) {
+ GGML_ASSERT(cparams.ctx_other != nullptr);
+ const auto * model_other = llama_get_model(cparams.ctx_other);
+ GGML_ASSERT(model_other->output != nullptr && "DSpark decoder requires the target model's output projection");
+ output = model_other->output;
+ output_s = model_other->output_s;
+ }
+
+ ggml_tensor * logits = build_lora_mm(output, cur, output_s);
+
+ // softcap does not apply to the markov bias
+ if (hparams.f_final_logit_softcapping) {
+ logits = ggml_scale(ctx0, logits, 1.0f / hparams.f_final_logit_softcapping);
+ logits = ggml_tanh(ctx0, logits);
+ logits = ggml_scale(ctx0, logits, hparams.f_final_logit_softcapping);
+ }
+ cb(logits, "result_output", -1);
+ res->t_logits = logits;
+
+ if (model.dspark_markov_w1) {
+ build_dspark_markov_head(*this, model, inp_tokens, cur);
+ }
+}
+
std::unique_ptr<llm_graph_context> llama_model_dflash::build_arch_graph(const llm_graph_params & params) const {
switch (params.gtype) {
case LLM_GRAPH_TYPE_ENCODER:
return std::make_unique<graph<true>>(*this, params);
case LLM_GRAPH_TYPE_DEFAULT:
case LLM_GRAPH_TYPE_DECODER:
+ if (hparams.n_lora_kv > 0) {
+ return std::make_unique<graph_mistral_large4>(*this, params);
+ }
if (hparams.dsv4_hc_mult > 0 && hparams.dsv4_hc_single_pass) {
return std::make_unique<graph_dsv41>(*this, params);
}
diff --git a/src/models/mistral-large4.cpp b/src/models/mistral-large4.cpp
new file mode 100644
index 000000000..9cfdf825b
--- /dev/null
+++ b/src/models/mistral-large4.cpp
@@ -0,0 +1,507 @@
+#include "models.h"
+
+#include "llama-kv-cache.h"
+#include "llama-kv-cache-dsa.h"
+
+void llama_model_mistral_large4::load_arch_hparams(llama_model_loader & ml) {
+ ml.get_key_or_arr(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp_arr, hparams.n_layer_all);
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps);
+ ml.get_key(LLM_KV_ATTENTION_LAYERNORM_EPS, hparams.f_norm_eps); // eps for the indexer layer norms
+
+ // MoE parameters
+ ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared);
+ ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false);
+ ml.get_key(LLM_KV_EXPERT_WEIGHTS_NORM, hparams.expert_weights_norm, false);
+ ml.get_key(LLM_KV_EXPERT_GATING_FUNC, hparams.expert_gating_func);
+
+ // MLA parameters
+ ml.get_key(LLM_KV_ATTENTION_Q_LORA_RANK, hparams.n_lora_q, false);
+ ml.get_key(LLM_KV_ATTENTION_KV_LORA_RANK, hparams.n_lora_kv);
+ ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH_MLA, hparams.n_embd_head_k_mla_impl);
+ ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH_MLA, hparams.n_embd_head_v_mla_impl);
+
+ // DSA parameters
+ ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head);
+ ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size);
+ ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k);
+
+ std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 1);
+ ml.get_key_or_arr(LLM_KV_ATTENTION_INDEXER_TYPES, hparams.is_indexer_full_impl, hparams.n_layer(), false);
+ if (!hparams.is_indexer_full(0)) {
+ throw std::runtime_error("the first layer must have a full indexer");
+ }
+
+ ml.get_key_or_arr(LLM_KV_ATTENTION_ROPE_PATTERN, hparams.rope_pattern, hparams.n_layer(), false);
+ ml.get_key_or_arr(LLM_KV_ATTENTION_KV_SOURCE_MAP, hparams.kv_source_impl, hparams.n_layer(), false);
+ for (uint32_t il = 0; il < hparams.n_layer(); ++il) {
+ const uint32_t il_src = hparams.kv_source(il);
+ if (il_src > il || !hparams.has_kv(il_src) || hparams.has_rope(il_src) != hparams.has_rope(il)) {
+ throw std::runtime_error(format("invalid KV source layer %u for layer %u", il_src, il));
+ }
+ }
+
+ hparams.f_final_logit_softcapping = 0.0f;
+ ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false);
+ ml.get_key(LLM_KV_ROPE_SCALING_YARN_LOG_MUL, hparams.rope_yarn_log_mul, false);
+
+ // SSMax uses log(pos + 1), computed via the attention temperature input as 1 + log(pos + 1)
+ hparams.n_attn_temp_floor_scale = 1;
+ hparams.f_attn_temp_scale = 1.0f;
+ hparams.f_attn_temp_offset = 0.0f;
+
+ switch (hparams.n_layer()) {
+ default: type = LLM_TYPE_UNKNOWN;
+ }
+}
+
+void llama_model_mistral_large4::load_arch_tensors(llama_model_loader &) {
+ LLAMA_LOAD_LOCALS;
+
+ if (!hparams.is_mla()) {
+ throw std::runtime_error("MISTRAL_LARGE4 architecture requires MLA");
+ }
+
+ const int64_t n_embd_head_k_mla = hparams.n_embd_head_k_mla();
+ const int64_t n_embd_head_v_mla = hparams.n_embd_head_v_mla();
+
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+
+ const int64_t n_ff_exp = hparams.n_ff_exp();
+ const int64_t n_expert_shared = hparams.n_expert_shared;
+
+ const int64_t n_indexer_head = hparams.indexer_n_head;
+ const int64_t n_embd_indexer_head = hparams.indexer_head_size;
+
+ tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0);
+
+ // output
+ output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0);
+ output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, 0);
+
+ for (int i = 0; i < n_layer; ++i) {
+ auto & layer = layers[i];
+
+ // NoPE layers use the full head for the non-rotary part
+ const int64_t n_rot_l = hparams.has_rope(i) ? n_rot : 0;
+ const int64_t n_embd_head_qk_nope = n_embd_head_k_mla - n_rot_l;
+
+ // learned rope frequencies (LeRoPE), shared by all layers
+ layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, i != 0 ? TENSOR_DUPLICATED : 0);
+
+ layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0);
+
+ layer.wq = create_tensor(tn(LLM_TENSOR_ATTN_Q, "weight", i), {n_embd, n_head * n_embd_head_k_mla}, 0);
+
+ if (hparams.has_kv(i)) {
+ layer.wkv_a_mqa = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_MQA, "weight", i), {n_embd, kv_lora_rank + n_rot_l}, 0);
+ layer.attn_kv_a_norm = create_tensor(tn(LLM_TENSOR_ATTN_KV_A_NORM, "weight", i), {kv_lora_rank}, 0);
+ }
+
+ layer.wk_b = create_tensor(tn(LLM_TENSOR_ATTN_K_B, "weight", i), {n_embd_head_qk_nope, kv_lora_rank, n_head}, 0);
+ layer.wv_b = create_tensor(tn(LLM_TENSOR_ATTN_V_B, "weight", i), {kv_lora_rank, n_embd_head_v_mla, n_head}, 0);
+
+ layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_head * n_embd_head_v_mla, n_embd}, 0);
+
+ layer.attn_ssmax_s = create_tensor(tn(LLM_TENSOR_ATTN_SSMAX, "weight", i), {n_head}, 0);
+ layer.attn_ssmax_b = create_tensor(tn(LLM_TENSOR_ATTN_SSMAX, "bias", i), {n_head}, 0);
+
+ // DSA indexer
+ if (hparams.is_indexer_full(i)) {
+ layer.indexer_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {n_embd}, 0);
+ layer.indexer_q_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "bias", i), {n_embd}, 0);
+ layer.indexer_attn_q_b = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_Q_B, "weight", i), {n_embd, n_indexer_head * n_embd_indexer_head}, 0);
+ layer.indexer_attn_k = create_tensor(tn(LLM_TENSOR_INDEXER_ATTN_K, "weight", i), {n_embd, n_embd_indexer_head}, 0);
+ layer.indexer_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {n_embd_indexer_head}, 0);
+ layer.indexer_k_norm_b = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "bias", i), {n_embd_indexer_head}, 0);
+ layer.indexer_proj = create_tensor(tn(LLM_TENSOR_INDEXER_PROJ, "weight", i), {n_embd, n_indexer_head}, 0);
+ layer.indexer_ssmax_s = create_tensor(tn(LLM_TENSOR_INDEXER_SSMAX, "weight", i), {n_indexer_head}, 0);
+ layer.indexer_ssmax_b = create_tensor(tn(LLM_TENSOR_INDEXER_SSMAX, "bias", i), {n_indexer_head}, 0);
+ }
+
+ layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0);
+
+ if (i < (int) hparams.n_layer_dense_lead) {
+ layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0);
+ layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0);
+ layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0);
+ } else {
+ if (n_expert == 0) {
+ throw std::runtime_error("n_expert must be > 0");
+ }
+ if (n_expert_used == 0) {
+ throw std::runtime_error("n_expert_used must be > 0");
+ }
+
+ layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0);
+
+ // MoE branch
+ layer.ffn_gate_exps = create_tensor(tn(LLM_TENSOR_FFN_GATE_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, 0);
+ layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), {n_ff_exp, n_embd, n_expert}, 0);
+ layer.ffn_up_exps = create_tensor(tn(LLM_TENSOR_FFN_UP_EXPS, "weight", i), { n_embd, n_ff_exp, n_expert}, 0);
+
+ // Shared expert branch
+ layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+ layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), { n_ff_exp * n_expert_shared, n_embd}, 0);
+ layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, n_ff_exp * n_expert_shared}, 0);
+ }
+ }
+}
+
+std::unique_ptr<llm_graph_context> llama_model_mistral_large4::build_arch_graph(const llm_graph_params & params) const {
+ return std::make_unique<graph>(*this, params);
+}
+
+llama_model_mistral_large4::graph::graph(const llama_model & model, const llm_graph_params & params) :
+ llm_graph_context(params) {
+ GGML_ASSERT(hparams.is_mla());
+
+ const int64_t n_embd_head_k = hparams.n_embd_head_k_mla();
+
+ const int64_t n_indexer_head = hparams.indexer_n_head;
+ const int64_t n_embd_indexer_head = hparams.indexer_head_size;
+ const uint32_t n_indexer_top_k = hparams.indexer_top_k;
+
+ // the indexer head layout is [rope | nope]
+ GGML_ASSERT(n_rot <= n_embd_indexer_head);
+
+ const int64_t kv_lora_rank = hparams.n_lora_kv;
+
+ // YaRN is used without magnitude scaling (apply_scale = false), so no mscale in kq_scale
+ const float kq_scale = 1.0f / sqrtf(float(n_embd_head_k));
+
+ ggml_tensor * cur;
+ ggml_tensor * inpL;
+
+ // {n_embd, n_tokens}
+ inpL = build_inp_embd(model.tok_embd);
+
+ ggml_tensor * inp_pos = build_inp_pos();
+
+ // SSMax position basis: log(pos + 1), {1, 1, n_tokens}
+ ggml_tensor * ssmax_pos = ggml_scale_bias(ctx0, build_inp_attn_scale(), 1.0f, -1.0f);
+ cb(ssmax_pos, "ssmax_pos", -1);
+
+ // per-head SSMax scale s * log(pos + 1) + b, {1, n_head, n_tokens}
+ auto build_ssmax = [&](ggml_tensor * s, ggml_tensor * b, int64_t n_head_l) {
+ ggml_tensor * pos = ggml_repeat_4d(ctx0, ssmax_pos, 1, n_head_l, n_tokens, 1);
+ ggml_tensor * res = ggml_mul(ctx0, pos, ggml_reshape_3d(ctx0, s, 1, n_head_l, 1));
+ return ggml_add(ctx0, res, ggml_reshape_3d(ctx0, b, 1, n_head_l, 1));
+ };
+
+ // NoPE layers keep the cache layout of rope layers by zero-filling the rope part
+ auto pad_rope = [&](ggml_tensor * x) {
+ ggml_tensor * zeros = ggml_new_tensor_3d(ctx0, x->type, n_rot, x->ne[1], x->ne[2]);
+ zeros = ggml_fill(ctx0, zeros, 0.0f);
+ return ggml_concat(ctx0, x, zeros, 0);
+ };
+
+ llm_graph_input_attn_k_dsa * inp_attn_dsa = build_attn_inp_k_dsa();
+
+ ggml_tensor * inp_out_ids = build_inp_out_ids();
+
+ // shared indexer layers reuse the top_k from the previous full indexer layer
+ ggml_tensor * prev_top_k = nullptr;
+
+ // DSpark reads last hidden state of all tokens, select output rows after the last layer
+ const bool extract_last = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer];
+
+ for (int il = 0; il < n_layer; ++il) {
+ const auto & layer = model.layers[il];
+
+ res->t_layer_inp[il] = inpL;
+
+ const bool use_rope = hparams.has_rope(il);
+
+ const int64_t n_rot_l = use_rope ? n_rot : 0;
+ const int64_t n_embd_head_qk_nope = n_embd_head_k - n_rot_l;
+
+ ggml_tensor * rope_factors = model.get_rope_factors(cparams, il);
+
+ ggml_tensor * inpSA = inpL;
+
+ cur = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "attn_norm", il);
+
+ // self_attention
+ {
+ ggml_tensor * top_k = nullptr;
+
+ // lightning indexer
+ if (hparams.is_indexer_full(il)) {
+ ggml_tensor * indexer_q = build_norm(cur, layer.indexer_q_norm, layer.indexer_q_norm_b, LLM_NORM, il);
+ cb(indexer_q, "indexer_q_norm", il);
+
+ indexer_q = ggml_mul_mat(ctx0, layer.indexer_attn_q_b, indexer_q);
+ cb(indexer_q, "indexer_q", il);
+
+ // {n_embd_indexer_head, n_indexer_head, n_tokens}
+ indexer_q = ggml_reshape_3d(ctx0, indexer_q, n_embd_indexer_head, n_indexer_head, n_tokens);
+ indexer_q = ggml_mul(ctx0, indexer_q, build_ssmax(layer.indexer_ssmax_s, layer.indexer_ssmax_b, n_indexer_head));
+ cb(indexer_q, "indexer_q_ssmax", il);
+
+ ggml_tensor * indexer_k = ggml_mul_mat(ctx0, layer.indexer_attn_k, cur);
+ cb(indexer_k, "indexer_k", il);
+
+ indexer_k = build_norm(indexer_k, layer.indexer_k_norm, layer.indexer_k_norm_b, LLM_NORM, il);
+ cb(indexer_k, "indexer_k", il);
+
+ // {n_embd_indexer_head, 1, n_tokens}
+ indexer_k = ggml_reshape_3d(ctx0, indexer_k, n_embd_indexer_head, 1, n_tokens);
+
+ if (use_rope) {
+ indexer_q = ggml_rope_ext(ctx0, indexer_q, inp_pos, rope_factors, n_rot,
+ rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(indexer_q, "indexer_q", il);
+
+ indexer_k = ggml_rope_ext(ctx0, indexer_k, inp_pos, rope_factors, n_rot,
+ rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(indexer_k, "indexer_k", il);
+ }
+
+ // perform Hadamard transform on indexer q and k
+ indexer_q = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_q);
+ cb(indexer_q, "indexer_q", il);
+ indexer_k = ggml_mul_mat(ctx0, inp_attn_dsa->self_k_rot_lid, indexer_k);
+ cb(indexer_k, "indexer_k", il);
+
+ // store indexer keys to KV cache
+ const auto * mctx_lid = inp_attn_dsa->mctx->get_lid();
+ const auto & k_idxs_lid = inp_attn_dsa->get_k_idxs_lid();
+ ggml_build_forward_expand(gf, mctx_lid->cpy_k(ctx0, indexer_k, k_idxs_lid, il));
+
+ ggml_tensor * indexer_weights = ggml_mul_mat(ctx0, layer.indexer_proj, cur);
+ cb(indexer_weights, "indexer_weights", il);
+
+ // get cached indexer keys
+ indexer_k = mctx_lid->get_k(ctx0, il);
+
+ // split the batch into streams if needed
+ const auto n_stream = indexer_k->ne[3];
+ indexer_q = ggml_view_4d(ctx0, indexer_q, indexer_q->ne[0], indexer_q->ne[1], indexer_q->ne[2]/n_stream, n_stream, indexer_q->nb[1], indexer_q->nb[2], indexer_q->nb[3]/n_stream, 0);
+ indexer_weights = ggml_view_4d(ctx0, indexer_weights, indexer_weights->ne[0], indexer_weights->ne[1]/n_stream, indexer_weights->ne[2], n_stream, indexer_weights->nb[1], indexer_weights->nb[2]/n_stream, indexer_weights->nb[3]/n_stream, 0);
+
+ // pre-scale weights to avoid scaling operations on huge indexer_score tensor
+ indexer_weights = ggml_scale(ctx0, indexer_weights, 1.0f / sqrtf(float(n_embd_indexer_head * n_indexer_head)));
+ cb(indexer_weights, "indexer_weights", il);
+
+ ggml_tensor * indexer_score = nullptr;
+ if (cparams.fused_lid) {
+ indexer_score = ggml_lightning_indexer(ctx0, indexer_q, indexer_k, indexer_weights, inp_attn_dsa->get_kq_mask_lid());
+ cb(indexer_score, "indexer_score", il);
+ res->add_fused_node({LLM_FUSED_OP_LIGHTNING_INDEXER, indexer_score, il});
+ } else {
+ // calculate indexer kq
+ indexer_q = ggml_permute(ctx0, indexer_q, 0, 2, 1, 3);
+ cb(indexer_q, "indexer_q", il);
+ indexer_k = ggml_permute(ctx0, indexer_k, 0, 2, 1, 3);
+ cb(indexer_k, "indexer_k", il);
+
+ ggml_tensor * indexer_kq = ggml_mul_mat(ctx0, indexer_k, indexer_q);
+ cb(indexer_kq, "indexer_kq", il);
+
+ // ReLU requires contiguous tensors
+ indexer_kq = ggml_cont(ctx0, ggml_permute(ctx0, indexer_kq, 2, 1, 0, 3));
+ cb(indexer_kq, "indexer_kq", il);
+
+ // apply ReLU
+ indexer_score = ggml_relu(ctx0, indexer_kq);
+ cb(indexer_score, "indexer_score", il);
+
+ // multiply scores by indexer weights
+ indexer_score = ggml_mul(ctx0, indexer_score, indexer_weights);
+ cb(indexer_score, "indexer_score", il);
+
+ // sum by q n_indexer_head dimension
+ indexer_score = ggml_sum_rows(ctx0, indexer_score);
+ cb(indexer_score, "indexer_score", il);
+
+ // permute result to match KQ mask
+ indexer_score = ggml_cont(ctx0, ggml_permute(ctx0, indexer_score, 2, 1, 0, 3));
+ cb(indexer_score, "indexer_score", il);
+
+ // mask indexer scores
+ ggml_tensor * indexer_kq_mask = inp_attn_dsa->get_kq_mask_lid();
+ indexer_score = ggml_add(ctx0, indexer_score, indexer_kq_mask);
+ cb(indexer_score, "indexer_score", il);
+ }
+
+ // get indices of top k indexer scores
+ uint32_t n_top_k = indexer_score->ne[0] < n_indexer_top_k ? indexer_score->ne[0] : n_indexer_top_k;
+ top_k = ggml_cont(ctx0, ggml_top_k(ctx0, indexer_score, n_top_k));
+ prev_top_k = top_k;
+ cb(top_k, "top_k", il);
+ } else {
+ GGML_ASSERT(prev_top_k != nullptr && "shared indexer layer must follow a previous full indexer layer");
+ top_k = prev_top_k;
+ cb(top_k, "top_k", il);
+ }
+
+ ggml_tensor * q = ggml_mul_mat(ctx0, layer.wq, cur);
+ cb(q, "q", il);
+
+ // {n_embd_head_qk_nope, n_head, n_tokens}
+ ggml_tensor * q_nope = ggml_view_3d(ctx0, q, n_embd_head_qk_nope, n_head, n_tokens,
+ ggml_row_size(q->type, n_embd_head_k),
+ ggml_row_size(q->type, n_embd_head_k) * n_head, 0);
+ cb(q_nope, "q_nope", il);
+
+ // {n_embd_head_qk_nope, n_tokens, n_head}
+ q_nope = ggml_permute(ctx0, q_nope, 0, 2, 1, 3);
+ cb(q_nope, "q_nope_perm", il);
+
+ // {n_embd_head_qk_nope, kv_lora_rank, n_head} x {n_embd_head_qk_nope, n_tokens, n_head}
+ ggml_tensor * q_nope_absorbed = ggml_mul_mat(ctx0, layer.wk_b, q_nope);
+ cb(q_nope_absorbed, "q_nope_absorbed", il);
+
+ // {kv_lora_rank, n_head, n_tokens}
+ q_nope_absorbed = ggml_permute(ctx0, q_nope_absorbed, 0, 2, 1, 3);
+ cb(q_nope_absorbed, "q_nope_absorbed_perm", il);
+
+ // {kv_lora_rank + n_rot, n_head, n_tokens}
+ // note: rope must go first for in-place context shifting in build_rope_shift()
+ ggml_tensor * Qcur = nullptr;
+ if (use_rope) {
+ ggml_tensor * q_pe = ggml_view_3d(ctx0, q, n_rot, n_head, n_tokens,
+ ggml_row_size(q->type, n_embd_head_k),
+ ggml_row_size(q->type, n_embd_head_k) * n_head,
+ ggml_row_size(q->type, n_embd_head_qk_nope));
+ q_pe = ggml_rope_ext(ctx0, q_pe, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(q_pe, "q_pe", il);
+
+ Qcur = ggml_concat(ctx0, q_nope_absorbed, q_pe, 0);
+ } else {
+ Qcur = pad_rope(q_nope_absorbed);
+ }
+ cb(Qcur, "Qcur", il);
+
+ Qcur = ggml_mul(ctx0, Qcur, build_ssmax(layer.attn_ssmax_s, layer.attn_ssmax_b, n_head));
+ cb(Qcur, "Qcur_ssmax", il);
+
+ // layers without own KV read the cache written by their source layer
+ ggml_tensor * Kcur = nullptr;
+ ggml_tensor * Vcur = nullptr;
+
+ if (hparams.has_kv(il)) {
+ ggml_tensor * kv_cmpr_pe = ggml_mul_mat(ctx0, layer.wkv_a_mqa, cur);
+ cb(kv_cmpr_pe, "kv_cmpr_pe", il);
+
+ // {kv_lora_rank, n_tokens}
+ ggml_tensor * kv_cmpr = ggml_view_2d(ctx0, kv_cmpr_pe, kv_lora_rank, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot_l), 0);
+ cb(kv_cmpr, "kv_cmpr", il);
+
+ kv_cmpr = build_norm(kv_cmpr, layer.attn_kv_a_norm, nullptr, LLM_NORM_RMS, il);
+ cb(kv_cmpr, "kv_cmpr", il);
+
+ kv_cmpr = ggml_reshape_3d(ctx0, kv_cmpr, kv_lora_rank, 1, n_tokens);
+ cb(kv_cmpr, "kv_cmpr_reshape", il);
+
+ if (use_rope) {
+ // {n_rot, 1, n_tokens}
+ ggml_tensor * k_pe = ggml_view_3d(ctx0, kv_cmpr_pe, n_rot, 1, n_tokens,
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank + n_rot),
+ ggml_row_size(kv_cmpr_pe->type, kv_lora_rank));
+ k_pe = ggml_rope_ext(ctx0, k_pe, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale,
+ ext_factor, attn_factor, beta_fast, beta_slow);
+ cb(k_pe, "k_pe", il);
+
+ Kcur = ggml_concat(ctx0, kv_cmpr, k_pe, 0);
+ } else {
+ Kcur = pad_rope(kv_cmpr);
+ }
+ cb(Kcur, "Kcur", il);
+
+ Vcur = kv_cmpr;
+ cb(Vcur, "Vcur", il);
+ }
+
+ // note: MLA with the absorption optimization converts into MQA (ie: GQA with 1 group)
+ cur = build_attn(inp_attn_dsa,
+ layer.wo, NULL, layer.wo_s,
+ Qcur, Kcur, Vcur, nullptr, nullptr, layer.wv_b, top_k, kq_scale, il);
+ }
+
+ if (il == n_layer - 1 && inp_out_ids && !extract_last) {
+ cur = ggml_get_rows(ctx0, cur, inp_out_ids);
+ inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids);
+ }
+
+ ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA);
+ cb(ffn_inp, "ffn_inp", il);
+
+ cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il);
+ cb(cur, "ffn_norm", il);
+
+ if ((uint32_t) il < hparams.n_layer_dense_lead) {
+ cur = build_ffn(cur,
+ layer.ffn_up, NULL, layer.ffn_up_s,
+ layer.ffn_gate, NULL, layer.ffn_gate_s,
+ layer.ffn_down, NULL, layer.ffn_down_s,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(cur, "ffn_out", il);
+ } else {
+ ggml_tensor * moe_out = build_moe_ffn(cur,
+ layer.ffn_gate_inp,
+ layer.ffn_up_exps,
+ layer.ffn_gate_exps,
+ layer.ffn_down_exps,
+ nullptr,
+ n_expert, n_expert_used,
+ LLM_FFN_SILU, hparams.expert_weights_norm,
+ hparams.expert_weights_scale,
+ (llama_expert_gating_func_type) hparams.expert_gating_func,
+ il,
+ nullptr,
+ layer.ffn_gate_up_exps,
+ layer.ffn_up_exps_s,
+ layer.ffn_gate_exps_s,
+ layer.ffn_down_exps_s);
+ cb(moe_out, "ffn_moe_out", il);
+
+ ggml_tensor * ffn_shexp = build_ffn(cur,
+ layer.ffn_up_shexp, NULL, layer.ffn_up_shexp_s,
+ layer.ffn_gate_shexp, NULL, layer.ffn_gate_shexp_s,
+ layer.ffn_down_shexp, NULL, layer.ffn_down_shexp_s,
+ NULL, LLM_FFN_SILU, LLM_FFN_PAR, il);
+ cb(ffn_shexp, "ffn_shexp", il);
+
+ cur = ggml_add(ctx0, moe_out, ffn_shexp);
+ cb(cur, "ffn_out", il);
+ }
+
+ cur = ggml_add(ctx0, cur, ffn_inp);
+
+ cur = build_cvec(cur, il);
+ cb(cur, "l_out", il);
+
+ inpL = cur;
+ }
+
+ if (extract_last) {
+ res->t_layer_inp[n_layer] = inpL;
+ if (inp_out_ids) {
+ inpL = ggml_get_rows(ctx0, inpL, inp_out_ids);
+ }
+ }
+
+ cur = build_norm(inpL, model.output_norm, NULL, LLM_NORM_RMS, -1);
+
+ cb(cur, "result_norm", -1);
+ res->t_embd = cur;
+
+ cur = ggml_mul_mat(ctx0, model.output, cur);
+
+ if (hparams.f_final_logit_softcapping) {
+ cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping);
+ cur = ggml_tanh(ctx0, cur);
+ cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping);
+ }
+
+ cb(cur, "result_output", -1);
+ res->t_logits = cur;
+}
diff --git a/src/models/models.h b/src/models/models.h
index 0fe40f4d0..9a8602b9c 100644
--- a/src/models/models.h
+++ b/src/models/models.h
@@ -1617,6 +1617,24 @@ struct llama_model_dflash : public llama_model_base {
graph_dsv41(const llama_model & model, const llm_graph_params & params);
};
+ // Mistral Large 4 DSpark: dense MLA stages, non-causal, rnn markov head
+ struct graph_mistral_large4 : public llm_graph_context {
+ graph_mistral_large4(const llama_model & model, const llm_graph_params & params);
+ };
+
+ std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
+};
+
+
+struct llama_model_mistral_large4 : public llama_model_base {
+ llama_model_mistral_large4(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_graph_context {
+ graph(const llama_model & model, const llm_graph_params & params);
+ };
+
std::unique_ptr<llm_graph_context> build_arch_graph(const llm_graph_params & params) const override;
};
diff --git a/tests/fusion/MTL.csv b/tests/fusion/MTL.csv
index 42fa5c912..e08d0c81b 100644
--- a/tests/fusion/MTL.csv
+++ b/tests/fusion/MTL.csv
@@ -194,6 +194,10 @@ minimax-m2 ,0 ,any ,RMS_NORM+MUL , 9
minimax-m3 ,0 ,any ,ADD+ADD , 1
minimax-m3 ,0 ,any ,MUL+ADD , 2
minimax-m3 ,0 ,any ,RMS_NORM+MUL , 11
+mistral-large4 ,0 ,any ,ADD+ADD , 3
+mistral-large4 ,0 ,any ,MUL+ADD , 6
+mistral-large4 ,0 ,any ,NORM+MUL+ADD , 6
+mistral-large4 ,0 ,any ,RMS_NORM+MUL , 12
mistral3 ,0 ,any ,RMS_NORM+MUL , 5
mistral3 ,0 ,any ,MUL+ADD , 4
mistral3 ,0 ,any ,RMS_NORM+MUL , 5
diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp
index 333551157..0479e1020 100644
--- a/tests/test-llama-archs.cpp
+++ b/tests/test-llama-archs.cpp
@@ -281,10 +281,14 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|| arch == LLM_ARCH_KIMI_K3
|| arch == LLM_ARCH_GLM5_NEXT
|| arch == LLM_ARCH_MISTRAL4
+ || arch == LLM_ARCH_MISTRAL_LARGE4
|| arch == LLM_ARCH_HY_V4) {
n_embd = 128;
n_head = 1;
n_ff = 192;
+ if (arch == LLM_ARCH_MISTRAL_LARGE4) {
+ n_layer = 4; // cover NoPE, shared indexer and KV sharing layers
+ }
} else if (arch == LLM_ARCH_NEMOTRON_H || arch == LLM_ARCH_NEMOTRON_H_MOE) {
n_layer = 3;
} else if (arch == LLM_ARCH_CHAMELEON) {
@@ -376,7 +380,8 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
|| arch == LLM_ARCH_KIMI_K3
|| arch == LLM_ARCH_GLM5_NEXT
|| arch == LLM_ARCH_HY_V4
- || arch == LLM_ARCH_MISTRAL4) {
+ || arch == LLM_ARCH_MISTRAL4
+ || arch == LLM_ARCH_MISTRAL_LARGE4) {
// GLM5 next MLA is nope only, the cache row is the compressed latent alone.
ms.add_kv(LLM_KV_ATTENTION_KEY_LENGTH, arch == LLM_ARCH_GLM5_NEXT ? uint32_t(512) : uint32_t(576));
ms.add_kv(LLM_KV_ATTENTION_VALUE_LENGTH, uint32_t(512));
@@ -539,6 +544,14 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) {
ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, indexer_types);
}
+ if (arch == LLM_ARCH_MISTRAL_LARGE4) {
+ // layer 3 reads the KV cache of layer 2, both are NoPE
+ ms.add_kv(LLM_KV_ATTENTION_ROPE_PATTERN, std::vector<uint32_t>({1, 1, 0, 0}));
+ ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, std::vector<uint32_t>({1, 0, 1, 1}));
+ ms.add_kv(LLM_KV_ATTENTION_KV_SOURCE_MAP, std::vector<uint32_t>({0, 1, 2, 2}));
+ ms.add_kv(LLM_KV_FINAL_LOGIT_SOFTCAPPING, 30.0f);
+ }
+
if (arch == LLM_ARCH_DEEPSEEK4) {
ms.add_kv(LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, uint32_t(8));
ms.add_kv(LLM_KV_ATTENTION_OUTPUT_LORA_RANK, uint32_t(32));
@@ -946,6 +959,7 @@ static bool moe_mandatory(const llm_arch arch) {
case LLM_ARCH_GLM5_NEXT:
case LLM_ARCH_STEP35:
case LLM_ARCH_MISTRAL4:
+ case LLM_ARCH_MISTRAL_LARGE4:
case LLM_ARCH_MELLUM:
case LLM_ARCH_LAGUNA:
case LLM_ARCH_MAPLE: