diff --git a/conversion/__init__.py b/conversion/__init__.py index 02ea6385208a..398f00ab3a59 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -155,6 +155,8 @@ "MiniCPMForCausalLM": "minicpm", "MiniCPMV4_6ForConditionalGeneration": "minicpm", "MiniMaxM2ForCausalLM": "minimax", + "MiniMaxM3SparseForCausalLM": "minimax", + "MiniMaxM3SparseForConditionalGeneration": "minimax", "Ministral3ForCausalLM": "mistral3", "Mistral3ForConditionalGeneration": "mistral3", "MistralForCausalLM": "llama", @@ -282,6 +284,7 @@ "LlavaForConditionalGeneration": "llava", "MERaLiON2ForConditionalGeneration": "ultravox", "MiMoV2ForCausalLM": "mimo", + "MiniMaxM3SparseForConditionalGeneration": "minimax", "MiniCPMV4_6ForConditionalGeneration": "minicpm", "Mistral3ForConditionalGeneration": "llava", "NemotronH_Nano_VL_V2": "nemotron", diff --git a/conversion/base.py b/conversion/base.py index 0421aa4bc4d3..fd4d71458841 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1154,7 +1154,7 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca or "projector." in name or "pre_mm_projector_norm" in name \ or "image_newline" in name or "view_seperator" in name \ or "patch_embed" in name or "patch_embedding" in name \ - or "patch_merger." in name or "model.connector." in name: + or "patch_merger." in name or "patch_merge_mlp." in name or "model.connector." in name: return None return super().filter_tensors(item) @@ -1201,7 +1201,7 @@ def set_gguf_parameters(self): self.gguf_writer.add_embedding_length(n_embd) logger.info(f"gguf: embedding length = {n_embd}") - if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None: + if (n_ff := self.find_hparam(["prefix_dense_intermediate_size", "dense_intermediate_size", "intermediate_size", "n_inner", "hidden_dim"], optional=True)) is not None: self.gguf_writer.add_feed_forward_length(n_ff) logger.info(f"gguf: feed forward length = {n_ff}") diff --git a/conversion/minimax.py b/conversion/minimax.py index 4857775cbfb9..4c5c9876a49e 100644 --- a/conversion/minimax.py +++ b/conversion/minimax.py @@ -7,7 +7,7 @@ if TYPE_CHECKING: from torch import Tensor -from .base import ModelBase, TextModel, gguf +from .base import ModelBase, TextModel, MmprojModel, gguf @ModelBase.register("MiniMaxM2ForCausalLM") @@ -23,7 +23,7 @@ def set_gguf_parameters(self): def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): # merge expert weights - if 'experts' in name: + if "block_sparse_moe.experts." in name: n_experts = self.find_hparam(["num_local_experts", "num_experts"]) assert bid is not None @@ -52,3 +52,107 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None): return yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("MiniMaxM3SparseForCausalLM", "MiniMaxM3SparseForConditionalGeneration") +class MiniMaxM3Model(MiniMaxM2Model): + model_arch = gguf.MODEL_ARCH.MINIMAXM3 + def set_gguf_parameters(self): + super().set_gguf_parameters() + + self.gguf_writer.add_expert_shared_count(self.find_hparam(["n_shared_experts"])) + 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"]) + 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"]) + n_dense = 0 + for v in moe_layer_freq: + if v == 0: + n_dense += 1 + else: + break + self.gguf_writer.add_leading_dense_block_count(n_dense) + + 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 + if name.endswith("norm.weight"): + data_torch = data_torch + 1.0 + + yield from super().modify_tensors(data_torch, name, bid) + + +@ModelBase.register("MiniMaxM3SparseForConditionalGeneration", "MiniMaxM3VLForConditionalGeneration") +class MiniMaxM3VisionModel(MmprojModel): + @classmethod + def filter_tensors(cls, item): + name, gen = item + # keep only the vision-side tensors; text / mtp / sparse-index are dropped + if not name.startswith(("vision_tower.", "multi_modal_projector.", "patch_merge_mlp.")): + return None + return super().filter_tensors((name, gen)) + + def set_gguf_parameters(self): + super().set_gguf_parameters() + assert self.hparams_vision is not None + + self.gguf_writer.add_clip_projector_type(gguf.VisionProjectorType.MINIMAXM3) + self.gguf_writer.add_vision_use_gelu(True) + + # the ViT carries its own LayerNorm eps (text tower uses a different one) + self.gguf_writer.add_vision_attention_layernorm_eps( + self.hparams_vision.get("layer_norm_eps", 1e-5) + ) + + comp = self.hparams_vision.get("img_token_compression_config", {}) + merge_size = comp.get("spatial_merge_size", 2) + self.gguf_writer.add_vision_spatial_merge_size(int(merge_size)) + + def modify_tensors(self, data_torch, name, bid): + assert self.hparams_vision is not None + + # Conv3d patch embed -> Conv2d slices + if name == "vision_tower.vision_model.embeddings.patch_embedding.weight": + if data_torch.ndim != 5: + raise ValueError(f"unexpected patch_embedding rank {data_torch.ndim} for {name}") + kt = data_torch.shape[2] + base = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.V_ENC_EMBD_PATCH] + for t in range(kt): + suffix = ".weight" if t == 0 else f".weight.{t}" + yield (base + suffix, data_torch[:, :, t, ...]) + return + + # Permute ViT q/k. HF [Ta Ha Wa | Tb Hb Wb | pad] reorder to [Ta Tb | Ha Hb | Wa Wb | pad]. + for new_name, tensor in super().modify_tensors(data_torch, name, bid): + if ".attn_q." in new_name or ".attn_k." in new_name: + tensor = self._permute_vit_qk(tensor, new_name) + yield new_name, tensor + + def _permute_vit_qk(self, t: "Tensor", new_name: str) -> "Tensor": + n_head = self.hparams_vision["num_attention_heads"] + d_head = t.shape[0] // n_head + axis_dim = 2 * ((2 * (d_head // 2) // 3) // 2) + ah = axis_dim // 2 + half = 3 * ah + perm = (list(range(0, ah)) + list(range(half, half + ah)) + + list(range(ah, 2 * ah)) + list(range(half + ah, half + 2*ah)) + + list(range(2 * ah, 3 * ah)) + list(range(half + 2*ah, half + 3*ah)) + + list(range(2 * half, d_head))) + + assert axis_dim % 2 == 0 + assert 3 * axis_dim <= d_head + assert len(perm) == d_head + assert sorted(perm) == list(range(d_head)), "perm is not a bijection of d_head" + assert t.shape[0] == n_head * d_head, f"{new_name}: {t.shape[0]} != {n_head}*{d_head}" + assert d_head == 80 + + idx = torch.tensor(perm, dtype=torch.long) + 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) diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index cd4cdef8991f..b8d9519907a1 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -200,6 +200,8 @@ class Indexer: HEAD_COUNT = "{arch}.attention.indexer.head_count" KEY_LENGTH = "{arch}.attention.indexer.key_length" TOP_K = "{arch}.attention.indexer.top_k" + BLOCK_SIZE = "{arch}.attention.indexer.block_size" #MSA + LOCAL_BLOCKS = "{arch}.attention.indexer.local_blocks" #MSA class HyperConnection: COUNT = "{arch}.hyper_connection.count" @@ -525,6 +527,7 @@ class MODEL_ARCH(IntEnum): APERTUS = auto() COGVLM = auto() MINIMAXM2 = auto() + MINIMAXM3 = auto() RND1 = auto() PANGU_EMBED = auto() MISTRAL3 = auto() @@ -771,6 +774,9 @@ class MODEL_TENSOR(IntEnum): INDEXER_PROJ = auto() INDEXER_ATTN_K = auto() INDEXER_ATTN_Q_B = auto() + INDEXER_Q_PROJ = auto() + INDEXER_K_PROJ = auto() + INDEXER_Q_NORM = auto() INDEXER_COMPRESSOR_WKV = auto() INDEXER_COMPRESSOR_WGATE = auto() INDEXER_COMPRESSOR_APE = auto() @@ -848,6 +854,8 @@ class MODEL_TENSOR(IntEnum): V_MM_UP = auto() # cogvlm V_MM_DOWN = auto() # cogvlm V_MM_GATE = auto() # cogvlm + V_MM_MERGE_FC1 = auto() # minimax-m3 (patch-merge MLP) + V_MM_MERGE_FC2 = auto() # minimax-m3 (patch-merge MLP) V_TOK_BOI = auto() # cogvlm V_TOK_EOI = auto() # cogvlm V_TOK_IMG_BEGIN = auto() # hunyuanvl @@ -1105,6 +1113,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.GROVEMOE: "grovemoe", MODEL_ARCH.APERTUS: "apertus", MODEL_ARCH.MINIMAXM2: "minimax-m2", + MODEL_ARCH.MINIMAXM3: "minimax-m3", MODEL_ARCH.COGVLM: "cogvlm", MODEL_ARCH.RND1: "rnd1", MODEL_ARCH.PANGU_EMBED: "pangu-embedded", @@ -1350,6 +1359,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.INDEXER_PROJ: "blk.{bid}.indexer.proj", MODEL_TENSOR.INDEXER_ATTN_K: "blk.{bid}.indexer.attn_k", MODEL_TENSOR.INDEXER_ATTN_Q_B: "blk.{bid}.indexer.attn_q_b", + 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_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", @@ -1426,6 +1438,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.V_MM_UP: "mm.up", MODEL_TENSOR.V_MM_DOWN: "mm.down", MODEL_TENSOR.V_MM_GATE: "mm.gate", + MODEL_TENSOR.V_MM_MERGE_FC1: "mm.merge.fc1", + MODEL_TENSOR.V_MM_MERGE_FC2: "mm.merge.fc2", MODEL_TENSOR.V_TOK_BOI: "v.boi", MODEL_TENSOR.V_TOK_EOI: "v.eoi", MODEL_TENSOR.V_MM_PRE_NORM: "mm.pre_norm", @@ -1622,6 +1636,8 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.V_RESMPL_QUERY, MODEL_TENSOR.V_TOK_EMBD_IMG_BREAK, MODEL_TENSOR.V_MM_PATCH_MERGER, + MODEL_TENSOR.V_MM_MERGE_FC1, + MODEL_TENSOR.V_MM_MERGE_FC2, MODEL_TENSOR.V_DS_NORM, MODEL_TENSOR.V_DS_FC1, MODEL_TENSOR.V_DS_FC2, @@ -4102,6 +4118,34 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_UP_EXP, MODEL_TENSOR.FFN_EXP_PROBS_B, ], + MODEL_ARCH.MINIMAXM3: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_EXP_PROBS_B, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + MODEL_TENSOR.INDEXER_Q_PROJ, + MODEL_TENSOR.INDEXER_K_PROJ, + MODEL_TENSOR.INDEXER_Q_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + ], MODEL_ARCH.COGVLM: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, @@ -4670,6 +4714,7 @@ class VisionProjectorType: YOUTUVL = "youtuvl" NEMOTRON_V2_VL = "nemotron_v2_vl" HUNYUANVL = "hunyuanvl" + MINIMAXM3 = "minimax_m3" MINICPMV4_6 = "minicpmv4_6" GRANITE_SPEECH = "granite_speech" # audio MIMOVL = "mimovl" diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 1e277f0687c5..12c2fa31fa2e 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -793,6 +793,12 @@ def add_indexer_key_length(self, length: int) -> None: def add_indexer_top_k(self, top_k: int) -> None: self.add_uint32(Keys.Attention.Indexer.TOP_K.format(arch=self.arch), top_k) + def add_indexer_block_size(self, block_size: int) -> None: + self.add_uint32(Keys.Attention.Indexer.BLOCK_SIZE.format(arch=self.arch), block_size) + + def add_indexer_local_blocks(self, local_blocks: int) -> None: + self.add_uint32(Keys.Attention.Indexer.LOCAL_BLOCKS.format(arch=self.arch), local_blocks) + def add_max_alibi_bias(self, bias: float) -> None: self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 9efb36f8a447..cfb810bfc7a9 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1263,7 +1263,8 @@ class TensorNameMap: ), MODEL_TENSOR.INDEXER_K_NORM: ( - "model.layers.{bid}.self_attn.indexer.k_norm", # DSA + "model.layers.{bid}.self_attn.indexer.k_norm", # DSA + "model.layers.{bid}.self_attn.index_k_norm", # MSA ), MODEL_TENSOR.INDEXER_PROJ: ( @@ -1278,6 +1279,19 @@ class TensorNameMap: "model.layers.{bid}.self_attn.indexer.wq_b", # DSA ), + MODEL_TENSOR.INDEXER_Q_PROJ: ( + "model.layers.{bid}.self_attn.index_q_proj", # MSA + ), + + MODEL_TENSOR.INDEXER_K_PROJ: ( + "model.layers.{bid}.self_attn.index_k_proj", # MSA + ), + + MODEL_TENSOR.INDEXER_Q_NORM: ( + "model.layers.{bid}.self_attn.index_q_norm", # MSA + ), + + ############################################################################ # TODO: these do not belong to block_mappings_cfg - move them to mappings_cfg MODEL_TENSOR.ENC_OUTPUT_NORM: ( @@ -1824,6 +1838,14 @@ class TensorNameMap: "visual.downsample", # glm4v ), + MODEL_TENSOR.V_MM_MERGE_FC1: ( + "patch_merge_mlp.linear_1", # minimax-m3 + ), + + MODEL_TENSOR.V_MM_MERGE_FC2: ( + "patch_merge_mlp.linear_2", # minimax-m3 + ), + MODEL_TENSOR.V_DS_NORM: ( "model.visual.deepstack_merger_list.{bid}.norm", # deepstack in qwen3vl ), diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index b890e66fcf6e..2baa67abd401 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -125,6 +125,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GROVEMOE, "grovemoe" }, { LLM_ARCH_APERTUS, "apertus" }, { LLM_ARCH_MINIMAX_M2, "minimax-m2" }, + { LLM_ARCH_MINIMAX_M3, "minimax-m3" }, { LLM_ARCH_COGVLM, "cogvlm" }, { LLM_ARCH_RND1, "rnd1" }, { LLM_ARCH_PANGU_EMBED, "pangu-embedded" }, @@ -251,6 +252,8 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, "%s.attention.indexer.head_count" }, { LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, "%s.attention.indexer.key_length" }, { LLM_KV_ATTENTION_INDEXER_TOP_K, "%s.attention.indexer.top_k" }, + { LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, "%s.attention.indexer.block_size" }, + { LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, "%s.attention.indexer.local_blocks" }, { LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, "%s.attention.output_group_count" }, { LLM_KV_ATTENTION_OUTPUT_LORA_RANK, "%s.attention.output_lora_rank" }, { LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, "%s.attention.compress_rope_freq_base" }, @@ -594,6 +597,9 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" }, { LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" }, { LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" }, + { 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_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" }, @@ -829,6 +835,9 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_INDEXER_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_K, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_INDEXER_ATTN_Q_B, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {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_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_ADD}}, @@ -998,6 +1007,7 @@ bool llm_arch_supports_sm_tensor(const llm_arch & arch) { case LLM_ARCH_LFM2: case LLM_ARCH_LFM2MOE: case LLM_ARCH_MINIMAX_M2: + case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_MISTRAL4: case LLM_ARCH_KIMI_LINEAR: return false; diff --git a/src/llama-arch.h b/src/llama-arch.h index a4f5091e7170..2b28337775fd 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -144,6 +144,7 @@ enum llm_arch { LLM_ARCH_TALKIE, LLM_ARCH_MELLUM, LLM_ARCH_EAGLE3, + LLM_ARCH_MINIMAX_M3, LLM_ARCH_DFLASH, LLM_ARCH_UNKNOWN, }; @@ -256,6 +257,8 @@ enum llm_kv { LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, LLM_KV_ATTENTION_INDEXER_TOP_K, + LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, + LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, LLM_KV_ATTENTION_OUTPUT_GROUP_COUNT, LLM_KV_ATTENTION_OUTPUT_LORA_RANK, LLM_KV_ATTENTION_COMPRESS_ROPE_FREQ_BASE, @@ -594,6 +597,9 @@ enum llm_tensor { LLM_TENSOR_INDEXER_PROJ, LLM_TENSOR_INDEXER_ATTN_K, LLM_TENSOR_INDEXER_ATTN_Q_B, + LLM_TENSOR_INDEXER_Q_PROJ, + LLM_TENSOR_INDEXER_K_PROJ, + LLM_TENSOR_INDEXER_Q_NORM, LLM_TENSOR_INDEXER_COMPRESSOR_WKV, LLM_TENSOR_INDEXER_COMPRESSOR_WGATE, LLM_TENSOR_INDEXER_COMPRESSOR_APE, diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 0465430df43a..9fbfdc0db27c 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2325,7 +2325,8 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_KIMI_LINEAR || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || - model.arch == LLM_ARCH_DEEPSEEK4) { + model.arch == LLM_ARCH_DEEPSEEK4 || + model.arch == LLM_ARCH_MINIMAX_M3) { return std::max(n_tokens * 40, 32u * model.n_tensors()); } uint32_t res = std::max(1024u, 8u*model.n_tensors()); diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 4c86e43c1f74..3a1de5cd22cf 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -1700,6 +1700,17 @@ ggml_tensor * llm_graph_context::build_ffn( cur = ggml_swiglu(ctx0, cur); cb(cur, "ffn_swiglu", il); } break; + case LLM_FFN_SWIGLU_OAI_MOE: + if (gate && type_gate == LLM_FFN_PAR) { + //Same constants as LLM_FFN_SWIGLU_OAI_MOE + const float alpha = 1.702f; + const float limit = 7.0f; + cur = ggml_swiglu_oai(ctx0, cur, tmp, alpha, limit); + cb(cur, "ffn_swiglu_oai", il); + type_gate = LLM_FFN_SEQ; + } else { + GGML_ABORT("LLM_FFN_SWIGLU_OAI requires a parallel gate"); + } break; case LLM_FFN_GEGLU: { cur = ggml_geglu(ctx0, cur); @@ -2629,7 +2640,8 @@ ggml_tensor * llm_graph_context::build_attn( ggml_tensor * sinks, ggml_tensor * v_mla, // TODO: remove float kq_scale, - int il) const { + int il, + ggml_tensor * kq_mask_override) const { GGML_ASSERT(v_mla == nullptr); if (inp->self_k_rot) { @@ -2659,7 +2671,7 @@ ggml_tensor * llm_graph_context::build_attn( ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); } - const auto & kq_mask = inp->get_kq_mask(); + ggml_tensor * kq_mask = kq_mask_override ? kq_mask_override : inp->get_kq_mask(); ggml_tensor * q = q_cur; ggml_tensor * k = mctx_cur->get_k(ctx0, il); diff --git a/src/llama-graph.h b/src/llama-graph.h index 4b5b75c632ab..e067aec47d61 100644 --- a/src/llama-graph.h +++ b/src/llama-graph.h @@ -1085,7 +1085,8 @@ struct llm_graph_context { ggml_tensor * sinks, // [n_head_q] ggml_tensor * v_mla, // [n_embd_head_v_mla, n_embd_head_v, n_head_v] // TODO: remove float kq_scale, - int il) const; + int il, + ggml_tensor * kq_mask_override = nullptr) const; llm_graph_input_attn_k * build_attn_inp_k() const; diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 9d0683d2fec4..0219fc65ef3e 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -180,6 +180,16 @@ uint32_t llama_hparams::n_embd_v_gqa_max() const { return val; } +uint32_t llama_hparams::n_embd_k_idx(uint32_t il) const { + if (indexer_head_size == 0) { + return 0; // arch without MSA + } + if (il < n_layer_dense_lead) { + return 0; // leading dense layers carry no indexer + } + return indexer_head_size; // 128 +} + uint32_t llama_hparams::n_embd_r() const { if (wkv_head_size != 0) { // for RWKV models diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 8be5f28f39e6..3174bd8e5026 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -226,6 +226,9 @@ struct llama_hparams { uint32_t indexer_n_head = 0; uint32_t indexer_head_size = 0; uint32_t indexer_top_k = 0; + // MSA + uint32_t indexer_block_size = 0; + uint32_t indexer_local_blocks = 0; // DeepSeek-V4 uint32_t dsv4_o_group_count = 0; @@ -344,6 +347,9 @@ struct llama_hparams { uint32_t n_embd_k_gqa_max() const; uint32_t n_embd_v_gqa_max() const; + // dimension of the single-head MSA indexer key stream + uint32_t n_embd_k_idx(uint32_t il = 0) const; + // dimension of the rolling state embeddings // corresponds to Mamba's conv_states size or RWKV's token_shift states size uint32_t n_embd_r() const; diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 12bf5c37914d..7e50fc4d9b8c 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -128,7 +128,7 @@ llama_kv_cache::llama_kv_cache( auto it = ctx_map.find(buft); if (it == ctx_map.end()) { ggml_init_params params = { - /*.mem_size =*/ size_t(2u*(1 + n_stream)*n_layer*ggml_tensor_overhead()), + /*.mem_size =*/ size_t(3u*(1 + n_stream)*n_layer*ggml_tensor_overhead()), /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; @@ -258,9 +258,24 @@ llama_kv_cache::llama_kv_cache( v_stream.push_back(has_v ? ggml_view_2d(ctx, v, n_embd_v_gqa, kv_size, v->nb[1], s*v->nb[2]) : nullptr); } + const uint32_t n_embd_k_idx = hparams.n_embd_k_idx(il); + ggml_tensor * k_idx = n_embd_k_idx > 0 + ? ggml_new_tensor_3d(ctx, GGML_TYPE_F32, n_embd_k_idx, kv_size, n_stream) + : nullptr; + if (k_idx) { + ggml_format_name(k_idx, "cache_k_idx_l%d", il); + } + + std::vector k_idx_stream; + for (uint32_t s = 0; s < n_stream; ++s) { + k_idx_stream.push_back(k_idx + ? ggml_view_2d(ctx, k_idx, n_embd_k_idx, kv_size, k_idx->nb[1], s*k_idx->nb[2]) + : nullptr); + } + map_layer_ids[il] = layers.size(); - layers.push_back({ il, k, v, k_stream, v_stream, }); + layers.push_back({ il, k, v, k_idx, k_stream, v_stream, k_idx_stream }); } if (reuse) { @@ -311,11 +326,21 @@ llama_kv_cache::llama_kv_cache( { const size_t memory_size_k = size_k_bytes(); const size_t memory_size_v = size_v_bytes(); + const size_t memory_size_k_idx = size_k_idx_bytes(); + const size_t memory_size_total = memory_size_k + memory_size_v + memory_size_k_idx; - LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__, - (float)(memory_size_k + memory_size_v) / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream, + if (memory_size_k_idx > 0) { + LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB, K_idx (%s): %7.2f MiB\n", __func__, + (float)memory_size_total / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream, + ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f), + ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f), + ggml_type_name(GGML_TYPE_F32), (float)memory_size_k_idx / (1024.0f * 1024.0f)); + } else { + LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u/%u seqs), K (%s): %7.2f MiB, V (%s): %7.2f MiB\n", __func__, + (float)memory_size_total / (1024.0f * 1024.0f), kv_size, (int) layers.size(), n_seq_max, n_stream, ggml_type_name(type_k), (float)memory_size_k / (1024.0f * 1024.0f), ggml_type_name(type_v), (float)memory_size_v / (1024.0f * 1024.0f)); + } } // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] @@ -862,6 +887,10 @@ bool llama_kv_cache::update(llama_context * lctx, bool do_shift, const stream_co if (layer.v_stream[ssrc]) { ggml_backend_tensor_copy(layer.v_stream[ssrc], layer.v_stream[sdst]); } + if (layer.k_idx_stream[ssrc]) { + GGML_ASSERT(layer.k_idx_stream[sdst]); + ggml_backend_tensor_copy(layer.k_idx_stream[ssrc], layer.k_idx_stream[sdst]); + } } } } @@ -1192,6 +1221,12 @@ bool llama_kv_cache::get_can_shift() const { if (hparams.n_pos_per_embd() > 1) { return false; } + // shifting would leave k_idx stale + for (const auto & layer : layers) { + if (layer.k_idx) { + return false; + } + } return true; } @@ -1308,6 +1343,23 @@ ggml_tensor * llama_kv_cache::get_v(ggml_context * ctx, int32_t il, uint32_t n_k ggml_row_size(v->type, kv_size*n_embd_v_gqa)*sinfo.s0); } +ggml_tensor * llama_kv_cache::get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const { + const int32_t ikv = map_layer_ids.at(il); + auto * k_idx = layers[ikv].k_idx; + GGML_ASSERT(k_idx); + + const uint64_t kv_size = get_size(); + const int64_t n_idx = k_idx->ne[0]; // 128 + const uint32_t ns = sinfo.s1 - sinfo.s0 + 1; + + return ggml_view_4d(ctx, k_idx, + n_idx, 1, n_kv, ns, + ggml_row_size(k_idx->type, n_idx), // nb1 (single head) + ggml_row_size(k_idx->type, n_idx), // nb2 (per cell) + ggml_row_size(k_idx->type, n_idx*kv_size), // nb3 (per stream) + ggml_row_size(k_idx->type, n_idx*kv_size)*sinfo.s0); +} + ggml_tensor * llama_kv_cache::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const { GGML_UNUSED(sinfo); @@ -1409,6 +1461,28 @@ ggml_tensor * llama_kv_cache::build_input_k_idxs(ggml_context * ctx, const llama return k_idxs; } +ggml_tensor * llama_kv_cache::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const { + GGML_UNUSED(sinfo); + const int32_t ikv = map_layer_ids.at(il); + ggml_tensor * k_idx = layers[ikv].k_idx; + GGML_ASSERT(k_idx && "cpy_k_idx on a layer with no indexer cache"); + + const int64_t n_embd_head = k_idx_cur->ne[0]; // 128 + const int64_t n_head = k_idx_cur->ne[1]; // 1 + const int64_t n_tokens = k_idx_cur->ne[2]; + const int64_t n_embd_gqa = n_embd_head*n_head; // 128 + + GGML_ASSERT(ggml_row_size(k_idx_cur->type, n_embd_head) == k_idx_cur->nb[1]); + k_idx_cur = ggml_view_2d(ctx, k_idx_cur, n_embd_gqa, n_tokens, k_idx_cur->nb[2], 0); + + const int64_t n_stream = k_idx->ne[2]; + if (n_stream > 1) { + const int64_t kv_size = get_size(); + k_idx = ggml_reshape_2d(ctx, k_idx, n_embd_gqa, kv_size*n_stream); + } + return ggml_set_rows(ctx, k_idx, k_idx_cur, k_idxs); // same k_idxs as the K store +} + ggml_tensor * llama_kv_cache::build_input_v_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { const uint32_t n_tokens = ubatch.n_tokens; @@ -1843,6 +1917,18 @@ size_t llama_kv_cache::size_v_bytes() const { return size_v_bytes; } +size_t llama_kv_cache::size_k_idx_bytes() const { + size_t size_k_idx_bytes = 0; + + for (const auto & layer : layers) { + if (layer.k_idx) { + size_k_idx_bytes += ggml_nbytes(layer.k_idx); + } + } + + return size_k_idx_bytes; +} + ggml_tensor * llama_kv_cache::build_rope_shift( const llama_cparams & cparams, ggml_context * ctx, @@ -2150,6 +2236,36 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t } } + if (size_k_idx_bytes() > 0) { + const uint32_t has_k_idx_u32 = 1; + io.write(&has_k_idx_u32, sizeof(has_k_idx_u32)); + + for (const auto & layer : layers) { + const uint32_t layer_has_k_idx = layer.k_idx ? 1 : 0; + io.write(&layer_has_k_idx, sizeof(layer_has_k_idx)); + + if (!layer_has_k_idx) { + continue; + } + + GGML_ASSERT(layer.k_idx_stream[cr.strm]); + + const int32_t k_idx_type_i = (int32_t) layer.k_idx->type; + io.write(&k_idx_type_i, sizeof(k_idx_type_i)); + + const uint64_t k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]); + io.write(&k_idx_size_row, sizeof(k_idx_size_row)); + + for (const auto & range : cr.data) { + const size_t range_size = range.second - range.first; + const size_t buf_size = range_size * k_idx_size_row; + const size_t offset = range.first * k_idx_size_row; + + io.write_tensor(layer.k_idx_stream[cr.strm], offset, buf_size); + } + } + } + if (!v_trans) { for (const auto & layer : layers) { const uint32_t il = layer.il; @@ -2398,6 +2514,68 @@ bool llama_kv_cache::state_read_data(llama_io_read_i & io, uint32_t strm, uint32 } } + if (size_k_idx_bytes() > 0) { + uint32_t has_k_idx_u32 = 0; + io.read(&has_k_idx_u32, sizeof(has_k_idx_u32)); + + if (has_k_idx_u32 != 1) { + LLAMA_LOG_ERROR("%s: missing k_idx data in KV cache state\n", __func__); + return false; + } + + for (const auto & layer : layers) { + uint32_t layer_has_k_idx = 0; + io.read(&layer_has_k_idx, sizeof(layer_has_k_idx)); + + const uint32_t expected_layer_has_k_idx = layer.k_idx ? 1 : 0; + + if (layer_has_k_idx != expected_layer_has_k_idx) { + LLAMA_LOG_ERROR( + "%s: mismatched k_idx state for layer: got %u, expected %u\n", + __func__, layer_has_k_idx, expected_layer_has_k_idx); + return false; + } + + if (!layer_has_k_idx) { + continue; + } + + GGML_ASSERT(layer.k_idx_stream[strm]); + + int32_t k_idx_type_i = -1; + io.read(&k_idx_type_i, sizeof(k_idx_type_i)); + + if (k_idx_type_i != (int32_t) layer.k_idx->type) { + LLAMA_LOG_ERROR( + "%s: mismatched k_idx type: got %d, expected %d\n", + __func__, k_idx_type_i, (int32_t) layer.k_idx->type); + return false; + } + + uint64_t k_idx_size_row = 0; + io.read(&k_idx_size_row, sizeof(k_idx_size_row)); + + const uint64_t expected_k_idx_size_row = ggml_row_size(layer.k_idx->type, layer.k_idx->ne[0]); + + if (k_idx_size_row != expected_k_idx_size_row) { + LLAMA_LOG_ERROR( + "%s: mismatched k_idx row size: got %zu, expected %zu\n", + __func__, (size_t) k_idx_size_row, (size_t) expected_k_idx_size_row); + return false; + } + + if (cell_count) { + if (sinfo.is_contiguous()) { + io.read_tensor(layer.k_idx_stream[strm], sinfo.head() * k_idx_size_row, cell_count * k_idx_size_row); + } else { + for (uint32_t i = 0; i < cell_count; ++i) { + io.read_tensor(layer.k_idx_stream[strm], sinfo.idxs[0][i] * k_idx_size_row, k_idx_size_row); + } + } + } + } + } + if (!this->v_trans) { for (const auto & layer : layers) { const uint32_t il = layer.il; @@ -2599,6 +2777,10 @@ ggml_tensor * llama_kv_cache_context::get_v(ggml_context * ctx, int32_t il) cons return kv->get_v(ctx, il, n_kv, sinfos[i_cur]); } +ggml_tensor * llama_kv_cache_context::get_k_idx(ggml_context * ctx, int32_t il) const { + return kv->get_k_idx(ctx, il, n_kv, sinfos[i_cur]); +} + ggml_tensor * llama_kv_cache_context::cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const { return kv->cpy_k(ctx, k_cur, k_idxs, il, sinfos[i_cur]); } @@ -2607,6 +2789,10 @@ ggml_tensor * llama_kv_cache_context::cpy_v(ggml_context * ctx, ggml_tensor * v_ return kv->cpy_v(ctx, v_cur, v_idxs, il, sinfos[i_cur]); } +ggml_tensor * llama_kv_cache_context::cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const { + return kv->cpy_k_idx(ctx, k_idx_cur, k_idxs, il, sinfos[i_cur]); +} + ggml_tensor * llama_kv_cache_context::build_input_k_idxs(ggml_context * ctx, const llama_ubatch & ubatch) const { return kv->build_input_k_idxs(ctx, ubatch); } diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 531d99dbdec1..e5afa20101be 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -173,10 +173,12 @@ class llama_kv_cache : public llama_memory_i { // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; + ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il, uint32_t n_kv, const slot_info & sinfo) const; // store k_cur and v_cur in the cache based on the provided head location ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const; ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il, const slot_info & sinfo) const; + ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il, const slot_info & sinfo) const; // // preparation API @@ -228,9 +230,11 @@ class llama_kv_cache : public llama_memory_i { ggml_tensor * k; ggml_tensor * v; + ggml_tensor * k_idx; // MSA single-head indexer keys, F32 std::vector k_stream; std::vector v_stream; + std::vector k_idx_stream; }; bool v_trans = true; // the value tensor is transposed @@ -291,6 +295,7 @@ class llama_kv_cache : public llama_memory_i { size_t size_k_bytes() const; size_t size_v_bytes() const; + size_t size_k_idx_bytes() const; ggml_tensor * build_rope_shift( const llama_cparams & cparams, @@ -370,6 +375,7 @@ class llama_kv_cache_context : public llama_memory_context_i { // get views of the current state of the cache ggml_tensor * get_k(ggml_context * ctx, int32_t il) const; ggml_tensor * get_v(ggml_context * ctx, int32_t il) const; + ggml_tensor * get_k_idx(ggml_context * ctx, int32_t il) const; // store k_cur and v_cur in the cache based on the provided head location // note: the heads in k_cur and v_cur should be laid out contiguously in memory @@ -379,6 +385,7 @@ class llama_kv_cache_context : public llama_memory_context_i { // - v_idxs [n_tokens] or [n_tokens*n_embd_v_gqa] depending if V cache is transposed ggml_tensor * cpy_k(ggml_context * ctx, ggml_tensor * k_cur, ggml_tensor * k_idxs, int32_t il) const; ggml_tensor * cpy_v(ggml_context * ctx, ggml_tensor * v_cur, ggml_tensor * v_idxs, int32_t il) const; + ggml_tensor * cpy_k_idx(ggml_context * ctx, ggml_tensor * k_idx_cur, ggml_tensor * k_idxs, int32_t il) const; // create destination indices for each head of the current batch for where it would be written in the KV cache // the indices address the global KV cache (not per stream) - this is not relevant for the user of this API, but diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d58ebac28b9b..3db980289e3a 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -280,6 +280,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_apertus(params); case LLM_ARCH_MINIMAX_M2: return new llama_model_minimax_m2(params); + case LLM_ARCH_MINIMAX_M3: + return new llama_model_minimax_m3(params); case LLM_ARCH_COGVLM: return new llama_model_cogvlm(params); case LLM_ARCH_PANGU_EMBED: @@ -803,6 +805,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_122B_A10B: return "122B.A10B"; case LLM_TYPE_196B_A11B: return "196B.A11B"; case LLM_TYPE_230B_A10B: return "230B.A10B"; + case LLM_TYPE_428B_A23B: return "428B.A23B"; case LLM_TYPE_235B_A22B: return "235B.A22B"; case LLM_TYPE_300B_A47B: return "300B.A47B"; case LLM_TYPE_310B_A15B: return "310B.A15B"; @@ -2515,6 +2518,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GROVEMOE: case LLM_ARCH_APERTUS: case LLM_ARCH_MINIMAX_M2: + case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_COGVLM: case LLM_ARCH_PANGU_EMBED: case LLM_ARCH_AFMOE: diff --git a/src/llama-model.h b/src/llama-model.h index 4800d2928c52..c2f75131f9d7 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -134,6 +134,7 @@ enum llm_type { LLM_TYPE_122B_A10B, // Qwen3.5 LLM_TYPE_196B_A11B, // Step3.5-Flash LLM_TYPE_230B_A10B, // Minimax M2 + LLM_TYPE_428B_A23B, // Minimax M3 LLM_TYPE_235B_A22B, LLM_TYPE_300B_A47B, // Ernie MoE big LLM_TYPE_310B_A15B, // /MiMo-V2-Flash @@ -515,6 +516,12 @@ struct llama_layer { struct ggml_tensor * indexer_attn_k = nullptr; struct ggml_tensor * indexer_attn_q_b = nullptr; // note: for lora a/b, not bias + // MSA + struct ggml_tensor * index_q_proj = nullptr; + struct ggml_tensor * index_k_proj = nullptr; + struct ggml_tensor * index_q_norm = nullptr; + struct ggml_tensor * index_k_norm = nullptr; + // gemma4 layer output scale, reused for talkie embedding skip scale struct ggml_tensor * out_scale = nullptr; diff --git a/src/models/minimax-m3.cpp b/src/models/minimax-m3.cpp new file mode 100644 index 000000000000..df36a2e4e0e9 --- /dev/null +++ b/src/models/minimax-m3.cpp @@ -0,0 +1,625 @@ +#include "models.h" +#include "llama-kv-cache.h" +#include +#include +#include +#include + +// MiniMax-M3: MiniMax-M2 style GQA (per-head QK-norm, partial rotary) with +// DeepSeek-V3 leading-dense + routed/shared experts (sigmoid gating, routed scaling) and +// swigluoai activation + Minimax Sparse Attention. MTP is dropped. + +void llama_model_minimax_m3::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp); + ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); + 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); + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, hparams.indexer_block_size, false); + ml.get_key(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, hparams.indexer_local_blocks, false); + msa_p = { (int) hparams.indexer_block_size, (int) hparams.indexer_top_k, (int) hparams.indexer_local_blocks }; + + switch (hparams.n_layer()) { + case 60: type = LLM_TYPE_428B_A23B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_minimax_m3::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + const int64_t n_expert_shared = hparams.n_expert_shared; + const int64_t n_ff_exp = hparams.n_ff_exp; + + 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]; + + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_gqa, n_embd_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), { n_embd_head_k * n_head, n_embd }, 0); + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + // per-head QK-norm: a single head_dim vector applied to every head + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), {n_embd_head_k}, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), {n_embd_head_k}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + 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); + 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 { + // routed experts + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_exp_probs_b = create_tensor(tn(LLM_TENSOR_FFN_EXP_PROBS_B, "bias", i), {n_expert}, 0); + 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 + 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); + + // 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); + } + } +} + +std::unique_ptr llama_model_minimax_m3::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// Per-step local-block force for the MSA decode path. bias[b, i] = +BIG iff block b is one of +// query i's local blocks (L = pos/blk, plus local-1 neighbours). Added to block scores before +// top_k so the local window is always selected, matching msa_select_*'s bs[qb-l]=INF. +class llm_graph_input_msa_local : public llm_graph_input_i { +public: + llm_graph_input_msa_local(int blk, int local, int64_t nblk) : blk(blk), local(local), nblk(nblk) {} + void set_input(const llama_ubatch * ubatch) override { + if (!bias || !ubatch->pos) return; + const int64_t n_tokens = ubatch->n_tokens; + std::vector data((size_t) nblk * n_tokens, 0.0f); + for (int64_t i = 0; i < n_tokens; ++i) { + const int L = (int) (ubatch->pos[i] / blk); + for (int l = 0; l < local && L - l >= 0; ++l) { + if (L - l < nblk) data[(size_t) i * nblk + (L - l)] = 1e30f; + } + } + ggml_backend_tensor_set(bias, data.data(), 0, data.size() * sizeof(float)); + } + ggml_tensor * bias = nullptr; // [nblk, n_tokens] f32 + int blk; + int local; + int64_t nblk; +}; + +static inline ggml_fp16_t msa_neg_val (ggml_fp16_t) { return ggml_fp32_to_fp16(-INFINITY); } +static inline float msa_neg_val (float) { return -INFINITY; } + +// ---- Decomposed MSA ---- + +static inline bool msa_score_masked(float x) { return x <= -1e30f; } // -inf or pool -FLT_MAX; not +inf + +// ---- MSA 4-way per-group selection ------------------------------------------------- +// dst = [n_kv, S, Hd, ns] f16/f32. Channel h = base causal mask (from src[1]) with the +// blocks NOT in group h's top-k forced to -inf. bs = [nblk, Hd, S] f32 (per-group block +// scores from the decomposed OP). Per (query i, group h) the nblk scores are contiguous +// at bs + i*nblk*Hd + h*nblk. No global memcpy: each (i,h) copies its base-mask column, +// so threads partitioned over i never race (channels within a column are disjoint). +template +static inline void msa_select_4way_t( + MT * dst, const MT * base_mask, + int64_t mask_skey, int64_t dst_squery, int64_t base_squery, int64_t chan_stride, + const float * bs_in, int Hd, int S, int64_t n_kv, + int blk, int topk_blocks, int local, + int ith, int nth, int nblk ) { + + const int topk = topk_blocks < nblk ? topk_blocks : nblk; + std::vector bs (nblk); + std::vector ord(nblk); + std::vector sel(nblk); + + for (int i = ith; i < S; i += nth) { + for (int h = 0; h < Hd; ++h) { + MT * out = dst + (int64_t) h*chan_stride + (int64_t) i*dst_squery; + const MT * src = base_mask + (int64_t) i*base_squery; + for (int64_t j = 0; j < n_kv; ++j) out[j*mask_skey] = src[j*mask_skey]; // base col i + + const float * bcol = bs_in + (size_t) i*nblk*Hd + (size_t) h*nblk; + for (int bk = 0; bk < nblk; ++bk) bs[bk] = bcol[bk]; + + // local block = highest non-masked block for this query; force-keep it (+ local-1 neighbors) + int qb = -1; + for (int bk = nblk - 1; bk >= 0; --bk) { if (!msa_score_masked(bs[bk])) { qb = bk; break; } } + for (int l = 0; l < local && qb - l >= 0; ++l) bs[qb - l] = INFINITY; + + for (int t = 0; t < nblk; ++t) ord[t] = t; + // top-k SET only; partial_sort keeps the "break at first empty" logic and + // at O(nblk log topk) instead of O(nblk log nblk). + std::partial_sort(ord.begin(), ord.begin() + topk, ord.end(), + [&](int a, int b){ return bs[a] > bs[b]; }); + + std::fill(sel.begin(), sel.end(), (char) 0); + for (int t = 0; t < topk; ++t) { + const int bk = ord[t]; + if (msa_score_masked(bs[bk])) break; // sorted desc: first empty => fewer than topk real + sel[bk] = 1; + } + + for (int bk = 0; bk < nblk; ++bk) { + if (sel[bk]) continue; + const int j0 = bk * blk; + const int j1 = (int) std::min((int64_t)(bk + 1) * blk, n_kv); + for (int j = j0; j < j1; ++j) out[(int64_t) j * mask_skey] = msa_neg_val(MT(0)); + } + } + } +} + +static void msa_mask_from_scores_4way_op(struct ggml_tensor * dst, int ith, int nth, void * userdata) { + const struct ggml_tensor * bs = dst->src[0]; // [nblk, Hd, S] f32 + const struct ggml_tensor * mask = dst->src[1]; // [n_kv, S, 1, ns] base causal mask + const msa_params * p = (const msa_params *) userdata; + + const int nblk = bs->ne[0]; + const int Hd = bs->ne[1]; + const int S = bs->ne[2]; + const int64_t n_kv = dst->ne[0]; + + GGML_ASSERT(bs->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(bs)); + GGML_ASSERT(ggml_is_contiguous(dst)); + GGML_ASSERT(dst->type == mask->type); + GGML_ASSERT(dst->ne[0] == mask->ne[0]); // n_kv + GGML_ASSERT(dst->ne[1] == mask->ne[1]); // S + GGML_ASSERT(dst->ne[2] == (int64_t) Hd); // channels = index heads + GGML_ASSERT(mask->ne[2] == 1); + GGML_ASSERT(mask->ne[3] == 1 && "MSA 4-way assumes single stream (-np 1)"); + GGML_ASSERT(n_kv % p->blk == 0); + GGML_ASSERT((int64_t) nblk == n_kv / p->blk); + + const int64_t ts = ggml_type_size(dst->type); + const int64_t mask_skey = 1; + const int64_t dst_squery = dst->nb[1] / ts; // = n_kv + const int64_t base_squery = mask->nb[1] / ts; + const int64_t chan_stride = dst->nb[2] / ts; // = n_kv * S + + if (dst->type == GGML_TYPE_F16) { + msa_select_4way_t( + (ggml_fp16_t *) dst->data, (const ggml_fp16_t *) mask->data, + mask_skey, dst_squery, base_squery, chan_stride, + (const float *) bs->data, Hd, S, n_kv, + p->blk, p->topk_blocks, p->local, ith, nth, nblk); + } else { + msa_select_4way_t( + (float *) dst->data, (const float *) mask->data, + mask_skey, dst_squery, base_squery, chan_stride, + (const float *) bs->data, Hd, S, n_kv, + p->blk, p->topk_blocks, p->local, ith, nth, nblk); + } +} + +ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_4way( + llm_graph_input_attn_kv * inp, + ggml_tensor * wo, ggml_tensor * wo_s, + ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, + ggml_tensor * msa_mask4, float kq_scale, int il) const { + + GGML_ASSERT(!inp->self_k_rot && !inp->self_v_rot && "MSA 4-way: attn-rot not supported"); + + // --- store K/V to cache (mirror build_attn) --- + ggml_build_forward_expand(gf, q_cur); + ggml_build_forward_expand(gf, v_cur); + ggml_build_forward_expand(gf, k_cur); + const auto * mctx_cur = inp->mctx; + { + const auto & k_idxs = inp->get_k_idxs(); + const auto & v_idxs = inp->get_v_idxs(); + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + } + + ggml_tensor * k = mctx_cur->get_k(ctx0, il); // [D, n_head_kv, n_kv, ns] + ggml_tensor * v = mctx_cur->get_v(ctx0, il); // [D, n_head_kv, n_kv, ns] (v_trans=false under FA) + + const int64_t D = k->ne[0]; + const int64_t HKV = k->ne[1]; + const int64_t n_kv = k->ne[2]; + const int64_t ns = k->ne[3]; + const int64_t HQ = q_cur->ne[1]; + const int64_t T = q_cur->ne[2]; + const int64_t Gp = HQ / HKV; // query heads per group (16) + GGML_ASSERT(HQ % HKV == 0); + GGML_ASSERT(ns == 1 && "MSA 4-way assumes single stream (-np 1)"); + GGML_ASSERT(msa_mask4->ne[2] == HKV); // one channel per group + + const bool v_trans = v->nb[1] > v->nb[2]; // false under FA + + ggml_tensor * acc = nullptr; + for (int g = 0; g < (int) HKV; ++g) { + // Q heads [Gp*g, Gp*g+Gp) -> [D, Gp, T, 1] -> permute -> [D, T, Gp, 1] + ggml_tensor * qg = ggml_view_4d(ctx0, q_cur, D, Gp, T, 1, + q_cur->nb[1], q_cur->nb[2], q_cur->nb[3], + (size_t) g * Gp * q_cur->nb[1]); + qg = ggml_permute(ctx0, qg, 0, 2, 1, 3); + + // K kv-head g -> [D, 1, n_kv, ns] -> permute -> [D, n_kv, 1, ns] + ggml_tensor * kg = ggml_view_4d(ctx0, k, D, 1, n_kv, ns, + k->nb[1], k->nb[2], k->nb[3], (size_t) g * k->nb[1]); + kg = ggml_permute(ctx0, kg, 0, 2, 1, 3); + + // V kv-head g (same head=ne[1] slice; works for both v_trans layouts) + ggml_tensor * vg = ggml_view_4d(ctx0, v, v->ne[0], 1, v->ne[2], ns, + v->nb[1], v->nb[2], v->nb[3], (size_t) g * v->nb[1]); + vg = ggml_permute(ctx0, vg, 0, 2, 1, 3); + if (v_trans) vg = ggml_transpose(ctx0, vg); // no-op under FA + + if (kg->type == GGML_TYPE_F32) kg = ggml_cast(ctx0, kg, GGML_TYPE_F16); + if (vg->type == GGML_TYPE_F32) vg = ggml_cast(ctx0, vg, GGML_TYPE_F16); + + // mask channel g -> [n_kv, S, 1] contiguous slice (FA broadcasts over the Gp heads) + ggml_tensor * mg = ggml_view_3d(ctx0, msa_mask4, msa_mask4->ne[0], msa_mask4->ne[1], 1, + msa_mask4->nb[1], msa_mask4->nb[2], + (size_t) g * msa_mask4->nb[2]); + + ggml_tensor * og = ggml_flash_attn_ext(ctx0, qg, kg, vg, mg, kq_scale, + hparams.f_max_alibi_bias, 0.0f); + ggml_flash_attn_ext_set_prec(og, GGML_PREC_F32); // [D, Gp, T, 1] + cb(og, LLAMA_TENSOR_NAME_FATTN, il); + + acc = acc ? ggml_concat(ctx0, acc, og, 1) : og; // concat along HEAD axis (ne[1]) + } + + // [D, HQ, T, 1] -> [n_embd, T] + ggml_tensor * cur = ggml_reshape_2d(ctx0, acc, acc->ne[0]*acc->ne[1], acc->ne[2]*acc->ne[3]); + cb(cur, "kqv_out", il); + + if (wo) cur = build_lora_mm(wo, cur, wo_s); // o_proj + return cur; +} + +ggml_tensor * llama_model_minimax_m3::graph::build_attn_msa_decode( + llm_graph_input_attn_kv * inp, + ggml_tensor * wo, ggml_tensor * wo_s, + ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, + ggml_tensor * bs, // [nblk, Hd, 1] per-group block scores (mask-added) from the front + ggml_tensor * local_bias, // [nblk, 1] f32, +BIG at local block(s) + ggml_tensor * kqm, // [n_kv, 1, 1, 1] causal mask (f16/f32), contiguous + int topk_blocks, float kq_scale, int il) const { + + GGML_ASSERT(!inp->self_k_rot && !inp->self_v_rot && "MSA decode: attn-rot not supported"); + GGML_ASSERT(q_cur->ne[2] == 1 && "MSA decode path is S==1 only"); + + // --- store K/V to cache--- + ggml_build_forward_expand(gf, q_cur); + ggml_build_forward_expand(gf, v_cur); + ggml_build_forward_expand(gf, k_cur); + const auto * mctx_cur = inp->mctx; + { + const auto & k_idxs = inp->get_k_idxs(); + const auto & v_idxs = inp->get_v_idxs(); + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + } + ggml_tensor * k = mctx_cur->get_k(ctx0, il); // [D, HKV, n_kv, 1] + ggml_tensor * v = mctx_cur->get_v(ctx0, il); // [D, HKV, n_kv, 1] (v_trans=false under FA) + + const int64_t D = k->ne[0]; + const int64_t HKV = k->ne[1]; + const int64_t n_kv = k->ne[2]; + const int64_t HQ = q_cur->ne[1]; + const int64_t Gp = HQ / HKV; + const int64_t nblk = bs->ne[0]; + const int64_t Hd = bs->ne[1]; + const int blk = (int) (n_kv / nblk); + const int K = topk_blocks < (int) nblk ? topk_blocks : (int) nblk; + GGML_ASSERT(Hd == HKV); + GGML_ASSERT((int64_t) blk * nblk == n_kv); + GGML_ASSERT(!(v->nb[1] > v->nb[2]) && "MSA decode assumes v_trans=false (FA on)"); + + // --- force local block(s): bs += local_bias (broadcast over Hd) --- + ggml_tensor * bsf = ggml_add(ctx0, bs, ggml_reshape_3d(ctx0, local_bias, nblk, 1, 1)); // [nblk,Hd,1] + bsf = ggml_reshape_2d(ctx0, bsf, nblk, Hd); // [nblk,Hd] + + // --- top-k block indices per group --- + ggml_tensor * idx = ggml_top_k(ctx0, bsf, K); // I32 [K, Hd] + + // --- expand block idx -> token idx: blk*idx + arange(blk) --- + ggml_tensor * idxf = ggml_scale(ctx0, ggml_cast(ctx0, idx, GGML_TYPE_F32), (float) blk); // [K,Hd] f32 + idxf = ggml_reshape_3d(ctx0, idxf, 1, K, Hd); // [1,K,Hd] + ggml_tensor * tgt = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk, K, Hd); + ggml_tensor * rep = ggml_repeat(ctx0, idxf, tgt); // [blk,K,Hd] + ggml_tensor * ar = ggml_reshape_3d(ctx0, ggml_arange(ctx0, 0.f, (float) blk, 1.f), blk, 1, 1); + ggml_tensor * tokf = ggml_add(ctx0, rep, ar); // [blk,K,Hd] + tokf = ggml_reshape_2d(ctx0, tokf, (int64_t) blk * K, Hd); // [blk*K,Hd] + ggml_tensor * tok = ggml_cast(ctx0, tokf, GGML_TYPE_I32); // [blk*K,Hd] I32 + + ggml_tensor * km1 = ggml_reshape_2d(ctx0, kqm, 1, n_kv); // [1, n_kv] for mask-gather (decode: kqm is [n_kv,1]) + + ggml_tensor * acc = nullptr; + for (int g = 0; g < (int) HKV; ++g) { + ggml_tensor * tg = ggml_cont(ctx0, ggml_view_2d(ctx0, tok, (int64_t) blk * K, 1, + tok->nb[1], (size_t) g * tok->nb[1])); // [blk*K,1] I32 + + // gather K/V for kv-head g from the strided [D,n_kv] head-slice (no cont of the cache) + ggml_tensor * Kg2 = ggml_view_2d(ctx0, k, D, n_kv, k->nb[2], (size_t) g * k->nb[1]); + ggml_tensor * Vg2 = ggml_view_2d(ctx0, v, D, n_kv, v->nb[2], (size_t) g * v->nb[1]); + ggml_tensor * Kg = ggml_get_rows(ctx0, Kg2, tg); // [D, blk*K] f32 + ggml_tensor * Vg = ggml_get_rows(ctx0, Vg2, tg); // [D, blk*K] f32 + ggml_tensor * mg = ggml_get_rows(ctx0, km1, tg); // [1, blk*K] f32 (causal/padding per key) + + // shape for FA: q_g [D,1,Gp,1]; gathered K/V [D, blk*K, 1, 1] f16; mask [blk*K,1,1,1] f16 + ggml_tensor * qg = ggml_view_3d(ctx0, q_cur, D, Gp, 1, + q_cur->nb[1], q_cur->nb[2], (size_t) g * Gp * q_cur->nb[1]); + qg = ggml_permute(ctx0, ggml_reshape_4d(ctx0, ggml_cont(ctx0, qg), D, Gp, 1, 1), 0, 2, 1, 3); + ggml_tensor * kgf = ggml_cast(ctx0, ggml_reshape_4d(ctx0, Kg, D, (int64_t) blk * K, 1, 1), GGML_TYPE_F16); + ggml_tensor * vgf = ggml_cast(ctx0, ggml_reshape_4d(ctx0, Vg, D, (int64_t) blk * K, 1, 1), GGML_TYPE_F16); + ggml_tensor * mgf = ggml_cast(ctx0, ggml_reshape_4d(ctx0, mg, (int64_t) blk * K, 1, 1, 1), GGML_TYPE_F16); + + ggml_tensor * og = ggml_flash_attn_ext(ctx0, qg, kgf, vgf, mgf, kq_scale, + hparams.f_max_alibi_bias, 0.0f); + ggml_flash_attn_ext_set_prec(og, GGML_PREC_F32); // [D, Gp, 1, 1] + cb(og, LLAMA_TENSOR_NAME_FATTN, il); + acc = acc ? ggml_concat(ctx0, acc, og, 1) : og; // concat along head axis + } + + ggml_tensor * cur = ggml_reshape_2d(ctx0, acc, acc->ne[0]*acc->ne[1], acc->ne[2]*acc->ne[3]); // [n_embd,1] + cb(cur, "kqv_out", il); + if (wo) cur = build_lora_mm(wo, cur, wo_s); + return cur; +} + +llama_model_minimax_m3::graph::graph(const llama_model & model, const llm_graph_params & params) : llm_graph_context(params) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + const auto & mm = static_cast(model); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + // partial rotary: head_dim != n_rot, so don't assert n_embd_head == n_rot + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + ggml_tensor * inp_pos = build_inp_pos(); + auto inp_attn = build_attn_inp_kv(); + + // MSA decode-only gather path is active for single-token batches; build its local-force input once. + // MSA's decode-gather AND 4-way paths both call ggml_flash_attn_ext directly and assume the + // non-transposed V layout that llama.cpp only provides when flash attention is enabled + // (v_trans = !cparams.flash_attn). So MSA as a whole requires FA; without it we fall back to + // incorrect dense build_attn, which handles the transposed-V / explicit-softmax path. + const bool fa_on = cparams.flash_attn; // resolved FA: supported AND enabled + const bool single_stream = cparams.n_seq_max == 1; // MSA block selection is anchored to absolute + // KV slots, valid only for single-sequence decode + const bool want_bypass = getenv("MSA_BYPASS"); + const bool msa_enabled = fa_on && single_stream && !want_bypass; + + static bool warned_no_fa = false; + if (!fa_on && !want_bypass && !warned_no_fa) { + LLAMA_LOG_WARN("%s: flash attention disabled; MSA requires it -> running DENSE attention " + "(no sparse selection). Enable flash attention for MSA.\n", __func__); + warned_no_fa = true; + } + static bool warned_multi_seq = false; + if (fa_on && !single_stream && !want_bypass && !warned_multi_seq) { + LLAMA_LOG_WARN("%s: n_seq_max > 1; MSA is single-stream only (-np 1) -> running DENSE " + "attention (no sparse selection). Use -np 1 to enable MSA.\n", __func__); + warned_multi_seq = true; + } + + const bool msa_decode = msa_enabled && (n_tokens == 1); + + + llm_graph_input_msa_local * msa_loc = nullptr; + if (msa_decode) { + ggml_tensor * kqm0 = inp_attn->get_kq_mask(); + const int64_t n_kv0 = kqm0->ne[0]; + GGML_ASSERT(n_kv0 % mm.msa_p.blk == 0 && + "MSA: KV/mask n_kv must be a multiple of indexer.block_size (128); " + "the flash-attention KV padding must be a multiple of the block size. " + "A non-multiple would silently drop the partial tail block."); + const int64_t nblk0 = n_kv0 / mm.msa_p.blk; + auto loc = std::make_unique(mm.msa_p.blk, mm.msa_p.local, nblk0); + loc->bias = ggml_new_tensor_2d(ctx0, GGML_TYPE_F32, nblk0, n_tokens); // [nblk,1] at decode + ggml_set_input(loc->bias); + msa_loc = (llm_graph_input_msa_local *) res->add_input(std::move(loc)); + } + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + ggml_tensor * inpSA = inpL; + + // self-attention + { + cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, n_head, n_head_kv, il); + + // per-head QK RMSNorm (weights already include Gemma's +1) + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, NULL, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, NULL, LLM_NORM_RMS, il); + cb(Kcur, "Kcur_normed", il); + + // partial rotary: only the first n_rot dims are rotated + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + ggml_tensor * msa_mask4 = nullptr; // [n_kv, S, Hd, ns] per-group (default) + bool msa_decode_done = false; + if (msa_enabled && il >= (int) hparams.n_layer_dense_lead) { // sparse layers == MoE layers for M3 + const int64_t n_idx_dim = hparams.indexer_head_size; // 128 + const int64_t n_idx_head = hparams.indexer_n_head; // 4 + + ggml_tensor * iq = build_lora_mm(model.layers[il].index_q_proj, cur); // [512, n_tokens] + ggml_tensor * ik = build_lora_mm(model.layers[il].index_k_proj, cur); // [128, n_tokens] + iq = ggml_reshape_3d(ctx0, iq, n_idx_dim, n_idx_head, n_tokens); // [128,4,T] + ik = ggml_reshape_3d(ctx0, ik, n_idx_dim, 1, n_tokens); // [128,1,T] + iq = build_norm(iq, model.layers[il].index_q_norm, NULL, LLM_NORM_RMS, il); // +1 baked + ik = build_norm(ik, model.layers[il].index_k_norm, NULL, LLM_NORM_RMS, il); + iq = ggml_rope_ext(ctx0, iq, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + ik = ggml_rope_ext(ctx0, ik, inp_pos, nullptr, n_rot, rope_type, n_ctx_orig, + freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + + const auto * mctx_cur = inp_attn->mctx; + const auto & k_idxs = inp_attn->get_k_idxs(); + ggml_build_forward_expand(gf, mctx_cur->cpy_k_idx(ctx0, ik, k_idxs, il)); + ggml_tensor * ik_kv = mctx_cur->get_k_idx(ctx0, il); // [128,1,n_kv,1] + + ggml_tensor * kqm = inp_attn->get_kq_mask(); // f16 (FA on) / f32 + const int64_t n_kv = kqm->ne[0]; + + // --- DEFAULT: per-group decomposed (no head-amax) --- + GGML_ASSERT(n_kv % mm.msa_p.blk == 0 && + "MSA: KV/mask n_kv must be a multiple of indexer.block_size (128); " + "the flash-attention KV padding must be a multiple of the block size. " + "A non-multiple would silently drop the partial tail block."); + const int blk = mm.msa_p.blk; + + ggml_tensor * ik2d = ggml_reshape_2d(ctx0, ik_kv, n_idx_dim, n_kv); + ik2d = ggml_cont(ctx0, ik2d); // force contiguous + ggml_tensor * iq2d = ggml_reshape_2d(ctx0, iq, n_idx_dim, n_idx_head*n_tokens); // [D, Hd*S] + ggml_tensor * sc = ggml_mul_mat(ctx0, ik2d, iq2d); // [n_kv, Hd*S] f32 + sc = ggml_reshape_3d(ctx0, sc, n_kv, n_idx_head, n_tokens); // [n_kv, Hd, S] + + // + causal mask, broadcast over Hd (dim 1). Single-stream: mask is [n_kv,S,1,1]. + ggml_tensor * mf = (kqm->type == GGML_TYPE_F32) ? kqm : ggml_cast(ctx0, kqm, GGML_TYPE_F32); + mf = ggml_reshape_3d(ctx0, mf, n_kv, 1, n_tokens); // [n_kv, 1, S] + sc = ggml_add(ctx0, sc, mf); // [n_kv, Hd, S] + + // block-amax over n_kv (dim 0), keep Hd -> [nblk, Hd, S] + ggml_tensor * bs = ggml_pool_2d(ctx0, sc, GGML_OP_POOL_MAX, blk, 1, blk, 1, 0, 0); + + if (msa_decode) { + // decode gather skip the CPU tail op + msa_mask4 entirely + cur = build_attn_msa_decode(inp_attn, + model.layers[il].wo, model.layers[il].wo_s, + Qcur, Kcur, Vcur, bs, msa_loc->bias, kqm, + mm.msa_p.topk_blocks, 1.0f/sqrtf(float(n_embd_head)), il); + msa_decode_done = true; // signal the dispatch below to skip + } else { + ggml_tensor * srcs[4] = { bs, kqm, nullptr, nullptr }; + int nsrc = 2; + msa_mask4 = ggml_custom_4d(ctx0, kqm->type, + n_kv, n_tokens, n_idx_head, kqm->ne[3], + srcs, nsrc, msa_mask_from_scores_4way_op, GGML_N_TASKS_MAX, + const_cast(&mm.msa_p)); + cb(msa_mask4, "msa_mask4", il); + } + + } + + if (msa_decode_done) { + // cur already computed by build_attn_msa_decode above + } else if (il >= (int) hparams.n_layer_dense_lead && msa_mask4) { + cur = build_attn_msa_4way(inp_attn, model.layers[il].wo, model.layers[il].wo_s, + Qcur, Kcur, Vcur, msa_mask4, 1.0f/sqrtf(float(n_embd_head)), il); + } else { + cur = build_attn(inp_attn, model.layers[il].wo, NULL, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, + 1.0f/sqrtf(float(n_embd_head)), il, nullptr); + } + } + + if (il == n_layer - 1 && inp_out_ids) { + 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, model.layers[il].ffn_norm, NULL, LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + if ((uint32_t) il < hparams.n_layer_dense_lead) { + // leading dense FFN (swigluoai) + cur = build_ffn(cur, + model.layers[il].ffn_up, NULL, NULL, + model.layers[il].ffn_gate, NULL, NULL, + model.layers[il].ffn_down, NULL, NULL, + NULL, + LLM_FFN_SWIGLU_OAI_MOE, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + } else { + // routed experts (swigluoai MoE) + ggml_tensor * moe_out = build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + model.layers[il].ffn_exp_probs_b, + n_expert, n_expert_used, + LLM_FFN_SWIGLU_OAI_MOE, hparams.expert_weights_norm, + hparams.expert_weights_scale, + (llama_expert_gating_func_type) hparams.expert_gating_func, + il); + cb(moe_out, "ffn_moe_out", il); + + // shared expert (swigluoai) + ggml_tensor * ffn_shexp = build_ffn(cur, + model.layers[il].ffn_up_shexp, NULL, NULL, + model.layers[il].ffn_gate_shexp, NULL, NULL, + model.layers[il].ffn_down_shexp, NULL, NULL, + NULL, + LLM_FFN_SWIGLU_OAI_MOE, 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); + + // input for next layer + inpL = cur; + } + + cur = inpL; + + cur = build_norm(cur, model.output_norm, NULL, LLM_NORM_RMS, -1); + cb(cur, "result_norm", -1); + res->t_embd = cur; + + // lm_head + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + diff --git a/src/models/models.h b/src/models/models.h index 7a52e7bc1ab7..0beda3185fe1 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1870,6 +1870,37 @@ struct llama_model_minimax_m2 : public llama_model_base { std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; }; +struct msa_params { + int blk; + int topk_blocks; + int local; +}; + +struct llama_model_minimax_m3 : public llama_model_base { + llama_model_minimax_m3(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; + msa_params msa_p; + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + + // MSA 4-way per-group FA split (exact per-group selection). msa_mask4 is + // [n_kv, S, Hd, ns] f16, channel g = group g's combined causal+block mask. + ggml_tensor * build_attn_msa_4way( + llm_graph_input_attn_kv * inp, + ggml_tensor * wo, ggml_tensor * wo_s, + ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, + ggml_tensor * msa_mask4, float kq_scale, int il) const; + // MSA decode-only gather path (S==1): per-group top-k block gather + FA. + ggml_tensor * build_attn_msa_decode( + llm_graph_input_attn_kv * inp, + ggml_tensor * wo, ggml_tensor * wo_s, + ggml_tensor * q_cur, ggml_tensor * k_cur, ggml_tensor * v_cur, + ggml_tensor * bs, ggml_tensor * local_bias, ggml_tensor * kqm, + int topk_blocks, float kq_scale, int il) const; + }; + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; struct llama_model_cogvlm : public llama_model_base { llama_model_cogvlm(const struct llama_model_params & params) : llama_model_base(params) {} diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index f39abe773fc6..2085f43fdec8 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -352,6 +352,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_LLADA_MOE: case LLM_ARCH_GROVEMOE: case LLM_ARCH_MINIMAX_M2: + case LLM_ARCH_MINIMAX_M3: case LLM_ARCH_RND1: case LLM_ARCH_PADDLEOCR: case LLM_ARCH_MIMO2: diff --git a/tools/mtmd/CMakeLists.txt b/tools/mtmd/CMakeLists.txt index ea684d9f156d..33d60f0c4cd1 100644 --- a/tools/mtmd/CMakeLists.txt +++ b/tools/mtmd/CMakeLists.txt @@ -40,6 +40,7 @@ add_library(mtmd models/paddleocr.cpp models/pixtral.cpp models/qwen2vl.cpp + models/minimax-m3.cpp models/qwen3vl.cpp models/mimovl.cpp models/qwen3a.cpp diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 5b413681f040..6f4fcc06f9d9 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -131,6 +131,8 @@ #define TN_MM_SOFT_EMB_N "mm.soft_emb_norm.weight" // gemma3 #define TN_MM_PROJECTOR "mm.model.fc.%s" // idefics3, deepseekocr #define TN_MM_PATCH_MERGER "mm.patch_merger.%s" // mistral small 3.1, glm4v +#define TN_MM_MERGE_FC1 "mm.merge.fc1.%s" // minimax-m3 patch-merge MLP +#define TN_MM_MERGE_FC2 "mm.merge.fc2.%s" #define TN_TOK_IMG_BREAK "v.token_embd.img_break" // pixtral #define TN_TOK_GLM_BOI "adapter.boi" // glm-edge (these embeddings are not in text model) #define TN_TOK_GLM_EOI "adapter.eoi" // glm-edge (these embeddings are not in text model) @@ -370,6 +372,7 @@ enum projector_type { PROJECTOR_TYPE_MINICPMV4_6, PROJECTOR_TYPE_GRANITE_SPEECH, PROJECTOR_TYPE_MIMOVL, + PROJECTOR_TYPE_MINIMAX_M3, PROJECTOR_TYPE_GRANITE4_VISION, PROJECTOR_TYPE_UNKNOWN, }; @@ -424,6 +427,7 @@ static std::map PROJECTOR_TYPE_NAMES = { { PROJECTOR_TYPE_MINICPMV4_6, "minicpmv4_6"}, { PROJECTOR_TYPE_GRANITE_SPEECH, "granite_speech"}, { PROJECTOR_TYPE_MIMOVL, "mimovl"}, + { PROJECTOR_TYPE_MINIMAX_M3, "minimax_m3"}, { PROJECTOR_TYPE_GRANITE4_VISION, "granite4_vision"}, }; diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index 46be39a641c1..608691e2887c 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -396,6 +396,10 @@ struct clip_model { ggml_tensor * mm_0_b = nullptr; ggml_tensor * mm_2_w = nullptr; ggml_tensor * mm_2_b = nullptr; + ggml_tensor * mm_merge_fc1_w = nullptr; // minimax-m3 + ggml_tensor * mm_merge_fc1_b = nullptr; + ggml_tensor * mm_merge_fc2_w = nullptr; + ggml_tensor * mm_merge_fc2_b = nullptr; ggml_tensor * image_newline = nullptr; ggml_tensor * view_seperator = nullptr; diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index d2226b3be1d7..fea41a81b35f 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -915,6 +915,10 @@ static std::unique_ptr clip_get_graph_builder(clip_ctx * ctx, const { builder = std::make_unique(ctx, img); } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + builder = std::make_unique(ctx, img); + } break; case PROJECTOR_TYPE_STEP3VL: { builder = std::make_unique(ctx, img); @@ -1467,6 +1471,16 @@ struct clip_model_loader { LOG_WRN("%s: more info: https://github.com/ggml-org/llama.cpp/issues/16842\n\n", __func__); } } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + hparams.n_merge = 2; // spatial_merge_size + hparams.image_resize_algo = RESIZE_ALGO_BICUBIC; + get_u32(KEY_SPATIAL_MERGE_SIZE, hparams.n_merge, false); + hparams.rope_theta = 10000.0f; // vision_config.rope_theta + // MiniMax-M3: max_pixels 451584 (=672^2) -> 576 merged tokens (image_seq_length) + hparams.set_limit_image_tokens(8, 576); + hparams.set_warmup_n_tokens(16*16); + } break; case PROJECTOR_TYPE_MIMOVL: { hparams.n_merge = 2; // spatial_merge_size @@ -2078,6 +2092,19 @@ struct clip_model_loader { model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight")); model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias"), false); } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + // per-patch MLP: mm.1 -> gelu -> mm.2 + model.mm_1_w = get_tensor(string_format(TN_LLAVA_PROJ, 1, "weight")); + model.mm_1_b = get_tensor(string_format(TN_LLAVA_PROJ, 1, "bias")); + model.mm_2_w = get_tensor(string_format(TN_LLAVA_PROJ, 2, "weight")); + model.mm_2_b = get_tensor(string_format(TN_LLAVA_PROJ, 2, "bias")); + // 2x2 merge MLP: mm.merge.fc1 -> gelu -> mm.merge.fc2 + model.mm_merge_fc1_w = get_tensor(string_format(TN_MM_MERGE_FC1, "weight")); + model.mm_merge_fc1_b = get_tensor(string_format(TN_MM_MERGE_FC1, "bias")); + model.mm_merge_fc2_w = get_tensor(string_format(TN_MM_MERGE_FC2, "weight")); + model.mm_merge_fc2_b = get_tensor(string_format(TN_MM_MERGE_FC2, "bias")); + } break; case PROJECTOR_TYPE_STEP3VL: { model.mm_0_w = get_tensor(string_format(TN_LLAVA_PROJ, 0, "weight")); @@ -3346,6 +3373,7 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { case PROJECTOR_TYPE_QWEN3VL: case PROJECTOR_TYPE_EXAONE4_5: case PROJECTOR_TYPE_MIMOVL: + case PROJECTOR_TYPE_MINIMAX_M3: case PROJECTOR_TYPE_GLM4V: case PROJECTOR_TYPE_YOUTUVL: { @@ -3845,6 +3873,26 @@ bool clip_image_batch_encode(clip_ctx * ctx, int n_threads, const clip_image_f32 set_input_i32("positions", positions); } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + const int gh = image_size_height / patch_size; + const int gw = image_size_width / patch_size; + std::vector pos_t, pos_h, pos_w; + pos_t.reserve(gh * gw); + pos_h.reserve(gh * gw); + pos_w.reserve(gh * gw); + for (int bh = 0; bh < gh / 2; bh++) + for (int bw = 0; bw < gw / 2; bw++) + for (int mh = 0; mh < 2; mh++) + for (int mw = 0; mw < 2; mw++) { + pos_t.push_back(0); + pos_h.push_back(bh * 2 + mh); + pos_w.push_back(bw * 2 + mw); + } + set_input_i32("minimax_pos_t", pos_t); + set_input_i32("minimax_pos_h", pos_h); + set_input_i32("minimax_pos_w", pos_w); + } break; case PROJECTOR_TYPE_DOTS_OCR: { const int pw = image_size_width / patch_size; @@ -4545,6 +4593,8 @@ int clip_n_mmproj_embd(const struct clip_ctx * ctx) { return ctx->model.mm_ffn_down_w->ne[1]; case PROJECTOR_TYPE_GLM_EDGE: return ctx->model.mm_model_mlp_3_w->ne[1]; + case PROJECTOR_TYPE_MINIMAX_M3: + return ctx->model.mm_merge_fc2_b->ne[0]; case PROJECTOR_TYPE_QWEN2VL: case PROJECTOR_TYPE_QWEN25VL: case PROJECTOR_TYPE_EXAONE4_5: diff --git a/tools/mtmd/models/minimax-m3.cpp b/tools/mtmd/models/minimax-m3.cpp new file mode 100644 index 000000000000..8b602e0fd437 --- /dev/null +++ b/tools/mtmd/models/minimax-m3.cpp @@ -0,0 +1,85 @@ +#include "models.h" + +// MiniMax-M3 vision graph + +ggml_tensor * clip_graph_minimax_m3::apply_rope( + ggml_tensor * x, ggml_tensor * pos_t, ggml_tensor * pos_h, ggml_tensor * pos_w) { + const int64_t Hn = x->ne[1]; + const int64_t P = x->ne[2]; + const size_t es = ggml_element_size(x); + const int dh = (int) x->ne[0]; + const int axd = 2 * ((2 * (dh / 2) / 3) / 2); + + GGML_ASSERT(x->nb[0] == es); + GGML_ASSERT(3 * axd <= dh); + + const float th = hparams.rope_theta; + auto sl = [&](int off, int n) { + return ggml_cont(ctx0, ggml_view_3d(ctx0, x, n, Hn, P, x->nb[1], x->nb[2], (size_t) off * es)); + }; + ggml_tensor * t = sl(0, axd); + ggml_tensor * h = sl(axd, axd); + ggml_tensor * w = sl(2 * axd, axd); + ggml_tensor * pad = sl(3 * axd, dh - 3 * axd); + + t = ggml_rope_ext(ctx0, t, pos_t, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + h = ggml_rope_ext(ctx0, h, pos_h, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + w = ggml_rope_ext(ctx0, w, pos_w, nullptr, axd, GGML_ROPE_TYPE_NEOX, 0, th, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f); + return ggml_concat(ctx0, ggml_concat(ctx0, ggml_concat(ctx0, t, h, 0), w, 0), pad, 0); +} + +ggml_cgraph * clip_graph_minimax_m3::build() { + GGML_ASSERT(model.patch_bias == nullptr); + GGML_ASSERT(model.class_embedding == nullptr); + GGML_ASSERT(model.patch_embeddings_0 && model.patch_embeddings_1); + GGML_ASSERT(model.mm_1_w && model.mm_2_w); + GGML_ASSERT(model.mm_merge_fc1_w && model.mm_merge_fc2_w); + + const int batch_size = 1; + const int n_pos = n_patches; + const int merge = 2; + + // patch embedding + ggml_tensor * inp_raw = build_inp_raw(); + ggml_tensor * inp = ggml_add(ctx0, + ggml_conv_2d(ctx0, model.patch_embeddings_0, inp_raw, patch_size, patch_size, 0, 0, 1, 1), + ggml_conv_2d(ctx0, model.patch_embeddings_1, inp_raw, patch_size, patch_size, 0, 0, 1, 1)); + + // spatial merge + { + inp = ggml_permute(ctx0, inp, 1, 2, 0, 3); + inp = ggml_cont_4d(ctx0, inp, n_embd * 2, n_patches_x / 2, n_patches_y, batch_size); + inp = ggml_reshape_4d(ctx0, inp, n_embd * 2, n_patches_x / 2, 2, batch_size * (n_patches_y / 2)); + inp = ggml_permute(ctx0, inp, 0, 2, 1, 3); + inp = ggml_cont_3d(ctx0, inp, n_embd, n_patches_x * n_patches_y, batch_size); + } + + ggml_tensor * pos_t = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); + ggml_set_name(pos_t, "minimax_pos_t"); ggml_set_input(pos_t); + ggml_tensor * pos_h = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); + ggml_set_name(pos_h, "minimax_pos_h"); ggml_set_input(pos_h); + ggml_tensor * pos_w = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_pos); + ggml_set_name(pos_w, "minimax_pos_w"); ggml_set_input(pos_w); + + ggml_tensor * inpL = build_vit( + inp, n_pos, NORM_TYPE_NORMAL, FFN_GELU_ERF, nullptr, + [&](ggml_tensor * c, const clip_layer &) { + return apply_rope(c, pos_t, pos_h, pos_w); + }); + + // projector + ggml_tensor * emb = inpL; + emb = build_ffn(emb, model.mm_1_w, model.mm_1_b, + nullptr, nullptr, + model.mm_2_w, model.mm_2_b, FFN_GELU_ERF, -1); + + const int64_t proj = emb->ne[0]; + emb = ggml_reshape_2d(ctx0, emb, proj * merge * merge, n_pos / (merge * merge)); + + emb = build_ffn(emb, model.mm_merge_fc1_w, model.mm_merge_fc1_b, + nullptr, nullptr, + model.mm_merge_fc2_w, model.mm_merge_fc2_b, FFN_GELU_ERF, -1); + + ggml_build_forward_expand(gf, emb); + return gf; +} diff --git a/tools/mtmd/models/models.h b/tools/mtmd/models/models.h index 12d5e6949320..93b0adc8f40e 100644 --- a/tools/mtmd/models/models.h +++ b/tools/mtmd/models/models.h @@ -40,6 +40,12 @@ struct clip_graph_qwen3vl : clip_graph_qwen2vl { ggml_cgraph * build() override; }; +struct clip_graph_minimax_m3 : clip_graph { + clip_graph_minimax_m3(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} + ggml_cgraph * build() override; + ggml_tensor * apply_rope(ggml_tensor * x, ggml_tensor * pos_t, ggml_tensor * pos_h, ggml_tensor * pos_w); +}; + struct clip_graph_mimovl : clip_graph { clip_graph_mimovl(clip_ctx * ctx, const clip_image_f32 & img) : clip_graph(ctx, img) {} ggml_cgraph * build() override; diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 724538b5857a..8b110f6c5c3f 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -463,6 +463,13 @@ struct mtmd_context { img_end = "<|vision_end|>"; image_preproc = std::make_unique(ctx_v); } break; + case PROJECTOR_TYPE_MINIMAX_M3: + { + // ]<]start of image[>[ ... (image embeddings) ... ]<]end of image[>[ + img_beg = "]<]start of image[>["; + img_end = "]<]end of image[>["; + image_preproc = std::make_unique(ctx_v); + } break; case PROJECTOR_TYPE_YOUTUVL: { // <|vision_start|> ... (image embeddings) ... <|vision_end|>