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_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_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_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_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_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_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_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_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 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 #include +#include // // 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( 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::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 { + 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(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(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 llama_model_dflash::build_arch_graph(const llm_graph_params & params) const { switch (params.gtype) { case LLM_GRAPH_TYPE_ENCODER: return std::make_unique>(*this, params); case LLM_GRAPH_TYPE_DEFAULT: case LLM_GRAPH_TYPE_DECODER: + if (hparams.n_lora_kv > 0) { + return std::make_unique(*this, params); + } if (hparams.dsv4_hc_mult > 0 && hparams.dsv4_hc_single_pass) { return std::make_unique(*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 llama_model_mistral_large4::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*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 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 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({1, 1, 0, 0})); + ms.add_kv(LLM_KV_ATTENTION_INDEXER_TYPES, std::vector({1, 0, 1, 1})); + ms.add_kv(LLM_KV_ATTENTION_KV_SOURCE_MAP, std::vector({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: