Add MiniMax-M2 and MiMo-V2-Flash model support

[whn09]

Description

Add NxDI contrib model adapters for MiniMax-M2 and MiMo-V2-Flash, two large-scale Mixture-of-Experts (MoE) models. Both require trn2.48xlarge with 64 logical cores (LNC=2) and feature custom MoE routing, partial RoPE, and architecture-specific attention mechanisms. Includes vLLM-neuron integration with benchmark scripts and performance results.

Model Information

Model Name: MiniMax-M2, MiMo-V2-Flash

Model Architecture: Decoder-only MoE transformer

Purpose: Text generation

Checklist

Required Components

• Accuracy Test (test/integration/test_model.py)

  • MiniMax-M2: 3/3 integration tests passed (import, config, neuron config)
  • MiMo-V2-Flash: 4/4 integration tests passed (import, config, neuron config, state dict converter)

• README.md with the following sections:

  • Usage Example: Code examples showing model compilation and inference
  • Compatibility Matrix: Tested on Trn2 (trn2.48xlarge) with Neuron SDK 2.22+
  • Example Checkpoints: Links to HuggingFace model hubs
  • Testing Instructions: pytest commands for integration tests

• Source Code (src/)

  • models/minimax_m2/ — MiniMax-M2: NKI attention kernel, Neuron-native QK norm, sigmoid router with e_score_correction_bias, fused_qkv, partial RoPE
  • models/mimo_v2/ — MiMo-V2-Flash: hybrid attention (full + sliding window), asymmetric head dims, attention sink bias
  • models/mimo_v2/conversion_script/ — FP8 to BF16 weight converter
  • utils/constants.py — Register both models in MODEL_TYPES for vllm-neuron discovery

Optional Components

• Unit Tests — Integration tests in test/integration/

Folder Structure

contrib/models/MiniMax-M2/
  README.md
  src/
    modeling_minimax_m2.py          # Re-exports from src/ (thin wrapper)
  test/
    integration/
      test_model.py

contrib/models/MiMo-V2-Flash/
  README.md
  src/
    modeling_mimo_v2.py             # Re-exports from src/ (thin wrapper)
  test/
    integration/
      test_model.py

src/neuronx_distributed_inference/
  models/
    minimax_m2/
      __init__.py
      config.json                     # Bundled HF config for unit tests
      configuration_minimax_m2.py     # Custom config class
      modeling_minimax_m2.py          # Full implementation (~1400 lines)
        - NKI attention kernel (nki-library attention_block_tkg with partial RoPE)
        - Neuron-native QK norm (RmsNorm.apply, per-rank, no all-reduce)
        - RouterTopKWithBias (e_score_correction_bias as nn.Parameter)
        - Fused QKV support with correct state dict key ordering
    mimo_v2/
      __init__.py
      modeling_mimo_v2.py             # Full implementation (~1333 lines)
        - Hybrid attention (9 full + 39 sliding window layers)
        - Asymmetric head dims (Q/K=192, V=128)
        - Attention sink bias
      conversion_script/
        preprocess_mimo_v2_fp8.py     # FP8 -> BF16 converter
  utils/
    constants.py                      # MODEL_TYPES entries for both models

perf_test/
  README.md                           # Benchmark plan and instructions
  0_setup.sh                          # vllm-neuron install + model download
  1_bench_mimo_v2_flash.sh            # MiMo benchmark (BS=1/32/128)
  2_bench_minimax_m2.sh               # MiniMax benchmark (BS=1/256)
  vllm-neuron-mimo-minimax.patch      # vllm-neuron patch for MiMo/MiniMax

Testing

How did you test this change?

Tested on trn2.48xlarge with TP=64, LNC=2, Neuron SDK 2.22+.

MiniMax-M2:

• 3/3 integration tests passed
• Architecture: 62 layers, 256 experts (top-8), hidden=3072, Q=48/KV=8 heads, partial RoPE (50%)
• Correctness verified: 4/4 sanity prompts answered correctly (math, knowledge, reasoning, generation)

Config	Concurrency	Throughput (tok/s)	TPOT (ms)	TTFT (ms)
BS=1, TP=64/EP=1, fused_qkv	1	40.56	12.89	1075
BS=256, TP=64/EP=64, CB	1	5.76	173.83	165
BS=256, TP=64/EP=64, CB	16	54.69	287.09	513
BS=256, TP=64/EP=64, CB	32	75.85	408.66	1066
BS=256, TP=64/EP=64, CB	128	106.72	1158.08	3950
BS=256, TP=64/EP=64, CB	256	128.94	1860.69	11263

MiMo-V2-Flash:

• 4/4 integration tests passed
• Architecture: 48 layers, 256 experts (top-8), hidden=4096, Q=64 heads, asymmetric Q/K=192 + V=128

Standalone NxDI (BF16, TP=64, EP=64):

Batch Size	Throughput (tok/s)
1	29.92
8	215.94
32	649.14

vLLM Serving (BS=32, TP=64/EP=64, CB, 900/90 tokens):

Concurrency	Throughput (tok/s)	TPOT (ms)	TTFT (ms)
1	27.98	33.65	222
16	224.57	64.95	570
32	302.61	90.23	1351

Note: Large MoE models require extended engine startup time (~30-50 min for compile+load). Set VLLM_ENGINE_READY_TIMEOUT_S=3600 before launching the vLLM server.

Compatibility

Tested with:

• Neuron SDK Version(s): 2.22+
• Instance Type(s): Trn2 (trn2.48xlarge)
• PyTorch Version: 2.9
• Python Version: 3.12

Additional Information

MiniMax-M2 key features:

• NKI attention kernel: nki-library attention_block_tkg with partial RoPE support (rotary_dim=64, head_dim=128) and flat QK RMSNorm fused into the kernel
• Neuron-native QK norm: Uses RmsNorm.apply (AwsNeuronRmsNorm custom call) instead of hand-rolled PyTorch ops, which compiled into different HLO in CE vs TG NEFFs
• Router bias: e_score_correction_bias preserved as nn.Parameter (not register_buffer) with non-uniform init to prevent XLA optimization from eliminating the add operation
• Fused QKV: convert_state_dict_to_fused_qkv runs before key rename to match expected state dict key paths
• Partial RoPE (50% of head_dim=128)
• Sigmoid router with learnable e_score_correction_bias

MiMo-V2-Flash key features:

• Hybrid attention: 9 full attention + 39 sliding window layers
• Asymmetric head dims (Q/K=192, V=128) — fused_qkv not supported
• Attention sink bias in sliding window layers
• EP hybrid sharding (EP=64 for prefill, EP=1 for token generation)
• Requires FP8→BF16 conversion (OCP block-wise FP8 not compatible with Neuron)

Related Issues

N/A

vLLM Integration

• This model/feature is intended for use with vLLM
• Documentation includes vLLM registration instructions (see READMEs and perf_test/ directory)
• vllm-neuron patch included (perf_test/vllm-neuron-mimo-minimax.patch)

---

By submitting this PR, I confirm that:

• I have read and followed the contributing guidelines
• This is a community contribution and may have limited testing compared to officially-supported models
• The code follows best practices and is well-documented
• All required components listed above are included

[whn09]

End-to-End Test Results on trn2.48xlarge

Both models have been compiled and tested on a trn2.48xlarge instance (32 NeuronCores, logical_nc_config=2 → 64 logical cores) using NxDI 2.22+ (PyTorch 2.9).

---

MiMo-V2-Flash (XiaomiMiMo/MiMo-V2-Flash)

Configuration:

• TP=64, EP=64 (hybrid sharding: CTE EP=64/TP=1, TKG EP=1/TP=64)
• logical_nc_config=2, fused_qkv=False (asymmetric Q/K=192, V=128)
• save_sharded_checkpoint=True (64 shards × ~19.3GB each)
• BF16 weights (dequantized from FP8 original)

Result: ✅ Passed — generates coherent text output

---

MiniMax-M2 (MiniMax/MiniMax-M2)

Configuration:

• TP=64, EP=64 (hybrid sharding: CTE EP=64/TP=1, TKG EP=1/TP=64)
• logical_nc_config=2, fused_qkv=True, sigmoid router
• save_sharded_checkpoint=True (64 shards, 857GB total)
• BF16 weights from s3://datalab/minimax/model_hf/MiniMax-M2-BF16/

Important: Must NOT enable use_shard_on_intermediate_dynamic_while in blockwise_matmul_config — intermediate_size=1536 with TP=64 yields I_TP=24, which gets padded to 256 (10.7× inflation), causing compiler OOM.

Result: ✅ Passed — generates correct, coherent output with chat template

Sample outputs:

Prompt: "What is 2+3? Answer briefly."
Response: <think>The user asks: "What is 2+3? Answer briefly."
This is a simple arithmetic question. The answer is 5.</think>
5

Prompt: "Explain what a neural network is in one sentence."
Response: A neural network is a computational model made up of interconnected
nodes (neurons) that process inputs via weighted connections and activation
functions to learn patterns and generate outputs such as predictions or
classifications.

Prompt: "Write a haiku about the ocean."
Response: Moonlit tide whispers / silver over midnight waves / shells hold ancient breath

---

Test Environment

• Instance: trn2.48xlarge (2TB RAM, 1.7TB NVMe)
• Software: NxDI from this branch (contrib/llm-models), installed via pip install -e .
• Compiler: neuronx-cc (NeuronSDK 2.22+)

[whn09]

Benchmark Results: MiMo-V2-Flash & MiniMax-M2 on trn2.48xlarge

Environment

• Instance: trn2.48xlarge (32 NeuronCores, logical_nc_config=2 → 64 logical cores)
• neuronx-cc: 2.24.5133.0+58f8de22
• torch-neuronx: 2.9.0.2.13.24727+8e870898
• vllm-neuron: release-0.5.0 (with hf_config passthrough + snapshot_download patch)
• NxDI: this PR branch

Configuration

• TP=64, EP=1, non-continuous-batching, BS=1
• Input: ~900 tokens, Output: ~90 tokens, 16 requests, concurrency=1
• MiMo: fused_qkv=false, async_mode=false, moe_mask_padded_tokens=true, blockwise_matmul_config={use_shard_on_intermediate_dynamic_while: false, skip_dma_token: true}
• MiniMax: fused_qkv=true, async_mode=false, moe_mask_padded_tokens=true, blockwise_matmul_config={use_shard_on_intermediate_dynamic_while: false, skip_dma_token: true}

Results

Metric	MiMo-V2-Flash	MiniMax-M2
Successful requests	16	16
Failed requests	0	0
Output throughput (tok/s)	46.17	39.35
Peak output throughput (tok/s)	91.00	74.00
Total token throughput (tok/s)	506.43	431.61
Mean TTFT (ms)	1027.72	1079.53
Median TTFT (ms)	1026.56	1079.03
P99 TTFT (ms)	1038.51	1084.67
Mean TPOT (ms)	10.39	13.61
Median TPOT (ms)	10.36	13.59
P99 TPOT (ms)	10.57	13.77
Mean ITL (ms)	10.39	13.61
P99 ITL (ms)	12.15	15.93

Notes

• MiMo-V2-Flash is ~17% faster on output throughput and ~24% lower TPOT compared to MiniMax-M2
• Both models load and serve successfully on trn2.48xlarge with TP=64/EP=1
• EP=64 (Expert Parallel) configurations are currently blocked — both models hit RuntimeError: Cannot concatenate arrays with different element types: S32 vs S64 in blockwise matmul during compilation. This prevents testing continuous batching / high-throughput (large batch size) configs. Likely a limitation in neuronx-cc 2.24 with these MoE architectures.
• Model weight loading takes ~10-20 min (MiMo ~570GB, MiniMax ~442GB); first-time compilation adds ~50 min per model