Skip to content

Commit 4bbbff4

Browse files
committed
[None][feat] Add AD custom model for MiniMax-M2 family
Replace the existing MiniMax-M2 MoE patch with a full custom model implementation using AD canonical ops. Covers both MiniMaxAI/MiniMax-M2 and MiniMaxAI/MiniMax-M2.5 (same architecture, model_type: minimax_m2). Key architecture features: - MoE with 256 experts, top-8, sigmoid routing + e_score_correction_bias - GQA (48 Q heads, 8 KV heads, head_dim=128) - Partial RoPE (rotary_dim=64 out of head_dim=128) - Per-layer QK normalization (RMSNorm on full Q/K before reshape) - FP8 block-wise quantized checkpoint Canonical ops used: torch_rmsnorm, torch_rope_with_explicit_cos_sin, torch_attention (GQA-native, no repeat_kv), torch_moe. Signed-off-by: Lucas Liebenwein <11156568+lucaslie@users.noreply.github.com>
1 parent 8e0ca19 commit 4bbbff4

5 files changed

Lines changed: 1183 additions & 79 deletions

File tree

examples/auto_deploy/model_registry/models.yaml

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -229,7 +229,7 @@ models:
229229
yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml']
230230
# --- MiniMax-M2.5 (Feb 2026) ---
231231
- name: MiniMaxAI/MiniMax-M2.5
232-
yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml']
232+
yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml', 'minimax_m2.yaml']
233233
# --- MiMo-V2-Flash (Feb 2026) ---
234234
- name: XiaomiMiMo/MiMo-V2-Flash
235235
yaml_extra: ['dashboard_default.yaml', 'world_size_8.yaml']

tensorrt_llm/_torch/auto_deploy/models/custom/__init__.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@
66
from .modeling_hunyuan_dense_v1 import HunYuanDenseV1ForCausalLM
77
from .modeling_kimi_k2 import KimiK2ForCausalLM, KimiK25ForConditionalGeneration
88
from .modeling_llama3 import Llama3ForCausalLM
9+
from .modeling_minimax_m2 import MiniMaxM2ForCausalLM
910
from .modeling_nemotron_flash import NemotronFlashForCausalLM, NemotronFlashPreTrainedTokenizerFast
1011
from .modeling_nemotron_h import NemotronHForCausalLM
1112
from .modeling_qwen3 import Qwen3ForCausalLM
@@ -22,6 +23,7 @@
2223
"KimiK2ForCausalLM",
2324
"KimiK25ForConditionalGeneration",
2425
"Llama3ForCausalLM",
26+
"MiniMaxM2ForCausalLM",
2527
"NemotronFlashForCausalLM",
2628
"NemotronFlashPreTrainedTokenizerFast",
2729
"NemotronHForCausalLM",

0 commit comments

Comments
 (0)