From a9a210a1ea6686960f828b3f0dbfcc8bc6987ec9 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 23 Oct 2026 06:31:02 -0700 Subject: [PATCH 1/3] model : MiniMax-M3.1 support (MXFP8 + NVFP4 conversion, MSA on dense-lead layers, DSpark draft, BPE NFC) - convert: MiniMax-M3.1 checkpoints (MXFP8 dense weights, NVFP4 routed experts). Experts pass through as native NVFP4 with a per-expert .scale tensor by default; --fp8-as-q8 and --outtype f32/bf16/f16/q8_0 dequantize exactly. Raw config.json is read for the sparse attention fields (transformers >= 5.15). - minimax-m3: a layer runs MSA iff it carries indexer tensors (M3.1 enables it on the dense-lead layers), NVFP4 expert scales are passed to build_moe_ffn, and the MSA batch path is chunked below 2^31 elements. - vocab: the BPE session applies tokenizer.ggml.normalizer.nfc, as the Unigram session already does. - DSpark: DSparkMiniMaxDraftModel converter, swigluoai in dflash, per-layer target taps in minimax-m3. --- conversion/__init__.py | 1 + conversion/minimax.py | 347 +++++++++++++++++++++++++++++++++++++- src/llama-model.cpp | 6 +- src/llama-vocab.cpp | 28 ++- src/models/dflash.cpp | 3 + src/models/minimax-m3.cpp | 149 ++++++++++------ src/models/models.h | 2 + 7 files changed, 482 insertions(+), 54 deletions(-) diff --git a/conversion/__init__.py b/conversion/__init__.py index 5f252b490d98..a203c5a2203a 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -70,6 +70,7 @@ "LingDSparkModel": "qwen", "DeepseekV4ForCausalLM": "deepseek", "DeepseekV4DSparkModel": "deepseek", + "DSparkMiniMaxDraftModel": "minimax", "DistilBertForMaskedLM": "bert", "DistilBertForSequenceClassification": "bert", "DistilBertModel": "bert", diff --git a/conversion/minimax.py b/conversion/minimax.py index aac340c61414..7e0788567bba 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -1,13 +1,18 @@ from __future__ import annotations -from typing import Iterable, Sequence, TYPE_CHECKING +import json +import math +import re +from typing import Callable, Iterable, Sequence, TYPE_CHECKING +import numpy as np import torch if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, TextModel, MmprojModel, gguf, logger +from .base import LazyTorchTensor, ModelBase, TextModel, MmprojModel, gguf, logger +from .qwen import DSparkModel @ModelBase.register("MiniMaxText01ForCausalLM") @@ -169,11 +174,197 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): class MiniMaxM3Model(MiniMaxM2Model): model_arch = gguf.MODEL_ARCH.MINIMAXM3 + # MiniMax-M3.1 checkpoints: quant_method "mxfp8" (attention / indexer / shared / dense MLP as + # F8_E4M3 + E8M0 weight_scale_inv per [1, 32]) with routed experts in compressed-tensors NVFP4 + # (weight_packed, weight_scale F8_E4M3 per 16, weight_global_scale, input_global_scale). + # --outtype auto keeps the experts NVFP4 (per-expert 1/global as ffn_*_exps.scale) and writes the + # MXFP8 tensors exactly as BF16 (or Q8_0 with --fp8-as-q8); an explicit --outtype dequantizes + # every tensor exactly through F32 first. + _NVFP4_PROJ = {"w1": gguf.MODEL_TENSOR.FFN_GATE_EXP, "w2": gguf.MODEL_TENSOR.FFN_DOWN_EXP, "w3": gguf.MODEL_TENSOR.FFN_UP_EXP} + + def __init__(self, dir_model, ftype, *args, **kwargs): + self._nvfp4_native = ftype == gguf.LlamaFileType.GUESSED + super().__init__(dir_model, ftype, *args, **kwargs) + def tensor_force_quant(self, name, new_name, bid, n_dims): if ".indexer." in new_name: return gguf.GGMLQuantizationType.F32 return super().tensor_force_quant(name, new_name, bid, n_dims) + @staticmethod + def _dequant_mxfp8(weight: Tensor, scale: Tensor) -> Tensor: + # exact: fp8 * 2^(e - 127) is representable in fp32 for every E8M0 byte except 255 (NaN) + rows, cols = weight.shape + e = scale.view(torch.uint8).to(torch.float64) + s = torch.exp2(e - 127.0).repeat_interleave(32, dim=1)[:, :cols] + return (weight.to(torch.float64) * s).to(torch.float32).reshape(rows, cols) + + _e2m1_pair_lut: Tensor | None = None + + @classmethod + def _dequant_nvfp4_eager(cls, packed: Tensor, scale: Tensor, global_scale: Tensor) -> Tensor: + # exact: e2m1 * e4m3 / global evaluated in fp64, rounded once to fp32 + if cls._e2m1_pair_lut is None: + v = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] + v16 = v + [-x for x in v] + # byte b holds element 2i in the low nibble and 2i+1 in the high nibble + cls._e2m1_pair_lut = torch.tensor([[v16[b & 15], v16[b >> 4]] for b in range(256)], dtype=torch.float64) + rows = packed.shape[0] + vals = cls._e2m1_pair_lut[packed.view(torch.uint8).to(torch.int64)].reshape(rows, -1) + s = scale.to(torch.float64).repeat_interleave(16, dim=1) + return (vals * s / global_scale.to(torch.float64)).to(torch.float32) + + @classmethod + def _dequant_nvfp4(cls, packed: Tensor, scale: Tensor, global_scale: Tensor) -> Tensor: + if not isinstance(packed, LazyTorchTensor): + return cls._dequant_nvfp4_eager(packed, scale, global_scale) + rows, packed_cols = packed.shape + # one lazy node around the eager decode (a lazy op chain costs ~10 s per expert tensor) + return LazyTorchTensor(meta=LazyTorchTensor.meta_with_dtype_and_shape(torch.float32, (rows, packed_cols * 2)), + args=(packed, scale, global_scale), func=cls._dequant_nvfp4_eager) + + @staticmethod + def _nvfp4_pack_np(packed: Tensor, scale: Tensor) -> np.ndarray: + # same bytes as ModelBase._nvfp4_pack (compressed-tensors packing -> ggml block_nvfp4), in numpy + p = packed.contiguous().view(torch.uint8).numpy() + d = scale.contiguous().view(torch.uint8).numpy() & 0x7F + rows, packed_cols = p.shape + n_blocks = packed_cols // 8 + p = p.reshape(rows, n_blocks, 8) + vals = np.stack((p & 0x0F, p >> 4), axis=-1).reshape(rows, n_blocks, 16) + qs = vals[:, :, :8] | (vals[:, :, 8:] << 4) + n_super = n_blocks // 4 + raw = np.concatenate([d.reshape(rows, n_super, 4), qs.reshape(rows, n_super, 32)], axis=-1) + return raw.reshape(rows, n_super * 36) + + # stacked [n_expert, rows, cols] routed-expert tensors are written straight to gguf_writer, built + # lazily (only the tensor being written is resident, cf. _mxfp4_expert_tensor) and per expert in a + # thread pool: the torch / numpy kernels release the GIL, a serial loop takes hours on the real model + _EXPERT_THREADS = 16 + _DEQUANT_RAW = { + gguf.LlamaFileType.ALL_F32: (gguf.GGMLQuantizationType.F32, torch.float32, 4), + gguf.LlamaFileType.MOSTLY_F16: (gguf.GGMLQuantizationType.F16, torch.float16, 2), + gguf.LlamaFileType.MOSTLY_BF16: (gguf.GGMLQuantizationType.BF16, torch.bfloat16, 2), + } + + def _expert_tensor(self, loaders: list[tuple[Callable[[], Tensor], ...]], native: bool): + rows, packed_cols = loaders[0][0]().shape + cols = packed_cols * 2 + if native: + if cols % 64 != 0: + raise ValueError(f"NVFP4 expert row has {cols} values, expected a multiple of 64") + row_bytes = (cols // 64) * 36 + else: + row_bytes = cols * self._DEQUANT_RAW[self.ftype][2] + byte_shape = (len(loaders), rows, row_bytes) + + def one(fns) -> np.ndarray: + t = [LazyTorchTensor.to_eager(f()) for f in fns] + if native: + return self._nvfp4_pack_np(t[0], t[1]) + w = self._dequant_nvfp4_eager(t[0], t[1], t[2]) + # torch's float -> bf16 / f16 casts round to nearest even, like gguf.quants + return w.to(self._DEQUANT_RAW[self.ftype][1]).contiguous().view(torch.uint8).numpy() + + def load(fns_list) -> np.ndarray: + from concurrent.futures import ThreadPoolExecutor + out = np.empty(byte_shape, dtype=np.uint8) + with ThreadPoolExecutor(self._EXPERT_THREADS) as ex: + for eid, arr in enumerate(ex.map(one, fns_list)): + out[eid] = arr + return out + + return gguf.LazyNumpyTensor( + meta=gguf.LazyNumpyTensor.meta_with_dtype_and_shape(np.uint8, byte_shape), + args=(loaders,), + func=load, + ) + + def dequant_model(self): + qc = self.hparams.get("quantization_config") or {} + if not isinstance(qc, dict) or qc.get("quant_method") != "mxfp8": + return super().dequant_model() + + if list(qc.get("weight_block_size", [1, 32])) != [1, 32]: + raise NotImplementedError(f"mxfp8 weight_block_size {qc.get('weight_block_size')!r} is not supported") + + remove: list[str] = [] + experts: dict[tuple[int, str], dict[int, str]] = {} + n_input_scale = n_input_scale_one = 0 + + for name in list(self.model_tensors.keys()): + if name.endswith(".weight_scale_inv"): + weight_name = name.removesuffix("_scale_inv") + if weight_name not in self.model_tensors: + raise KeyError(f"{name} has no matching {weight_name}") + w, s = self.model_tensors[weight_name], self.model_tensors[name] + self.model_tensors[weight_name] = lambda w=w, s=s: self._dequant_mxfp8(w(), s()) + remove.append(name) + if self._fp8_as_q8: + self._fp8_dequantized.add(weight_name) + elif name.endswith(".weight_packed"): + if qc.get("moe_quant_format") != "nvfp4-pack-quantized" or int(qc.get("group_size", 16)) != 16: + raise NotImplementedError(f"packed weight {name} with moe_quant_format {qc.get('moe_quant_format')!r}") + base = name.removesuffix("_packed") # ...experts.E.wX.weight + module = base.removesuffix(".weight") + for aux in (base + "_scale", base + "_global_scale"): + if aux not in self.model_tensors: + raise KeyError(f"{name}: missing {aux}") + # input_global_scale is the trained W4A4 activation scale; llama.cpp computes its own + # activation scales (or none), so it is dropped + if (module + ".input_global_scale") in self.model_tensors: + n_input_scale += 1 + if float(LazyTorchTensor.to_eager(self.model_tensors[module + ".input_global_scale"]()).float()) == 1.0: + n_input_scale_one += 1 + remove.append(module + ".input_global_scale") + m = re.search(r"layers\.(\d+)\.block_sparse_moe\.experts\.(\d+)\.(w[123])$", module) + if m is not None and (self._nvfp4_native or self.ftype in self._DEQUANT_RAW): + experts.setdefault((int(m.group(1)), m.group(3)), {})[int(m.group(2))] = base + continue + p = self.model_tensors[name] + s = self.model_tensors[base + "_scale"] + g = self.model_tensors[base + "_global_scale"] + self.model_tensors[base] = lambda p=p, s=s, g=g: MiniMaxM3Model._dequant_nvfp4(p(), s(), g()) + remove += [name, base + "_scale", base + "_global_scale"] + + if n_input_scale: + logger.info(f"dropped {n_input_scale} NVFP4 input_global_scale tensors ({n_input_scale_one} equal to 1.0): " + "activation quantization is not part of the converted model") + + n_experts = self.find_hparam(["num_local_experts", "num_experts"]) + for (bid, wid), by_eid in sorted(experts.items()): + if sorted(by_eid) != list(range(n_experts)): + raise KeyError(f"layer {bid} {wid}: {len(by_eid)} of {n_experts} NVFP4 experts present") + loaders = [] + inv_global: list[float] = [] + for eid in range(n_experts): + base = by_eid[eid] + loaders.append((self.model_tensors[base + "_packed"], self.model_tensors[base + "_scale"], + self.model_tensors[base + "_global_scale"])) + g = float(LazyTorchTensor.to_eager(self.model_tensors[base + "_global_scale"]()).float()) + if not (g > 0.0 and math.isfinite(g)): + raise ValueError(f"{base}_global_scale = {g}") + inv_global.append(1.0 / g) + remove += [base + "_packed", base + "_scale", base + "_global_scale"] + new_name = self.format_tensor_name(self._NVFP4_PROJ[wid], bid) + qtype = gguf.GGMLQuantizationType.NVFP4 if self._nvfp4_native else self._DEQUANT_RAW[self.ftype][0] + data = self._expert_tensor(loaders, self._nvfp4_native) + shape = gguf.quant_shape_from_byte_shape(data.shape, qtype) + logger.info(f"{new_name}: {n_experts} experts {'kept NVFP4' if self._nvfp4_native else 'dequantized to ' + qtype.name}, " + f"shape = {{{', '.join(str(n) for n in reversed(shape))}}}") + self.gguf_writer.add_tensor(new_name, data, raw_dtype=qtype) + if self._nvfp4_native: + # per-expert 1/global, applied to the matmul output by build_moe_ffn + self.gguf_writer.add_tensor(new_name.replace(".weight", ".scale"), np.array(inv_global, dtype=np.float32)) + + for name in remove: + self.model_tensors.pop(name, None) + + def _raw_text_config(self) -> dict: + with open(self.dir_model / "config.json", "r", encoding="utf-8") as f: + raw = json.load(f) + return raw.get("text_config", raw) + def set_gguf_parameters(self): super().set_gguf_parameters() @@ -181,14 +372,44 @@ def set_gguf_parameters(self): self.gguf_writer.add_expert_weights_scale(self.find_hparam(["routed_scaling_factor"])) self.gguf_writer.add_expert_weights_norm(True) - sac = self.find_hparam(["sparse_attention_config"]) + # transformers >= 5.15 resolves minimax_m3_vl natively and folds sparse_attention_config / + # moe_layer_freq into index_* / layer_types / mlp_layer_types, dropping the fields it does not + # consume (sparse_score_type, sparse_init_block, sparse_disable_index_value): read the raw file + raw = self._raw_text_config() + sac = raw.get("sparse_attention_config") or self.find_hparam(["sparse_attention_config"]) + for flat, legacy in (("index_n_heads", "sparse_num_index_heads"), ("index_head_dim", "sparse_index_dim"), + ("index_topk_blocks", "sparse_topk_blocks"), ("index_block_size", "sparse_block_size"), + ("index_local_blocks", "sparse_local_block")): + if flat in self.hparams and self.hparams[flat] != sac[legacy]: + raise ValueError(f"resolved {flat}={self.hparams[flat]} != config {legacy}={sac[legacy]}") + # only the MSA variant every released config uses is implemented + if sac.get("sparse_score_type", "max") != "max": + raise NotImplementedError(f"sparse_score_type {sac['sparse_score_type']!r}") + if int(sac.get("sparse_init_block", 0)) != 0: + raise NotImplementedError(f"sparse_init_block {sac['sparse_init_block']}") + # which layers run MSA is carried by the presence of blk.N.indexer.* (loader side), so the + # indexer tensors must exist exactly on the layers sparse_attention_freq marks sparse + freq = sac.get("sparse_attention_freq") + n_layer = self.find_hparam(["num_hidden_layers"]) + if freq is not None: + if len(freq) != n_layer: + raise ValueError(f"sparse_attention_freq has {len(freq)} entries, expected {n_layer}") + disable_v = sac.get("sparse_disable_index_value") or [1] * n_layer + for il, f in enumerate(freq): + has_idx = any(f"layers.{il}.self_attn.index_k_proj." in n for n in self.model_tensors) + if bool(f) != has_idx: + raise ValueError(f"layer {il}: sparse_attention_freq={f} but indexer tensors present={has_idx}") + if f and not disable_v[il]: + raise NotImplementedError(f"layer {il}: MSA index value path (sparse_disable_index_value=0)") self.gguf_writer.add_indexer_head_count(sac["sparse_num_index_heads"]) self.gguf_writer.add_indexer_key_length(sac["sparse_index_dim"]) self.gguf_writer.add_indexer_top_k(sac["sparse_topk_blocks"]) self.gguf_writer.add_indexer_block_size(sac["sparse_block_size"]) self.gguf_writer.add_indexer_local_blocks(sac["sparse_local_block"]) - moe_layer_freq = self.find_hparam(["moe_layer_freq"]) + moe_layer_freq = raw.get("moe_layer_freq") or self.find_hparam(["moe_layer_freq"]) + if "mlp_layer_types" in self.hparams: + assert [t == "sparse" for t in self.hparams["mlp_layer_types"]] == [bool(v) for v in moe_layer_freq] n_dense = 0 for v in moe_layer_freq: if v == 0: @@ -279,3 +500,121 @@ def _permute_vit_qk(self, t: "Tensor", new_name: str) -> "Tensor": if t.ndim == 2: return t.reshape(n_head, d_head, t.shape[1])[:, idx, :].reshape(t.shape) return t.reshape(n_head, d_head)[:, idx].reshape(t.shape) + + +@ModelBase.register("DSparkMiniMaxDraftModel") +@ModelBase.example("MiniMaxAI/MiniMax-M3.1-Flash") +class DSparkMiniMaxModel(DSparkModel): + # MiniMax-M3.1 DSpark draft (the `dspark/` folder of the checkpoint): a DFlash block drafter with + # MiniMax-M3 dense layers (Gemma norms, per-head Gemma QK norms, partial NeoX RoPE, SwiGLU-OAI MLP) + # plus the vanilla Markov head and the confidence head. Embeddings and lm_head come from the target. + model_arch = gguf.MODEL_ARCH.DFLASH + + _PREFIX = "language_model.model.dspark." + + def __init__(self, dir_model, *args, **kwargs): + hparams = kwargs.pop("hparams", None) + if hparams is None: + # read config.json as is: transformers >= 5.15 resolves minimax_m3_vl to its own config class, + # which drops the DSpark keys and reports the default hidden_act + with open(dir_model / "config.json", "r", encoding="utf-8") as f: + hparams = json.load(f) + hparams = {**hparams, **hparams.get("text_config", {})} + hparams.pop("text_config", None) + + if not hparams.get("dspark_use_dense_ffn", True): + raise ValueError("MiniMax DSpark: only the dense-FFN draft (dspark_use_dense_ffn) is supported") + if hparams.get("dspark_hybrid_context_fusion", False): + raise ValueError("MiniMax DSpark: hybrid_context_fusion is not supported") + if hparams.get("hidden_act") != "swigluoai": + raise ValueError(f"MiniMax DSpark: unexpected hidden_act {hparams.get('hidden_act')!r}") + # llama.cpp's LLM_FFN_SWIGLU_OAI_MOE hard-codes the gpt-oss constants + if float(hparams.get("swiglu_alpha", 1.702)) != 1.702 or float(hparams.get("swiglu_limit", 7.0)) != 7.0: + raise ValueError("MiniMax DSpark: swiglu_alpha/limit differ from 1.702/7.0") + if hparams.get("attention_output_gate", False): + raise ValueError("MiniMax DSpark: attention_output_gate is not supported") + if hparams.get("qk_norm_type", "per_head") != "per_head" or not hparams.get("use_qk_norm", True): + raise ValueError("MiniMax DSpark: expects per-head QK norm") + if not hparams.get("use_gemma_norm", True): + raise ValueError("MiniMax DSpark: expects Gemma (1 + w) norms") + if hparams.get("dspark_markov_head_type", hparams.get("markov_head_type", "vanilla")) != "vanilla": + raise ValueError("MiniMax DSpark: only the vanilla Markov head is supported") + layer_types = hparams.get("layer_types") or hparams.get("dflash_config", {}).get("layer_types") or [] + if any(lt != "sliding_attention" for lt in layer_types): + raise ValueError(f"MiniMax DSpark: expects uniform sliding_attention layers, got {layer_types}") + + # the draft's FFN is the dense one (dspark_ffn_hidden_size), not the MoE expert width + hparams["dense_intermediate_size"] = hparams.get("dspark_ffn_hidden_size", hparams["dense_intermediate_size"]) + for k in ("num_local_experts", "num_experts_per_tok", "n_shared_experts", "shared_intermediate_size", + "moe_intermediate_size", "num_experts", "scoring_func", "use_routing_bias", "routed_scaling_factor"): + hparams.pop(k, None) + + dflash_config = dict(hparams.get("dflash_config", {})) + dflash_config.setdefault("block_size", hparams["dspark_block_size"]) + dflash_config.setdefault("mask_token_id", hparams["dspark_noise_token_id"]) + dflash_config.setdefault("target_layer_ids", hparams["dspark_target_layer_ids"]) + if dflash_config["mask_token_id"] != hparams["dspark_noise_token_id"]: + raise ValueError("MiniMax DSpark: dflash_config.mask_token_id != dspark_noise_token_id") + if list(dflash_config["target_layer_ids"]) != list(hparams["dspark_target_layer_ids"]): + raise ValueError("MiniMax DSpark: dflash_config.target_layer_ids != dspark_target_layer_ids") + if dflash_config.get("context_source", "target_layers") != "target_layers": + raise ValueError("MiniMax DSpark: only context_source=target_layers is supported") + hparams["dflash_config"] = dflash_config + hparams["markov_rank"] = hparams.get("dspark_markov_rank", hparams.get("markov_rank")) + hparams["use_sliding_window"] = True + hparams["layer_types"] = layer_types + + super().__init__(dir_model, *args, hparams=hparams, **kwargs) + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + if not name.startswith(cls._PREFIX): + raise ValueError(f"MiniMax DSpark: unexpected tensor {name!r}") + name = name.removeprefix(cls._PREFIX).replace(".decoder_layer.", ".") + if name.startswith("final_norm."): + name = "norm." + name.removeprefix("final_norm.") + return super().filter_tensors((name, gen)) + + def dequant_model(self): + # MXFP8 (quant_method "mxfp8", block [1, 32]): F8_E4M3 weight * 2^(E8M0 - 127) per 32 inputs + quant_config = self.hparams.get("quantization_config") or {} + if quant_config.get("quant_method") == "mxfp8": + if list(quant_config.get("weight_block_size", [1, 32])) != [1, 32]: + raise ValueError(f"MiniMax DSpark: unexpected MXFP8 block {quant_config.get('weight_block_size')}") + + def dequant_mxfp8(w: Tensor, s: Tensor) -> Tensor: + if w.dtype != torch.float8_e4m3fn or s.dtype != torch.uint8: + raise ValueError(f"MiniMax DSpark: MXFP8 expects F8_E4M3 + U8, got {w.dtype} + {s.dtype}") + if w.shape[-1] != s.shape[-1] * 32 or w.shape[:-1] != s.shape[:-1]: + raise ValueError(f"MiniMax DSpark: MXFP8 shape mismatch {tuple(w.shape)} vs {tuple(s.shape)}") + scale = torch.exp2(s.to(torch.float32) - 127.0) + out = w.to(torch.float32).reshape(*w.shape[:-1], s.shape[-1], 32) * scale.unsqueeze(-1) + return out.reshape(w.shape) + + for name in list(self.model_tensors.keys()): + if name.endswith(".weight_scale_inv"): + weight_name = name.removesuffix("_scale_inv") + if weight_name not in self.model_tensors: + raise ValueError(f"MiniMax DSpark: scale {name!r} has no weight") + w = self.model_tensors[weight_name] + s = self.model_tensors.pop(name) + self.model_tensors[weight_name] = lambda w=w, s=s: dequant_mxfp8(w(), s()) + self.hparams = {k: v for k, v in self.hparams.items() if k != "quantization_config"} + super().dequant_model() + + def set_gguf_parameters(self): + super().set_gguf_parameters() + self.gguf_writer.add_hidden_act("swigluoai") + + def tensor_force_quant(self, name, new_name, bid, n_dims): + # confidence head is a tiny F32 projection in the reference (params_dtype=float32) + if "confidence_head" in name or new_name.startswith("conf_proj"): + return gguf.GGMLQuantizationType.F32 + return super().tensor_force_quant(name, new_name, bid, n_dims) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # Gemma-style (1 + w) RMSNorm, computed in fp32 by the reference: add 1 in fp32, not in bf16 + if name.endswith("norm.weight"): + data_torch = data_torch.to(torch.float32) + 1.0 + yield from super().modify_tensors(data_torch, name, bid) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 4cb04097ef95..e1d21984c1ae 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2815,9 +2815,11 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } break; case LLM_ARCH_MINIMAX_M3: { - // sparse (MSA) layers carry an indexer key cache, but leading dense layers do not + // sparse (MSA) layers carry an indexer key cache; a layer is sparse iff it has indexer + // tensors (MiniMax-M3: the MoE layers, MiniMax-M3.1: every layer) + const auto * mm3 = static_cast(this); llama_kv_cache::layer_filter_cb filter_idx = - [&](int32_t il) { return (uint32_t) il >= hparams.n_layer_dense_lead; }; + [mm3](int32_t il) { return mm3->is_msa_layer(il); }; res = new llama_kv_cache_msa( *this, diff --git a/src/llama-vocab.cpp b/src/llama-vocab.cpp index f72fbd2215cb..286358f4029c 100644 --- a/src/llama-vocab.cpp +++ b/src/llama-vocab.cpp @@ -637,7 +637,33 @@ struct llm_tokenizer_bpe_session { } } - virtual void tokenize(const std::string & text, std::vector & output) { + // HF tokenizers apply the normalizer (e.g. {"type": "NFC"}) before the pre-tokenizer + static std::string normalize_nfc(const std::string & text) { + bool ascii = true; + for (const unsigned char c : text) { + if (c >= 0x80) { + ascii = false; + break; + } + } + if (ascii) { + return text; // ASCII is already in NFC + } + try { + std::string canonical; + canonical.reserve(text.size()); + for (const uint32_t cpt : unicode_cpts_normalize_nfc(unicode_cpts_from_utf8(text))) { + canonical += unicode_cpt_to_utf8(cpt); + } + return canonical; + } catch (const std::invalid_argument &) { + return text; // malformed UTF-8: keep the input as is + } + } + + virtual void tokenize(const std::string & raw_text, std::vector & output) { + const std::string nfc_text = vocab.get_normalizer_opts().nfc ? normalize_nfc(raw_text) : std::string(); + const std::string & text = vocab.get_normalizer_opts().nfc ? nfc_text : raw_text; auto tok_pre = vocab.get_pre_type(); std::vector word_collection; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index e15af6afdc74..7bd79b8af93a 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -15,6 +15,9 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { if (ml.get_key(LLM_KV_HIDDEN_ACT, hidden_act, false)) { if (hidden_act == "gelu" || hidden_act == "gelu_pytorch_tanh") { hparams.llm_ffn_op = LLM_FFN_GELU; + } else if (hidden_act == "swigluoai") { + // MiniMax-M3 DSpark: clamped SwiGLU, alpha = 1.702, limit = 7.0 (checked by the converter) + hparams.llm_ffn_op = LLM_FFN_SWIGLU_OAI_MOE; } else if (hidden_act != "silu") { throw std::runtime_error("unsupported DFlash hidden activation: " + hidden_act); } diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp index 35e7b1a815c7..b20f27821c49 100644 --- a/src/models/minimax-m3.cpp +++ b/src/models/minimax-m3.cpp @@ -1,6 +1,8 @@ #include "models.h" #include "llama-kv-cache-msa.h" +#include #include +#include #include #include @@ -62,6 +64,22 @@ void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) { layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + // indexer: required on the MoE layers (MiniMax-M3); optional on the leading dense layers, which + // run MSA when it is present (MiniMax-M3.1 sets sparse_attention_freq on every layer). A layer + // is sparse iff it carries indexer tensors, see llama_model_minimax_m3::is_msa_layer() + { + const int idx_flags = i < (int) hparams.n_layer_dense_lead ? TENSOR_NOT_REQUIRED : 0; + layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, hparams.indexer_n_head * hparams.indexer_head_size}, idx_flags); + layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, hparams.indexer_head_size}, idx_flags); + layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {hparams.indexer_head_size}, idx_flags); + layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, idx_flags); + const int n_idx = (layer.index_q_proj != nullptr) + (layer.index_k_proj != nullptr) + + (layer.index_q_norm != nullptr) + (layer.index_k_norm != nullptr); + if (n_idx != 0 && n_idx != 4) { + throw std::runtime_error(format("%s: layer %d has %d of 4 indexer tensors", __func__, i, n_idx)); + } + } + if (i < (int) hparams.n_layer_dense_lead) { // leading dense layers layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); @@ -80,15 +98,24 @@ void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) { 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); - // indexer - layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, hparams.indexer_n_head * hparams.indexer_head_size}, 0); - layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", i), {n_embd, hparams.indexer_head_size}, 0); - layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", i), {hparams.indexer_head_size}, 0); - layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", i), {hparams.indexer_head_size}, 0); } } } +// tokens per MSA batch chunk so that per_token_elems * tokens stays below 2^31 elements +// (MSA_MAX_ELEMS overrides the cap, for testing the chunked path at small contexts) +static int64_t msa_tokens_per_chunk(int64_t per_token_elems, int64_t n_tokens) { + static const int64_t max_elems = [] { + const char * e = getenv("MSA_MAX_ELEMS"); + return e ? std::max(1, atoll(e)) : (int64_t) INT32_MAX; + }(); + return std::max(1, std::min(n_tokens, max_elems / std::max(1, per_token_elems))); +} + +bool llama_model_minimax_m3::is_msa_layer(int il) const { + return layers[il].index_k_proj != nullptr; +} + std::unique_ptr llama_model_minimax_m3::build_arch_graph(const llm_graph_params & params) const { return std::make_unique(*this, params); } @@ -253,7 +280,6 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ // hoisted per-graph MSA state (shared by every sparse layer) llm_graph_input_msa * msa = nullptr; ggml_tensor * msa_kqm = nullptr; - ggml_tensor * msa_mf = nullptr; // F32 copy of the FA mask for the final mask add int64_t n_kv = 0, n_ps = 0, nblk = 0, ns = 1, n_tps = 0; bool msa_decode = false; // gather (1 token per stream) vs mask const int blk = mm.msa_p.blk; @@ -291,8 +317,6 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ } else { inp->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, ns); ggml_set_input(inp->cell_blk); - - msa_mf = ggml_cast(ctx0, msa_kqm, GGML_TYPE_F32); } msa = (llm_graph_input_msa *) res->add_input(std::move(inp)); @@ -300,7 +324,13 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * inp_out_ids = build_inp_out_ids(); + // DFlash/DSpark drafts read the residual stream entering layer il (index n_layer = the last layer's output, + // before the final norm); that one needs every row, so the output-row crop moves after the loop + const bool tap_last = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]; + for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; + ggml_tensor * inpSA = inpL; // self-attention @@ -331,7 +361,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); - const bool is_sparse = msa_enabled && il >= (int) hparams.n_layer_dense_lead; + const bool is_sparse = msa_enabled && mm.is_msa_layer(il); if (!is_sparse) { cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s, @@ -404,6 +434,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * bsf = ggml_add(ctx0, bs, ggml_reshape_4d(ctx0, msa->bias, nblk, 1, 1, ns)); ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // position blocks + cb(idx, "msa_topk", il); // pos idx: tj[t,k,h,s] = blk*idx[k,h,s] + t (positions - mask gather) // cell idx: cs[t,k,h,s] = pos_slot[tj] (pos -> cell translation) @@ -446,84 +477,97 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ cur = build_attn_msa_fa(Qcur, kfa, vfa, mfa, Gp, kq_scale, il); } else { - // batch: per-stream loop - std::vector outs(ns); + // batch: per stream, in token chunks. The largest intermediates (indexer scores + // [n_ps, Hd, n], block mask expanded to cells [n_kv, n, Hd]) grow with context x tokens; + // chunking keeps every one of them below 2^31 elements (int32 indexing in backend + // kernels, e.g. CUDA pool2d, overflows past that: 270K ctx at -ub 2048). + const int64_t n_chunk = msa_tokens_per_chunk(std::max(n_ps, n_kv)*Hd, n_tps); + std::vector outs; for (int64_t st = 0; st < ns; ++st) { - ggml_tensor * iq_s = ggml_view_3d(ctx0, iq, n_idx_dim, Hd, n_tps, - iq->nb[1], iq->nb[2], st*n_tps*iq->nb[2]); ggml_tensor * ik_s = ggml_view_2d(ctx0, ik_kv, n_idx_dim, n_kv, ik_kv->nb[2], st*ik_kv->nb[3]); ggml_tensor * psl_s = ggml_view_1d(ctx0, msa->pos_slot_i, n_ps, st*msa->pos_slot_i->nb[1]); - ggml_tensor * pm_s = ggml_view_3d(ctx0, msa->pos_mask, n_ps, 1, n_tps, - msa->pos_mask->nb[1], msa->pos_mask->nb[1], st*n_tps*msa->pos_mask->nb[1]); ggml_tensor * cb_s = ggml_view_1d(ctx0, msa->cell_blk, n_kv, st*msa->cell_blk->nb[1]); - ggml_tensor * mf_s = ggml_view_3d(ctx0, msa_mf, n_kv, n_tps, 1, - msa_mf->nb[1], msa_mf->nb[3], st*msa_mf->nb[3]); - ggml_tensor * bias_s = ggml_view_3d(ctx0, msa->bias, nblk, 1, n_tps, - msa->bias->nb[1], msa->bias->nb[1], st*n_tps*msa->bias->nb[1]); - ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, n_tps, - Qcur->nb[1], Qcur->nb[2], st*n_tps*Qcur->nb[2]); ggml_tensor * k_s = ggml_view_4d(ctx0, k, D, HKV, n_kv, 1, k->nb[1], k->nb[2], k->nb[3], st*k->nb[3]); ggml_tensor * v_s = ggml_view_4d(ctx0, v, D, HKV, n_kv, 1, v->nb[1], v->nb[2], v->nb[3], st*v->nb[3]); - // block scores: the indexer keys are gathered through the pos -> cell map first - // scores are unscaled, only the top-k ordering matters + // the indexer keys are gathered through the pos -> cell map once per stream ggml_tensor * ikp = ggml_get_rows(ctx0, ik_s, psl_s); // [n_idx_dim, n_ps] + + // cache views with groups on ne[3]; + ggml_tensor * kfa = ggml_permute(ctx0, k_s, 0, 3, 1, 2); + ggml_tensor * vfa = ggml_permute(ctx0, v_s, 0, 3, 1, 2); + + for (int64_t c0 = 0; c0 < n_tps; c0 += n_chunk) { + const int64_t nt = std::min(n_chunk, n_tps - c0); // tokens in this chunk + const int64_t t0 = st*n_tps + c0; // first token (stream-grouped) + + ggml_tensor * iq_s = ggml_view_3d(ctx0, iq, n_idx_dim, Hd, nt, + iq->nb[1], iq->nb[2], t0*iq->nb[2]); + ggml_tensor * pm_s = ggml_view_3d(ctx0, msa->pos_mask, n_ps, 1, nt, + msa->pos_mask->nb[1], msa->pos_mask->nb[1], t0*msa->pos_mask->nb[1]); + // F32 copy of this chunk's rows of the FA (f16) mask for the final mask add + ggml_tensor * mf_s = ggml_cast(ctx0, ggml_view_3d(ctx0, msa_kqm, n_kv, nt, 1, + msa_kqm->nb[1], msa_kqm->nb[3], st*msa_kqm->nb[3] + c0*msa_kqm->nb[1]), GGML_TYPE_F32); + ggml_tensor * bias_s = ggml_view_3d(ctx0, msa->bias, nblk, 1, nt, + msa->bias->nb[1], msa->bias->nb[1], t0*msa->bias->nb[1]); + ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, nt, + Qcur->nb[1], Qcur->nb[2], t0*Qcur->nb[2]); + + // block scores, unscaled: only the top-k ordering matters ggml_tensor * sc = ggml_mul_mat(ctx0, ikp, - ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*n_tps)); + ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*nt)); // indexer scores run in F32 ggml_prec_set_acc(sc, GGML_PREC_F32); - sc = ggml_reshape_3d(ctx0, sc, n_ps, Hd, n_tps); + sc = ggml_reshape_3d(ctx0, sc, n_ps, Hd, nt); // unmapped positions (holes, padding, empty cells) come out -inf sc = ggml_add_inplace(ctx0, sc, pm_s); ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0); cb(bs, "msa_bs", il); // bias the scores so locally-forced blocks always rank first - ggml_tensor * bsf = ggml_add(ctx0, bs, bias_s); // [nblk, Hd, n_tps] + ggml_tensor * bsf = ggml_add(ctx0, bs, bias_s); // [nblk, Hd, nt] cb(bsf, "msa_bsf", il); - ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // [K, Hd, n_tps] i32 + ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // [K, Hd, nt] i32 + cb(idx, "msa_topk", il); ggml_tensor * ninf = ggml_cast(ctx0, ggml_scale_bias(ctx0, bias_s, 0.0f, -1e30f), - GGML_TYPE_F16); // [nblk, 1, n_tps] - ninf = ggml_repeat_4d(ctx0, ninf, nblk, Hd, n_tps, 1); + GGML_TYPE_F16); // [nblk, 1, nt] + ninf = ggml_repeat_4d(ctx0, ninf, nblk, Hd, nt, 1); ggml_tensor * zero = ggml_scale(ctx0, ggml_cast(ctx0, idx, GGML_TYPE_F32), 0.0f); ggml_tensor * bm = ggml_set_rows(ctx0, - ggml_reshape_3d(ctx0, ninf, 1, nblk, Hd*n_tps), - ggml_reshape_3d(ctx0, zero, 1, K, Hd*n_tps), - ggml_reshape_2d(ctx0, idx, K, Hd*n_tps)); - bm = ggml_reshape_3d(ctx0, bm, nblk, Hd, n_tps); - bm = ggml_cont(ctx0, ggml_permute(ctx0, bm, 0, 2, 1, 3)); // [nblk, n_tps, Hd] + ggml_reshape_3d(ctx0, ninf, 1, nblk, Hd*nt), + ggml_reshape_3d(ctx0, zero, 1, K, Hd*nt), + ggml_reshape_2d(ctx0, idx, K, Hd*nt)); + bm = ggml_reshape_3d(ctx0, bm, nblk, Hd, nt); + bm = ggml_cont(ctx0, ggml_permute(ctx0, bm, 0, 2, 1, 3)); // [nblk, nt, Hd] cb(bm, "msa_block_mask", il); // expand block -> cell granularity through the cell -> position block // map, then combine with the causal mask. empty cells are masked by the causal mask. ggml_tensor * bm2 = ggml_cont(ctx0, ggml_transpose(ctx0, - ggml_reshape_2d(ctx0, bm, nblk, n_tps*Hd))); // [n_tps*Hd, nblk] - ggml_tensor * bmc = ggml_get_rows(ctx0, bm2, cb_s); // [n_tps*Hd, n_kv] F32 + ggml_reshape_2d(ctx0, bm, nblk, nt*Hd))); // [nt*Hd, nblk] + ggml_tensor * bmc = ggml_get_rows(ctx0, bm2, cb_s); // [nt*Hd, n_kv] F32 ggml_tensor * bmx = ggml_cont(ctx0, ggml_transpose(ctx0, bmc)); - bmx = ggml_reshape_3d(ctx0, bmx, n_kv, n_tps, Hd); + bmx = ggml_reshape_3d(ctx0, bmx, n_kv, nt, Hd); ggml_tensor * mask4 = ggml_add_inplace(ctx0, bmx, mf_s); mask4 = ggml_cast(ctx0, - ggml_reshape_4d(ctx0, mask4, n_kv, n_tps, 1, Hd), GGML_TYPE_F16); + ggml_reshape_4d(ctx0, mask4, n_kv, nt, 1, Hd), GGML_TYPE_F16); cb(mask4, "msa_mask4", il); - // cache views with groups on ne[3]; - ggml_tensor * kfa = ggml_permute(ctx0, k_s, 0, 3, 1, 2); - ggml_tensor * vfa = ggml_permute(ctx0, v_s, 0, 3, 1, 2); - - outs[st] = build_attn_msa_fa(q_s, kfa, vfa, mask4, Gp, kq_scale, il); + outs.push_back(build_attn_msa_fa(q_s, kfa, vfa, mask4, Gp, kq_scale, il)); + } } cur = outs[0]; - for (int64_t st = 1; st < ns; ++st) { - cur = ggml_concat(ctx0, cur, outs[st], 1); + for (size_t i = 1; i < outs.size(); ++i) { + cur = ggml_concat(ctx0, cur, outs[i], 1); } } if (inp_attn->self_v_rot) { @@ -536,7 +580,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ } } - if (il == n_layer - 1 && inp_out_ids) { + if (il == n_layer - 1 && inp_out_ids && !tap_last) { cur = ggml_get_rows(ctx0, cur, inp_out_ids); inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); } @@ -568,7 +612,11 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ LLM_FFN_SWIGLU_OAI_MOE, hparams.expert_weights_norm, hparams.expert_weights_scale, (llama_expert_gating_func_type) hparams.expert_gating_func, - il); + il, + nullptr, nullptr, + model.layers[il].ffn_up_exps_s, + model.layers[il].ffn_gate_exps_s, + model.layers[il].ffn_down_exps_s); // NVFP4 experts: per-expert 1/global scale cb(moe_out, "ffn_moe_out", il); // shared expert (swigluoai) @@ -595,6 +643,13 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ cur = inpL; + if (tap_last) { + res->t_layer_inp[n_layer] = cur; + if (inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + } + } + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); cb(cur, "result_norm", -1); res->t_embd = cur; diff --git a/src/models/models.h b/src/models/models.h index bdcc1a4faef9..78d21f5ef9b1 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2419,6 +2419,8 @@ struct llama_model_minimax_m3 : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; msa_params msa_p; + // a layer runs MSA iff it carries the indexer tensors (M3: MoE layers only; M3.1: every layer) + bool is_msa_layer(int il) const; struct graph : public llm_graph_context { graph(const llama_model & model, const llm_graph_params & params); From ef8f94ccb0d24216cd361924aeec2fe0ab436d72 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Fri, 23 Oct 2026 07:14:30 -0700 Subject: [PATCH 2/3] convert : minimax-m3 add the Gemma-norm +1 in fp32 Upcast the norm weight to F32 before adding 1 instead of relying on the incoming dtype. prepare_tensors already turns BF16 tensors into F32 before modify_tensors, so BF16 checkpoints convert byte-identically (M3.1 and MiniMax-M3 layout fixtures, DSpark draft, all outtypes); the change keeps the add exact for F16 sources and for callers that pass a raw BF16 tensor, and matches the DSpark draft converter which already adds 1 in F32. --- conversion/minimax.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/conversion/minimax.py b/conversion/minimax.py index 7e0788567bba..51fe20e3f3b1 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -420,8 +420,9 @@ def set_gguf_parameters(self): def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): # Gemma-style (1 + w) RMSNorm: bake the +1 in so llama.cpp can use plain RMSNorm + # (in fp32 whatever the incoming dtype: adding 1.0 in F16 / BF16 rounds away the low bits of w) if name.endswith("norm.weight"): - data_torch = data_torch + 1.0 + data_torch = data_torch.to(torch.float32) + 1.0 yield from super().modify_tensors(data_torch, name, bid) From 1040caef41bf22a1d97804697ef8999ceb3ef2c1 Mon Sep 17 00:00:00 2001 From: Daniel Han Date: Mon, 26 Oct 2026 04:35:05 -0700 Subject: [PATCH 3/3] minimax-m3 : trim comments --- conversion/minimax.py | 50 +++++++++++++-------------------------- src/llama-model.cpp | 3 +-- src/models/dflash.cpp | 2 +- src/models/minimax-m3.cpp | 20 ++++------------ src/models/models.h | 2 +- 5 files changed, 25 insertions(+), 52 deletions(-) diff --git a/conversion/minimax.py b/conversion/minimax.py index 51fe20e3f3b1..c1a4bd2faf16 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -174,12 +174,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): class MiniMaxM3Model(MiniMaxM2Model): model_arch = gguf.MODEL_ARCH.MINIMAXM3 - # MiniMax-M3.1 checkpoints: quant_method "mxfp8" (attention / indexer / shared / dense MLP as - # F8_E4M3 + E8M0 weight_scale_inv per [1, 32]) with routed experts in compressed-tensors NVFP4 - # (weight_packed, weight_scale F8_E4M3 per 16, weight_global_scale, input_global_scale). - # --outtype auto keeps the experts NVFP4 (per-expert 1/global as ffn_*_exps.scale) and writes the - # MXFP8 tensors exactly as BF16 (or Q8_0 with --fp8-as-q8); an explicit --outtype dequantizes - # every tensor exactly through F32 first. + # --outtype auto keeps experts NVFP4 and writes MXFP8 tensors exactly as BF16 (or Q8_0 with --fp8-as-q8) _NVFP4_PROJ = {"w1": gguf.MODEL_TENSOR.FFN_GATE_EXP, "w2": gguf.MODEL_TENSOR.FFN_DOWN_EXP, "w3": gguf.MODEL_TENSOR.FFN_UP_EXP} def __init__(self, dir_model, ftype, *args, **kwargs): @@ -193,7 +188,7 @@ def tensor_force_quant(self, name, new_name, bid, n_dims): @staticmethod def _dequant_mxfp8(weight: Tensor, scale: Tensor) -> Tensor: - # exact: fp8 * 2^(e - 127) is representable in fp32 for every E8M0 byte except 255 (NaN) + # exact in fp32 for every E8M0 byte except 255 (NaN) rows, cols = weight.shape e = scale.view(torch.uint8).to(torch.float64) s = torch.exp2(e - 127.0).repeat_interleave(32, dim=1)[:, :cols] @@ -203,11 +198,11 @@ def _dequant_mxfp8(weight: Tensor, scale: Tensor) -> Tensor: @classmethod def _dequant_nvfp4_eager(cls, packed: Tensor, scale: Tensor, global_scale: Tensor) -> Tensor: - # exact: e2m1 * e4m3 / global evaluated in fp64, rounded once to fp32 + # exact: evaluated in fp64, rounded once to fp32 if cls._e2m1_pair_lut is None: v = [0.0, 0.5, 1.0, 1.5, 2.0, 3.0, 4.0, 6.0] v16 = v + [-x for x in v] - # byte b holds element 2i in the low nibble and 2i+1 in the high nibble + # low nibble = element 2i, high nibble = 2i+1 cls._e2m1_pair_lut = torch.tensor([[v16[b & 15], v16[b >> 4]] for b in range(256)], dtype=torch.float64) rows = packed.shape[0] vals = cls._e2m1_pair_lut[packed.view(torch.uint8).to(torch.int64)].reshape(rows, -1) @@ -219,13 +214,13 @@ def _dequant_nvfp4(cls, packed: Tensor, scale: Tensor, global_scale: Tensor) -> if not isinstance(packed, LazyTorchTensor): return cls._dequant_nvfp4_eager(packed, scale, global_scale) rows, packed_cols = packed.shape - # one lazy node around the eager decode (a lazy op chain costs ~10 s per expert tensor) + # a lazy op chain costs ~10 s per expert tensor return LazyTorchTensor(meta=LazyTorchTensor.meta_with_dtype_and_shape(torch.float32, (rows, packed_cols * 2)), args=(packed, scale, global_scale), func=cls._dequant_nvfp4_eager) @staticmethod def _nvfp4_pack_np(packed: Tensor, scale: Tensor) -> np.ndarray: - # same bytes as ModelBase._nvfp4_pack (compressed-tensors packing -> ggml block_nvfp4), in numpy + # same bytes as ModelBase._nvfp4_pack p = packed.contiguous().view(torch.uint8).numpy() d = scale.contiguous().view(torch.uint8).numpy() & 0x7F rows, packed_cols = p.shape @@ -237,9 +232,7 @@ def _nvfp4_pack_np(packed: Tensor, scale: Tensor) -> np.ndarray: raw = np.concatenate([d.reshape(rows, n_super, 4), qs.reshape(rows, n_super, 32)], axis=-1) return raw.reshape(rows, n_super * 36) - # stacked [n_expert, rows, cols] routed-expert tensors are written straight to gguf_writer, built - # lazily (only the tensor being written is resident, cf. _mxfp4_expert_tensor) and per expert in a - # thread pool: the torch / numpy kernels release the GIL, a serial loop takes hours on the real model + # experts are built lazily and in a thread pool: a serial loop takes hours on the real model _EXPERT_THREADS = 16 _DEQUANT_RAW = { gguf.LlamaFileType.ALL_F32: (gguf.GGMLQuantizationType.F32, torch.float32, 4), @@ -263,7 +256,6 @@ def one(fns) -> np.ndarray: if native: return self._nvfp4_pack_np(t[0], t[1]) w = self._dequant_nvfp4_eager(t[0], t[1], t[2]) - # torch's float -> bf16 / f16 casts round to nearest even, like gguf.quants return w.to(self._DEQUANT_RAW[self.ftype][1]).contiguous().view(torch.uint8).numpy() def load(fns_list) -> np.ndarray: @@ -310,8 +302,7 @@ def dequant_model(self): for aux in (base + "_scale", base + "_global_scale"): if aux not in self.model_tensors: raise KeyError(f"{name}: missing {aux}") - # input_global_scale is the trained W4A4 activation scale; llama.cpp computes its own - # activation scales (or none), so it is dropped + # trained activation scale; llama.cpp computes its own if (module + ".input_global_scale") in self.model_tensors: n_input_scale += 1 if float(LazyTorchTensor.to_eager(self.model_tensors[module + ".input_global_scale"]()).float()) == 1.0: @@ -354,7 +345,7 @@ def dequant_model(self): f"shape = {{{', '.join(str(n) for n in reversed(shape))}}}") self.gguf_writer.add_tensor(new_name, data, raw_dtype=qtype) if self._nvfp4_native: - # per-expert 1/global, applied to the matmul output by build_moe_ffn + # per-expert 1/global, applied by build_moe_ffn self.gguf_writer.add_tensor(new_name.replace(".weight", ".scale"), np.array(inv_global, dtype=np.float32)) for name in remove: @@ -372,9 +363,7 @@ def set_gguf_parameters(self): self.gguf_writer.add_expert_weights_scale(self.find_hparam(["routed_scaling_factor"])) self.gguf_writer.add_expert_weights_norm(True) - # transformers >= 5.15 resolves minimax_m3_vl natively and folds sparse_attention_config / - # moe_layer_freq into index_* / layer_types / mlp_layer_types, dropping the fields it does not - # consume (sparse_score_type, sparse_init_block, sparse_disable_index_value): read the raw file + # transformers >= 5.15 folds the sparse config and drops fields: read the raw file raw = self._raw_text_config() sac = raw.get("sparse_attention_config") or self.find_hparam(["sparse_attention_config"]) for flat, legacy in (("index_n_heads", "sparse_num_index_heads"), ("index_head_dim", "sparse_index_dim"), @@ -387,8 +376,7 @@ def set_gguf_parameters(self): raise NotImplementedError(f"sparse_score_type {sac['sparse_score_type']!r}") if int(sac.get("sparse_init_block", 0)) != 0: raise NotImplementedError(f"sparse_init_block {sac['sparse_init_block']}") - # which layers run MSA is carried by the presence of blk.N.indexer.* (loader side), so the - # indexer tensors must exist exactly on the layers sparse_attention_freq marks sparse + # the loader marks MSA layers by the presence of blk.N.indexer.* freq = sac.get("sparse_attention_freq") n_layer = self.find_hparam(["num_hidden_layers"]) if freq is not None: @@ -420,7 +408,6 @@ def set_gguf_parameters(self): def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): # Gemma-style (1 + w) RMSNorm: bake the +1 in so llama.cpp can use plain RMSNorm - # (in fp32 whatever the incoming dtype: adding 1.0 in F16 / BF16 rounds away the low bits of w) if name.endswith("norm.weight"): data_torch = data_torch.to(torch.float32) + 1.0 @@ -506,9 +493,7 @@ def _permute_vit_qk(self, t: "Tensor", new_name: str) -> "Tensor": @ModelBase.register("DSparkMiniMaxDraftModel") @ModelBase.example("MiniMaxAI/MiniMax-M3.1-Flash") class DSparkMiniMaxModel(DSparkModel): - # MiniMax-M3.1 DSpark draft (the `dspark/` folder of the checkpoint): a DFlash block drafter with - # MiniMax-M3 dense layers (Gemma norms, per-head Gemma QK norms, partial NeoX RoPE, SwiGLU-OAI MLP) - # plus the vanilla Markov head and the confidence head. Embeddings and lm_head come from the target. + # MiniMax-M3.1 DSpark draft; embeddings and lm_head come from the target model_arch = gguf.MODEL_ARCH.DFLASH _PREFIX = "language_model.model.dspark." @@ -516,8 +501,7 @@ class DSparkMiniMaxModel(DSparkModel): def __init__(self, dir_model, *args, **kwargs): hparams = kwargs.pop("hparams", None) if hparams is None: - # read config.json as is: transformers >= 5.15 resolves minimax_m3_vl to its own config class, - # which drops the DSpark keys and reports the default hidden_act + # raw config.json: transformers >= 5.15 drops the DSpark keys with open(dir_model / "config.json", "r", encoding="utf-8") as f: hparams = json.load(f) hparams = {**hparams, **hparams.get("text_config", {})} @@ -544,7 +528,7 @@ def __init__(self, dir_model, *args, **kwargs): if any(lt != "sliding_attention" for lt in layer_types): raise ValueError(f"MiniMax DSpark: expects uniform sliding_attention layers, got {layer_types}") - # the draft's FFN is the dense one (dspark_ffn_hidden_size), not the MoE expert width + # dense FFN width, not the MoE expert width hparams["dense_intermediate_size"] = hparams.get("dspark_ffn_hidden_size", hparams["dense_intermediate_size"]) for k in ("num_local_experts", "num_experts_per_tok", "n_shared_experts", "shared_intermediate_size", "moe_intermediate_size", "num_experts", "scoring_func", "use_routing_bias", "routed_scaling_factor"): @@ -578,7 +562,7 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca return super().filter_tensors((name, gen)) def dequant_model(self): - # MXFP8 (quant_method "mxfp8", block [1, 32]): F8_E4M3 weight * 2^(E8M0 - 127) per 32 inputs + # MXFP8: F8_E4M3 weight * 2^(E8M0 - 127) per 32 inputs quant_config = self.hparams.get("quantization_config") or {} if quant_config.get("quant_method") == "mxfp8": if list(quant_config.get("weight_block_size", [1, 32])) != [1, 32]: @@ -609,13 +593,13 @@ def set_gguf_parameters(self): self.gguf_writer.add_hidden_act("swigluoai") def tensor_force_quant(self, name, new_name, bid, n_dims): - # confidence head is a tiny F32 projection in the reference (params_dtype=float32) + # F32 in the reference if "confidence_head" in name or new_name.startswith("conf_proj"): return gguf.GGMLQuantizationType.F32 return super().tensor_force_quant(name, new_name, bid, n_dims) def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: - # Gemma-style (1 + w) RMSNorm, computed in fp32 by the reference: add 1 in fp32, not in bf16 + # add 1 in fp32: bf16 rounding loses the low bits of w if name.endswith("norm.weight"): data_torch = data_torch.to(torch.float32) + 1.0 yield from super().modify_tensors(data_torch, name, bid) diff --git a/src/llama-model.cpp b/src/llama-model.cpp index 758e2e005131..e1ca8caadedc 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -2824,8 +2824,7 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, } break; case LLM_ARCH_MINIMAX_M3: { - // sparse (MSA) layers carry an indexer key cache; a layer is sparse iff it has indexer - // tensors (MiniMax-M3: the MoE layers, MiniMax-M3.1: every layer) + // a layer is sparse (MSA) iff it has indexer tensors const auto * mm3 = static_cast(this); llama_kv_cache::layer_filter_cb filter_idx = [mm3](int32_t il) { return mm3->is_msa_layer(il); }; diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index b7e20e5e7160..7aa7ea6c994d 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -16,7 +16,7 @@ void llama_model_dflash::load_arch_hparams(llama_model_loader & ml) { if (hidden_act == "gelu" || hidden_act == "gelu_pytorch_tanh") { hparams.llm_ffn_op = LLM_FFN_GELU; } else if (hidden_act == "swigluoai") { - // MiniMax-M3 DSpark: clamped SwiGLU, alpha = 1.702, limit = 7.0 (checked by the converter) + // alpha and limit are checked by the converter hparams.llm_ffn_op = LLM_FFN_SWIGLU_OAI_MOE; } else if (hidden_act != "silu") { throw std::runtime_error("unsupported DFlash hidden activation: " + hidden_act); diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp index 42fa9f33af4d..522c55044f6e 100644 --- a/src/models/minimax-m3.cpp +++ b/src/models/minimax-m3.cpp @@ -64,9 +64,7 @@ void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) { layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); - // indexer: required on the MoE layers (MiniMax-M3); optional on the leading dense layers, which - // run MSA when it is present (MiniMax-M3.1 sets sparse_attention_freq on every layer). A layer - // is sparse iff it carries indexer tensors, see llama_model_minimax_m3::is_msa_layer() + // indexer optional on dense layers: a layer is sparse iff it carries indexer tensors { const int idx_flags = i < (int) hparams.n_layer_dense_lead ? TENSOR_NOT_REQUIRED : 0; layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", i), {n_embd, hparams.indexer_n_head * hparams.indexer_head_size}, idx_flags); @@ -102,8 +100,7 @@ void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) { } } -// tokens per MSA batch chunk so that per_token_elems * tokens stays below 2^31 elements -// (MSA_MAX_ELEMS overrides the cap, for testing the chunked path at small contexts) +// keeps per_token_elems * tokens below 2^31 (MSA_MAX_ELEMS overrides, for testing) static int64_t msa_tokens_per_chunk(int64_t per_token_elems, int64_t n_tokens) { static const int64_t max_elems = [] { const char * e = getenv("MSA_MAX_ELEMS"); @@ -324,8 +321,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * inp_out_ids = build_inp_out_ids(); - // DFlash/DSpark drafts read the residual stream entering layer il (index n_layer = the last layer's output, - // before the final norm); that one needs every row, so the output-row crop moves after the loop + // DFlash/DSpark drafts need every row, so the output-row crop moves after the loop const bool tap_last = (size_t) n_layer < cparams.embeddings_layer_inp.size() && cparams.embeddings_layer_inp[n_layer]; for (int il = 0; il < n_layer; ++il) { @@ -477,10 +473,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ cur = build_attn_msa_fa(Qcur, kfa, vfa, mfa, Gp, kq_scale, il); } else { - // batch: per stream, in token chunks. The largest intermediates (indexer scores - // [n_ps, Hd, n], block mask expanded to cells [n_kv, n, Hd]) grow with context x tokens; - // chunking keeps every one of them below 2^31 elements (int32 indexing in backend - // kernels, e.g. CUDA pool2d, overflows past that: 270K ctx at -ub 2048). + // chunk tokens: backend kernels index with int32 and overflow past 2^31 elements const int64_t n_chunk = msa_tokens_per_chunk(std::max(n_ps, n_kv)*Hd, n_tps); std::vector outs; for (int64_t st = 0; st < ns; ++st) { @@ -495,10 +488,8 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * v_s = ggml_view_4d(ctx0, v, D, HKV, n_kv, 1, v->nb[1], v->nb[2], v->nb[3], st*v->nb[3]); - // the indexer keys are gathered through the pos -> cell map once per stream ggml_tensor * ikp = ggml_get_rows(ctx0, ik_s, psl_s); // [n_idx_dim, n_ps] - // cache views with groups on ne[3]; ggml_tensor * kfa = ggml_permute(ctx0, k_s, 0, 3, 1, 2); ggml_tensor * vfa = ggml_permute(ctx0, v_s, 0, 3, 1, 2); @@ -510,7 +501,6 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ iq->nb[1], iq->nb[2], t0*iq->nb[2]); ggml_tensor * pm_s = ggml_view_3d(ctx0, msa->pos_mask, n_ps, 1, nt, msa->pos_mask->nb[1], msa->pos_mask->nb[1], t0*msa->pos_mask->nb[1]); - // F32 copy of this chunk's rows of the FA (f16) mask for the final mask add ggml_tensor * mf_s = ggml_cast(ctx0, ggml_view_3d(ctx0, msa_kqm, n_kv, nt, 1, msa_kqm->nb[1], msa_kqm->nb[3], st*msa_kqm->nb[3] + c0*msa_kqm->nb[1]), GGML_TYPE_F32); ggml_tensor * bias_s = ggml_view_3d(ctx0, msa->bias, nblk, 1, nt, @@ -518,7 +508,7 @@ llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * q_s = ggml_view_3d(ctx0, Qcur, D, n_head, nt, Qcur->nb[1], Qcur->nb[2], t0*Qcur->nb[2]); - // block scores, unscaled: only the top-k ordering matters + // unscaled: only the top-k ordering matters ggml_tensor * sc = ggml_mul_mat(ctx0, ikp, ggml_reshape_2d(ctx0, iq_s, n_idx_dim, Hd*nt)); // indexer scores run in F32 diff --git a/src/models/models.h b/src/models/models.h index 1fa7228411eb..30ed6ffd1e3a 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -2495,7 +2495,7 @@ struct llama_model_minimax_m3 : public llama_model_base { void load_arch_hparams(llama_model_loader & ml) override; void load_arch_tensors(llama_model_loader & ml) override; msa_params msa_p; - // a layer runs MSA iff it carries the indexer tensors (M3: MoE layers only; M3.1: every layer) + // a layer runs MSA iff it carries the indexer tensors bool is_msa_layer(int il) const; struct graph : public llm_graph_context { graph(const llama_model & model, const llm_graph_params & params);