From 05b1e5e4ccd531ee3941bcaf140331e6e61799c7 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 06:18:29 +0800 Subject: [PATCH 01/27] Add contrib MiMo-V2.5 as a rename of MiMo-V2-Flash XiaomiMiMo/MiMo-V2.5 supersedes MiMo-V2-Flash: same architecture (48-layer MoE, 256 experts, hybrid full+SWA attention, partial RoPE, sink bias, sigmoid + noaux_tc routing, attention_value_scale=0.707) with new tokenizer id, larger vocab, and multimodal heads (vision + audio) that the NxDI path does not use. Copies the Flash tree and renames: - contrib/models/MiMo-V2-Flash -> contrib/models/MiMo-V2.5 - preprocess_mimo_v2_flash_fp8.py -> preprocess_mimo_v2_5_fp8.py - bench_mimo_v2_flash.sh -> bench_mimo_v2_5.sh - smoke_{compile,generate}_mimo_v2_flash.py -> ..._mimo_v2_5.py - MiMoV2FlashForCausalLM -> MiMoV2ForCausalLM (HF arch name in V2.5) - NXDI_CONTRIB_MIMO_V2_FLASH_SRC -> NXDI_CONTRIB_MIMO_V2_5_SRC - MODEL_TYPES key "mimov2flash" -> "mimov2" The unused legacy preprocess_mimo_v2_fp8.py (Jim's first version, superseded by the streaming variant) is dropped. Preprocess adjustments for V2.5's published FP8 checkpoint layout: - LazyWeightMap aliases legacy `model_N-00001-of-00002.safetensors` filenames referenced by safetensors.index.json to the actual shard names on disk (`model_pp0_epN_shardM.safetensors`). V2.5 ships both naming conventions inconsistently: HF Hub stores the latter while the index still references the former. Setup script: - 0_setup.sh downloads from HuggingFace directly (V2.5 is a public repo), drops the S3 fallback and the stale "BF16" path. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/README.md | 364 ++++ contrib/models/MiMo-V2.5/perf_test/0_setup.sh | 65 + .../MiMo-V2.5/perf_test/bench_mimo_v2_5.sh | 228 +++ .../MiMo-V2.5/perf_test/run_bench_single.sh | 76 + .../MiMo-V2.5/perf_test/sanity_check.sh | 59 + .../perf_test/smoke_compile_mimo_v2_5.py | 195 ++ .../perf_test/smoke_generate_mimo_v2_5.py | 171 ++ .../perf_test/vllm-neuron-patch.patch | 107 ++ contrib/models/MiMo-V2.5/src/__init__.py | 0 .../preprocess_mimo_v2_5_fp8.py | 510 +++++ .../models/MiMo-V2.5/src/modeling_mimo_v2.py | 1677 +++++++++++++++++ contrib/models/MiMo-V2.5/test/__init__.py | 0 .../MiMo-V2.5/test/integration/__init__.py | 0 .../MiMo-V2.5/test/integration/test_model.py | 53 + .../models/MiMo-V2.5/test/unit/__init__.py | 0 15 files changed, 3505 insertions(+) create mode 100644 contrib/models/MiMo-V2.5/README.md create mode 100755 contrib/models/MiMo-V2.5/perf_test/0_setup.sh create mode 100755 contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh create mode 100755 contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh create mode 100755 contrib/models/MiMo-V2.5/perf_test/sanity_check.sh create mode 100755 contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py create mode 100755 contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py create mode 100644 contrib/models/MiMo-V2.5/perf_test/vllm-neuron-patch.patch create mode 100644 contrib/models/MiMo-V2.5/src/__init__.py create mode 100644 contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py create mode 100644 contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py create mode 100644 contrib/models/MiMo-V2.5/test/__init__.py create mode 100644 contrib/models/MiMo-V2.5/test/integration/__init__.py create mode 100644 contrib/models/MiMo-V2.5/test/integration/test_model.py create mode 100644 contrib/models/MiMo-V2.5/test/unit/__init__.py diff --git a/contrib/models/MiMo-V2.5/README.md b/contrib/models/MiMo-V2.5/README.md new file mode 100644 index 00000000..7011ef15 --- /dev/null +++ b/contrib/models/MiMo-V2.5/README.md @@ -0,0 +1,364 @@ +# Contrib Model: MiMo-V2.5 + +NeuronX Distributed Inference implementation of [XiaomiMiMo/MiMo-V2.5](https://huggingface.co/XiaomiMiMo/MiMo-V2.5). + +## Model Information + +- **HuggingFace ID:** `XiaomiMiMo/MiMo-V2.5` +- **Model Type:** Decoder-only MoE transformer with hybrid attention +- **Architecture:** Custom MoE with full + sliding window attention +- **License:** Check HuggingFace model card + +## Architecture Details + +| Parameter | Value | +|-----------|-------| +| Hidden Size | 4096 | +| Layers | 48 | +| Attention Heads | 64 Q | +| KV Heads (full attn) | 4 | +| KV Heads (sliding window) | 8 | +| Q/K Head Dim | 192 | +| V Head Dim | 128 | +| Experts | 256 (top-8 routing) | +| Expert Intermediate | 2048 | +| Vocab Size | 151,936 | +| RoPE | Partial (34% of dims), theta=5M (full), 10K (SWA) | +| Sliding Window | 128 | +| Max Position | 262,144 | +| Total Params | ~143B (FP8 native) / ~286B (BF16 upcast) | + +Key features: +- **Hybrid Attention**: 9 full attention layers (0, 5, 11, 17, 23, 29, 35, 41, 47) + 39 sliding window layers +- **Asymmetric Head Dims**: Q/K use 192, V uses 128 (fused_qkv not supported) +- **Attention Sink Bias**: Learnable per-head bias on sliding window layers only +- **Sigmoid Router + noaux_tc**: e_score_correction_bias added to sigmoid scores before top-k selection; unbiased scores become affinity weights +- **attention_value_scale = 0.707**: HF MiMo-V2 multiplies `value_states` by this before the attention softmax × V (NOT applied to attn_output); the NxDI model matches + +## Prerequisites + +- **Instance**: trn2.48xlarge (32 NeuronCores, logical_nc_config=2 → 64 logical cores) +- **Neuron SDK**: 2.29 (Python 3.12, PyTorch 2.9) +- **Venvs**: `/opt/aws_neuronx_venv_pytorch_2_9_nxd_inference` (for preprocess + NxDI direct smoke), `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (for vLLM serving). Both ship with the DLAMI. +- **Disk**: ~700 GB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~290 GB, the Neuron-FP8 preprocessed output is ~310 GB, and `save_sharded_checkpoint=true` writes another ~300 GB per compiled config). + +## Quick Start (FP8 on Trn2) + +End-to-end recipe to go from a fresh trn2.48xlarge to a working vLLM OpenAI server serving MiMo-V2.5 FP8. First-time compile takes ~45-60 minutes; subsequent runs hit the neuronx-cc cache and start in a few minutes. + +```bash +# 1. Clone this repo on the Trn2 instance +cd $HOME +git clone /neuronx-distributed-inference.git +cd neuronx-distributed-inference +git checkout contrib/MiMo-V2.5 # the branch this README lives on + +# 2. Download the HuggingFace FP8 checkpoint (~290 GB). Any HF-compatible +# downloader works; huggingface-cli example: +huggingface-cli download XiaomiMiMo/MiMo-V2.5 \ + --local-dir /opt/dlami/nvme/models/MiMo-V2.5 + +# 3. Preprocess HF FP8 -> Neuron FP8 (~20 min, ~24 GB peak RAM) +source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate +python contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py \ + --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5 \ + --save_path /opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8 \ + --tp_degree 64 + +# 4. (Optional) sanity-check the Neuron-FP8 checkpoint without vLLM +# ~45 min first compile; subsequent runs ~30s to load the pre-sharded NEFF. +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate +python contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py # compile +python contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py # 20-token generate + +# 5. Install vllm-neuron with the contrib registration patch +bash contrib/models/MiMo-V2.5/perf_test/0_setup.sh + +# 6. Start vLLM serving MiMo-V2.5 FP8 (first compile ~60 min; subsequent ~3 min) +bash contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh +``` + +The bench script runs two configurations (BS=32 and BS=128, both +`moe_tp_degree=1 / moe_ep_degree=64`) and logs results under +`/tmp/bench_results/mimo_v2_5/`. + +For a quick `curl` sanity check while the server is up: + +```bash +curl -s http://localhost:8000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{"model": "/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8", + "messages": [{"role": "user", "content": "Hello! Introduce yourself in one sentence."}], + "max_tokens": 64, "temperature": 0.0}' | python3 -m json.tool +``` + +If you get fluent sentence-ending output on a 30+ token generation, the +FP8 path is working correctly. If you see repetition collapse +("helpful helpful helpful..."), double-check that `moe_tp_degree=1`, +`moe_ep_degree=64`, `batch_size>=32`, and that you are loading the +preprocessed Neuron-FP8 checkpoint (not the raw HF FP8 directory). + +## Checkpoint Preparation + +The HuggingFace checkpoint ships as block-wise OCP FP8 (E4M3, ±448 range), which is not directly compatible with Neuron FP8 (IEEE-754 E4M3, ±240 range). Two preprocess scripts are provided: + +### Recommended: FP8 → Neuron-FP8 (streaming) + +`src/conversion_script/preprocess_mimo_v2_5_fp8.py` performs a per-layer streaming rescale from OCP FP8 to Neuron FP8 (per-row scales for attention Q/K/V and layer-0 dense MLP; blockwise scales for MoE experts). `o_proj` is listed in HF's `quantization_config.ignored_layers` and is kept BF16 on the Neuron side (it binds to a plain `RowParallelLinear`, not `QuantizedRowParallel`). Output is ~310 GB across 48 per-layer safetensors shards. + +```bash +python contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py \ + --hf_model_path /path/to/MiMo-V2.5 \ + --save_path /path/to/MiMo-V2.5-Neuron-FP8 \ + --tp_degree 64 +``` + +Peak RAM during preprocessing is ~24 GB; total runtime ~20 minutes on a trn2.48xlarge instance. + +### Fallback: FP8 → BF16 + +`src/conversion_script/preprocess_mimo_v2_fp8.py` dequantizes the entire checkpoint to BF16. Output is ~290 GB; BF16 is numerically equivalent to the published HF FP8 weights and is useful as a known-good reference. Throughput is ~2× worse than the FP8 path because every attention/MLP matmul operates on full BF16 weights. + +## Usage + +```python +import sys +from pathlib import Path + +# Make this contrib package's src/ importable (flat, per upstream contrib convention). +sys.path.insert(0, str(Path("contrib/models/MiMo-V2.5/src").resolve())) + +import torch +from transformers import AutoConfig, AutoTokenizer +from neuronx_distributed_inference.models.config import MoENeuronConfig, OnDeviceSamplingConfig +from neuronx_distributed_inference.utils.hf_adapter import load_pretrained_config, HuggingFaceGenerationAdapter + +from modeling_mimo_v2 import NeuronMiMoV2ForCausalLM, MiMoV2InferenceConfig + +model_path = "/path/to/MiMo-V2.5-Neuron-FP8/" +compiled_path = "/path/to/compiled/" + +# Recommended FP8 recipe: +# moe_tp_degree = 1, moe_ep_degree = 64 +# See "FP8 Configuration Notes" below for why other moe_tp/ep ratios collapse. +neuron_config = MoENeuronConfig( + tp_degree=64, + ep_degree=1, # keep outer EP = 1; only MoE-internal EP varies + moe_tp_degree=1, + moe_ep_degree=64, + batch_size=32, # must be >= num_experts / top_k = 256 / 8 = 32 + max_batch_size=32, + ctx_batch_size=1, + tkg_batch_size=32, + seq_len=1024, + n_active_tokens=128, + torch_dtype=torch.bfloat16, + logical_nc_config=2, + capacity_factor=1.0, + glu_mlp=True, + fused_qkv=False, # required: asymmetric Q/K (192) vs V (128) head dims + router_config={"act_fn": "sigmoid", "dtype": "float32"}, + blockwise_matmul_config={ + "use_shard_on_block_dynamic_while": True, + "block_sharding_strategy": "PING_PONG", + }, + save_sharded_checkpoint=True, + quantized=True, + quantized_checkpoints_path=model_path, + quantization_dtype="f8e4m3", + quantization_type="blockwise_symmetric", + quantization_block_axis=[1, 2], + quantization_block_size=[128, 128], + modules_to_not_convert=[ + "embed_tokens", "lm_head", "norm", "router", "o_proj", + ], + on_device_sampling_config=OnDeviceSamplingConfig( + do_sample=True, temperature=0.6, top_k=20, top_p=0.95, + ), +) + +# trust_remote_code is required by MiMo-V2's HF config; pre-load via AutoConfig +# and pass to NxDI so load_pretrained_config does not re-load without the flag. +hf_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) +config = MiMoV2InferenceConfig( + neuron_config, load_config=load_pretrained_config(hf_config=hf_config), +) + +model = NeuronMiMoV2ForCausalLM(model_path, config) +model.compile(compiled_path) +model.load(compiled_path) + +tokenizer = AutoTokenizer.from_pretrained(model_path, trust_remote_code=True) +adapter = HuggingFaceGenerationAdapter(model) +inputs = tokenizer(["Hello, how are you?"] * 32, return_tensors="pt", padding=True) +output = adapter.generate( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + max_new_tokens=128, +) +``` + +For a minimal end-to-end smoke test that bypasses vLLM, see: + +- `perf_test/smoke_compile_mimo_v2_5.py` — compile + load (STAGE=instantiate|compile|load|all, DRY_RUN, SKIP_WARMUP) +- `perf_test/smoke_generate_mimo_v2_5.py` — 20-token generation via HuggingFaceGenerationAdapter + +Both default to the recommended FP8 recipe (`moe_tp=1`, `moe_ep=64`). + +## FP8 Configuration Notes + +### moe_tp_degree = 1, moe_ep_degree = 64 + +**Why**: at `moe_tp_degree=64` each rank owns 1/64 of the intermediate dim, which for MiMo-V2.5 (MoE intermediate = 2048) is 32 rows — **below the 128-row blockwise scale block**. NxDI's `_setup_for_scale` detects `weight_shape[axis] < block_size` and collapses the per-rank scale dim to 1, losing per-channel FP8 scale granularity. The resulting drift compounds across MiMo-V2.5's 47 MoE layers and manifests as output collapse ("helpful helpful helpful ...") after roughly 30 decode tokens. + +`moe_tp_degree=1, moe_ep_degree=64` keeps each expert's weights and blockwise scales intact on a single rank (4 experts per rank), which preserves per-channel scale and produces correct output even on long multi-turn prompts. + +Intermediate ratios (`moe_tp=32/ep=2` or `moe_tp=16/ep=4`) have been empirically tested and still produce gibberish, so this is the only currently-supported moe_tp/ep combination for MiMo-V2.5 FP8. + +### batch_size >= 32 + +NxDI's TKG (token generation) path refuses Expert Parallelism when `batch_size < num_experts / top_k`. For MiMo-V2.5 that is 256 / 8 = 32, so the smallest working BS on the FP8 path is 32. BS=1 latency demos are not currently possible on FP8; use the BF16 checkpoint with `moe_tp=64, moe_ep=1, batch_size=1` for single-stream latency measurements. + +### outer ep_degree = 1 + +`MoENeuronConfig.ep_degree` is the **full-model** expert-parallel factor. Setting it to anything > 1 multiplies `world_size` to `tp_degree * ep_degree`, which on a 64-NC Trn2 overflows the device (ranks beyond 63 have no backing hardware, sharded-checkpoint size grows linearly, and load fails). The MoE-internal expert parallelism is controlled exclusively by `moe_ep_degree` — keep `ep_degree=1` at the outer level. + +## vLLM Integration + +MiMo-V2.5 can be served via [vllm-neuron](https://github.com/aws-neuron/vllm-neuron). A contrib registration patch is required to plug the NxDI modeling code into vllm-neuron's lookup tables. + +### Setup + +```bash +# The setup script clones vllm-project/vllm-neuron at release-0.5.0, applies +# the contrib registration patch, installs it editable, and downloads MiMo-V2.5 +# weights (BF16 by default; set MIMO_V2_5_PATH to override). +bash contrib/models/MiMo-V2.5/perf_test/0_setup.sh +``` + +The patch (`perf_test/vllm-neuron-patch.patch`) is 40 lines and only touches `vllm_neuron/__init__.py`. It adds a `_register_contrib_models()` hook that, when `NXDI_CONTRIB_MIMO_V2_5_SRC` is set, registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under the key `mimo_v2_5` **and** registers the `MiMoV2ForCausalLM` architecture into vLLM's `ModelRegistry`. No upstream vLLM or NxDI source is modified. + +### Serving (FP8, recommended) + +```bash +export NXDI_CONTRIB_MIMO_V2_5_SRC=/path/to/neuronx-distributed-inference/contrib/models/MiMo-V2.5/src +export MIMO_V2_5_PATH=/path/to/MiMo-V2.5-Neuron-FP8 +# First-time compile of MiMo-V2.5's 256-expert MoE takes 30-60 minutes. +export VLLM_ENGINE_READY_TIMEOUT_S=7200 +# Optional: isolate compile cache per config so parallel MiMo-V2.5/Pro/etc. compiles +# don't race on the default /var/tmp/neuron-compile-cache lock files. +export NEURON_COMPILED_ARTIFACTS=/path/to/compiled/mimo_v2_5_bs32_moetp1_ep64_fp8 + +python3 -m vllm.entrypoints.openai.api_server \ + --model "$MIMO_V2_5_PATH" \ + --tensor-parallel-size 64 \ + --max-model-len 1024 \ + --max-num-seqs 32 \ + --no-enable-chunked-prefill \ + --no-enable-prefix-caching \ + --trust_remote_code \ + --additional-config '{ + "override_neuron_config": { + "tp_degree": 64, + "logical_nc_config": 2, + "fused_qkv": false, + "sequence_parallel_enabled": false, + "glu_mlp": true, + "normalize_top_k_affinities": true, + "save_sharded_checkpoint": true, + "router_config": {"act_fn": "sigmoid", "dtype": "float32"}, + "quantized": true, + "quantized_checkpoints_path": "/path/to/MiMo-V2.5-Neuron-FP8", + "quantization_dtype": "f8e4m3", + "quantization_type": "blockwise_symmetric", + "quantization_block_axis": [1, 2], + "quantization_block_size": [128, 128], + "modules_to_not_convert": ["embed_tokens", "lm_head", "norm", "router", "o_proj"], + "blockwise_matmul_config": {"use_shard_on_block_dynamic_while": true, "block_sharding_strategy": "PING_PONG"}, + "moe_tp_degree": 1, + "moe_ep_degree": 64, + "batch_size": 32, + "ctx_batch_size": 1, + "tkg_batch_size": 32, + "max_context_length": 1024, + "seq_len": 1024, + "is_continuous_batching": true, + "enable_bucketing": true, + "context_encoding_buckets": [1024], + "token_generation_buckets": [1024], + "async_mode": true, + "on_device_sampling_config": { + "do_sample": true, "temperature": 0.6, "top_k": 20, "top_p": 0.95 + } + } + }' +``` + +See `perf_test/bench_mimo_v2_5.sh` for the full benchmark recipe at BS=32 and BS=128. + +### vllm-neuron patch summary + +The patch is applied to vllm-neuron 0.5.0 and: + +- Maps the `MiMoV2ForCausalLM` architecture to MiMo-V2.5's model loader (reusing the Qwen2-family loader path, which MiMo-V2.5's tokenizer inherits from). +- Passes `hf_config` from vLLM into `load_pretrained_config` so NxDI does not re-load the config without `trust_remote_code=True`. +- Replaces vllm-neuron's internal `AutoModelForCausalLM.from_pretrained` call with `huggingface_hub.snapshot_download`, which is the only path that works for `trust_remote_code=True` models when no GPU is available for HF's CUDA-gated FP8 quantizer. + +## Performance + +> These numbers are from the earlier BF16 recipe (pre-FP8 rollout). FP8 numbers will be added once a stable bench run completes on the new recipe; preliminary single-stream qualitative tests show fluent multi-sentence output on long Chinese chat prompts with `moe_tp=1, moe_ep=64, batch_size=32`. + +### Standalone NxDI (trn2.48xlarge, BF16, TP=64, EP=64) + +| Batch Size | Throughput (tok/s) | +|------------|-------------------| +| 1 | 29.92 | +| 8 | 215.94 | +| 32 | 649.14 | + +### vLLM Serving (trn2.48xlarge, BF16, BS=32, TP=64/EP=64, CB) + +Input/output: 900/90 tokens (random dataset) + +| 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 | + +> **Compile time:** the first MiMo-V2.5 compile on SDK 2.29 is ~30-60 minutes for the TKG NEFF and similar for the CTE NEFF. Subsequent runs with the same `override_neuron_config` hit the neuronx-cc cache and start in ~1-2 minutes. `save_sharded_checkpoint=true` additionally persists per-rank FP8 shards under `/weights/`, letting future `load()` calls skip the ~10-minute shard_checkpoint pass. + +## Compatibility Matrix + +| Instance | Neuron SDK 2.29+ (PyTorch 2.9) | 2.21 and earlier | +|----------|--------------------------------|------------------| +| Trn2 (trn2.48xlarge) | Tested | Not tested | +| Trn1 | Not supported (requires 64 logical cores via logical_nc_config=2) | Not supported | +| Inf2 | Not supported | Not supported | + +## Testing + +```bash +pytest contrib/models/MiMo-V2.5/test/integration/test_model.py -v +``` + +## Key Implementation Notes + +1. **Hybrid Attention**: `hybrid_layer_pattern` list determines full vs sliding window per layer; the modeling code constructs one `NeuronMiMoV2Attention` per layer with the correct `is_sliding_window` flag and rope_theta. +2. **CONVERT_TO_MHA**: When `tp_degree > num_kv_heads` (64 > 4 full / 64 > 8 SWA), K/V are replicated to `num_attention_heads` (64) during state-dict conversion; this applies to both `.weight` and the per-row `.scale` on the FP8 path. +3. **Attention Sink Bias**: Learnable per-head bias added as an extra "sink" column to attention scores in sliding window layers (not added in full-attention layers). Per-rank slicing of the bias happens inside `forward()` based on `parallel_state.get_tensor_model_parallel_rank()`. +4. **FP8 Path Caveats**: + - Must use `moe_tp_degree=1, moe_ep_degree=64` (see "FP8 Configuration Notes" above). + - Must use `batch_size >= 32` (NxDI EP>1 requirement). + - Must keep outer `ep_degree=1` (only `moe_ep_degree` should vary). + - Several runtime monkey-patches (router bias, blockwise scale stride, 2D per-channel, EP scale handling) are installed automatically in `NeuronMiMoV2ForCausalLM.__init__` when `quantized=True`; the BF16 path is untouched. + +## Example Checkpoints + +* [XiaomiMiMo/MiMo-V2.5](https://huggingface.co/XiaomiMiMo/MiMo-V2.5) — HF FP8 source checkpoint + +## Maintainer + +Henan Wan (whn09) + +**Last Updated:** 2026-04-25 diff --git a/contrib/models/MiMo-V2.5/perf_test/0_setup.sh b/contrib/models/MiMo-V2.5/perf_test/0_setup.sh new file mode 100755 index 00000000..7dfec5ca --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/0_setup.sh @@ -0,0 +1,65 @@ +#!/bin/bash +# Setup for MiMo-V2.5 vLLM benchmarking on Trn2. +# +# Clones upstream vllm-project/vllm-neuron at release-0.5.0 and applies +# vllm-neuron-patch.patch, which adds a runtime registration hook so the +# contrib NeuronMiMoV2ForCausalLM is plugged into both NxDI's MODEL_TYPES +# (under the key "mimov2") and vLLM's ModelRegistry (as +# MiMoV2ForCausalLM) at vllm-neuron plugin init time. +# +# Then downloads XiaomiMiMo/MiMo-V2.5 from HuggingFace (FP8 blockwise, ~320 GB). +set -e + +echo "==========================================" +echo "Setup: vllm-neuron + MiMo-V2.5 weights" +echo "==========================================" + +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate + +PATCH_FILE="$(cd "$(dirname "$0")" && pwd)/vllm-neuron-patch.patch" + +echo "" +echo "[1/2] Installing vllm-neuron (release-0.5.0) with the contrib registration patch..." + +if [ ! -d $HOME/vllm-neuron ]; then + git clone --branch release-0.5.0 https://github.com/vllm-project/vllm-neuron.git $HOME/vllm-neuron +fi + +cd $HOME/vllm-neuron + +# Apply patch (idempotent via `git apply --check` first). +if git apply --check "$PATCH_FILE" 2>/dev/null; then + git apply "$PATCH_FILE" + echo " Applied $PATCH_FILE" +else + echo " Patch already applied or conflicts; continuing." +fi + +pip install --extra-index-url=https://pip.repos.neuron.amazonaws.com -e . + +python3 -c "import vllm_neuron; print('vllm-neuron installed:', vllm_neuron.__file__)" + +echo "" +echo "[2/2] Downloading MiMo-V2.5 FP8 weights from HuggingFace..." + +MIMO_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5}" +if [ -d "$MIMO_PATH" ] && [ "$(ls "$MIMO_PATH"/*.safetensors 2>/dev/null | wc -l)" -gt 0 ]; then + echo " MiMo-V2.5 weights already exist at $MIMO_PATH, skipping download" +else + mkdir -p "$MIMO_PATH" + huggingface-cli download XiaomiMiMo/MiMo-V2.5 --local-dir "$MIMO_PATH" --max-workers 16 + echo " Download complete: $(du -sh $MIMO_PATH | cut -f1)" +fi + +CONTRIB_SRC="$(cd "$(dirname "$0")/.." && pwd)/src" + +echo "" +echo "Next, preprocess the FP8 checkpoint for Neuron (~15 min, ~15 GB peak RAM):" +echo " python $CONTRIB_SRC/conversion_script/preprocess_mimo_v2_5_fp8.py \\" +echo " --hf_model_path $MIMO_PATH \\" +echo " --save_path ${MIMO_PATH}-Neuron-FP8 \\" +echo " --tp_degree 64" +echo "" +echo "Then before running the benchmark, export:" +echo " export MIMO_V2_5_PATH=${MIMO_PATH}-Neuron-FP8" +echo " export NXDI_CONTRIB_MIMO_V2_5_SRC=$CONTRIB_SRC" diff --git a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh new file mode 100755 index 00000000..45089d62 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh @@ -0,0 +1,228 @@ +#!/bin/bash +set -e + +# MiMo-V2.5 FP8 vLLM benchmark on Trn2. +# +# Requires a Neuron-FP8 preprocessed checkpoint (see +# `src/conversion_script/preprocess_mimo_v2_5_fp8.py`). The configs below +# all use moe_tp_degree=1 / moe_ep_degree=64 (experts sharded by expert +# parallelism only, no intra-expert TP split) because moe_tp_degree=64 collapses +# the per-rank FP8 blockwise scale to a singleton — per-rank expert +# intermediate is 32 rows, below the 128-row blockwise block, so +# NxDI's `_setup_for_scale` drops per-channel scale granularity. The resulting +# drift compounds across 47 MoE layers and gives repetition / output collapse. +# Using moe_ep_degree=64 keeps all of each expert's weight + scale on one rank +# (4 experts per rank), which preserves the blockwise scale intact. +# +# NxDI's TKG path refuses Expert Parallelism with BS < num_experts/top_k +# (256 / 8 = 32 for MiMo-V2.5), so the smallest working batch size here is 32. +# If you want BS=1 behaviour, the FP8 path is not currently supported on +# this model on Trn2 — use the BF16 checkpoint with the old bench recipe +# (`moe_tp_degree=64, moe_ep_degree=1, batch_size=1`). + +source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate + +MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8}" +# The NxDI contrib MiMo-V2.5 modeling code is registered into vLLM / +# NxDI lookup tables by vllm-neuron's register() hook using this env var. +# Default to this contrib package's own src/ relative to the script. +: "${NXDI_CONTRIB_MIMO_V2_5_SRC:=$(cd "$(dirname "$0")/.." && pwd)/src}" +export NXDI_CONTRIB_MIMO_V2_5_SRC + +# First-time MiMo-V2.5 FP8 compile takes 30-60 minutes; extend vLLM's ready +# timeout and the compiler's environment variables for FP8 numerics. +export VLLM_ENGINE_READY_TIMEOUT_S=7200 + +PORT=8000 +RESULTS_DIR="/tmp/bench_results/mimo_v2_5" +mkdir -p "$RESULTS_DIR" + +# Common neuron config shared across all MiMo-V2.5 FP8 configs. +# save_sharded_checkpoint=true persists per-rank sharded weights to +# /weights/tp{N}_sharded_checkpoint.safetensors during compile; +# load() then reads those directly (~30s) instead of re-sharding the entire +# checkpoint on every vllm-neuron startup (~10+ min). +COMMON_MIMO_CONFIG='"tp_degree": 64, + "logical_nc_config": 2, + "fused_qkv": false, + "sequence_parallel_enabled": false, + "glu_mlp": true, + "normalize_top_k_affinities": true, + "save_sharded_checkpoint": true, + "router_config": {"act_fn": "sigmoid", "dtype": "float32"}, + "quantized": true, + "quantized_checkpoints_path": "'"$MODEL_PATH"'", + "quantization_dtype": "f8e4m3", + "quantization_type": "blockwise_symmetric", + "quantization_block_axis": [1, 2], + "quantization_block_size": [128, 128], + "modules_to_not_convert": ["embed_tokens", "lm_head", "norm", "router", "o_proj"], + "blockwise_matmul_config": {"use_shard_on_block_dynamic_while": true, "block_sharding_strategy": "PING_PONG"}' + +# Helper: wait for vLLM server to be ready. First-time compilation of a +# 256-expert MoE model takes 30-90 minutes, so we poll for up to 2 hours. +wait_for_server() { + echo " Waiting for vLLM server to be ready (up to 2h for first compile)..." + local interval=10 + local max_attempts=720 # 720 * 10s = 7200s = 2h + local start=$SECONDS + for i in $(seq 1 $max_attempts); do + if curl -s http://localhost:$PORT/health > /dev/null 2>&1; then + echo " Server ready! (waited $((SECONDS - start))s)" + return 0 + fi + # Show a progress blip every minute so the user knows we're alive + if [ $((i % 6)) -eq 0 ]; then + echo " ...still waiting ($((SECONDS - start))s elapsed)" + fi + sleep $interval + done + echo " ERROR: Server did not start within $((max_attempts * interval))s" + return 1 +} + +# Helper: run benchmark +run_bench() { + local config_name=$1 + local concurrency=$2 + local num_prompts=$3 + + echo " Benchmark: concurrency=$concurrency, prompts=$num_prompts" + vllm bench serve \ + --backend vllm \ + --model "$MODEL_PATH" \ + --tokenizer "$MODEL_PATH" \ + --endpoint /v1/completions \ + --dataset-name random \ + --num-prompts "$num_prompts" \ + --random-input-len 900 \ + --random-output-len 90 \ + --random-range-ratio 0.03 \ + --max-concurrency "$concurrency" \ + 2>&1 | tee "$RESULTS_DIR/${config_name}_c${concurrency}.txt" + echo "" +} + +# Helper: stop server +stop_server() { + echo " Stopping vLLM server..." + pkill -f "vllm.entrypoints.openai.api_server" 2>/dev/null || true + sleep 5 +} + +# Helper: quick sanity check +sanity_check() { + echo " Running sanity check..." + curl -s http://localhost:$PORT/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d '{ + "messages": [{"role": "user", "content": "What is 1+1? Answer briefly."}], + "model": "'"$MODEL_PATH"'", + "max_tokens": 64, + "temperature": 0.0, + "stream": false + }' | python3 -c "import sys,json; r=json.load(sys.stdin); print(' Sanity:', r['choices'][0]['message']['content'][:100])" 2>/dev/null || echo " Sanity check: could not parse response" +} + +echo "==========================================" +echo "MiMo-V2.5 FP8 Performance Benchmark" +echo "==========================================" +echo "Model: $MODEL_PATH" +echo "Results: $RESULTS_DIR" +echo "" + +############################################################################### +# Config 1: BS=32, TP=64 + moe_tp=1/moe_ep=64, CB + bucketing (smallest BS +# that satisfies NxDI's Expert-Parallel BS >= num_experts/top_k requirement). +############################################################################### +CONFIG_NAME="bs32_tp64_moetp1_ep64" +echo "--- Config 1: BS=32, moe_tp=1/moe_ep=64, CB + bucketing ---" + +python3 -m vllm.entrypoints.openai.api_server \ + --model "$MODEL_PATH" \ + --tokenizer "$MODEL_PATH" \ + --tensor-parallel-size 64 \ + --max-model-len 1024 \ + --max-num-seqs 32 \ + --no-enable-chunked-prefill \ + --no-enable-prefix-caching \ + --port $PORT \ + --trust_remote_code \ + --additional-config '{ + "override_neuron_config": { + '"$COMMON_MIMO_CONFIG"', + "moe_tp_degree": 1, + "moe_ep_degree": 64, + "batch_size": 32, + "ctx_batch_size": 1, + "tkg_batch_size": 32, + "max_context_length": 1024, + "seq_len": 1024, + "is_continuous_batching": true, + "enable_bucketing": true, + "context_encoding_buckets": [1024], + "token_generation_buckets": [1024], + "async_mode": true, + "on_device_sampling_config": { + "do_sample": true, "temperature": 0.6, "top_k": 20, "top_p": 0.95 + } + } + }' & + +wait_for_server +sanity_check +run_bench "$CONFIG_NAME" 1 16 +run_bench "$CONFIG_NAME" 16 128 +run_bench "$CONFIG_NAME" 32 128 +stop_server + +############################################################################### +# Config 2: BS=128, TP=64 + moe_tp=1/moe_ep=64, CB + bucketing (throughput). +############################################################################### +CONFIG_NAME="bs128_tp64_moetp1_ep64" +echo "--- Config 2: BS=128, moe_tp=1/moe_ep=64, CB + bucketing ---" + +python3 -m vllm.entrypoints.openai.api_server \ + --model "$MODEL_PATH" \ + --tokenizer "$MODEL_PATH" \ + --tensor-parallel-size 64 \ + --max-model-len 1024 \ + --max-num-seqs 128 \ + --no-enable-chunked-prefill \ + --no-enable-prefix-caching \ + --port $PORT \ + --trust_remote_code \ + --additional-config '{ + "override_neuron_config": { + '"$COMMON_MIMO_CONFIG"', + "moe_tp_degree": 1, + "moe_ep_degree": 64, + "batch_size": 128, + "ctx_batch_size": 1, + "tkg_batch_size": 128, + "max_context_length": 1024, + "seq_len": 1024, + "is_continuous_batching": true, + "enable_bucketing": true, + "context_encoding_buckets": [1024], + "token_generation_buckets": [1024], + "async_mode": true, + "on_device_sampling_config": { + "do_sample": true, "temperature": 0.6, "top_k": 20, "top_p": 0.95 + } + } + }' & + +wait_for_server +sanity_check +run_bench "$CONFIG_NAME" 1 16 +run_bench "$CONFIG_NAME" 16 128 +run_bench "$CONFIG_NAME" 32 128 +run_bench "$CONFIG_NAME" 128 512 +stop_server + +echo "==========================================" +echo "MiMo-V2.5 FP8 benchmarks complete!" +echo "Results saved to: $RESULTS_DIR" +echo "==========================================" +ls -la "$RESULTS_DIR" diff --git a/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh new file mode 100755 index 00000000..a3c5ecb6 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh @@ -0,0 +1,76 @@ +#!/bin/bash +# Run a single vllm-bench-serve pass against an already-running vLLM server. +# +# Unlike bench_mimo_v2_5.sh this script does NOT launch or kill the vLLM +# server — you bring your own. That makes it convenient when the bench driver +# in bench_mimo_v2_5.sh times out during first-time compilation: the server +# keeps running, and once it's ready you can collect numbers with this. +# +# Usage: +# bash run_bench_single.sh # defaults: c=1, 16 prompts +# CONCURRENCY=16 NUM_PROMPTS=128 bash run_bench_single.sh +# CONFIG_NAME=bs32_tp1_ep64_opt CONCURRENCY=16 NUM_PROMPTS=128 bash run_bench_single.sh +# +# Environment knobs: +# PORT vLLM server port (default 8000) +# MIMO_V2_5_PATH Path to the BF16 checkpoint (default +# /opt/dlami/nvme/models/MiMo-V2.5-BF16) +# CONCURRENCY --max-concurrency (default 1) +# NUM_PROMPTS --num-prompts (default 16) +# INPUT_LEN --random-input-len (default 900) +# OUTPUT_LEN --random-output-len (default 90) +# RANGE_RATIO --random-range-ratio (default 0.03) +# CONFIG_NAME Used in the output filename (default bs1_tp64_ep1) +# RESULTS_DIR Where to dump per-run log (default /tmp/bench_results/mimo_v2_5) + +set -e + +source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate + +MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-BF16}" +PORT="${PORT:-8000}" +CONCURRENCY="${CONCURRENCY:-1}" +NUM_PROMPTS="${NUM_PROMPTS:-16}" +INPUT_LEN="${INPUT_LEN:-900}" +OUTPUT_LEN="${OUTPUT_LEN:-90}" +RANGE_RATIO="${RANGE_RATIO:-0.03}" +CONFIG_NAME="${CONFIG_NAME:-bs1_tp64_ep1}" +RESULTS_DIR="${RESULTS_DIR:-/tmp/bench_results/mimo_v2_5}" + +mkdir -p "$RESULTS_DIR" + +echo "==========================================" +echo "MiMo-V2.5 single-run benchmark" +echo "==========================================" +echo " Model: $MODEL_PATH" +echo " Port: $PORT" +echo " Config: $CONFIG_NAME" +echo " Concurrency: $CONCURRENCY" +echo " Prompts: $NUM_PROMPTS" +echo " Input len: $INPUT_LEN Output len: $OUTPUT_LEN" +echo " Results: $RESULTS_DIR/${CONFIG_NAME}_c${CONCURRENCY}.txt" +echo "" + +# Quick health check +if ! curl -sf "http://localhost:$PORT/health" > /dev/null; then + echo "ERROR: vLLM server is not responding on http://localhost:$PORT" + echo "Start it first (e.g., bench_mimo_v2_5.sh) and wait until" + echo "'Application startup complete.' is printed." + exit 1 +fi + +vllm bench serve \ + --backend vllm \ + --model "$MODEL_PATH" \ + --tokenizer "$MODEL_PATH" \ + --endpoint /v1/completions \ + --dataset-name random \ + --num-prompts "$NUM_PROMPTS" \ + --random-input-len "$INPUT_LEN" \ + --random-output-len "$OUTPUT_LEN" \ + --random-range-ratio "$RANGE_RATIO" \ + --max-concurrency "$CONCURRENCY" \ + 2>&1 | tee "$RESULTS_DIR/${CONFIG_NAME}_c${CONCURRENCY}.txt" + +echo "" +echo "Saved to: $RESULTS_DIR/${CONFIG_NAME}_c${CONCURRENCY}.txt" diff --git a/contrib/models/MiMo-V2.5/perf_test/sanity_check.sh b/contrib/models/MiMo-V2.5/perf_test/sanity_check.sh new file mode 100755 index 00000000..8dca4c7c --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/sanity_check.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# Quick sanity check against an already-running vLLM server. +# +# Assumes vLLM is already listening on $PORT (default 8000) with MiMo-V2.5 +# loaded. Sends a single chat completion and prints the model's reply. +# +# Usage: +# bash sanity_check.sh # uses defaults +# PORT=8001 bash sanity_check.sh # custom port +# PROMPT="..." bash sanity_check.sh # custom prompt + +set -e + +MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-BF16}" +PORT="${PORT:-8000}" +PROMPT="${PROMPT:-What is 1+1? Answer briefly.}" +MAX_TOKENS="${MAX_TOKENS:-64}" + +echo "Sanity check: POST /v1/chat/completions on port $PORT" +echo " Model: $MODEL_PATH" +echo " Prompt: $PROMPT" +echo " Max tokens: $MAX_TOKENS" +echo "" + +# Health check first — fail fast if server isn't up. +if ! curl -sf "http://localhost:$PORT/health" > /dev/null; then + echo "ERROR: vLLM server is not responding on http://localhost:$PORT" + echo "Start it with 'bash bench_mimo_v2_5.sh' or your own launcher first." + exit 1 +fi + +RESPONSE=$(curl -s "http://localhost:$PORT/v1/chat/completions" \ + -H 'Content-Type: application/json' \ + -d "$(cat </dev/null || echo "$RESPONSE" +echo "" + +# Extract the model's reply for a human-friendly one-liner summary. +REPLY=$(echo "$RESPONSE" | python3 -c " +import json, sys +try: + r = json.load(sys.stdin) + print(r['choices'][0]['message']['content'].strip()) +except Exception as e: + print(f'(could not parse reply: {e})') +" 2>/dev/null) + +echo "Model reply: $REPLY" diff --git a/contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py b/contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py new file mode 100755 index 00000000..7113e611 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python3 +"""Minimal compile+load smoke test for MiMo-V2.5 FP8 on Trn2. + +Bypasses vLLM entirely so we can iterate on the preprocessed Neuron-FP8 +checkpoint without paying vllm-neuron's startup cost. Builds the MiMo-V2.5 BS=1 +recipe (TP=64, EP=1, blockwise FP8 for routed experts), compiles to a temp +dir, then loads. EP=1 lets the TKG path enter forward_selective_loading +legally so BS=1 compiles — with EP>1 NxDI raises NotImplementedError and +forces BS>=num_experts/top_k = 32. + +STAGE controls how far we go: + instantiate | compile | load | all (default: all) + +DRY_RUN=1 does HLO-only compile (no torch.jit.save + shard). Fastest sanity +check for the preprocessed checkpoint. SKIP_WARMUP=1 on load() skips the +forward pass that allocates the shared scratchpad — useful when HBM is +tight. + +Run under /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16 (same venv used +by the bench script). +""" + +import os +import sys +import time +import traceback + +MODEL_PATH = os.environ.get( + "MIMO_V2_5_MODEL_PATH", + "/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8", +) +COMPILED_PATH = os.environ.get( + "MIMO_V2_5_COMPILED_PATH", + "/opt/dlami/nvme/compiled/mimo_v2_5_tp64_moetp1_ep64_fp8/", +) + +TP_DEGREE = int(os.environ.get("TP_DEGREE", "64")) +SEQ_LEN = int(os.environ.get("SEQ_LEN", "1024")) +BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "1")) +CTX_BATCH_SIZE = int(os.environ.get("CTX_BATCH_SIZE", "1")) +# Default to moe_tp=1 / moe_ep=64. Under FP8 + moe_tp=64 (our old default) +# each rank's MoE expert intermediate slice is 32 rows (<128, the scale +# block size), which collapses the per-rank scale to a singleton in +# NxDI's `_setup_for_scale` — losing per-channel FP8 scale granularity +# and producing a BF16-accumulator drift that compounds into output +# collapse after ~30 decode tokens. moe_tp=1/moe_ep=64 keeps every expert +# on a single rank (4 full experts per rank), so each expert's scale +# survives intact. Override via MOE_TP / MOE_EP env vars for other recipes. +MOE_TP = int(os.environ.get("MOE_TP", "1")) +MOE_EP = int(os.environ.get("MOE_EP", "64")) + +STAGE = os.environ.get("STAGE", "all").lower() + +os.makedirs(COMPILED_PATH, exist_ok=True) + +# NxDI's model builder uses a per-process temp workdir for HLO/NEFF staging +# (BASE_COMPILE_WORK_DIR, default "/tmp/nxd_model/"). If two compiles run in +# parallel with the same default, they silently overwrite each other's +# .hlo_module.pb files and one or both compilations crash with +# "neuronx-cc returned non-zero exit status 70". Pin the workdir to a +# unique per-COMPILED_PATH subdir to stay safe under any parallel invocation. +os.environ.setdefault( + "BASE_COMPILE_WORK_DIR", + os.path.join("/tmp/nxd_model", os.path.basename(COMPILED_PATH.rstrip("/"))), +) + + +def main(): + from neuronx_distributed_inference.models.config import MoENeuronConfig + from neuronx_distributed_inference.utils.hf_adapter import load_pretrained_config + + # Import the contrib wrapper (sibling src dir). + contrib_src = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "src", + ) + sys.path.insert(0, os.path.abspath(contrib_src)) + + from modeling_mimo_v2 import ( + MiMoV2InferenceConfig, + NeuronMiMoV2ForCausalLM, + ) + + print(f"[smoke] MODEL_PATH={MODEL_PATH}") + print(f"[smoke] COMPILED_PATH={COMPILED_PATH}") + print(f"[smoke] TP_DEGREE={TP_DEGREE}, SEQ_LEN={SEQ_LEN}, BS={BATCH_SIZE}") + print(f"[smoke] MOE_TP={MOE_TP}, MOE_EP={MOE_EP}") + print(f"[smoke] STAGE={STAGE}") + + print("[smoke] Building MoENeuronConfig (quantized FP8 MoE, blockwise_symmetric)...") + # NOTE: ep_degree at the top level controls the OUTER (full model) + # expert-parallel factor, which multiplies world_size to + # tp_degree * ep_degree and duplicates non-MoE weights per replica. + # At world_size > 64 on a 64-NC Trn2, sharded weights grow accordingly + # (e.g. tp=64 + ep=4 -> 256 ranks -> 4x the sharded checkpoint size, + # and at runtime the model doesn't fit on the device). For MoE-only + # EP we want ep_degree=1 at the outer level and the per-MoE split + # controlled solely by moe_ep_degree (which Pro's working benches + # also do). Keep ep_degree=1 unconditionally. + neuron_config = MoENeuronConfig( + tp_degree=TP_DEGREE, + ep_degree=1, + logical_nc_config=2, + batch_size=BATCH_SIZE, + max_batch_size=BATCH_SIZE, + ctx_batch_size=CTX_BATCH_SIZE, + tkg_batch_size=BATCH_SIZE, + seq_len=SEQ_LEN, + n_active_tokens=128, + torch_dtype="bfloat16", + capacity_factor=1.0, + glu_mlp=True, + moe_ep_degree=MOE_EP, + moe_tp_degree=MOE_TP, + context_encoding_buckets=[SEQ_LEN], + router_config={"act_fn": "sigmoid", "dtype": "float32"}, + # SDK 2.29 ships only bwmm_shard_on_block / bwmm_shard_on_intermediate; + # default routes to _call_shard_hidden_kernel which is missing, so we + # take the shard-on-block path via this flag. + blockwise_matmul_config={ + "use_shard_on_block_dynamic_while": True, + "block_sharding_strategy": "PING_PONG", + }, + # Persist sharded FP8 weights to disk so subsequent load()s skip the + # ~10-minute shard_checkpoint step (writes weights/tp{0..63}_*.safetensors + # on NVMe; NxDI load() reads these directly when present). + save_sharded_checkpoint=True, + # FP8 blockwise for routed experts (Kimi-K2 recipe). + quantized=True, + quantized_checkpoints_path=MODEL_PATH, + quantization_dtype="f8e4m3", + quantization_type="blockwise_symmetric", + quantization_block_axis=[1, 2], + quantization_block_size=[128, 128], + modules_to_not_convert=[ + "embed_tokens", + "lm_head", + "norm", + "router", + "o_proj", + ], + ) + + print("[smoke] Building MiMoV2InferenceConfig...") + from transformers import AutoConfig + hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) + config = MiMoV2InferenceConfig( + neuron_config, load_config=load_pretrained_config(hf_config=hf_config) + ) + print(f"[smoke] config.hidden_size={config.hidden_size}") + print(f"[smoke] config.num_hidden_layers={config.num_hidden_layers}") + print(f"[smoke] config.n_routed_experts={config.n_routed_experts}") + print(f"[smoke] config.num_experts_per_tok={config.num_experts_per_tok}") + print(f"[smoke] config.layer_uses_moe[:5]={config.layer_uses_moe[:5]}") + print(f"[smoke] config.layer_attention_types[:5]={config.layer_attention_types[:5]}") + + print("[smoke] Instantiating NeuronMiMoV2ForCausalLM (build model-on-cpu)...") + t0 = time.time() + model = NeuronMiMoV2ForCausalLM(MODEL_PATH, config) + print(f"[smoke] Instantiated in {time.time() - t0:.1f}s") + + if STAGE == "instantiate": + print("[smoke] STAGE=instantiate only, skipping compile/load.") + return + + DRY_RUN = os.environ.get("DRY_RUN", "0") == "1" + if STAGE in ("compile", "all"): + label = "Dry-run compile (HLO only)" if DRY_RUN else "Full compile" + print(f"[smoke] {label} -> {COMPILED_PATH}") + t0 = time.time() + try: + model.compile(COMPILED_PATH, dry_run=DRY_RUN) + print(f"[smoke] {label} OK in {time.time() - t0:.1f}s") + except Exception: + print(f"[smoke] {label} FAILED:") + traceback.print_exc() + raise + + if STAGE in ("load", "all") and not DRY_RUN: + SKIP_WARMUP = os.environ.get("SKIP_WARMUP", "1") == "1" + print(f"[smoke] Loading compiled model from {COMPILED_PATH} (skip_warmup={SKIP_WARMUP})") + t0 = time.time() + model.load(COMPILED_PATH, skip_warmup=SKIP_WARMUP) + print(f"[smoke] Loaded in {time.time() - t0:.1f}s") + + print("[smoke] Done.") + + +if __name__ == "__main__": + try: + main() + except Exception: + traceback.print_exc() + sys.exit(1) diff --git a/contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py b/contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py new file mode 100755 index 00000000..65dfa61e --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py @@ -0,0 +1,171 @@ +#!/usr/bin/env python3 +"""Minimal generate smoke test for MiMo-V2.5 FP8 on Trn2. + +Assumes the compiled NEFF already exists at MIMO_V2_5_COMPILED_PATH +(from smoke_compile_mimo_v2_5.py). Rebuilds the same MoENeuronConfig / +MiMo-V2.5 wrapper, loads with skip_warmup=False, and generates 20 tokens for a +single prompt via HuggingFaceGenerationAdapter. Purpose: sanity-check that +the FP8 MoE + preprocessed scales actually produce coherent tokens. + +Run under /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16. +""" + +import os +import sys +import time +import traceback + +MODEL_PATH = os.environ.get( + "MIMO_V2_5_MODEL_PATH", + "/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8", +) +COMPILED_PATH = os.environ.get( + "MIMO_V2_5_COMPILED_PATH", + "/opt/dlami/nvme/compiled/mimo_v2_5_tp64_moetp1_ep64_fp8/", +) + +# Must match smoke_compile_mimo_v2_5.py exactly, else load() sees a +# mismatched NEFF. +TP_DEGREE = int(os.environ.get("TP_DEGREE", "64")) +SEQ_LEN = int(os.environ.get("SEQ_LEN", "1024")) +BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "1")) +CTX_BATCH_SIZE = int(os.environ.get("CTX_BATCH_SIZE", "1")) +MOE_TP = int(os.environ.get("MOE_TP", "1")) +MOE_EP = int(os.environ.get("MOE_EP", "64")) + +PROMPT = os.environ.get( + "MIMO_V2_5_PROMPT", + "Hello! Please introduce yourself in one sentence.", +) +MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "20")) + +# Keep the per-compile BASE_COMPILE_WORK_DIR in sync with +# smoke_compile_mimo_v2_5.py so load() under the same COMPILED_PATH +# doesn't collide with a concurrent compile or reuse a stale workdir. +os.environ.setdefault( + "BASE_COMPILE_WORK_DIR", + os.path.join("/tmp/nxd_model", os.path.basename(COMPILED_PATH.rstrip("/"))), +) + + +def main(): + from transformers import AutoConfig, AutoTokenizer, GenerationConfig + + from neuronx_distributed_inference.models.config import MoENeuronConfig + from neuronx_distributed_inference.utils.hf_adapter import ( + HuggingFaceGenerationAdapter, + load_pretrained_config, + ) + + contrib_src = os.path.join( + os.path.dirname(os.path.abspath(__file__)), + "..", + "src", + ) + sys.path.insert(0, os.path.abspath(contrib_src)) + + from modeling_mimo_v2 import ( + MiMoV2InferenceConfig, + NeuronMiMoV2ForCausalLM, + ) + + print(f"[gen] MODEL_PATH={MODEL_PATH}") + print(f"[gen] COMPILED_PATH={COMPILED_PATH}") + print(f"[gen] TP={TP_DEGREE}, SEQ={SEQ_LEN}, BS={BATCH_SIZE}") + + # Outer ep_degree must match the compile-time value (kept at 1 so + # world_size = tp_degree; see smoke_compile_mimo_v2_5.py comment). + neuron_config = MoENeuronConfig( + tp_degree=TP_DEGREE, + ep_degree=1, + logical_nc_config=2, + batch_size=BATCH_SIZE, + max_batch_size=BATCH_SIZE, + ctx_batch_size=CTX_BATCH_SIZE, + tkg_batch_size=BATCH_SIZE, + seq_len=SEQ_LEN, + n_active_tokens=128, + torch_dtype="bfloat16", + capacity_factor=1.0, + glu_mlp=True, + moe_ep_degree=MOE_EP, + moe_tp_degree=MOE_TP, + context_encoding_buckets=[SEQ_LEN], + router_config={"act_fn": "sigmoid", "dtype": "float32"}, + blockwise_matmul_config={ + "use_shard_on_block_dynamic_while": True, + "block_sharding_strategy": "PING_PONG", + }, + save_sharded_checkpoint=True, + quantized=True, + quantized_checkpoints_path=MODEL_PATH, + quantization_dtype="f8e4m3", + quantization_type="blockwise_symmetric", + quantization_block_axis=[1, 2], + quantization_block_size=[128, 128], + modules_to_not_convert=[ + "embed_tokens", + "lm_head", + "norm", + "router", + "o_proj", + ], + ) + + hf_config = AutoConfig.from_pretrained(MODEL_PATH, trust_remote_code=True) + config = MiMoV2InferenceConfig( + neuron_config, load_config=load_pretrained_config(hf_config=hf_config) + ) + + print("[gen] Instantiating model...") + t0 = time.time() + model = NeuronMiMoV2ForCausalLM(MODEL_PATH, config) + print(f"[gen] Instantiated in {time.time() - t0:.1f}s") + + # skip_warmup=False so generate() hits a primed graph (the warmup forward + # allocates the shared scratchpad the generation path needs). + print(f"[gen] Loading from {COMPILED_PATH} (skip_warmup=False)") + t0 = time.time() + model.load(COMPILED_PATH, skip_warmup=False) + print(f"[gen] Loaded in {time.time() - t0:.1f}s") + + tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) + adapter = HuggingFaceGenerationAdapter(model) + + inputs = tokenizer([PROMPT] * BATCH_SIZE, return_tensors="pt", padding=True) + gen_config = GenerationConfig( + max_new_tokens=MAX_NEW_TOKENS, + min_new_tokens=MAX_NEW_TOKENS, + do_sample=False, + pad_token_id=getattr(tokenizer, "pad_token_id", None) or tokenizer.eos_token_id, + ) + + print(f"[gen] prompt: {PROMPT!r}") + print(f"[gen] input_ids.shape={tuple(inputs['input_ids'].shape)}") + t0 = time.time() + output_ids = adapter.generate( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + generation_config=gen_config, + ) + dt = time.time() - t0 + + prompt_len = inputs["input_ids"].shape[1] + new_tokens = output_ids[0, prompt_len:] + decoded = tokenizer.decode(new_tokens, skip_special_tokens=True) + full = tokenizer.decode(output_ids[0], skip_special_tokens=True) + + print(f"[gen] generated {new_tokens.numel()} tokens in {dt:.2f}s " + f"({new_tokens.numel() / dt:.2f} tok/s)") + print(f"[gen] new token ids: {new_tokens.tolist()}") + print(f"[gen] new text : {decoded!r}") + print(f"[gen] full text : {full!r}") + print("[gen] Done.") + + +if __name__ == "__main__": + try: + main() + except Exception: + traceback.print_exc() + sys.exit(1) diff --git a/contrib/models/MiMo-V2.5/perf_test/vllm-neuron-patch.patch b/contrib/models/MiMo-V2.5/perf_test/vllm-neuron-patch.patch new file mode 100644 index 00000000..67a88579 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/vllm-neuron-patch.patch @@ -0,0 +1,107 @@ +diff --git a/vllm_neuron/worker/neuronx_distributed_model_loader.py b/vllm_neuron/worker/neuronx_distributed_model_loader.py +index d2099eb..0c162e4 100644 +--- a/vllm_neuron/worker/neuronx_distributed_model_loader.py ++++ b/vllm_neuron/worker/neuronx_distributed_model_loader.py +@@ -922,6 +922,94 @@ def _camel_to_kebab(name: str) -> str: + return re.sub("([a-z0-9])([A-Z])", r"\1-\2", s1).lower() + + ++ ++def _patch_autoconfig_trust_remote_code(): ++ """Monkey-patch ``AutoConfig.from_pretrained`` to default ``trust_remote_code=True``. ++ ++ NxDI's ``hf_adapter.load_config`` calls ``AutoConfig.from_pretrained(path)`` ++ without ``trust_remote_code``. Contrib models like MiMo-V2.5 that ++ ship a ``configuration_*.py`` with the checkpoint require custom code ++ execution, so the default behaviour crashes with ``ValueError: The ++ repository ... contains custom code which must be executed``. ++ ++ vLLM's top-level ``--trust-remote-code`` flag only affects vLLM's own ++ config load, not NxDI's. Patching here is cheap and idempotent. ++ """ ++ try: ++ from transformers import AutoConfig ++ except ImportError: ++ return ++ if getattr(AutoConfig, "_nxdi_contrib_patched", False): ++ return ++ _orig = AutoConfig.from_pretrained ++ ++ def _patched(*args, **kwargs): ++ kwargs.setdefault("trust_remote_code", True) ++ return _orig(*args, **kwargs) ++ ++ AutoConfig.from_pretrained = _patched ++ AutoConfig._nxdi_contrib_patched = True ++ ++ ++def _register_contrib_models(): ++ """Lazy-register NxDI contrib models on each process that calls the loader. ++ ++ Driven by env vars: ++ NXDI_CONTRIB_MIMO_V2_5_SRC -> path to contrib MiMo-V2.5 src/ ++ NXDI_CONTRIB_MINIMAX_M2_SRC -> path to contrib MiniMax-M2 src/ ++ ++ Registers the contrib model class into NxDI's MODEL_TYPES and, where ++ vLLM does not already know the architecture, registers it into vLLM's ++ ModelRegistry. Runs every time _get_neuron_model_cls is called so that ++ vLLM's spawn'd EngineCore workers (which don't inherit the parent's ++ module-level state) pick up the registration too. Registration is ++ idempotent. ++ """ ++ import os as _os ++ import sys as _sys ++ import warnings as _w ++ ++ _patch_autoconfig_trust_remote_code() ++ ++ mimo_src = _os.environ.get("NXDI_CONTRIB_MIMO_V2_5_SRC") ++ if mimo_src and _os.path.isdir(mimo_src) and "mimov2" not in MODEL_TYPES: ++ if mimo_src not in _sys.path: ++ _sys.path.insert(0, mimo_src) ++ try: ++ from modeling_mimo_v2 import NeuronMiMoV2ForCausalLM ++ MODEL_TYPES.setdefault( ++ "mimov2", {"causal-lm": NeuronMiMoV2ForCausalLM} ++ ) ++ try: ++ from vllm.model_executor.models.registry import ModelRegistry ++ if "MiMoV2ForCausalLM" not in ModelRegistry.get_supported_archs(): ++ ModelRegistry.register_model( ++ "MiMoV2ForCausalLM", NeuronMiMoV2ForCausalLM ++ ) ++ except ImportError: ++ pass ++ except Exception as e: ++ _w.warn( ++ f"Failed to register MiMo-V2.5 contrib model: {e}", ++ category=UserWarning, ++ ) ++ ++ minimax_src = _os.environ.get("NXDI_CONTRIB_MINIMAX_M2_SRC") ++ if minimax_src and _os.path.isdir(minimax_src) and "minimaxm2" not in MODEL_TYPES: ++ if minimax_src not in _sys.path: ++ _sys.path.insert(0, minimax_src) ++ try: ++ from modeling_minimax_m2 import NeuronMiniMaxM2ForCausalLM ++ MODEL_TYPES.setdefault( ++ "minimaxm2", {"causal-lm": NeuronMiniMaxM2ForCausalLM} ++ ) ++ except Exception as e: ++ _w.warn( ++ f"Failed to register MiniMax-M2 contrib model: {e}", ++ category=UserWarning, ++ ) ++ ++ + def _get_neuron_model_cls(architecture: str): + """ + Get Neuron model class from architecture string. +@@ -941,6 +1029,7 @@ def _get_neuron_model_cls(architecture: str): + _get_neuron_model_cls("NeuronLlamaForCausalLM") + + """ ++ _register_contrib_models() + # Handle Neuron class name (starts with "Neuron") - strip prefix + if architecture.startswith("Neuron") and "For" in architecture: + original_architecture = architecture diff --git a/contrib/models/MiMo-V2.5/src/__init__.py b/contrib/models/MiMo-V2.5/src/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py new file mode 100644 index 00000000..60cf77f8 --- /dev/null +++ b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py @@ -0,0 +1,510 @@ +""" +Preprocess MiMo-V2.5 FP8 checkpoint for Neuron inference. + +This is a streaming (per-layer) rewrite of preprocess_mimo_v2_fp8.py. The +original preprocess loaded the entire ~290 GB FP8 checkpoint into RAM via +load_state_dict(); that peaks well over 600 GB after dequantize/requantize +copies and is fragile. This version keeps a single safe_open handle live +at a time and emits per-layer safetensors shards, capping peak RAM at +~24 GB and finishing in ~20 minutes. + +MiMo-V2.5 checkpoint layout: + - q_proj, k_proj, v_proj are stored *separately* in the HF checkpoint + (not pre-fused). No split_qkv_fused needed. + - o_proj is BF16 (listed in quantization_config.ignored_layers); kept + as BF16 on the Neuron side (RowParallelLinear, not QuantizedRowParallel). + - Layer 0 is a dense MLP (moe_layer_freq[0] == 0) with intermediate_size + 16384; layers 1..47 are MoE with 256 experts each. + - Hybrid attention: 9 "full" layers (hybrid_layer_pattern[i] == 0) and + 39 "sliding window" layers (== 1). SWA layers carry + attention_sink_bias (add_swa_attention_sink_bias=True in the config; + add_full_attention_sink_bias=False, so full layers do NOT get it). + +Neuron-side rescaling (same as MiMo-V2 siblings): + - OCP FP8 e4m3 (±448) -> Neuron FP8 e4m3 (±240) with FP8_SCALING_FACTOR=448/240. + - Per-row scales for attention/dense-mlp projections (q/k/v/o, gate/up/down + of the dense layer). + - Blockwise (128x128) scales kept for MoE expert weights; per-expert weights + are transposed and fused to match ExpertFusedRowParallelLinear's packed + layout (gate_up_proj: [num_experts, H, 2*IM]; down_proj: [num_experts, IM, H]). + +Output layout: + save_path/ + config.json, tokenizer.*, chat_template.jinja if present + configuration_mimo_v2.py, modeling_mimo_v2.py (trust_remote_code) + model.safetensors.index.json (regenerated) + model_extras.safetensors (embed_tokens, norm, lm_head) + model_layer{N}.safetensors (one per decoder layer, N=0..47) + +Usage: + python preprocess_mimo_v2_5_fp8.py \\ + --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5 \\ + --save_path /opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8 \\ + --tp_degree 64 +""" + +import argparse +import gc +import json +import os +import shutil +import time +from typing import Dict, List, Optional, Tuple + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + + +FP8_SCALING_FACTOR = 448.0 / 240.0 +NEURON_FP8_MAX = 240.0 + + +# --------------------------------------------------------------------------- +# Quantization primitives +# --------------------------------------------------------------------------- + +def convert_bf16_to_fp8_per_row( + weight: torch.Tensor, +) -> Tuple[torch.Tensor, torch.Tensor]: + """BF16 [out, in] -> Neuron FP8 per-row (scales shape [out, 1]).""" + weight_float = weight.float() + row_max_abs = weight_float.abs().max(dim=1, keepdim=True)[0] + scales = torch.clamp(row_max_abs / NEURON_FP8_MAX, min=1e-10) + quantized = (weight_float / scales).to(torch.float8_e4m3fn) + return quantized, scales.to(torch.float32) + + +def rescale_fp8_to_per_row( + weight: torch.Tensor, scale: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Block-wise FP8 + blockwise scale -> Neuron per-row FP8. + + Dequantize to float32 using block broadcast, then per-row requantize. + """ + out_features, in_features = weight.shape + scale_h, scale_w = scale.shape + + block_h = (out_features + scale_h - 1) // scale_h + block_w = (in_features + scale_w - 1) // scale_w + + weight_float = weight.float() + dequantized = torch.zeros(out_features, in_features, dtype=torch.float32) + for i in range(scale_h): + for j in range(scale_w): + h0, h1 = i * block_h, min((i + 1) * block_h, out_features) + w0, w1 = j * block_w, min((j + 1) * block_w, in_features) + dequantized[h0:h1, w0:w1] = ( + weight_float[h0:h1, w0:w1] * scale[i, j].item() + ) + + row_max_abs = dequantized.abs().max(dim=1, keepdim=True)[0] + scales = torch.clamp(row_max_abs / NEURON_FP8_MAX, min=1e-10) + quantized = (dequantized / scales).to(torch.float8_e4m3fn) + return quantized, scales.to(torch.float32) + + +def rescale_fp8_weight_blockwise( + weight: torch.Tensor, scale: torch.Tensor +) -> Tuple[torch.Tensor, torch.Tensor]: + """Keep blockwise scales, just rescale into Neuron FP8 range. + + MoE expert weights stay block-quantized; only the dtype range changes. + """ + weight_bf16 = weight.bfloat16() + rescaled = (weight_bf16 / FP8_SCALING_FACTOR).to(torch.float8_e4m3fn) + neuron_scale = scale.float() * FP8_SCALING_FACTOR + return rescaled, neuron_scale.to(torch.float32) + + +# --------------------------------------------------------------------------- +# Streaming weight access (one open safetensors handle at a time) +# --------------------------------------------------------------------------- + +class LazyWeightMap: + """Lazily fetch tensors from sharded safetensors, keeping one handle live.""" + + def __init__(self, model_dir: str, weight_map: Dict[str, str]): + self.model_dir = model_dir + self.weight_map = weight_map + self._cur_filename: Optional[str] = None + self._cur_handle = None + # MiMo-V2.5's published index.json still references the legacy + # `model_N-00001-of-00002.safetensors` names, but the actual shards + # on disk are `model_pp0_epN_shardM.safetensors`. Build an alias + # table: legacy -> real filename, for all shards present locally. + self._filename_alias: Dict[str, str] = {} + import re + actual_files = os.listdir(model_dir) + legacy_re = re.compile(r"model_(\d+)-0000([12])-of-00002\.safetensors") + for name in weight_map.values(): + if name in self._filename_alias or name in actual_files: + continue + m = legacy_re.match(name) + if not m: + continue + ep_idx, shard_one_based = m.group(1), m.group(2) + shard_zero_based = str(int(shard_one_based) - 1) + candidate = f"model_pp0_ep{ep_idx}_shard{shard_zero_based}.safetensors" + if candidate in actual_files: + self._filename_alias[name] = candidate + + def _open(self, filename: str): + filename = self._filename_alias.get(filename, filename) + if self._cur_filename == filename: + return self._cur_handle + if self._cur_handle is not None: + self._cur_handle.__exit__(None, None, None) + self._cur_handle = None + path = os.path.join(self.model_dir, filename) + self._cur_handle = safe_open(path, framework="pt", device="cpu") + self._cur_handle.__enter__() + self._cur_filename = filename + return self._cur_handle + + def get(self, key: str) -> Optional[torch.Tensor]: + filename = self.weight_map.get(key) + if filename is None: + return None + return self._open(filename).get_tensor(key) + + def has(self, key: str) -> bool: + return key in self.weight_map + + def close(self): + if self._cur_handle is not None: + self._cur_handle.__exit__(None, None, None) + self._cur_handle = None + self._cur_filename = None + + +# --------------------------------------------------------------------------- +# Per-tensor helper +# --------------------------------------------------------------------------- + +def _maybe_fp8_to_neuron_per_row( + weight: torch.Tensor, scale: Optional[torch.Tensor] +) -> Tuple[torch.Tensor, Optional[torch.Tensor]]: + """FP8 blockwise -> per-row, or BF16 -> FP8 per-row. Pass-through otherwise.""" + if weight.dtype == torch.float8_e4m3fn and scale is not None: + return rescale_fp8_to_per_row(weight, scale) + if weight.dtype == torch.bfloat16: + return convert_bf16_to_fp8_per_row(weight) + return weight, scale + + +# --------------------------------------------------------------------------- +# Per-layer processing +# --------------------------------------------------------------------------- + +def process_layer( + layer_idx: int, + lazy: LazyWeightMap, + config: dict, + is_dense: bool, + is_swa: bool, +) -> Dict[str, torch.Tensor]: + out: Dict[str, torch.Tensor] = {} + prefix = f"model.layers.{layer_idx}." + out_prefix = f"layers.{layer_idx}." + + # --- Layer norms (BF16, untouched) --- + for name in ("input_layernorm", "post_attention_layernorm"): + t = lazy.get(f"{prefix}{name}.weight") + if t is not None: + out[f"{out_prefix}{name}.weight"] = t.detach().clone() + + # --- Attention: q/k/v/o are stored separately in MiMo-V2.5 --- + # q/k/v: rescale to Neuron FP8 per-row. + for proj in ("q_proj", "k_proj", "v_proj"): + w = lazy.get(f"{prefix}self_attn.{proj}.weight") + if w is None: + continue + s = lazy.get(f"{prefix}self_attn.{proj}.weight_scale_inv") + w2, s2 = _maybe_fp8_to_neuron_per_row(w, s) + out[f"{out_prefix}self_attn.{proj}.weight"] = w2 + if s2 is not None: + out[f"{out_prefix}self_attn.{proj}.scale"] = s2 + + # o_proj is listed in HF quantization_config.ignored_layers and ships as + # BF16; on Neuron it binds to a plain RowParallelLinear (see + # modeling_mimo_v2.py: self.o_proj = RowParallelLinear(...)), NOT a + # QuantizedRowParallel. Writing FP8 + .scale here would silently be + # reinterpreted as BF16 bytes at load time and produce garbage outputs. + # Keep BF16, never emit .scale. + o_w = lazy.get(f"{prefix}self_attn.o_proj.weight") + o_s = lazy.get(f"{prefix}self_attn.o_proj.weight_scale_inv") + if o_w is not None: + if o_w.dtype == torch.float8_e4m3fn: + # Defensive: if a future checkpoint FP8-quantizes o_proj, dequant + # blockwise back to BF16 (no per-row requant; RowParallelLinear has + # no .scale parameter). + assert o_s is not None, "FP8 o_proj requires weight_scale_inv" + out_features, in_features = o_w.shape + scale_h, scale_w = o_s.shape + block_h = (out_features + scale_h - 1) // scale_h + block_w = (in_features + scale_w - 1) // scale_w + wf = o_w.float() + tmp = torch.zeros(out_features, in_features, dtype=torch.float32) + for i in range(scale_h): + for j in range(scale_w): + h0, h1 = i * block_h, min((i + 1) * block_h, out_features) + w0, w1 = j * block_w, min((j + 1) * block_w, in_features) + tmp[h0:h1, w0:w1] = wf[h0:h1, w0:w1] * o_s[i, j].item() + o_bf16 = tmp.to(torch.bfloat16) + else: + o_bf16 = o_w.to(torch.bfloat16) + out[f"{out_prefix}self_attn.o_proj.weight"] = o_bf16.detach().clone() + + # --- attention_sink_bias: present only on SWA layers in MiMo-V2.5. + # config.add_swa_attention_sink_bias=True, add_full_attention_sink_bias=False. + if is_swa and config.get("add_swa_attention_sink_bias", False): + sink = lazy.get(f"{prefix}self_attn.attention_sink_bias") + if sink is not None: + out[f"{out_prefix}self_attn.attention_sink_bias"] = sink.detach().clone() + elif not is_swa and config.get("add_full_attention_sink_bias", False): + sink = lazy.get(f"{prefix}self_attn.attention_sink_bias") + if sink is not None: + out[f"{out_prefix}self_attn.attention_sink_bias"] = sink.detach().clone() + + # --- MLP: dense vs MoE --- + if is_dense: + # Dense MLP: gate_proj, up_proj, down_proj (FP8 blockwise in MiMo-V2.5 layer 0). + for proj in ("gate_proj", "up_proj", "down_proj"): + w = lazy.get(f"{prefix}mlp.{proj}.weight") + if w is None: + continue + s = lazy.get(f"{prefix}mlp.{proj}.weight_scale_inv") + w2, s2 = _maybe_fp8_to_neuron_per_row(w, s) + out[f"{out_prefix}mlp.{proj}.weight"] = w2 + if s2 is not None: + out[f"{out_prefix}mlp.{proj}.scale"] = s2 + return out + + # --- MoE --- + # Router: mlp.gate -> mlp.router.linear_router + router_w = lazy.get(f"{prefix}mlp.gate.weight") + if router_w is not None: + out[f"{out_prefix}mlp.router.linear_router.weight"] = router_w.detach().clone() + router_bias = lazy.get(f"{prefix}mlp.gate.e_score_correction_bias") + if router_bias is not None: + out[f"{out_prefix}mlp.router.e_score_correction_bias"] = router_bias.detach().clone() + + num_experts = config["n_routed_experts"] + + # Peek expert 0 to learn shapes/dtypes. + e0_gw = lazy.get(f"{prefix}mlp.experts.0.gate_proj.weight") + if e0_gw is None: + return out # no experts (shouldn't happen for MoE layers, but be safe) + e0_gs = lazy.get(f"{prefix}mlp.experts.0.gate_proj.weight_scale_inv") + + if e0_gw.dtype == torch.float8_e4m3fn and e0_gs is not None: + sample_w, sample_s = rescale_fp8_weight_blockwise(e0_gw, e0_gs) + elif e0_gw.dtype == torch.bfloat16: + # Should not happen for MiMo-V2.5 (experts ship in FP8); flag loudly. + raise NotImplementedError( + f"Layer {layer_idx} expert 0 gate_proj is BF16; MiMo-V2.5 expects FP8." + ) + else: + sample_w, sample_s = e0_gw, e0_gs + + intermediate_size, hidden_size = sample_w.shape # [IM, H] + # Packed transpose layout: [num_experts, H, 2*IM] for gate_up. + gate_up_proj = torch.empty( + num_experts, hidden_size, 2 * intermediate_size, dtype=sample_w.dtype + ) + i_blocks, h_blocks = sample_s.shape # [IM_blocks, H_blocks] + gate_up_scale = torch.empty( + num_experts, h_blocks, 2 * i_blocks, dtype=sample_s.dtype + ) + + e0_dw = lazy.get(f"{prefix}mlp.experts.0.down_proj.weight") + e0_ds = lazy.get(f"{prefix}mlp.experts.0.down_proj.weight_scale_inv") + if e0_dw.dtype == torch.float8_e4m3fn and e0_ds is not None: + sample_dw, sample_ds = rescale_fp8_weight_blockwise(e0_dw, e0_ds) + else: + raise NotImplementedError( + f"Layer {layer_idx} expert 0 down_proj dtype {e0_dw.dtype} not handled." + ) + d_h_blocks, d_i_blocks = sample_ds.shape # [H_blocks, IM_blocks] + down_proj = torch.empty( + num_experts, intermediate_size, hidden_size, dtype=sample_dw.dtype + ) + down_scale = torch.empty( + num_experts, d_i_blocks, d_h_blocks, dtype=sample_ds.dtype + ) + + # Slot expert 0 (already rescaled above). + gate_up_proj[0, :, :intermediate_size] = sample_w.T + gate_up_scale[0, :, :i_blocks] = sample_s.T + e0_uw = lazy.get(f"{prefix}mlp.experts.0.up_proj.weight") + e0_us = lazy.get(f"{prefix}mlp.experts.0.up_proj.weight_scale_inv") + up_w0, up_s0 = rescale_fp8_weight_blockwise(e0_uw, e0_us) + gate_up_proj[0, :, intermediate_size:] = up_w0.T + gate_up_scale[0, :, i_blocks:] = up_s0.T + down_proj[0] = sample_dw.T + down_scale[0] = sample_ds.T + del e0_gw, e0_gs, e0_uw, e0_us, e0_dw, e0_ds + del sample_w, sample_s, sample_dw, sample_ds, up_w0, up_s0 + + for e in range(1, num_experts): + gw = lazy.get(f"{prefix}mlp.experts.{e}.gate_proj.weight") + gs = lazy.get(f"{prefix}mlp.experts.{e}.gate_proj.weight_scale_inv") + uw = lazy.get(f"{prefix}mlp.experts.{e}.up_proj.weight") + us = lazy.get(f"{prefix}mlp.experts.{e}.up_proj.weight_scale_inv") + dw = lazy.get(f"{prefix}mlp.experts.{e}.down_proj.weight") + ds = lazy.get(f"{prefix}mlp.experts.{e}.down_proj.weight_scale_inv") + g_w, g_s = rescale_fp8_weight_blockwise(gw, gs) + u_w, u_s = rescale_fp8_weight_blockwise(uw, us) + d_w, d_s = rescale_fp8_weight_blockwise(dw, ds) + gate_up_proj[e, :, :intermediate_size] = g_w.T + gate_up_proj[e, :, intermediate_size:] = u_w.T + gate_up_scale[e, :, :i_blocks] = g_s.T + gate_up_scale[e, :, i_blocks:] = u_s.T + down_proj[e] = d_w.T + down_scale[e] = d_s.T + del gw, gs, uw, us, dw, ds, g_w, g_s, u_w, u_s, d_w, d_s + + out[f"{out_prefix}mlp.expert_mlps.mlp_op.gate_up_proj.weight"] = gate_up_proj + out[f"{out_prefix}mlp.expert_mlps.mlp_op.gate_up_proj.scale"] = gate_up_scale + out[f"{out_prefix}mlp.expert_mlps.mlp_op.down_proj.weight"] = down_proj + out[f"{out_prefix}mlp.expert_mlps.mlp_op.down_proj.scale"] = down_scale + return out + + +# --------------------------------------------------------------------------- +# Shard saving / index +# --------------------------------------------------------------------------- + +def save_shard( + tensors: Dict[str, torch.Tensor], + save_path: str, + filename: str, + weight_map: Dict[str, str], +) -> int: + """Save a sub-state-dict; clone tensors so safetensors doesn't complain + about views of mmapped storage. Returns bytes written.""" + path = os.path.join(save_path, filename) + materialized: Dict[str, torch.Tensor] = {} + total_bytes = 0 + for k, v in tensors.items(): + if not v.is_contiguous(): + v = v.contiguous() + v = v.detach().clone() + materialized[k] = v + total_bytes += v.numel() * v.element_size() + save_file(materialized, path) + for k in materialized.keys(): + weight_map[k] = filename + del materialized + return total_bytes + + +# --------------------------------------------------------------------------- +# Main driver +# --------------------------------------------------------------------------- + +def process_flash_checkpoint(hf_model_path: str, save_path: str, tp_degree: int): + os.makedirs(save_path, exist_ok=True) + + with open(os.path.join(hf_model_path, "model.safetensors.index.json")) as f: + weight_map_in = json.load(f)["weight_map"] + + with open(os.path.join(hf_model_path, "config.json")) as f: + config = json.load(f) + + num_layers = config["num_hidden_layers"] + hybrid = config.get("hybrid_layer_pattern", [0] * num_layers) + moe_freq = config.get("moe_layer_freq", [1] * num_layers) + + print( + f"Processing {num_layers} decoder layers" + f" (full={sum(1 for v in hybrid if v == 0)}," + f" swa={sum(1 for v in hybrid if v == 1)}," + f" dense={sum(1 for v in moe_freq if v == 0)}," + f" moe={sum(1 for v in moe_freq if v == 1)})", + flush=True, + ) + + lazy = LazyWeightMap(hf_model_path, weight_map_in) + weight_map_out: Dict[str, str] = {} + + try: + for li in range(num_layers): + t0 = time.time() + is_dense = moe_freq[li] == 0 + is_swa = hybrid[li] == 1 + layer_sd = process_layer(li, lazy, config, is_dense=is_dense, is_swa=is_swa) + filename = f"model_layer{li}.safetensors" + size = save_shard(layer_sd, save_path, filename, weight_map_out) + del layer_sd + gc.collect() + tag = "dense" if is_dense else "moe " + attn = "swa " if is_swa else "full" + print( + f" layer {li:2d} [{tag} {attn}] {size/1e9:6.2f} GB in {time.time()-t0:5.1f}s", + flush=True, + ) + + print("Processing embed_tokens, norm, lm_head ...", flush=True) + extras: Dict[str, torch.Tensor] = {} + for src, dst in ( + ("model.embed_tokens.weight", "embed_tokens.weight"), + ("model.norm.weight", "norm.weight"), + ("lm_head.weight", "lm_head.weight"), + ): + t = lazy.get(src) + if t is not None: + extras[dst] = t.detach().clone() + else: + print(f" WARNING: missing {src}", flush=True) + if "lm_head.weight" not in extras and "embed_tokens.weight" in extras: + # Tied embeddings + extras["lm_head.weight"] = extras["embed_tokens.weight"].detach().clone() + save_shard(extras, save_path, "model_extras.safetensors", weight_map_out) + del extras + finally: + lazy.close() + + # --- Index file --- + total_size = 0 + for f in set(weight_map_out.values()): + total_size += os.path.getsize(os.path.join(save_path, f)) + index = { + "metadata": {"total_size": total_size}, + "weight_map": weight_map_out, + } + with open(os.path.join(save_path, "model.safetensors.index.json"), "w") as f: + json.dump(index, f, indent=2) + + # --- Copy auxiliary files (config.json, tokenizer, chat template, + # and crucially the trust_remote_code modules the HF config references). + for name in sorted(os.listdir(hf_model_path)): + if name.endswith(".safetensors"): + continue + if name == "model.safetensors.index.json": + continue + src = os.path.join(hf_model_path, name) + if os.path.isfile(src): + shutil.copy(src, os.path.join(save_path, name)) + + print(f"\nPreprocess complete. total_size={total_size/1e9:.2f} GB", flush=True) + print(f" tensors written: {len(weight_map_out)}", flush=True) + print(f" output dir: {save_path}", flush=True) + + +def main(): + parser = argparse.ArgumentParser( + description="Preprocess MiMo-V2.5 FP8 checkpoint for Neuron inference" + ) + parser.add_argument("--hf_model_path", required=True) + parser.add_argument("--save_path", required=True) + parser.add_argument("--tp_degree", type=int, default=64, + help="Tensor parallelism (currently informational only; " + "the framework does the TP sharding at load time).") + args = parser.parse_args() + process_flash_checkpoint(args.hf_model_path, args.save_path, args.tp_degree) + + +if __name__ == "__main__": + main() diff --git a/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py new file mode 100644 index 00000000..947dfe2c --- /dev/null +++ b/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py @@ -0,0 +1,1677 @@ +# coding=utf-8 +# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# This implementation is based on the MiMo-V2.5 model from Xiaomi. +# Reference: https://huggingface.co/XiaomiMiMo/MiMo-V2.5 + +"""MiMo-V2.5 model for NXD inference.""" + +import gc +import math +import warnings +from typing import Any, Dict, List, Optional, Tuple, Type, Union + +import torch +import torch.nn.functional as F +from torch import nn, Tensor + +from neuronx_distributed.parallel_layers import parallel_state +from neuronx_distributed.parallel_layers.layers import ( + ColumnParallelLinear, + ParallelEmbedding, + RowParallelLinear, +) +from neuronx_distributed.parallel_layers.mappings import ( + gather_from_sequence_parallel_region, + gather_from_tensor_model_parallel_region_with_dim, +) +from neuronx_distributed.utils import cpu_mode + +from neuronx_distributed_inference.utils.distributed import ( + split_along_dim, + get_cp_rank, +) +from neuronx_distributed_inference.modules.attention.attention_process_groups import ( + get_context_parallel_attention_cp_group, +) + +from neuronx_distributed_inference.models.config import ( + InferenceConfig, + MoENeuronConfig, +) +from neuronx_distributed_inference.models.model_base import ( + NeuronBaseForCausalLM, + NeuronBaseModel, +) +from neuronx_distributed_inference.models.model_wrapper import ( + CONTEXT_ENCODING_MODEL_TAG, + TOKEN_GENERATION_MODEL_TAG, +) +from neuronx_distributed_inference.modules.attention.attention_base import ( + NeuronAttentionBase, +) +from neuronx_distributed_inference.modules.attention.utils import RotaryEmbedding +from neuronx_distributed_inference.modules.custom_calls import CustomRMSNorm +from neuronx_distributed_inference.modules.moe_v2 import initialize_moe_module + +try: + from neuronxcc.nki._private_kernels.attention import attention_isa_kernel +except ImportError: + from neuronxcc.nki.kernels.attention import attention_isa_kernel + +from torch_neuronx.xla_impl.ops import nki_jit + +_flash_fwd_call = nki_jit()(attention_isa_kernel) + + +def get_rmsnorm_cls(): + """Get appropriate RMSNorm class based on execution environment.""" + return MiMoV2RMSNorm if cpu_mode() else CustomRMSNorm + + +class MiMoV2RMSNorm(nn.Module): + """RMSNorm implementation for CPU mode.""" + + def __init__(self, hidden_size: int, eps: float = 1e-5): + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +class MiMoV2RotaryEmbedding(nn.Module): + """Rotary Position Embedding for MiMo-V2.5. + + Supports partial rotary embedding where only a fraction of dimensions + use rotary position encoding. + """ + + def __init__( + self, + dim: int, + max_position_embeddings: int = 262144, + base: float = 5000000.0, + partial_rotary_factor: float = 1.0, + ): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + self.partial_rotary_factor = partial_rotary_factor + + # Calculate the actual dimension used for rotary embedding + self.rope_dim = int(dim * partial_rotary_factor) + # Ensure rope_dim is even + self.rope_dim = self.rope_dim - (self.rope_dim % 2) + + inv_freq = 1.0 / ( + self.base ** (torch.arange(0, self.rope_dim, 2, dtype=torch.float32) / self.rope_dim) + ) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + @torch.no_grad() + def forward( + self, x: torch.Tensor, position_ids: torch.Tensor + ) -> Tuple[torch.Tensor, torch.Tensor]: + """Compute rotary embeddings. + + Args: + x: Input tensor of shape (batch_size, seq_len, hidden_size) + position_ids: Position indices of shape (batch_size, seq_len) + + Returns: + Tuple of (cos, sin) tensors for rotary embedding + """ + inv_freq_expanded = self.inv_freq[None, :, None].float().expand( + position_ids.shape[0], -1, 1 + ) + position_ids_expanded = position_ids[:, None, :].float() + + device_type = x.device.type if isinstance(x.device.type, str) else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded @ position_ids_expanded).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +def rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotate half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +def apply_rotary_pos_emb( + q: torch.Tensor, + k: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + position_ids: Optional[torch.Tensor] = None, + unsqueeze_dim: int = 1, +) -> Tuple[torch.Tensor, torch.Tensor]: + """Apply rotary position embedding to query and key tensors.""" + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +class MiMoV2InferenceConfig(InferenceConfig): + """Configuration class for MiMo-V2.5 inference on Neuron.""" + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + + # MoE configuration + self.num_local_experts = self.n_routed_experts + self.n_shared_experts = 0 # MiMo-V2.5 has no shared experts + + # Set intermediate_size for MoE layers + self.intermediate_size = self.moe_intermediate_size + + # Check and pad intermediate size if needed + self.maybe_pad_intermediate() + + # Router configuration + self.neuron_config.router_config.dtype = torch.float32 + self.neuron_config.router_config.act_fn = "sigmoid" # MiMo uses sigmoid + + # Disable numeric CC token as workaround + self.neuron_config.disable_numeric_cc_token = True + + # MiMo normalizes top-k affinities + self.neuron_config.normalize_top_k_affinities = True + + # Parse hybrid layer pattern + self._parse_hybrid_pattern() + + def _parse_hybrid_pattern(self): + """Parse hybrid layer pattern to determine attention types.""" + if hasattr(self, 'hybrid_layer_pattern') and self.hybrid_layer_pattern: + self.layer_attention_types = [ + "sliding_window" if p == 1 else "full" + for p in self.hybrid_layer_pattern + ] + else: + self.layer_attention_types = ["full"] * self.num_hidden_layers + + # Parse MoE layer frequency + if hasattr(self, 'moe_layer_freq') and self.moe_layer_freq: + self.layer_uses_moe = [bool(f) for f in self.moe_layer_freq] + else: + self.layer_uses_moe = [True] * self.num_hidden_layers + + def maybe_pad_intermediate(self): + """Pad intermediate size if required for efficient computation.""" + from neuronx_distributed_inference.models.config import ( + SHARD_ON_INTERMEDIATE_DIMENSION_PER_TP, + ) + + moe_tp_degree = self.neuron_config.moe_tp_degree + I_TP = self.moe_intermediate_size // moe_tp_degree + + if getattr( + self.neuron_config.blockwise_matmul_config, + "use_shard_on_intermediate_dynamic_while", + False, + ): + if I_TP % SHARD_ON_INTERMEDIATE_DIMENSION_PER_TP != 0: + padded_size = ( + math.ceil(I_TP / SHARD_ON_INTERMEDIATE_DIMENSION_PER_TP) + * SHARD_ON_INTERMEDIATE_DIMENSION_PER_TP + * moe_tp_degree + ) + self.moe_intermediate_pad_size = max( + padded_size - self.moe_intermediate_size, 0 + ) + self.moe_intermediate_size = padded_size + + def get_required_attributes(self) -> List[str]: + return [ + "attention_bias", + "head_dim", + "hidden_act", + "hidden_size", + "hybrid_layer_pattern", + "layernorm_epsilon", + "max_position_embeddings", + "moe_intermediate_size", + "moe_layer_freq", + "n_routed_experts", + "norm_topk_prob", + "num_attention_heads", + "num_experts_per_tok", + "num_hidden_layers", + "num_key_value_heads", + "partial_rotary_factor", + "rope_theta", + "scoring_func", + "sliding_window", + "swa_head_dim", + "swa_num_attention_heads", + "swa_num_key_value_heads", + "swa_rope_theta", + "swa_v_head_dim", + "tie_word_embeddings", + "v_head_dim", + "vocab_size", + ] + + @classmethod + def get_neuron_config_cls(cls) -> Type[MoENeuronConfig]: + return MoENeuronConfig + + +class NeuronMiMoV2Attention(NeuronAttentionBase): + """MiMo-V2.5 Attention implementation supporting hybrid attention patterns. + + Supports both full attention and sliding window attention with different + head dimensions for Q/K vs V. + """ + + def __init__( + self, + config: MiMoV2InferenceConfig, + layer_idx: int, + is_sliding_window: bool = False, + ): + self.layer_idx = layer_idx + self.is_sliding_window = is_sliding_window + + # Select parameters based on attention type + if is_sliding_window: + self.attn_head_dim = config.swa_head_dim + self.attn_v_head_dim = config.swa_v_head_dim + self.attn_num_heads = config.swa_num_attention_heads + self.attn_num_kv_heads = config.swa_num_key_value_heads + rope_theta = getattr(config, 'swa_rope_theta', 10000.0) + self.sliding_window_size = config.sliding_window + else: + self.attn_head_dim = config.head_dim + self.attn_v_head_dim = config.v_head_dim + self.attn_num_heads = config.num_attention_heads + self.attn_num_kv_heads = config.num_key_value_heads + rope_theta = config.rope_theta + self.sliding_window_size = None + + # Calculate partial rotary dimensions + self.partial_rotary_factor = config.partial_rotary_factor + self.rope_dim = int(self.attn_head_dim * self.partial_rotary_factor) + self.rope_dim = self.rope_dim - (self.rope_dim % 2) # Ensure even + self.nope_dim = self.attn_head_dim - self.rope_dim + + # Create rotary embedding + rotary_emb = MiMoV2RotaryEmbedding( + dim=self.attn_head_dim, + max_position_embeddings=config.max_position_embeddings, + base=rope_theta, + partial_rotary_factor=self.partial_rotary_factor, + ) + + # Initialize base attention + # NOTE: We pass v_head_dim to base class, but MiMo uses asymmetric Q/K (192) vs V (128). + # We override init_gqa_properties() to prevent the base class from creating + # incompatible projection layers (which cause crashes when CP > 1). + super().__init__( + config=config, + hidden_size=config.hidden_size, + num_attention_heads=self.attn_num_heads, + num_key_value_heads=self.attn_num_kv_heads, + head_dim=self.attn_v_head_dim, # Use v_head_dim for base class + rotary_emb=rotary_emb, + rms_norm_eps=config.layernorm_epsilon, + use_qk_norm=False, + ) + + # Initialize MiMo-specific projections with correct dimensions + self._init_projections(config) + + # Scaling factor + self.scaling = self.attn_head_dim ** -0.5 + # HF MiMoV2Attention (modeling_mimo_v2.py) multiplies value_states + # by config.attention_value_scale (0.707 for MiMo-V2) right after the V + # projection, before attention softmax*V. Matching that here — applied + # to value_states in forward() rather than to attn_output. + self.value_scale = float(getattr(config, "attention_value_scale", 1.0)) + + # Store cache KV heads for cache compatibility + # With CONVERT_TO_MHA, all layers have num_attention_heads KV heads + # Otherwise, use max of full and sliding window kv heads + tp_degree = config.neuron_config.tp_degree + if self.use_gqa_convert_to_mha: + # CONVERT_TO_MHA: cache stores num_attention_heads (same as Q heads) + self.cache_num_kv_heads = self.attn_num_heads + self.local_cache_kv_heads = self.local_num_heads + else: + # Standard GQA: cache uses max of full and sliding window kv heads + self.cache_num_kv_heads = max( + config.num_key_value_heads, + getattr(config, 'swa_num_key_value_heads', config.num_key_value_heads) + ) + self.local_cache_kv_heads = max(1, self.cache_num_kv_heads // tp_degree) + + def init_gqa_properties(self): + """Override base class to prevent creating incompatible QKV projections. + + MiMo-V2.5 has asymmetric Q/K head_dim (192) vs V head_dim (128), + which is incompatible with the base class's GroupQueryAttention_QKV. + MiMo uses its own custom projections via _init_projections() instead. + + When CP > 1, the base class would create cte_qkv_proj/tkg_qkv_proj with + wrong head_dim=128, causing compilation crashes. This no-op prevents that. + """ + pass + + def _init_projections(self, config: MiMoV2InferenceConfig): + """Initialize projection layers with correct dimensions. + + When CONVERT_TO_MHA is needed (tp_degree > num_kv_heads), K/V projections + are sized for num_attention_heads (not original num_kv_heads). The checkpoint + weights are replicated in preshard_hook before loading. + """ + dtype = config.neuron_config.torch_dtype + tp_degree = config.neuron_config.tp_degree + + # Check if we need GQA CONVERT_TO_MHA (when tp_degree > num_kv_heads) + self.use_gqa_convert_to_mha = tp_degree > self.attn_num_kv_heads + + # Store source heads for preshard_hook + self._src_num_kv_heads = self.attn_num_kv_heads + self._kv_replication_factor = self.attn_num_heads // self.attn_num_kv_heads if self.use_gqa_convert_to_mha else 1 + + if self.use_gqa_convert_to_mha: + # CONVERT_TO_MHA: K and V use num_attention_heads for proper TP splitting + k_num_heads = self.attn_num_heads + v_num_heads = self.attn_num_heads + else: + k_num_heads = self.attn_num_kv_heads + v_num_heads = self.attn_num_kv_heads + + # Q/K use head_dim, V uses v_head_dim + q_hidden_size = self.attn_num_heads * self.attn_head_dim + k_hidden_size = k_num_heads * self.attn_head_dim + v_hidden_size = v_num_heads * self.attn_v_head_dim + o_hidden_size = self.attn_num_heads * self.attn_v_head_dim + + if parallel_state.model_parallel_is_initialized(): + tp_group = parallel_state.get_tensor_model_parallel_group() + + # Q projection + self.q_proj = ColumnParallelLinear( + config.hidden_size, + q_hidden_size, + bias=config.attention_bias, + gather_output=False, + dtype=dtype, + tensor_model_parallel_group=tp_group, + ) + + # K projection + self.k_proj = ColumnParallelLinear( + config.hidden_size, + k_hidden_size, + bias=config.attention_bias, + gather_output=False, + dtype=dtype, + tensor_model_parallel_group=tp_group, + ) + + # V projection + self.v_proj = ColumnParallelLinear( + config.hidden_size, + v_hidden_size, + bias=config.attention_bias, + gather_output=False, + dtype=dtype, + tensor_model_parallel_group=tp_group, + ) + + # Output projection - with sequence parallel to scatter output + self.o_proj = RowParallelLinear( + o_hidden_size, + config.hidden_size, + bias=False, + input_is_parallel=True, + dtype=dtype, + tensor_model_parallel_group=tp_group, + sequence_parallel_enabled=self.sequence_parallel_enabled, + sequence_dimension=1 if self.sequence_parallel_enabled else None, + ) + + # Calculate local dimensions after TP split + self.local_num_heads = self.attn_num_heads // tp_degree + if self.use_gqa_convert_to_mha: + # With CONVERT_TO_MHA, local KV heads = local Q heads + self.local_num_kv_heads = self.local_num_heads + else: + self.local_num_kv_heads = max(1, self.attn_num_kv_heads // tp_degree) + else: + self.q_proj = nn.Linear(config.hidden_size, q_hidden_size, bias=config.attention_bias) + self.k_proj = nn.Linear(config.hidden_size, k_hidden_size, bias=config.attention_bias) + self.v_proj = nn.Linear(config.hidden_size, v_hidden_size, bias=config.attention_bias) + self.o_proj = nn.Linear(o_hidden_size, config.hidden_size, bias=False) + + self.local_num_heads = self.attn_num_heads + self.local_num_kv_heads = k_num_heads + + # Override base class attributes that were computed with wrong head_dim + # The base class init_gqa_properties() uses head_dim=v_head_dim which is wrong for Q/K + # We need to override these to ensure correct computation + self.num_heads = self.local_num_heads + self.num_key_value_heads = self.local_num_kv_heads + self.num_key_value_groups = self.local_num_heads // self.local_num_kv_heads + self.head_dim = self.attn_head_dim # Override to use actual Q/K head_dim (192) + + # Remove qkv_proj from base class if exists (we use separate q_proj, k_proj, v_proj) + if hasattr(self, 'qkv_proj'): + self.qkv_proj = None + + # Attention sink bias for attention layers (following HF implementation) + # This is a learnable parameter that allows attention to "sink" to an extra position + add_full_attention_sink_bias = getattr(config, 'add_full_attention_sink_bias', False) + add_swa_attention_sink_bias = getattr(config, 'add_swa_attention_sink_bias', True) + + # Determine if this layer uses sink bias based on config + self._use_sink_bias = (add_full_attention_sink_bias and not self.is_sliding_window) or \ + (add_swa_attention_sink_bias and self.is_sliding_window) + + if self._use_sink_bias: + # Shape: [num_attention_heads] - will be split across TP ranks + # The weight is loaded from checkpoint with shape [num_attention_heads] + # and will be sliced to [local_num_heads] during forward + self.attention_sink_bias = nn.Parameter( + torch.zeros(self.attn_num_heads, dtype=dtype), requires_grad=False + ) + else: + self.attention_sink_bias = None + + def preshard_hook(self, model_state_dict: dict, prefix: str) -> bool: + """Pre-shard hook to replicate K/V weights for CONVERT_TO_MHA. + + NOTE: This method is NOT currently called because NeuronMiMoV2Attention + is not a BaseGroupQueryAttention subclass. K/V weight replication is + instead done in convert_mimo_v2_hf_to_neuron_state_dict(). + + This method is kept for reference and potential future use. + """ + # This hook is not called - see note above + return False + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[torch.Tensor] = None, + cos_cache: Optional[torch.Tensor] = None, + sin_cache: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]: + """Forward pass for MiMo-V2.5 attention with Context Parallelism support.""" + + # Context Parallelism: only active during context encoding (no past_key_value) + is_context_parallel = past_key_value is None and self.cp_degree > 1 + cp_rank = None + + if is_context_parallel: + cp_rank = get_cp_rank( + self.rank_util.get_rank(), self.tp_degree, + self.cp_degree, self.neuron_config.switch_cc, + ) + # Split attention_mask (dim=2 = Q rows) and position_ids (dim=1 = seq) + attention_mask = split_along_dim( + attention_mask, dim=2, rank=cp_rank, num_partitions=self.cp_degree + ) + # Keep full position_ids for RoPE computation on full-length K/V + local_position_ids = split_along_dim( + position_ids, dim=1, rank=cp_rank, num_partitions=self.cp_degree + ) + + # Handle sequence parallel + if self.sequence_parallel_enabled and parallel_state.model_parallel_is_initialized(): + hidden_states = gather_from_sequence_parallel_region( + hidden_states, + self.sequence_dimension, + process_group=parallel_state.get_tensor_model_parallel_group(), + ) + + # Context Parallelism without sequence parallel: split hidden_states + if is_context_parallel and not self.sequence_parallel_enabled: + hidden_states = split_along_dim( + hidden_states, dim=1, rank=cp_rank, num_partitions=self.cp_degree + ) + + bsz, q_len, _ = hidden_states.size() + + # Determine if this is token generation (past_key_value is not None) + is_token_gen = past_key_value is not None + + # Project Q, K, V + query_states = self.q_proj(hidden_states) + key_states = self.k_proj(hidden_states) + value_states = self.v_proj(hidden_states) + + # HF MiMoV2Attention scales V by attention_value_scale (0.707 for MiMo-V2) + # right after v_proj, before the attention softmax*V. Earlier revisions + # of this file applied it post-attention or not at all; both produce + # gibberish for prompts longer than ~20 tokens. + if self.value_scale != 1.0: + value_states = value_states * self.value_scale + + # Reshape for multi-head attention: [bsz, num_heads, seq_len, head_dim] + query_states = query_states.view(bsz, q_len, self.local_num_heads, self.attn_head_dim).transpose(1, 2) + key_states = key_states.view(bsz, q_len, self.local_num_kv_heads, self.attn_head_dim).transpose(1, 2) + value_states = value_states.view(bsz, q_len, self.local_num_kv_heads, self.attn_v_head_dim).transpose(1, 2) + + # Split into rope and non-rope parts + query_rope = query_states[..., :self.rope_dim] + query_nope = query_states[..., self.rope_dim:] + key_rope = key_states[..., :self.rope_dim] + key_nope = key_states[..., self.rope_dim:] + + # Compute rotary embeddings + # IMPORTANT: Always compute for this layer because different layer types + # (full vs sliding window) use different rope_theta values. + # Full attention: rope_theta = 5000000 + # Sliding window: rope_theta = 10000 + # We cannot reuse cached cos/sin from other layers! + # + # For CP with sequence_parallel: Q/K/V have full S, use full position_ids for RoPE. + # For CP without sequence_parallel: Q/K/V have S/CP, use local_position_ids for RoPE + # (local_position_ids contain the correct global positions for this CP rank). + if is_context_parallel and not self.sequence_parallel_enabled: + rope_position_ids = local_position_ids + else: + rope_position_ids = position_ids + cos_cache, sin_cache = self.rotary_emb(value_states, rope_position_ids) + + # Apply rotary position embedding to rope parts only + query_rope, key_rope = apply_rotary_pos_emb( + query_rope, key_rope, cos_cache, sin_cache, rope_position_ids + ) + + # Concatenate rope and non-rope parts + query_states = torch.cat([query_rope, query_nope], dim=-1) + key_states = torch.cat([key_rope, key_nope], dim=-1) + + # Context Parallelism: split Q and save local KV for cache + if is_context_parallel: + if self.sequence_parallel_enabled: + # Q/K/V have full S. Split Q to local portion, save local KV for cache. + # Use split_along_dim (torch.index_select) instead of Python slicing + # because XLA tracing doesn't support dynamic tensor indices in slice notation. + query_states = split_along_dim(query_states, dim=2, rank=cp_rank, num_partitions=self.cp_degree) + key_states_for_cache = split_along_dim(key_states, dim=2, rank=cp_rank, num_partitions=self.cp_degree) + value_states_for_cache = split_along_dim(value_states, dim=2, rank=cp_rank, num_partitions=self.cp_degree) + q_len = q_len // self.cp_degree + # K/V stay at full S for attention computation + else: + # Q/K/V have S/CP. Save local KV for cache, then all-gather K/V. + key_states_for_cache = key_states + value_states_for_cache = value_states + key_states = gather_from_tensor_model_parallel_region_with_dim( + key_states, gather_dim=2, + process_group=get_context_parallel_attention_cp_group(), + ) + value_states = gather_from_tensor_model_parallel_region_with_dim( + value_states, gather_dim=2, + process_group=get_context_parallel_attention_cp_group(), + ) + # Q stays at S/CP + else: + # Store key/value states BEFORE GQA repeat for KV cache + key_states_for_cache = key_states + value_states_for_cache = value_states + + # WORKAROUND 1: Pad V from v_head_dim (128) to head_dim (192) for KV cache compatibility + if self.attn_v_head_dim < self.attn_head_dim: + pad_size = self.attn_head_dim - self.attn_v_head_dim + value_states_for_cache = F.pad(value_states_for_cache, (0, pad_size), value=0.0) + + # WORKAROUND 2: Pad KV heads if layer has fewer than cache expects + # Only needed when NOT using CONVERT_TO_MHA (standard GQA mode) + if not self.use_gqa_convert_to_mha and self.local_num_kv_heads < self.local_cache_kv_heads: + # Pad KV heads by repeating + repeat_factor = self.local_cache_kv_heads // self.local_num_kv_heads + key_states_for_cache = key_states_for_cache.repeat(1, repeat_factor, 1, 1) + value_states_for_cache = value_states_for_cache.repeat(1, repeat_factor, 1, 1) + + # Repeat KV heads for GQA (only needed without CONVERT_TO_MHA) + # With CONVERT_TO_MHA, K/V already have num_attention_heads + num_key_value_groups = self.local_num_heads // self.local_num_kv_heads + if num_key_value_groups > 1: + key_states = key_states.repeat_interleave(num_key_value_groups, dim=1) + value_states = value_states.repeat_interleave(num_key_value_groups, dim=1) + + if is_token_gen: + # Token generation: use decomposed attention with prior (cached) and active (current) KV + # past_key_value[0] = cached K, shape [bsz, cache_kv_heads, kv_seq_len, head_dim] + # past_key_value[1] = cached V, shape [bsz, cache_kv_heads, kv_seq_len, head_dim] (padded) + K_prior = past_key_value[0] + V_prior = past_key_value[1] + + # WORKAROUND 1: Slice KV heads if cache has more than layer needs + # Only needed when NOT using CONVERT_TO_MHA (standard GQA mode) + # With CONVERT_TO_MHA, cache and layer have same num_kv_heads + if not self.use_gqa_convert_to_mha and self.local_num_kv_heads < self.local_cache_kv_heads: + # Cache has repeated heads, just take the first local_num_kv_heads + K_prior = K_prior[:, :self.local_num_kv_heads, :, :] + V_prior = V_prior[:, :self.local_num_kv_heads, :, :] + + # WORKAROUND 2: Slice V_prior back to v_head_dim (128) from head_dim (192) + if self.attn_v_head_dim < self.attn_head_dim: + V_prior = V_prior[..., :self.attn_v_head_dim] + + # Repeat cached KV for GQA (only needed without CONVERT_TO_MHA) + # With CONVERT_TO_MHA, cached K/V already have num_attention_heads + if num_key_value_groups > 1: + K_prior = K_prior.repeat_interleave(num_key_value_groups, dim=1) + V_prior = V_prior.repeat_interleave(num_key_value_groups, dim=1) + + # Compute attention on prior (cached) KV + # K_prior shape: [bsz, num_heads, kv_seq_len, head_dim] + prior_scores = torch.matmul(query_states, K_prior.transpose(-2, -1)) * self.scaling + + # Apply attention mask to prior scores + if attention_mask is not None: + # Convert boolean mask to additive mask if needed + if attention_mask.dtype == torch.bool: + prior_scores = prior_scores.masked_fill(~attention_mask, float('-inf')) + else: + prior_scores = prior_scores + attention_mask + + # Apply sliding window mask for SWA layers + if self.is_sliding_window and self.sliding_window_size is not None and position_ids is not None: + kv_seq_len = prior_scores.size(-1) + current_pos = position_ids[0, 0] + pos_indices = torch.arange(kv_seq_len, device=prior_scores.device) + sliding_mask = pos_indices >= (current_pos - self.sliding_window_size + 1) + sliding_mask = sliding_mask[None, None, None, :] + prior_scores = prior_scores.masked_fill(~sliding_mask, float('-inf')) + + prior_scores = prior_scores.to(torch.float32) + + # Compute attention on active (current) KV + active_scores = torch.matmul(query_states, key_states.transpose(-2, -1)) * self.scaling + active_scores = active_scores.to(torch.float32) + + # Combined softmax over prior and active scores + all_scores = torch.cat([prior_scores, active_scores], dim=-1) + + # Add attention sink bias (following HF implementation) + # This must be applied to token generation as well! + use_sink = self._use_sink_bias and self.attention_sink_bias is not None + if use_sink: + tp_rank = parallel_state.get_tensor_model_parallel_rank() if parallel_state.model_parallel_is_initialized() else 0 + local_sink = self.attention_sink_bias[tp_rank * self.local_num_heads:(tp_rank + 1) * self.local_num_heads] + sink_bias = local_sink.reshape(1, -1, 1, 1).expand(bsz, -1, q_len, 1) + all_scores = torch.cat([all_scores, sink_bias], dim=-1) + + # Numerical stability: subtract max before softmax + all_scores = all_scores - all_scores.max(dim=-1, keepdim=True).values + attn_weights = F.softmax(all_scores, dim=-1, dtype=torch.float32) + + # Drop the sink column after softmax + if use_sink: + attn_weights = attn_weights[..., :-1] + + # Split attention weights back + prior_weights = attn_weights[..., :-q_len].to(V_prior.dtype) + active_weights = attn_weights[..., -q_len:].to(value_states.dtype) + + # Compute attention outputs + attn_prior = torch.matmul(prior_weights, V_prior) + attn_active = torch.matmul(active_weights, value_states) + attn_output = attn_prior + attn_active + else: + # Context encoding: standard attention + # With CP: Q is local [B, H, S/CP, D], K/V are full [B, H, S, D] + # Without CP: Q/K/V all have same seq_len + attn_weights = torch.matmul(query_states, key_states.transpose(-2, -1)) * self.scaling + + # Apply attention mask (additive mask: 0 = attend, -inf = mask out) + # The framework creates boolean masks, so we need to convert them + # With CP: attention_mask is already split to [B, 1, S/CP, S] (local Q rows, full K cols) + if attention_mask is not None: + # Convert boolean mask to additive mask if needed + if attention_mask.dtype == torch.bool: + # True = attend (0), False = mask (-inf) + additive_mask = torch.zeros_like(attn_weights) + additive_mask = additive_mask.masked_fill(~attention_mask, float('-inf')) + attn_weights = attn_weights + additive_mask + else: + # Already additive mask + attn_weights = attn_weights + attention_mask + + # Apply sliding window mask for SWA layers + if self.is_sliding_window and self.sliding_window_size is not None: + kv_seq_len = attn_weights.size(-1) + if is_context_parallel: + # With CP: Q has local seq len, K has full seq len. + # Use local_position_ids for correct global Q positions. + row_idx = local_position_ids[0].unsqueeze(1).to(attn_weights.device) + else: + row_idx = torch.arange(kv_seq_len, device=attn_weights.device).unsqueeze(1) + col_idx = torch.arange(kv_seq_len, device=attn_weights.device).unsqueeze(0) + # Causal: col <= row, and within window: col >= row - window_size + 1 + sliding_mask = (col_idx <= row_idx) & (col_idx >= row_idx - self.sliding_window_size + 1) + sliding_mask = sliding_mask[None, None, :, :] + # Convert to additive mask + attn_weights = attn_weights.masked_fill(~sliding_mask, float('-inf')) + + # Add attention sink bias (following HF implementation) + # This adds an extra "sink" column to attention weights + use_sink = self._use_sink_bias and self.attention_sink_bias is not None + if use_sink: + # Get local portion of sink bias for this TP rank + tp_rank = parallel_state.get_tensor_model_parallel_rank() if parallel_state.model_parallel_is_initialized() else 0 + local_sink = self.attention_sink_bias[tp_rank * self.local_num_heads:(tp_rank + 1) * self.local_num_heads] + # Reshape and expand: [local_num_heads] -> [bsz, local_num_heads, q_len, 1] + sink_bias = local_sink.reshape(1, -1, 1, 1).expand(bsz, -1, q_len, 1) + attn_weights = torch.cat([attn_weights, sink_bias], dim=-1) + + # Numerical stability: subtract max before softmax (like HF implementation) + attn_weights = attn_weights - attn_weights.max(dim=-1, keepdim=True).values + + # Softmax + attn_weights = F.softmax(attn_weights, dim=-1, dtype=torch.float32) + + # Drop the sink column after softmax + if use_sink: + attn_weights = attn_weights[..., :-1] + + attn_weights = attn_weights.to(value_states.dtype) + + # Apply attention to values + attn_output = torch.matmul(attn_weights, value_states) + + # Reshape and project output + attn_output = attn_output.transpose(1, 2).contiguous() + attn_output = attn_output.reshape(bsz, q_len, self.local_num_heads * self.attn_v_head_dim) + + # Context Parallelism: gather output across CP ranks BEFORE o_proj. + # With SP enabled, o_proj scatters along seq dim. The input must have full S + # (not S/CP), otherwise the SP-scattered output won't match the residual. + # Without SP, gather after o_proj to restore full seq_len for residual. + if is_context_parallel: + attn_output = gather_from_tensor_model_parallel_region_with_dim( + attn_output, gather_dim=1, + process_group=get_context_parallel_attention_cp_group(), + ) + + attn_output = self.o_proj(attn_output) + + # Prepare KV cache output - return as tuple for KV cache manager + # Return LOCAL key/value states for cache (each CP rank stores its portion) + new_key_value = (key_states_for_cache, value_states_for_cache) + + return attn_output, new_key_value, cos_cache, sin_cache + + +class MiMoV2MLP(nn.Module): + """Standard MLP for non-MoE layers in MiMo-V2.5.""" + + def __init__(self, config: MiMoV2InferenceConfig): + super().__init__() + self.hidden_size = config.hidden_size + # Use the dense intermediate size for non-MoE layers + self.intermediate_size = getattr(config, 'dense_intermediate_size', config.intermediate_size * 8) + + dtype = config.neuron_config.torch_dtype + + if parallel_state.model_parallel_is_initialized(): + tp_group = parallel_state.get_tensor_model_parallel_group() + + self.gate_proj = ColumnParallelLinear( + self.hidden_size, + self.intermediate_size, + bias=False, + gather_output=False, + dtype=dtype, + tensor_model_parallel_group=tp_group, + ) + self.up_proj = ColumnParallelLinear( + self.hidden_size, + self.intermediate_size, + bias=False, + gather_output=False, + dtype=dtype, + tensor_model_parallel_group=tp_group, + ) + self.down_proj = RowParallelLinear( + self.intermediate_size, + self.hidden_size, + bias=False, + input_is_parallel=True, + dtype=dtype, + tensor_model_parallel_group=tp_group, + ) + else: + self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False) + self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False) + + self.act_fn = F.silu + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x)) + + +class NeuronMiMoV2DecoderLayer(nn.Module): + """MiMo-V2.5 Decoder Layer with hybrid attention and conditional MoE.""" + + def __init__(self, config: MiMoV2InferenceConfig, layer_idx: int): + super().__init__() + self.hidden_size = config.hidden_size + self.layer_idx = layer_idx + + # Determine attention type for this layer + is_sliding_window = config.layer_attention_types[layer_idx] == "sliding_window" + self.attention_type = "sliding_window" if is_sliding_window else "full" + + # Create attention module + self.self_attn = NeuronMiMoV2Attention( + config=config, + layer_idx=layer_idx, + is_sliding_window=is_sliding_window, + ) + + # Determine if this layer uses MoE + self.uses_moe = config.layer_uses_moe[layer_idx] + + # Create MLP/MoE module + if self.uses_moe: + self.mlp = initialize_moe_module(config=config) + else: + self.mlp = MiMoV2MLP(config) + + # Layer norms + self.input_layernorm = get_rmsnorm_cls()( + config.hidden_size, + eps=config.layernorm_epsilon, + ) + self.post_attention_layernorm = get_rmsnorm_cls()( + config.hidden_size, + eps=config.layernorm_epsilon, + ) + + # Config flags + self.sequence_parallel_enabled = config.neuron_config.sequence_parallel_enabled + + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_value: Optional[Tuple[torch.Tensor]] = None, + padding_mask: Optional[torch.Tensor] = None, + cos_cache: Optional[torch.Tensor] = None, + sin_cache: Optional[torch.Tensor] = None, + **kwargs, + ) -> Tuple[torch.FloatTensor, ...]: + """Forward pass for decoder layer.""" + + # Self attention with residual + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + + hidden_states, present_key_value, cos_cache, sin_cache = self.self_attn( + hidden_states=hidden_states, + attention_mask=attention_mask, + position_ids=position_ids, + past_key_value=past_key_value, + cos_cache=cos_cache, + sin_cache=sin_cache, + **kwargs, + ) + hidden_states = residual + hidden_states + + # MLP/MoE with residual + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + + if self.uses_moe: + hidden_states = self.mlp(hidden_states, padding_mask)[0] + else: + hidden_states = self.mlp(hidden_states) + + hidden_states = residual + hidden_states + + outputs = (hidden_states, present_key_value, cos_cache, sin_cache, None) + return outputs + + +class NeuronMiMoV2Model(NeuronBaseModel): + """MiMo-V2.5 Model for NXD inference.""" + + def setup_attr_for_model(self, config: MiMoV2InferenceConfig): + self.on_device_sampling = config.neuron_config.on_device_sampling_config is not None + self.tp_degree = config.neuron_config.tp_degree + self.hidden_size = config.hidden_size + self.num_attention_heads = config.num_attention_heads + + # Check if we need GQA CONVERT_TO_MHA mode + # When tp_degree > num_kv_heads, we replicate K/V to match num_attention_heads + min_kv_heads = min( + config.num_key_value_heads, + getattr(config, 'swa_num_key_value_heads', config.num_key_value_heads) + ) + self.use_gqa_convert_to_mha = self.tp_degree > min_kv_heads + + if self.use_gqa_convert_to_mha: + # With CONVERT_TO_MHA, KV cache stores num_attention_heads (same as Q) + self.num_key_value_heads = config.num_attention_heads + else: + # Standard GQA: use the maximum num_kv_heads for KV cache + # (handles hybrid full/sliding window attention) + self.num_key_value_heads = max( + config.num_key_value_heads, + getattr(config, 'swa_num_key_value_heads', config.num_key_value_heads) + ) + + self.max_batch_size = config.neuron_config.max_batch_size + self.buckets = config.neuron_config.buckets + + # MiMo has hybrid attention (full + sliding window) + # NOTE: Do NOT set self.sliding_window here because it affects KV cache size globally. + # MiMo handles sliding window per-layer in the attention module itself. + # Setting has_mixed_attn = True enables proper mask creation without affecting cache size. + self.has_mixed_attn = True + + def init_model(self, config: MiMoV2InferenceConfig): + self.padding_idx = getattr(config, 'pad_token_id', None) + self.vocab_size = config.vocab_size + + self.embed_tokens = ParallelEmbedding( + config.vocab_size, + config.hidden_size, + self.padding_idx, + dtype=config.neuron_config.torch_dtype, + shard_across_embedding=True, + ) + + self.layers = nn.ModuleList([ + NeuronMiMoV2DecoderLayer(config, layer_idx) + for layer_idx in range(config.num_hidden_layers) + ]) + + self.norm = get_rmsnorm_cls()( + config.hidden_size, + eps=config.layernorm_epsilon, + ) + + self.lm_head = ColumnParallelLinear( + config.hidden_size, + config.vocab_size, + gather_output=not self.on_device_sampling, + bias=False, + ) + + +def _replicate_kv_weights_for_convert_to_mha( + tensor: torch.Tensor, + source_heads: int, + target_heads: int, + head_dim: int, +) -> torch.Tensor: + """Replicate K/V weights from source_heads to target_heads for CONVERT_TO_MHA. + + Args: + tensor: Weight tensor of shape [source_heads * head_dim, hidden_size] + source_heads: Number of source KV heads + target_heads: Number of target heads (num_attention_heads) + head_dim: Head dimension + + Returns: + Replicated tensor of shape [target_heads * head_dim, hidden_size] + """ + if tensor is None or source_heads >= target_heads: + return tensor + + repeats = target_heads // source_heads + + # Reshape to [source_heads, head_dim, hidden_size] + original_shape = tensor.shape + tensor = tensor.view(source_heads, head_dim, -1) + + # Repeat along head dimension + tensor = tensor.repeat_interleave(repeats, dim=0) + + # Reshape back to [num_heads * head_dim, hidden_size] + tensor = tensor.view(-1, original_shape[-1]) + + return tensor + + +def convert_mimo_v2_hf_to_neuron_state_dict( + neuron_state_dict: Dict[str, Any], + config: MiMoV2InferenceConfig, +) -> Dict[str, Any]: + """Convert HuggingFace MiMo-V2.5 weights to Neuron format. + + This handles: + 1. Router weight renaming + 2. Expert weight concatenation and transposition + 3. FP8 dequantization if needed + 4. K/V weight replication for CONVERT_TO_MHA mode + """ + + assert config.neuron_config.glu_mlp is True, "Only GLU MLP is supported" + + # Dequantize layers if needed + _maybe_dequantize_layer(neuron_state_dict, config) + + # Add rank utility tensors + neuron_state_dict["rank_util.rank"] = torch.arange( + 0, config.neuron_config.tp_degree, dtype=torch.int32 + ) + + # Determine if CONVERT_TO_MHA is needed + tp_degree = config.neuron_config.tp_degree + num_attention_heads = config.num_attention_heads + + # MiMo-V2.5 has different KV heads for full and sliding window attention + full_num_kv_heads = config.num_key_value_heads # 4 + swa_num_kv_heads = config.swa_num_key_value_heads # 8 + + # Check if we need to replicate K/V weights + full_use_convert_to_mha = tp_degree > full_num_kv_heads + swa_use_convert_to_mha = tp_degree > swa_num_kv_heads + + for layer_idx in range(config.num_hidden_layers): + # Add rank utility for attention + neuron_state_dict[f"layers.{layer_idx}.self_attn.rank_util.rank"] = torch.arange( + 0, config.neuron_config.tp_degree, dtype=torch.int32 + ) + + # Determine attention type for this layer + is_sliding_window = config.layer_attention_types[layer_idx] == "sliding_window" + + if is_sliding_window: + src_num_kv_heads = swa_num_kv_heads + use_convert_to_mha = swa_use_convert_to_mha + head_dim = config.swa_head_dim # 192 + v_head_dim = config.swa_v_head_dim # 128 + else: + src_num_kv_heads = full_num_kv_heads + use_convert_to_mha = full_use_convert_to_mha + head_dim = config.head_dim # 192 + v_head_dim = config.v_head_dim # 128 + + # Replicate K/V weights if CONVERT_TO_MHA is needed + if use_convert_to_mha: + k_proj_key = f"layers.{layer_idx}.self_attn.k_proj.weight" + v_proj_key = f"layers.{layer_idx}.self_attn.v_proj.weight" + + if k_proj_key in neuron_state_dict: + neuron_state_dict[k_proj_key] = _replicate_kv_weights_for_convert_to_mha( + neuron_state_dict[k_proj_key], + src_num_kv_heads, + num_attention_heads, + head_dim, + ) + + if v_proj_key in neuron_state_dict: + neuron_state_dict[v_proj_key] = _replicate_kv_weights_for_convert_to_mha( + neuron_state_dict[v_proj_key], + src_num_kv_heads, + num_attention_heads, + v_head_dim, + ) + + # FP8 path: replicate per-row scales ([src_heads*head_dim, 1]) in + # lockstep with the weights. Without this the shard_weights step + # rejects the scale shape mismatch (e.g. [12,1] vs expected [192,1]). + # BF16 has no .scale key, so this loop is a no-op there. + for proj, hd in (("k_proj", head_dim), ("v_proj", v_head_dim)): + scale_key = f"layers.{layer_idx}.self_attn.{proj}.scale" + if scale_key in neuron_state_dict: + neuron_state_dict[scale_key] = _replicate_kv_weights_for_convert_to_mha( + neuron_state_dict[scale_key], + src_num_kv_heads, + num_attention_heads, + hd, + ) + + # Only convert MoE layers + if not config.layer_uses_moe[layer_idx]: + continue + + # Check if this layer has MoE weights + gate_key = f"layers.{layer_idx}.mlp.gate.weight" + if gate_key not in neuron_state_dict: + continue + + # Rename router weights + neuron_state_dict[f"layers.{layer_idx}.mlp.router.linear_router.weight"] = ( + neuron_state_dict[gate_key].detach().clone() + ) + del neuron_state_dict[gate_key] + + # Get dimensions from first expert + expert_0_gate = f"layers.{layer_idx}.mlp.experts.0.gate_proj.weight" + if expert_0_gate not in neuron_state_dict: + continue + + intermediate_size, hidden_size = neuron_state_dict[expert_0_gate].shape + device = neuron_state_dict[expert_0_gate].device + dtype = neuron_state_dict[expert_0_gate].dtype + + num_experts = config.n_routed_experts + + # Concatenate gate and up projections + gate_up_proj = torch.empty( + num_experts, + hidden_size, + 2 * intermediate_size, + dtype=dtype, + device=device, + ) + + for e in range(num_experts): + gate_proj_weights = neuron_state_dict[ + f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight" + ].T.detach().clone() + up_proj_weights = neuron_state_dict[ + f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight" + ].T.detach().clone() + + gate_up_proj[e, :, :intermediate_size] = gate_proj_weights + gate_up_proj[e, :, intermediate_size:] = up_proj_weights + + del neuron_state_dict[f"layers.{layer_idx}.mlp.experts.{e}.gate_proj.weight"] + del neuron_state_dict[f"layers.{layer_idx}.mlp.experts.{e}.up_proj.weight"] + + # Pad if needed + pad_size = getattr(config, "moe_intermediate_pad_size", 0) + if pad_size > 0: + gate_up_proj = gate_up_proj.reshape(num_experts, hidden_size, 2, -1) + gate_up_proj = F.pad(gate_up_proj, (0, pad_size)) + gate_up_proj = gate_up_proj.reshape(num_experts, hidden_size, -1) + + neuron_state_dict[f"layers.{layer_idx}.mlp.expert_mlps.mlp_op.gate_up_proj.weight"] = gate_up_proj + + # Convert down projections + down_proj = torch.empty( + num_experts, + intermediate_size, + hidden_size, + dtype=dtype, + device=device, + ) + + for e in range(num_experts): + down_proj_weights = neuron_state_dict[ + f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight" + ].T.detach().clone() + down_proj[e] = down_proj_weights + del neuron_state_dict[f"layers.{layer_idx}.mlp.experts.{e}.down_proj.weight"] + + # Pad if needed + if pad_size > 0: + down_proj = F.pad(down_proj, (0, 0, 0, pad_size)) + + neuron_state_dict[f"layers.{layer_idx}.mlp.expert_mlps.mlp_op.down_proj.weight"] = down_proj + + gc.collect() + + # --- Expand MoE blockwise scales along the TP-partitioned dim (FP8 only). --- + # NxDI's shard_checkpoint splits the scale on its partition dim into + # `per_partition_size = dim_size / tp_degree`. At TP=64 both projections + # have per-rank "intermediate" smaller than the 128-wide scale block, so + # several ranks share one scale block — we need to replicate scale entries + # along that dim. Adjacent ranks whose weight falls inside the same + # 128-wide block genuinely share that block's scale. No-op when the + # .scale keys are absent (BF16 path). + if getattr(config.neuron_config, "quantized", False): + # IMPORTANT: MoE expert weights are sharded by moe_tp_degree (not the + # top-level tp_degree — attention uses tp_degree, MoE can use a + # different split). At moe_tp=64 the per-rank intermediate is 32 (<128) + # so we had to expand the scale to make the shard layout match; at + # moe_tp=16 per-rank intermediate is 128 (>=128) and no expansion is + # needed. + moe_tp = getattr(config.neuron_config, "moe_tp_degree", None) or config.neuron_config.tp_degree + for layer_idx in range(config.num_hidden_layers): + if not config.layer_uses_moe[layer_idx]: + continue + + # down_proj (RowParallel on intermediate dim). Scale: [E, I_blocks, H_blocks] + dp_key = f"layers.{layer_idx}.mlp.expert_mlps.mlp_op.down_proj.scale" + if dp_key in neuron_state_dict: + s = neuron_state_dict[dp_key] + i_blocks = s.shape[1] + h_blocks = s.shape[2] + intermediate = i_blocks * 128 + i_per_rank = intermediate // moe_tp + if i_per_rank < 128: + ranks_per_block = 128 // i_per_rank + s_exp = s.unsqueeze(2).expand(-1, -1, ranks_per_block, -1) + s_exp = s_exp.reshape(s.shape[0], i_blocks * ranks_per_block, h_blocks) + assert s_exp.shape[1] == moe_tp, ( + f"down_proj.scale expansion produced {s_exp.shape[1]} rows, " + f"expected moe_tp={moe_tp}" + ) + neuron_state_dict[dp_key] = s_exp.contiguous() + + # gate_up_proj (ColumnParallel on 2*intermediate dim, gate|up fused + # along last axis). Scale: [E, H_blocks, 2*I_blocks] stored as + # [gate_half | up_half]. Module parameter has per-rank last-dim=1 + # (via _apply_blockwise_scale_stride_fix patch forcing + # partition_stride=1), so the full scale must have last-dim=moe_tp + # with gate entries 0..moe_tp/2 and up entries moe_tp/2..moe_tp. + # Expand each half independently to preserve the gate/up boundary + # when NxD does `split(per_partition=2*I/moe_tp, dim=-1)`. + gu_key = f"layers.{layer_idx}.mlp.expert_mlps.mlp_op.gate_up_proj.scale" + if gu_key in neuron_state_dict: + s = neuron_state_dict[gu_key] + h_blocks = s.shape[1] + two_i_blocks = s.shape[2] + assert two_i_blocks % 2 == 0, ( + f"gate_up_proj.scale last dim must be 2*i_blocks, got {two_i_blocks}" + ) + i_blocks = two_i_blocks // 2 + intermediate = i_blocks * 128 + out_per_rank = (2 * intermediate) // moe_tp + if out_per_rank < 128: + assert moe_tp % 2 == 0, f"moe_tp={moe_tp} must be even for gate/up scale split" + ranks_per_half = moe_tp // 2 + assert ranks_per_half % i_blocks == 0, ( + f"ranks_per_half={ranks_per_half} must be divisible by " + f"i_blocks={i_blocks}" + ) + ranks_per_block = ranks_per_half // i_blocks + gate_half = s[..., :i_blocks] # [E, H_blocks, i_blocks] + up_half = s[..., i_blocks:] + gate_exp = ( + gate_half.unsqueeze(-1) + .expand(-1, -1, -1, ranks_per_block) + .reshape(s.shape[0], h_blocks, ranks_per_half) + ) + up_exp = ( + up_half.unsqueeze(-1) + .expand(-1, -1, -1, ranks_per_block) + .reshape(s.shape[0], h_blocks, ranks_per_half) + ) + s_exp = torch.cat([gate_exp, up_exp], dim=-1) + assert s_exp.shape[-1] == moe_tp, ( + f"gate_up_proj.scale expansion produced {s_exp.shape[-1]} " + f"entries, expected moe_tp={moe_tp}" + ) + neuron_state_dict[gu_key] = s_exp.contiguous() + + return neuron_state_dict + + +def _maybe_dequantize_layer( + neuron_state_dict: Dict[str, Any], + config: MiMoV2InferenceConfig, +): + """Dequantize FP8 layers if present.""" + scale_layers = [] + + for layer_key in list(neuron_state_dict.keys()): + if "_scale_inv" in layer_key: + scales = neuron_state_dict[layer_key] + scale_layers.append(layer_key) + + fp8_layer_name = layer_key.replace("_scale_inv", "") + if fp8_layer_name not in neuron_state_dict: + continue + + fp8_layer = neuron_state_dict[fp8_layer_name] + + # Get block size from config if available + if hasattr(config, 'quantization_config') and config.quantization_config: + block_size = config.quantization_config.get("weight_block_size", [128, 128]) + else: + block_size = [128, 128] + + # Expand scales and dequantize + scales_expanded = scales.repeat_interleave(block_size[0], dim=0) + scales_expanded = scales_expanded.repeat_interleave(block_size[1], dim=1) + + # Ensure shapes match + if scales_expanded.shape != fp8_layer.shape: + scales_expanded = scales_expanded[:fp8_layer.shape[0], :fp8_layer.shape[1]] + + scaled_layer = fp8_layer.to(torch.float32) * scales_expanded.to(torch.float32) + neuron_state_dict[fp8_layer_name] = scaled_layer.to(config.neuron_config.torch_dtype) + + # Remove scale layers + for scale_layer in scale_layers: + del neuron_state_dict[scale_layer] + + +class NeuronMiMoV2ForCausalLM(NeuronBaseForCausalLM): + """MiMo-V2.5 for Causal Language Modeling on Neuron.""" + + _model_cls = NeuronMiMoV2Model + + def __init__(self, *args, **kwargs): + # Install FP8 monkey-patches BEFORE super().__init__ so the patched + # RouterTopK.__init__ and quantization layer classes are in effect + # when NxDI builds the decoder (and instantiates routers). Harnesses + # that drive us via model.compile()/model.load() (e.g. vllm-neuron) + # call those methods AFTER construction, so patching from inside + # compile()/load() is too late — RouterTopK instances would already + # lack our e_score_correction_bias parameter, silently routing tokens + # to wrong experts and producing gibberish output. + # + # _install_fp8_patches() reads self.neuron_config, which needs to + # exist; grab it from the args or the config arg the same way the + # base class does. + ncfg = kwargs.get("config") or (args[1] if len(args) > 1 else None) + if ncfg is not None and getattr(getattr(ncfg, "neuron_config", None), "quantized", False): + self._apply_ep_scale_fix() + self._apply_blockwise_scale_stride_fix() + self._apply_2d_per_channel_fix() + self._apply_router_noaux_tc_fix() + super().__init__(*args, **kwargs) + + @staticmethod + def load_hf_model(model_path: str, **kwargs): + """Load HuggingFace model. + + Note: MiMo-V2.5 uses custom code, so we need trust_remote_code=True + """ + from transformers import AutoModelForCausalLM + return AutoModelForCausalLM.from_pretrained( + model_path, + trust_remote_code=True, + **kwargs, + ) + + @classmethod + def get_config_cls(cls) -> Type[MiMoV2InferenceConfig]: + return MiMoV2InferenceConfig + + @staticmethod + def convert_hf_to_neuron_state_dict( + state_dict: Dict[str, Any], + config: MiMoV2InferenceConfig, + ) -> Dict[str, Any]: + return convert_mimo_v2_hf_to_neuron_state_dict(state_dict, config) + + def enable_context_encoding(self): + self.compile_tag = CONTEXT_ENCODING_MODEL_TAG + super().enable_context_encoding() + + def enable_token_generation(self): + self.compile_tag = TOKEN_GENERATION_MODEL_TAG + super().enable_token_generation() + + # ------------------------------------------------------------------ + # FP8 quantized-inference monkey-patches (no-op unless quantized=True). + # + # Reconcile the preprocessed Neuron-FP8 checkpoint (blockwise-MoE + + # per-row-attn) with NxDI's global blockwise_symmetric q_config. All + # four are gated by self.neuron_config.quantized so the BF16 path is + # completely untouched. + # ------------------------------------------------------------------ + + @staticmethod + def _apply_ep_scale_fix(): + """Skip per-channel `scale` params when marking expert-parallel + weights; they have shape [1, 1, W] and cannot be EP-sharded.""" + from neuronx_distributed.modules.moe.moe_parallel_layers import ( + ExpertFusedLinear, + ) + + if getattr(ExpertFusedLinear, "_mimo_v2_ep_scale_patched", False): + return + + def _patched_mark( + self_inner, + iterable=None, + expert_parallel_group_size=None, + is_prefill=True, + expert_distribution=None, + ): + from neuronx_distributed.parallel_layers.parallel_state import ( + get_expert_model_parallel_size, + ) + + if expert_parallel_group_size is None: + expert_parallel_group_size = get_expert_model_parallel_size() + + if expert_parallel_group_size > 1: + if iterable is None: + params_to_mark = [] + for name, p in self_inner.named_parameters(): + if name == "scale" and p.shape[0] == 1: + continue + params_to_mark.append(p) + iterable = params_to_mark + + for p in iterable: + p.expert_model_parallel = True + if is_prefill: + p.is_prefill = True + p.expert_distribution = expert_distribution + + ExpertFusedLinear._mark_expert_parallel_weights = _patched_mark + ExpertFusedLinear._mimo_v2_ep_scale_patched = True + + @staticmethod + def _apply_blockwise_scale_stride_fix(): + """Force scale.partition_stride=1 for BLOCKWISE_SYMMETRIC quantization + — stride>1 causes strided-splitting failures when per-rank weight size + is smaller than a block.""" + from neuronx_distributed.quantization.quantization_config import ( + QuantizationType, + ) + from neuronx_distributed.quantization.quantization_layers import ( + BaseQuantizeParallelLinear, + ) + + if getattr(BaseQuantizeParallelLinear, "_mimo_v2_blockwise_stride_patched", False): + return + + _original_setup = BaseQuantizeParallelLinear._setup_for_scale + + def _patched_setup(self_inner, *args, **kwargs): + _original_setup(self_inner, *args, **kwargs) + if ( + hasattr(self_inner, "quantization_type") + and self_inner.quantization_type == QuantizationType.BLOCKWISE_SYMMETRIC + and hasattr(self_inner, "scale") + and hasattr(self_inner.scale, "partition_stride") + and self_inner.scale.partition_stride > 1 + ): + self_inner.scale.partition_stride = 1 + + BaseQuantizeParallelLinear._setup_for_scale = _patched_setup + BaseQuantizeParallelLinear._mimo_v2_blockwise_stride_patched = True + + @staticmethod + def _apply_2d_per_channel_fix(): + """Route 2D self_attn + layer-0 dense-MLP swaps through per_channel_symmetric. + + MiMo-V2.5's preprocess writes: + - MoE experts: 3D weights with (E, out//128, in//128) blockwise scales. + - self_attn q/k/v + layer-0 mlp gate/up/down: 2D weights with + (out, 1) per-row scales. + + NxDI's q_config is global blockwise_symmetric (to satisfy the MoE). + Feeding that into the 2D classes triggers + `block axis cannot be < 0 or > 2, received 2` in _setup_for_scale + (block axes [1, 2] exceed rank-2 weight_shape). This wraps the 2D + classes' from_float to override q_config on the fly. + """ + from neuronx_distributed.quantization.quantization_config import ( + QuantizationType, + ) + from neuronx_distributed.quantization.quantization_layers import ( + QuantizedColumnParallel, + QuantizedRowParallel, + ) + + def _wrap(cls): + if getattr(cls, "_mimo_v2_2d_patched", False): + return + original_from_float = cls.from_float + + def _patched_from_float(klass, mod, q_config=None, _orig=original_from_float): + if q_config is not None and q_config.get("quantization_type") == \ + QuantizationType.BLOCKWISE_SYMMETRIC: + q_config = dict(q_config) + q_config["quantization_type"] = QuantizationType.PER_CHANNEL_SYMMETRIC + q_config["quantization_per_channel_axis"] = 0 + q_config.pop("block_axis", None) + q_config.pop("block_size", None) + if q_config is None: + return _orig(mod) + return _orig(mod, q_config) + + cls.from_float = classmethod(_patched_from_float) + cls._mimo_v2_2d_patched = True + + _wrap(QuantizedColumnParallel) + _wrap(QuantizedRowParallel) + + @staticmethod + def _apply_router_noaux_tc_fix(): + """Register e_score_correction_bias on NxD RouterTopK and fold it into + top-k selection so MiMo-V2's noaux_tc routing matches HF reference. + + MiMo-V2's HF config uses topk_method='noaux_tc': each expert score is + `sigmoid(logits) + e_score_correction_bias`, top-k indices are chosen + from THAT biased score; the returned expert weights (affinities) + come from the UNBIASED sigmoid(logits). NxD's stock RouterTopK is + plain topk with no bias slot, so without this the bias is silently + dropped and ~all tokens route to wrong experts. + """ + from neuronx_distributed.modules.moe.routing import RouterTopK + + if getattr(RouterTopK, "_mimo_v2_noaux_tc_patched", False): + return + + original_init = RouterTopK.__init__ + + def _patched_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + # CRITICAL: dtype + init value both matter for XLA tracing. + # + # 1) dtype=torch.bfloat16: the NxDI checkpoint loader casts router + # bias from FP32 -> BF16 ("Found torch.float32 weights in + # checkpoint ... Will convert to torch.bfloat16"). If the traced + # NEFF expects FP32 but the checkpoint supplies BF16, the + # LayoutTransformation silently drops the weight and keeps the + # trace-time init values — so the bias at runtime is whatever + # we init here, not the checkpoint values. + # + # 2) init=arange, NOT zeros: if every entry is identical (all + # zeros), the `+ bias` op does not change the relative ordering + # of topk, so XLA's constant-folding passes can prove the add + # is a no-op and eliminate it entirely — dropping the bias + # parameter from the HLO. At that point checkpoint loading has + # nothing to bind to and the real bias is silently discarded. + # Using arange guarantees distinct per-expert values, forcing + # the compiler to keep the add as a runtime op with a live + # parameter. Source: Jim Burtoft's MiniMax-M2 fix notes + # (jimburtoft/neuronx-distributed-inference@49f8e164). + self.e_score_correction_bias = nn.Parameter( + torch.arange(self.num_experts, dtype=torch.bfloat16), + requires_grad=False, + ) + + def _patched_forward(self, hidden_states): + router_logits = self.get_router_logits(hidden_states) + expert_affinities = self.apply_activation_fn(router_logits) + + # MiMo (and MiniMax-M2) uses topk_method='noaux_tc': the bias is + # added ONLY for top-k selection, but the unbiased sigmoid scores + # remain as the expert-affinity weights passed to the experts. + scores_for_choice = ( + expert_affinities.float() + self.e_score_correction_bias.unsqueeze(0) + ) + _, expert_index = torch.topk(scores_for_choice, self.top_k, dim=-1) + + expert_affinities = expert_affinities.to(dtype=hidden_states.dtype) + expert_index = expert_index.detach().to(dtype=torch.long) + return router_logits, expert_affinities, expert_index + + RouterTopK.__init__ = _patched_init + RouterTopK.forward = _patched_forward + RouterTopK._mimo_v2_noaux_tc_patched = True + + def _install_fp8_patches(self): + """Install all FP8-specific runtime patches. No-op for BF16.""" + if not getattr(self.neuron_config, "quantized", False): + return + self._apply_ep_scale_fix() + self._apply_blockwise_scale_stride_fix() + self._apply_2d_per_channel_fix() + self._apply_router_noaux_tc_fix() + + def compile(self, *args, **kwargs): + # save_sharded_checkpoint=True serializes shards during compile() and + # that code path reads scale.partition_stride — patches must be live. + self._install_fp8_patches() + return super().compile(*args, **kwargs) + + def load(self, *args, **kwargs): + self._install_fp8_patches() + return super().load(*args, **kwargs) + + @classmethod + def save_quantized_state_dict(cls, model_path, config): + """MiMo-V2.5 ships pre-quantized FP8 safetensors via our preprocess script. + The base implementation calls AutoModelForCausalLM.from_pretrained to + re-quantize, which requires a CUDA GPU (finegrained_fp8 gate) and + materializes an ~600 GB BF16 copy. Skip if the checkpoint directory + already contains a Neuron-FP8 index produced by preprocess.""" + import os as _os + qpath = ( + getattr(config.neuron_config, "quantized_checkpoints_path", None) + or model_path + ) + if qpath and _os.path.isdir(qpath): + index = _os.path.join(qpath, "model.safetensors.index.json") + if _os.path.isfile(index): + return + return super().save_quantized_state_dict(model_path, config) + + def get_compiler_args(self) -> str: + """Get compiler arguments optimized for MiMo-V2.5.""" + if self.compile_tag == CONTEXT_ENCODING_MODEL_TAG: + optimization_level = "-O1" + elif self.compile_tag == TOKEN_GENERATION_MODEL_TAG: + optimization_level = "-O3" if self.neuron_config.moe_ep_degree > 1 else "-O1" + else: + optimization_level = "-O1" + + compiler_args = ( + f"--enable-saturate-infinity " + f"--enable-mixed-precision-accumulation " + f"--model-type transformer " + f"{optimization_level}" + ) + + # Add CC overlap optimization + compiler_args += ( + " --tensorizer-options='--enable-ccop-compute-overlap " + "--cc-pipeline-tiling-factor=2'" + ) + + compiler_args += " --auto-cast=none" + + # Enable vector-offset DGE + compiler_args += " --internal-enable-dge-levels vector_dynamic_offsets" + compiler_args += " --internal-hlo2tensorizer-options='--verify-hlo=true'" + + if self.neuron_config.scratchpad_page_size: + compiler_args += f" --hbm-scratchpad-page-size={self.neuron_config.scratchpad_page_size}" + + return compiler_args diff --git a/contrib/models/MiMo-V2.5/test/__init__.py b/contrib/models/MiMo-V2.5/test/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/contrib/models/MiMo-V2.5/test/integration/__init__.py b/contrib/models/MiMo-V2.5/test/integration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/contrib/models/MiMo-V2.5/test/integration/test_model.py b/contrib/models/MiMo-V2.5/test/integration/test_model.py new file mode 100644 index 00000000..d52fabb6 --- /dev/null +++ b/contrib/models/MiMo-V2.5/test/integration/test_model.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""Integration tests for MiMo-V2.5 NeuronX implementation.""" + +import pytest +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).parent.parent.parent / "src")) + + +def test_config_import(): + """Test that config class can be imported.""" + from modeling_mimo_v2 import MiMoV2InferenceConfig, NeuronMiMoV2ForCausalLM + assert MiMoV2InferenceConfig is not None + assert NeuronMiMoV2ForCausalLM is not None + print("PASS: Config and model classes imported successfully") + + +def test_required_attributes(): + """Test that required attributes are defined.""" + from modeling_mimo_v2 import MiMoV2InferenceConfig + # Check get_required_attributes without instantiation (requires many params) + required = MiMoV2InferenceConfig.get_required_attributes(MiMoV2InferenceConfig) + assert "hidden_size" in required + assert "n_routed_experts" in required + assert "num_experts_per_tok" in required + assert "hybrid_layer_pattern" in required + assert "v_head_dim" in required + assert "swa_head_dim" in required + print(f"PASS: {len(required)} required attributes defined") + + +def test_neuron_config_cls(): + """Test that MoENeuronConfig is returned.""" + from modeling_mimo_v2 import MiMoV2InferenceConfig + from neuronx_distributed_inference.models.config import MoENeuronConfig + assert MiMoV2InferenceConfig.get_neuron_config_cls() == MoENeuronConfig + print("PASS: MoENeuronConfig returned") + + +def test_state_dict_converter(): + """Test that state dict converter function exists.""" + from modeling_mimo_v2 import NeuronMiMoV2ForCausalLM + assert hasattr(NeuronMiMoV2ForCausalLM, "convert_hf_to_neuron_state_dict") + print("PASS: State dict converter exists") + + +if __name__ == "__main__": + test_config_import() + test_required_attributes() + test_neuron_config_cls() + test_state_dict_converter() + print("\nAll tests passed!") diff --git a/contrib/models/MiMo-V2.5/test/unit/__init__.py b/contrib/models/MiMo-V2.5/test/unit/__init__.py new file mode 100644 index 00000000..e69de29b From b830a78edf442651d3b6b458a0e960b2f4158a56 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 06:28:54 +0800 Subject: [PATCH 02/27] preprocess: rebuild weight_map from actual shard files for V2.5 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MiMo-V2.5's published model.safetensors.index.json references legacy shard filenames like `model_N-00001-of-00002.safetensors`, but the LFS objects on HuggingFace Hub (and therefore on disk after download) are named `model_pp0_epN_shardM.safetensors`. The N values between the two namings are not aligned either, so a mechanical legacy->new rewrite doesn't work — for example model.layers.0.input_layernorm is mapped to `model_1-00002-of-00002` in the index but actually lives in `model_pp0_ep0_shard1`. Rather than reverse-engineer the ep-index permutation, scan the on-disk shards once and rebuild weight_map directly from each safetensors file's manifest. This is a one-time O(num_shards) open at startup and avoids any heuristic filename mapping. Preserves the fast-path for pre-V2.5 checkpoints (where the index filenames match the on-disk names): if any overlap is detected the provided weight_map is used as-is. Co-Authored-By: Claude Opus 4.7 --- .../preprocess_mimo_v2_5_fp8.py | 40 +++++++++---------- 1 file changed, 19 insertions(+), 21 deletions(-) diff --git a/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py index 60cf77f8..75ee6169 100644 --- a/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py +++ b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py @@ -126,31 +126,29 @@ class LazyWeightMap: def __init__(self, model_dir: str, weight_map: Dict[str, str]): self.model_dir = model_dir - self.weight_map = weight_map self._cur_filename: Optional[str] = None self._cur_handle = None - # MiMo-V2.5's published index.json still references the legacy - # `model_N-00001-of-00002.safetensors` names, but the actual shards - # on disk are `model_pp0_epN_shardM.safetensors`. Build an alias - # table: legacy -> real filename, for all shards present locally. - self._filename_alias: Dict[str, str] = {} - import re - actual_files = os.listdir(model_dir) - legacy_re = re.compile(r"model_(\d+)-0000([12])-of-00002\.safetensors") - for name in weight_map.values(): - if name in self._filename_alias or name in actual_files: - continue - m = legacy_re.match(name) - if not m: - continue - ep_idx, shard_one_based = m.group(1), m.group(2) - shard_zero_based = str(int(shard_one_based) - 1) - candidate = f"model_pp0_ep{ep_idx}_shard{shard_zero_based}.safetensors" - if candidate in actual_files: - self._filename_alias[name] = candidate + # V2.5's published model.safetensors.index.json references filenames + # like `model_N-00001-of-00002.safetensors`, but the shards on disk + # are `model_pp0_epN_shardM.safetensors` and the ep/N numbers don't + # line up. Rather than try to reverse-engineer the mapping, scan the + # on-disk shards and rebuild weight_map by reading each file's + # manifest directly. Falls back to the provided weight_map when the + # shard files match the names on disk (pre-V2.5 checkpoints). + actual_files = sorted(f for f in os.listdir(model_dir) if f.endswith(".safetensors")) + names_in_weight_map = set(weight_map.values()) + if actual_files and not (names_in_weight_map & set(actual_files)): + rebuilt: Dict[str, str] = {} + for fname in actual_files: + path = os.path.join(model_dir, fname) + with safe_open(path, framework="pt", device="cpu") as fp: + for k in fp.keys(): + rebuilt[k] = fname + self.weight_map = rebuilt + else: + self.weight_map = weight_map def _open(self, filename: str): - filename = self._filename_alias.get(filename, filename) if self._cur_filename == filename: return self._cur_handle if self._cur_handle is not None: From d3816012930ac28ffb53fcce2cee72f319080e99 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 07:30:32 +0800 Subject: [PATCH 03/27] preprocess: split MiMo-V2.5 fused qkv_proj back into q/k/v for NxDI MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MiMo-V2.5's HF checkpoint stores attention as a fused self_attn.qkv_proj.weight tensor (shape [Q_dim+K_dim+V_dim, hidden]) even though its safetensors.index.json advertises separate q_proj/k_proj/v_proj keys. The actual LFS objects on the Hub carry only the fused form; HF's modeling code slices on the fly. NxDI's MiMoV2Attention hard-codes separate q_proj/k_proj/v_proj ColumnParallelLinear modules and actively deletes any qkv_proj attribute inherited from the base class, so the preprocess must produce split tensors. Slice the fused tensor along the output dim into Q / K / V chunks using the config's per-head dims (swa vs full), then run each through the same per-row FP8 rescale used for non-fused checkpoints. Slices the blockwise (128×128) scale along the output dim the same way — all Q/K/V output dims on V2.5 are multiples of 128, so the block boundaries line up. Any trailing block rows beyond Q+K+V (HF pads full-attention layers' scale to 108 blocks but the weight only has 106 blocks of content) are dropped with the unused weight rows. Falls back to the pre-V2.5 split-qkv path when qkv_proj.weight is absent, so Flash checkpoints still preprocess correctly. Co-Authored-By: Claude Opus 4.7 --- .../preprocess_mimo_v2_5_fp8.py | 76 ++++++++++++++++--- 1 file changed, 65 insertions(+), 11 deletions(-) diff --git a/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py index 75ee6169..079f9234 100644 --- a/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py +++ b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py @@ -212,17 +212,71 @@ def process_layer( if t is not None: out[f"{out_prefix}{name}.weight"] = t.detach().clone() - # --- Attention: q/k/v/o are stored separately in MiMo-V2.5 --- - # q/k/v: rescale to Neuron FP8 per-row. - for proj in ("q_proj", "k_proj", "v_proj"): - w = lazy.get(f"{prefix}self_attn.{proj}.weight") - if w is None: - continue - s = lazy.get(f"{prefix}self_attn.{proj}.weight_scale_inv") - w2, s2 = _maybe_fp8_to_neuron_per_row(w, s) - out[f"{out_prefix}self_attn.{proj}.weight"] = w2 - if s2 is not None: - out[f"{out_prefix}self_attn.{proj}.scale"] = s2 + # --- Attention: q/k/v --- + # MiMo-V2.5 ships the QKV projection *fused* into a single + # self_attn.qkv_proj.weight (shape [Q_dim+K_dim+V_dim, hidden]) with + # blockwise-128 FP8 scaling. The NxDI modeling code hard-codes separate + # q_proj / k_proj / v_proj, so slice the fused tensor back into three + # per-proj tensors. All per-head dimensions here are multiples of 128, + # so the blockwise scale slices cleanly along the fused output dim. + qkv_w = lazy.get(f"{prefix}self_attn.qkv_proj.weight") + qkv_s = lazy.get(f"{prefix}self_attn.qkv_proj.weight_scale_inv") + if qkv_w is not None: + num_heads = config["swa_num_attention_heads" if is_swa else "num_attention_heads"] + num_kv_heads = config["swa_num_key_value_heads" if is_swa else "num_key_value_heads"] + qk_head_dim = config["swa_head_dim" if is_swa else "head_dim"] + v_head_dim = config["swa_v_head_dim" if is_swa else "v_head_dim"] + q_dim = num_heads * qk_head_dim + k_dim = num_kv_heads * qk_head_dim + v_dim = num_kv_heads * v_head_dim + assert qkv_w.shape[0] == q_dim + k_dim + v_dim, ( + f"Layer {layer_idx} fused qkv_proj out_dim {qkv_w.shape[0]} != " + f"{q_dim}+{k_dim}+{v_dim}={q_dim+k_dim+v_dim} " + f"(is_swa={is_swa}, heads={num_heads}, kv_heads={num_kv_heads}, " + f"qk_hd={qk_head_dim}, v_hd={v_head_dim})" + ) + for block_size in (128,): + for name, dim in (("q_proj", q_dim), ("k_proj", k_dim), ("v_proj", v_dim)): + assert dim % block_size == 0, ( + f"{name} dim {dim} is not a multiple of {block_size}; " + f"cannot slice fused qkv blockwise scale" + ) + q_w_slice = qkv_w[:q_dim] + k_w_slice = qkv_w[q_dim : q_dim + k_dim] + v_w_slice = qkv_w[q_dim + k_dim :] + if qkv_s is not None: + q_blocks = q_dim // 128 + k_blocks = k_dim // 128 + q_s_slice = qkv_s[:q_blocks] + k_s_slice = qkv_s[q_blocks : q_blocks + k_blocks] + v_s_slice = qkv_s[q_blocks + k_blocks :] + else: + q_s_slice = k_s_slice = v_s_slice = None + for name, w, s in ( + ("q_proj", q_w_slice, q_s_slice), + ("k_proj", k_w_slice, k_s_slice), + ("v_proj", v_w_slice, v_s_slice), + ): + # `s` may still reference the parent qkv scale tensor; clone it so + # rescale_fp8 sees a contiguous view and can call .view/.contiguous + # on a real tensor. + s = s.contiguous().clone() if s is not None else None + w = w.contiguous().clone() + w2, s2 = _maybe_fp8_to_neuron_per_row(w, s) + out[f"{out_prefix}self_attn.{name}.weight"] = w2 + if s2 is not None: + out[f"{out_prefix}self_attn.{name}.scale"] = s2 + else: + # Fallback path for checkpoints that ship split q/k/v (pre-V2.5). + for proj in ("q_proj", "k_proj", "v_proj"): + w = lazy.get(f"{prefix}self_attn.{proj}.weight") + if w is None: + continue + s = lazy.get(f"{prefix}self_attn.{proj}.weight_scale_inv") + w2, s2 = _maybe_fp8_to_neuron_per_row(w, s) + out[f"{out_prefix}self_attn.{proj}.weight"] = w2 + if s2 is not None: + out[f"{out_prefix}self_attn.{proj}.scale"] = s2 # o_proj is listed in HF quantization_config.ignored_layers and ships as # BF16; on Neuron it binds to a plain RowParallelLinear (see From eb6b2aad6f577281e6444c8122054a800a04e231 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 09:20:59 +0800 Subject: [PATCH 04/27] preprocess: fix V2.5 fused qkv split to use 4 interleaved groups MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit MiMo-V2.5's fused qkv_proj.weight is NOT a simple [all_Q | all_K | all_V] concatenation — the first naive slicing approach produced garbled outputs because Q/K/V rows are physically interleaved in num_groups=4 per-layer groups: group g (g = 0..3): rows [g*R : g*R + qg] = Q heads [g*hpg : (g+1)*hpg] rows [g*R + qg : g*R + qg + kg] = K heads [g*kpg : (g+1)*kpg] rows [g*R + qg + kg : g*R + R] = V heads [g*kpg : (g+1)*kpg] The group count (4) is a model-level constant equal to the full- attention num_key_value_heads. SWA layers with num_kv_heads=8 pack kpg=2 K/V heads per group, which is why their fused weight row count is 14848 (= 4 * (8*192 + 2*192 + 2*128)) rather than the ~27136 one would expect from an 8-group layout. Full-attention layers with num_kv_heads=4 pack kpg=1 K/V head per group, giving 13568 rows. Scale rows also follow the per-group layout with phantom padding: full attention's K has kg=192 rows but consumes 2 scale blocks (the last half of the second block is unused), giving 4*(24+2+1)=108 total scale rows against 106 real blocks. Implementation ported from MiMo-V2.5-Pro's split_qkv_fused with the num_kv_heads/num_groups axis decoupled so it works for V2.5's asymmetric (num_kv_heads=4 full / 8 swa) config. Verified empirically by Q/K/V scale-magnitude probes — Q/K/V bands have distinct scale distributions that match the claimed slice boundaries on both full and SWA layers. Falls back to the pre-V2.5 per-proj path when qkv_proj.weight is absent, so Flash/other split-qkv checkpoints still preprocess. Co-Authored-By: Claude Opus 4.7 --- .../preprocess_mimo_v2_5_fp8.py | 223 ++++++++++++++---- 1 file changed, 177 insertions(+), 46 deletions(-) diff --git a/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py index 079f9234..80838239 100644 --- a/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py +++ b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py @@ -117,6 +117,160 @@ def rescale_fp8_weight_blockwise( return rescaled, neuron_scale.to(torch.float32) +def _requantize_per_row(dequant: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]: + """BF16/FP32 -> Neuron FP8 per-row.""" + row_max_abs = dequant.abs().max(dim=1, keepdim=True)[0] + scales = row_max_abs / NEURON_FP8_MAX + scales = torch.clamp(scales, min=1e-10) + quantized = (dequant / scales).to(torch.float8_e4m3fn) + return quantized, scales.to(torch.float32) + + +def split_qkv_fused( + qkv_weight: torch.Tensor, + qkv_scale: Optional[torch.Tensor], + num_q_heads: int, + num_kv_heads_this_layer: int, + num_groups: int, + head_dim: int, + v_head_dim: int, +) -> Dict[str, Tuple[torch.Tensor, Optional[torch.Tensor]]]: + """Split V2.5's fused qkv_proj into q/k/v. + + Layout — validated empirically via per-group Q/K/V magnitude probes: + + The weight is NOT `[all_Q | all_K | all_V]`. It is ``num_groups`` + interleaved groups, each packing ``hpg`` Q heads, ``kpg`` K heads, and + ``kpg`` V heads contiguously: + + group g (g = 0 .. num_groups-1): + rows [g*R : g*R + qg] = Q heads [g*hpg : (g+1)*hpg] + rows [g*R + qg : g*R + qg + kg] = K heads [g*kpg : (g+1)*kpg] + rows [g*R + qg + kg : g*R + R] = V heads [g*kpg : (g+1)*kpg] + where + hpg = num_q_heads / num_groups + kpg = num_kv_heads_this_layer / num_groups + qg = hpg * head_dim + kg = kpg * head_dim + vg = kpg * v_head_dim + R = qg + kg + vg + + ``num_groups`` is a model-level constant (= the full-attention + ``num_key_value_heads``, 4 for V2.5). It is the same for full and SWA + layers, so a SWA layer with ``num_kv_heads_this_layer=8`` packs + ``kpg=2`` K heads + 2 V heads per group. + + Scale layout: each group holds ``q_blk + k_blk + v_blk`` scale rows + where ``q_blk = qg // 128`` (exact; qg is always 128-aligned), + ``k_blk = ceil(kg / 128)`` (may add 64 rows of "phantom" padding on + V2.5 full layers where kg = 192 and the last half-block is unused), + ``v_blk = ceil(vg / 128)``. When a phantom half-block appears, the + physical weight rows stop before the phantom rows do — we recover the + correct dequant by padding each group's weight out to + ``per_group_scale * 128`` rows, broadcasting the scale, then stripping + the phantom rows. + """ + in_features = qkv_weight.shape[1] + assert num_q_heads % num_groups == 0 and num_kv_heads_this_layer % num_groups == 0, ( + f"num_q_heads={num_q_heads} and num_kv_heads_this_layer=" + f"{num_kv_heads_this_layer} must both be divisible by " + f"num_groups={num_groups}" + ) + hpg = num_q_heads // num_groups + kpg = num_kv_heads_this_layer // num_groups + qg_rows = hpg * head_dim + kg_rows = kpg * head_dim + vg_rows = kpg * v_head_dim + real_rows_per_group = qg_rows + kg_rows + vg_rows + total_real_rows = num_groups * real_rows_per_group + + BLOCK = 128 + q_scale_rows_per_group = qg_rows // BLOCK # exact + k_scale_rows_per_group = (kg_rows + BLOCK - 1) // BLOCK + v_scale_rows_per_group = (vg_rows + BLOCK - 1) // BLOCK + scale_rows_per_group = ( + q_scale_rows_per_group + k_scale_rows_per_group + v_scale_rows_per_group + ) + padded_rows_per_group = scale_rows_per_group * BLOCK + + assert qkv_weight.shape[0] == total_real_rows, ( + f"qkv_proj.weight row count {qkv_weight.shape[0]} != expected " + f"{total_real_rows} (num_groups={num_groups}, hpg={hpg}, kpg={kpg}, " + f"R={real_rows_per_group})" + ) + + if qkv_weight.dtype != torch.float8_e4m3fn or qkv_scale is None: + # BF16 path: no scale to worry about. + w = qkv_weight.view(num_groups, real_rows_per_group, in_features) + q_w = ( + w[:, :qg_rows, :] + .reshape(num_groups * qg_rows, in_features) + .contiguous() + ) + k_w = ( + w[:, qg_rows : qg_rows + kg_rows, :] + .reshape(num_groups * kg_rows, in_features) + .contiguous() + ) + v_w = ( + w[:, qg_rows + kg_rows :, :] + .reshape(num_groups * vg_rows, in_features) + .contiguous() + ) + q_w2, q_s2 = convert_bf16_to_fp8_per_row(q_w) + k_w2, k_s2 = convert_bf16_to_fp8_per_row(k_w) + v_w2, v_s2 = convert_bf16_to_fp8_per_row(v_w) + return {"q_proj": (q_w2, q_s2), "k_proj": (k_w2, k_s2), "v_proj": (v_w2, v_s2)} + + # FP8 + blockwise scale path. + expected_scale_rows = num_groups * scale_rows_per_group + expected_scale_cols = (in_features + BLOCK - 1) // BLOCK + assert qkv_scale.shape == (expected_scale_rows, expected_scale_cols), ( + f"qkv scale shape {tuple(qkv_scale.shape)} != expected " + f"({expected_scale_rows}, {expected_scale_cols}) for " + f"num_groups={num_groups}, per_group={scale_rows_per_group}" + ) + + w = qkv_weight.to(torch.float32).view( + num_groups, real_rows_per_group, in_features + ) + w_padded = torch.zeros( + num_groups, padded_rows_per_group, in_features, dtype=torch.float32 + ) + w_padded[:, :real_rows_per_group, :] = w + + s = qkv_scale.to(torch.float32).view( + num_groups, scale_rows_per_group, expected_scale_cols + ) + s_exp = s.repeat_interleave(BLOCK, dim=1).repeat_interleave(BLOCK, dim=2) + s_exp = s_exp[:, :padded_rows_per_group, :in_features] + + deq_padded = w_padded * s_exp + deq = deq_padded[:, :real_rows_per_group, :] + + q_deq = ( + deq[:, :qg_rows, :] + .reshape(num_groups * qg_rows, in_features) + .contiguous() + ) + k_deq = ( + deq[:, qg_rows : qg_rows + kg_rows, :] + .reshape(num_groups * kg_rows, in_features) + .contiguous() + ) + v_deq = ( + deq[:, qg_rows + kg_rows :, :] + .reshape(num_groups * vg_rows, in_features) + .contiguous() + ) + + q_w2, q_s2 = _requantize_per_row(q_deq) + k_w2, k_s2 = _requantize_per_row(k_deq) + v_w2, v_s2 = _requantize_per_row(v_deq) + + return {"q_proj": (q_w2, q_s2), "k_proj": (k_w2, k_s2), "v_proj": (v_w2, v_s2)} + + # --------------------------------------------------------------------------- # Streaming weight access (one open safetensors handle at a time) # --------------------------------------------------------------------------- @@ -213,59 +367,36 @@ def process_layer( out[f"{out_prefix}{name}.weight"] = t.detach().clone() # --- Attention: q/k/v --- - # MiMo-V2.5 ships the QKV projection *fused* into a single - # self_attn.qkv_proj.weight (shape [Q_dim+K_dim+V_dim, hidden]) with - # blockwise-128 FP8 scaling. The NxDI modeling code hard-codes separate - # q_proj / k_proj / v_proj, so slice the fused tensor back into three - # per-proj tensors. All per-head dimensions here are multiples of 128, - # so the blockwise scale slices cleanly along the fused output dim. + # MiMo-V2.5 ships QKV *fused* into a single self_attn.qkv_proj.weight, + # with an interleaved-group layout (see split_qkv_fused for details). + # The NxDI modeling code expects separate q_proj / k_proj / v_proj, so + # split the fused tensor back out. Falls back to per-proj tensors if + # the checkpoint is already split (pre-V2.5). qkv_w = lazy.get(f"{prefix}self_attn.qkv_proj.weight") qkv_s = lazy.get(f"{prefix}self_attn.qkv_proj.weight_scale_inv") if qkv_w is not None: num_heads = config["swa_num_attention_heads" if is_swa else "num_attention_heads"] - num_kv_heads = config["swa_num_key_value_heads" if is_swa else "num_key_value_heads"] + num_kv_heads_this = config[ + "swa_num_key_value_heads" if is_swa else "num_key_value_heads" + ] qk_head_dim = config["swa_head_dim" if is_swa else "head_dim"] - v_head_dim = config["swa_v_head_dim" if is_swa else "v_head_dim"] - q_dim = num_heads * qk_head_dim - k_dim = num_kv_heads * qk_head_dim - v_dim = num_kv_heads * v_head_dim - assert qkv_w.shape[0] == q_dim + k_dim + v_dim, ( - f"Layer {layer_idx} fused qkv_proj out_dim {qkv_w.shape[0]} != " - f"{q_dim}+{k_dim}+{v_dim}={q_dim+k_dim+v_dim} " - f"(is_swa={is_swa}, heads={num_heads}, kv_heads={num_kv_heads}, " - f"qk_hd={qk_head_dim}, v_hd={v_head_dim})" + v_hd = config["swa_v_head_dim" if is_swa else "v_head_dim"] + # num_groups is a model-level constant = full-attention num_kv_heads. + # SWA layers with num_kv_heads=8 still use 4 groups (2 K heads per group). + num_groups = config["num_key_value_heads"] + split = split_qkv_fused( + qkv_w, + qkv_s, + num_q_heads=num_heads, + num_kv_heads_this_layer=num_kv_heads_this, + num_groups=num_groups, + head_dim=qk_head_dim, + v_head_dim=v_hd, ) - for block_size in (128,): - for name, dim in (("q_proj", q_dim), ("k_proj", k_dim), ("v_proj", v_dim)): - assert dim % block_size == 0, ( - f"{name} dim {dim} is not a multiple of {block_size}; " - f"cannot slice fused qkv blockwise scale" - ) - q_w_slice = qkv_w[:q_dim] - k_w_slice = qkv_w[q_dim : q_dim + k_dim] - v_w_slice = qkv_w[q_dim + k_dim :] - if qkv_s is not None: - q_blocks = q_dim // 128 - k_blocks = k_dim // 128 - q_s_slice = qkv_s[:q_blocks] - k_s_slice = qkv_s[q_blocks : q_blocks + k_blocks] - v_s_slice = qkv_s[q_blocks + k_blocks :] - else: - q_s_slice = k_s_slice = v_s_slice = None - for name, w, s in ( - ("q_proj", q_w_slice, q_s_slice), - ("k_proj", k_w_slice, k_s_slice), - ("v_proj", v_w_slice, v_s_slice), - ): - # `s` may still reference the parent qkv scale tensor; clone it so - # rescale_fp8 sees a contiguous view and can call .view/.contiguous - # on a real tensor. - s = s.contiguous().clone() if s is not None else None - w = w.contiguous().clone() - w2, s2 = _maybe_fp8_to_neuron_per_row(w, s) - out[f"{out_prefix}self_attn.{name}.weight"] = w2 + for proj, (w2, s2) in split.items(): + out[f"{out_prefix}self_attn.{proj}.weight"] = w2 if s2 is not None: - out[f"{out_prefix}self_attn.{name}.scale"] = s2 + out[f"{out_prefix}self_attn.{proj}.scale"] = s2 else: # Fallback path for checkpoints that ship split q/k/v (pre-V2.5). for proj in ("q_proj", "k_proj", "v_proj"): From 94d001d451ce9251d786b37775e7876ca870c6d7 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 11:35:16 +0800 Subject: [PATCH 05/27] README/smoke/bench: V2.5-specific content + move scratch off /tmp MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit README rewritten to match what actually works today: - Fixed stale Flash-era facts (vocab 151,936 -> 152,576, param counts removed, wrong "fused_qkv not supported" claim updated). - New "Fused QKV on disk, split on Neuron" key feature note and a dedicated "V2.5-specific: fused qkv_proj split into 4 interleaved groups" subsection explaining the per-group layout and scale phantom-row handling that the preprocess implements. - "weight_map rebuild" note: V2.5's index.json references legacy shard filenames that don't exist on disk, and the preprocess scans actual files instead. - Dropped the "FP8 -> BF16 fallback" doc paragraph — that script never existed on this branch. - Mount instructions in Prerequisites: the DLAMI creates /dev/md0 as a 6.9 TB RAID0 but does not add it to /etc/fstab, so after a reboot /opt/dlami/nvme is empty until remounted. Document the `sudo mount /dev/md0 /opt/dlami/nvme` fix. - Updated timing numbers (preprocess 16 min / 15 GB peak RAM, first compile 30 min dominated by 27 min shard_checkpoint). - Dropped the stale BF16 benchmark numbers; FP8 numbers pending. Scratch locations off /tmp: - smoke_compile / smoke_generate: default BASE_COMPILE_WORK_DIR from /tmp/nxd_model/ to /opt/dlami/nvme/tmp/nxd_model/, so HLO/NEFF staging survives the nightly Trn2 reboot. - bench_mimo_v2_5.sh, run_bench_single.sh: RESULTS_DIR default from /tmp/bench_results/mimo_v2_5 to /opt/dlami/nvme/logs/bench_results/ mimo_v2_5. Why: the Trn2 instance reboots daily around 00:07 UTC and /tmp is wiped on reboot. A long-running compile that straddles the reboot loses all its intermediate files under /tmp. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/README.md | 118 +++++++++--------- .../MiMo-V2.5/perf_test/bench_mimo_v2_5.sh | 2 +- .../MiMo-V2.5/perf_test/run_bench_single.sh | 4 +- .../perf_test/smoke_compile_mimo_v2_5.py | 7 +- .../perf_test/smoke_generate_mimo_v2_5.py | 6 +- 5 files changed, 73 insertions(+), 64 deletions(-) diff --git a/contrib/models/MiMo-V2.5/README.md b/contrib/models/MiMo-V2.5/README.md index 7011ef15..3c77d200 100644 --- a/contrib/models/MiMo-V2.5/README.md +++ b/contrib/models/MiMo-V2.5/README.md @@ -1,12 +1,11 @@ # Contrib Model: MiMo-V2.5 -NeuronX Distributed Inference implementation of [XiaomiMiMo/MiMo-V2.5](https://huggingface.co/XiaomiMiMo/MiMo-V2.5). +NeuronX Distributed Inference implementation of [XiaomiMiMo/MiMo-V2.5](https://huggingface.co/XiaomiMiMo/MiMo-V2.5). MiMo-V2.5 supersedes the earlier MiMo-V2-Flash release with the same decoder-only MoE architecture, an updated tokenizer, and a multimodal (vision + audio) head that the NxDI language path does not use. ## Model Information - **HuggingFace ID:** `XiaomiMiMo/MiMo-V2.5` -- **Model Type:** Decoder-only MoE transformer with hybrid attention -- **Architecture:** Custom MoE with full + sliding window attention +- **Model Type:** Decoder-only MoE transformer with hybrid (full + SWA) attention - **License:** Check HuggingFace model card ## Architecture Details @@ -14,37 +13,44 @@ NeuronX Distributed Inference implementation of [XiaomiMiMo/MiMo-V2.5](https://h | Parameter | Value | |-----------|-------| | Hidden Size | 4096 | -| Layers | 48 | -| Attention Heads | 64 Q | +| Layers | 48 (layer 0 dense, layers 1–47 MoE) | +| Q Heads | 64 | | KV Heads (full attn) | 4 | | KV Heads (sliding window) | 8 | | Q/K Head Dim | 192 | | V Head Dim | 128 | | Experts | 256 (top-8 routing) | | Expert Intermediate | 2048 | -| Vocab Size | 151,936 | -| RoPE | Partial (34% of dims), theta=5M (full), 10K (SWA) | +| Vocab Size | 152,576 | +| RoPE | Partial (64 of 192 head dims = 33.4%), theta=5M (full) / 10K (SWA) | | Sliding Window | 128 | | Max Position | 262,144 | -| Total Params | ~143B (FP8 native) / ~286B (BF16 upcast) | Key features: -- **Hybrid Attention**: 9 full attention layers (0, 5, 11, 17, 23, 29, 35, 41, 47) + 39 sliding window layers -- **Asymmetric Head Dims**: Q/K use 192, V uses 128 (fused_qkv not supported) -- **Attention Sink Bias**: Learnable per-head bias on sliding window layers only -- **Sigmoid Router + noaux_tc**: e_score_correction_bias added to sigmoid scores before top-k selection; unbiased scores become affinity weights -- **attention_value_scale = 0.707**: HF MiMo-V2 multiplies `value_states` by this before the attention softmax × V (NOT applied to attn_output); the NxDI model matches +- **Hybrid Attention**: 9 full attention layers (0, 5, 11, 17, 23, 29, 35, 41, 47) + 39 sliding window layers (positions driven by `hybrid_layer_pattern`). +- **Asymmetric Head Dims**: Q/K use 192, V uses 128. Plus asymmetric `num_kv_heads` between full (4) and SWA (8) layers. +- **Fused QKV on disk, split on Neuron**: the HF checkpoint ships `qkv_proj.weight` fused (`attention_projection_layout="fused_qkv"`); the NxDI modeling code keeps separate `q_proj`/`k_proj`/`v_proj` linears, so the preprocess script slices the fused tensor back into three per-proj tensors (see "Checkpoint Preparation"). +- **Attention Sink Bias**: Learnable per-head bias on sliding window layers only (`add_swa_attention_sink_bias=true`, `add_full_attention_sink_bias=false`). +- **Sigmoid Router + noaux_tc**: `e_score_correction_bias` added to sigmoid scores before top-k selection; unbiased scores become the affinity weights. +- **attention_value_scale = 0.707**: HF MiMo-V2 multiplies `value_states` by this before the attention softmax × V (NOT applied to attn_output); the NxDI model matches. ## Prerequisites - **Instance**: trn2.48xlarge (32 NeuronCores, logical_nc_config=2 → 64 logical cores) - **Neuron SDK**: 2.29 (Python 3.12, PyTorch 2.9) -- **Venvs**: `/opt/aws_neuronx_venv_pytorch_2_9_nxd_inference` (for preprocess + NxDI direct smoke), `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (for vLLM serving). Both ship with the DLAMI. -- **Disk**: ~700 GB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~290 GB, the Neuron-FP8 preprocessed output is ~310 GB, and `save_sharded_checkpoint=true` writes another ~300 GB per compiled config). +- **Venv**: `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (ships with the DLAMI; has NxDI, vllm-neuron, and `huggingface_hub`/`s5cmd`). +- **Disk**: ~900 GB free under `/opt/dlami/nvme` (HF FP8 checkpoint ~295 GB, Neuron-FP8 preprocessed output ~310 GB, and `save_sharded_checkpoint=true` writes another ~300 GB of per-rank sharded weights per compiled config). The DLAMI creates a 6.9 TB RAID0 at `/dev/md0` across the instance-store NVMes but does **not** add it to `/etc/fstab`, so it is not mounted automatically after a reboot. Before running any of the steps below, remount it if needed: + + ```bash + # If /opt/dlami/nvme appears empty after an overnight reboot, the md0 array + # is still intact and just needs to be remounted: + mount | grep -q /opt/dlami/nvme || sudo mount /dev/md0 /opt/dlami/nvme + df -h /opt/dlami/nvme # should show ~6.9 TB + ``` ## Quick Start (FP8 on Trn2) -End-to-end recipe to go from a fresh trn2.48xlarge to a working vLLM OpenAI server serving MiMo-V2.5 FP8. First-time compile takes ~45-60 minutes; subsequent runs hit the neuronx-cc cache and start in a few minutes. +End-to-end recipe to go from a fresh trn2.48xlarge to a working vLLM OpenAI server serving MiMo-V2.5 FP8. First-time compile takes ~30 minutes; subsequent runs hit the neuronx-cc cache and start in a few minutes. ```bash # 1. Clone this repo on the Trn2 instance @@ -53,34 +59,33 @@ git clone /neuronx-distributed-inference.git cd neuronx-distributed-inference git checkout contrib/MiMo-V2.5 # the branch this README lives on -# 2. Download the HuggingFace FP8 checkpoint (~290 GB). Any HF-compatible -# downloader works; huggingface-cli example: +# 2. Download the HuggingFace FP8 checkpoint (~295 GB). +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate huggingface-cli download XiaomiMiMo/MiMo-V2.5 \ - --local-dir /opt/dlami/nvme/models/MiMo-V2.5 + --local-dir /opt/dlami/nvme/models/MiMo-V2.5 --max-workers 16 -# 3. Preprocess HF FP8 -> Neuron FP8 (~20 min, ~24 GB peak RAM) -source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate +# 3. Preprocess HF FP8 -> Neuron FP8 (~16 min, ~15 GB peak RAM) python contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py \ --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5 \ --save_path /opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8 \ --tp_degree 64 # 4. (Optional) sanity-check the Neuron-FP8 checkpoint without vLLM -# ~45 min first compile; subsequent runs ~30s to load the pre-sharded NEFF. -source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate -python contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py # compile +# ~30 min first compile (priority HLO + CE HLO + 27 min shard_checkpoint +# for 64 ranks); subsequent runs ~30s to load the pre-sharded NEFF. +python contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py # compile + shard python contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py # 20-token generate # 5. Install vllm-neuron with the contrib registration patch bash contrib/models/MiMo-V2.5/perf_test/0_setup.sh -# 6. Start vLLM serving MiMo-V2.5 FP8 (first compile ~60 min; subsequent ~3 min) +# 6. Start vLLM serving MiMo-V2.5 FP8 bash contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh ``` The bench script runs two configurations (BS=32 and BS=128, both `moe_tp_degree=1 / moe_ep_degree=64`) and logs results under -`/tmp/bench_results/mimo_v2_5/`. +`/opt/dlami/nvme/logs/bench_results/mimo_v2_5/`. For a quick `curl` sanity check while the server is up: @@ -100,11 +105,7 @@ preprocessed Neuron-FP8 checkpoint (not the raw HF FP8 directory). ## Checkpoint Preparation -The HuggingFace checkpoint ships as block-wise OCP FP8 (E4M3, ±448 range), which is not directly compatible with Neuron FP8 (IEEE-754 E4M3, ±240 range). Two preprocess scripts are provided: - -### Recommended: FP8 → Neuron-FP8 (streaming) - -`src/conversion_script/preprocess_mimo_v2_5_fp8.py` performs a per-layer streaming rescale from OCP FP8 to Neuron FP8 (per-row scales for attention Q/K/V and layer-0 dense MLP; blockwise scales for MoE experts). `o_proj` is listed in HF's `quantization_config.ignored_layers` and is kept BF16 on the Neuron side (it binds to a plain `RowParallelLinear`, not `QuantizedRowParallel`). Output is ~310 GB across 48 per-layer safetensors shards. +The HuggingFace checkpoint ships as block-wise OCP FP8 (E4M3, ±448 range), which is not directly compatible with Neuron FP8 (IEEE-754 E4M3, ±240 range). `src/conversion_script/preprocess_mimo_v2_5_fp8.py` performs a per-layer streaming rescale: per-row scales for attention Q/K/V (after fused-qkv split) and the layer-0 dense MLP; blockwise 128×128 scales for MoE experts. `o_proj` is listed in HF's `quantization_config.ignored_layers` and is kept BF16 on the Neuron side (it binds to a plain `RowParallelLinear`, not `QuantizedRowParallel`). Output is ~310 GB across 48 per-layer safetensors shards. ```bash python contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py \ @@ -113,11 +114,26 @@ python contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.p --tp_degree 64 ``` -Peak RAM during preprocessing is ~24 GB; total runtime ~20 minutes on a trn2.48xlarge instance. +Peak RAM during preprocessing is ~15 GB; total runtime ~16 minutes on a trn2.48xlarge instance. + +### V2.5-specific: fused qkv_proj split into 4 interleaved groups + +The HF checkpoint advertises `q_proj.weight` / `k_proj.weight` / `v_proj.weight` in its safetensors index, but the actual LFS objects on the Hub only carry a single fused `self_attn.qkv_proj.weight` tensor. NxDI's MiMoV2Attention hard-codes separate Q/K/V `ColumnParallelLinear` modules, so the preprocess script splits the fused tensor back into three per-proj tensors. + +The fused layout is **not** `[all_Q | all_K | all_V]`. It is **4 interleaved groups** (the group count equals the full-attention `num_key_value_heads = 4`), each packing `hpg` Q heads, `kpg` K heads, and `kpg` V heads contiguously: -### Fallback: FP8 → BF16 + group g (g = 0..3): + rows [g*R : g*R + qg] = Q heads [g*hpg : (g+1)*hpg] + rows [g*R + qg : g*R + qg + kg] = K heads [g*kpg : (g+1)*kpg] + rows [g*R + qg + kg : g*R + R] = V heads [g*kpg : (g+1)*kpg] -`src/conversion_script/preprocess_mimo_v2_fp8.py` dequantizes the entire checkpoint to BF16. Output is ~290 GB; BF16 is numerically equivalent to the published HF FP8 weights and is useful as a known-good reference. Throughput is ~2× worse than the FP8 path because every attention/MLP matmul operates on full BF16 weights. + where hpg = num_q_heads / 4, kpg = num_kv_heads / 4, + qg = hpg * 192, kg = kpg * 192, vg = kpg * 128, + R = qg + kg + vg + +For **full-attention layers** this gives `hpg=16, kpg=1, R=3392, total=13568` rows with 108 scale blocks (includes 2 phantom rows from `ceil(192/128)=2`). For **SWA layers** (`num_kv_heads=8`), `hpg=16, kpg=2, R=3712, total=14848` rows with 116 scale blocks (no phantom, since `kg=384` is 128-aligned). Layer 0 (dense) is still attention-FP8 and follows the full-layer layout. + +Any preprocess approach that treats the fused tensor as a plain `[Q|K|V]` concatenation produces garbled outputs — Q/K/V rows land in the wrong per-head slots after the split. ## Usage @@ -231,12 +247,12 @@ MiMo-V2.5 can be served via [vllm-neuron](https://github.com/aws-neuron/vllm-neu ```bash # The setup script clones vllm-project/vllm-neuron at release-0.5.0, applies -# the contrib registration patch, installs it editable, and downloads MiMo-V2.5 -# weights (BF16 by default; set MIMO_V2_5_PATH to override). +# the contrib registration patch, installs it editable, and downloads +# MiMo-V2.5 FP8 weights from HuggingFace (~295 GB; skipped if already present). bash contrib/models/MiMo-V2.5/perf_test/0_setup.sh ``` -The patch (`perf_test/vllm-neuron-patch.patch`) is 40 lines and only touches `vllm_neuron/__init__.py`. It adds a `_register_contrib_models()` hook that, when `NXDI_CONTRIB_MIMO_V2_5_SRC` is set, registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under the key `mimo_v2_5` **and** registers the `MiMoV2ForCausalLM` architecture into vLLM's `ModelRegistry`. No upstream vLLM or NxDI source is modified. +`perf_test/vllm-neuron-patch.patch` adds a `_register_contrib_models()` hook to `vllm_neuron/worker/neuronx_distributed_model_loader.py`. When `NXDI_CONTRIB_MIMO_V2_5_SRC` is set, it registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under the key `mimov2` **and** registers the `MiMoV2ForCausalLM` architecture into vLLM's `ModelRegistry`. The hook also patches `AutoConfig.from_pretrained` to default `trust_remote_code=True` so NxDI's `load_pretrained_config` can read the V2.5 config. No upstream vLLM or NxDI source is modified. ### Serving (FP8, recommended) @@ -306,27 +322,9 @@ The patch is applied to vllm-neuron 0.5.0 and: ## Performance -> These numbers are from the earlier BF16 recipe (pre-FP8 rollout). FP8 numbers will be added once a stable bench run completes on the new recipe; preliminary single-stream qualitative tests show fluent multi-sentence output on long Chinese chat prompts with `moe_tp=1, moe_ep=64, batch_size=32`. - -### Standalone NxDI (trn2.48xlarge, BF16, TP=64, EP=64) - -| Batch Size | Throughput (tok/s) | -|------------|-------------------| -| 1 | 29.92 | -| 8 | 215.94 | -| 32 | 649.14 | - -### vLLM Serving (trn2.48xlarge, BF16, BS=32, TP=64/EP=64, CB) - -Input/output: 900/90 tokens (random dataset) - -| 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 | +> Benchmark numbers will be added once a stable bench run completes on the FP8 recipe. Preliminary single-stream sanity test produces fluent MiMo self-introduction output on the recipe below (`moe_tp=1, moe_ep=64, batch_size=32`). -> **Compile time:** the first MiMo-V2.5 compile on SDK 2.29 is ~30-60 minutes for the TKG NEFF and similar for the CTE NEFF. Subsequent runs with the same `override_neuron_config` hit the neuronx-cc cache and start in ~1-2 minutes. `save_sharded_checkpoint=true` additionally persists per-rank FP8 shards under `/weights/`, letting future `load()` calls skip the ~10-minute shard_checkpoint pass. +> **Compile time:** the first MiMo-V2.5 compile on SDK 2.29 is ~30 minutes (TKG + CE HLO compilation, weight layout optimization, then `shard_checkpoint` for 64 ranks which dominates at ~27 minutes). Subsequent runs with the same `override_neuron_config` hit the neuronx-cc cache and the NEFF loads in ~1 minute. `save_sharded_checkpoint=true` persists per-rank FP8 shards under `/weights/`, letting future `load()` calls skip the `shard_checkpoint` pass entirely. ## Compatibility Matrix @@ -347,7 +345,9 @@ pytest contrib/models/MiMo-V2.5/test/integration/test_model.py -v 1. **Hybrid Attention**: `hybrid_layer_pattern` list determines full vs sliding window per layer; the modeling code constructs one `NeuronMiMoV2Attention` per layer with the correct `is_sliding_window` flag and rope_theta. 2. **CONVERT_TO_MHA**: When `tp_degree > num_kv_heads` (64 > 4 full / 64 > 8 SWA), K/V are replicated to `num_attention_heads` (64) during state-dict conversion; this applies to both `.weight` and the per-row `.scale` on the FP8 path. 3. **Attention Sink Bias**: Learnable per-head bias added as an extra "sink" column to attention scores in sliding window layers (not added in full-attention layers). Per-rank slicing of the bias happens inside `forward()` based on `parallel_state.get_tensor_model_parallel_rank()`. -4. **FP8 Path Caveats**: +4. **Fused qkv split in preprocess**: V2.5's HF checkpoint stores `self_attn.qkv_proj.weight` as 4 interleaved Q/K/V groups (see "Checkpoint Preparation" above). The preprocess script must slice these groups — naïve `[Q|K|V]` concat slicing produces garbage outputs. +5. **weight_map rebuild**: V2.5's `model.safetensors.index.json` references legacy `model_N-00001-of-00002.safetensors` filenames that do not match the actual `model_pp0_epN_shardM.safetensors` objects on disk. `LazyWeightMap` scans the on-disk shards at startup and rebuilds `weight_map` directly from each file's manifest; the inconsistent index is ignored. +6. **FP8 Path Caveats**: - Must use `moe_tp_degree=1, moe_ep_degree=64` (see "FP8 Configuration Notes" above). - Must use `batch_size >= 32` (NxDI EP>1 requirement). - Must keep outer `ep_degree=1` (only `moe_ep_degree` should vary). @@ -361,4 +361,4 @@ pytest contrib/models/MiMo-V2.5/test/integration/test_model.py -v Henan Wan (whn09) -**Last Updated:** 2026-04-25 +**Last Updated:** 2026-04-28 diff --git a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh index 45089d62..6606cb5d 100755 --- a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh +++ b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh @@ -34,7 +34,7 @@ export NXDI_CONTRIB_MIMO_V2_5_SRC export VLLM_ENGINE_READY_TIMEOUT_S=7200 PORT=8000 -RESULTS_DIR="/tmp/bench_results/mimo_v2_5" +RESULTS_DIR="/opt/dlami/nvme/logs/bench_results/mimo_v2_5" mkdir -p "$RESULTS_DIR" # Common neuron config shared across all MiMo-V2.5 FP8 configs. diff --git a/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh index a3c5ecb6..54455731 100755 --- a/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh @@ -21,7 +21,7 @@ # OUTPUT_LEN --random-output-len (default 90) # RANGE_RATIO --random-range-ratio (default 0.03) # CONFIG_NAME Used in the output filename (default bs1_tp64_ep1) -# RESULTS_DIR Where to dump per-run log (default /tmp/bench_results/mimo_v2_5) +# RESULTS_DIR Where to dump per-run log (default /opt/dlami/nvme/logs/bench_results/mimo_v2_5) set -e @@ -35,7 +35,7 @@ INPUT_LEN="${INPUT_LEN:-900}" OUTPUT_LEN="${OUTPUT_LEN:-90}" RANGE_RATIO="${RANGE_RATIO:-0.03}" CONFIG_NAME="${CONFIG_NAME:-bs1_tp64_ep1}" -RESULTS_DIR="${RESULTS_DIR:-/tmp/bench_results/mimo_v2_5}" +RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/logs/bench_results/mimo_v2_5}" mkdir -p "$RESULTS_DIR" diff --git a/contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py b/contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py index 7113e611..136dba6a 100755 --- a/contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py +++ b/contrib/models/MiMo-V2.5/perf_test/smoke_compile_mimo_v2_5.py @@ -59,9 +59,14 @@ # .hlo_module.pb files and one or both compilations crash with # "neuronx-cc returned non-zero exit status 70". Pin the workdir to a # unique per-COMPILED_PATH subdir to stay safe under any parallel invocation. +# Default under /opt/dlami/nvme rather than /tmp so the HLO/NEFF artifacts +# survive the nightly Trn2 reboot. os.environ.setdefault( "BASE_COMPILE_WORK_DIR", - os.path.join("/tmp/nxd_model", os.path.basename(COMPILED_PATH.rstrip("/"))), + os.path.join( + "/opt/dlami/nvme/tmp/nxd_model", + os.path.basename(COMPILED_PATH.rstrip("/")), + ), ) diff --git a/contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py b/contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py index 65dfa61e..45f4bce0 100755 --- a/contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py +++ b/contrib/models/MiMo-V2.5/perf_test/smoke_generate_mimo_v2_5.py @@ -42,9 +42,13 @@ # Keep the per-compile BASE_COMPILE_WORK_DIR in sync with # smoke_compile_mimo_v2_5.py so load() under the same COMPILED_PATH # doesn't collide with a concurrent compile or reuse a stale workdir. +# Default under /opt/dlami/nvme so artifacts survive the nightly Trn2 reboot. os.environ.setdefault( "BASE_COMPILE_WORK_DIR", - os.path.join("/tmp/nxd_model", os.path.basename(COMPILED_PATH.rstrip("/"))), + os.path.join( + "/opt/dlami/nvme/tmp/nxd_model", + os.path.basename(COMPILED_PATH.rstrip("/")), + ), ) From 1516d7a110aaea8f1ddecbf70179d0578749cece Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 11:39:41 +0800 Subject: [PATCH 06/27] README: add 'Testing the vLLM server' subsection + fix NC count - Step-by-step curl probes for both short sanity check and a longer generation, calling out what the outputs should look like when FP8 is working and what collapse symptoms to watch for. - Note that request-level temperature is ignored because on_device_sampling_config is baked into the NEFF at compile time. - Fix Prerequisites: trn2.48xlarge has 128 physical NeuronCores (not 32); with logical_nc_config=2 they appear as 64 logical cores. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/README.md | 28 +++++++++++++++++++++++++++- 1 file changed, 27 insertions(+), 1 deletion(-) diff --git a/contrib/models/MiMo-V2.5/README.md b/contrib/models/MiMo-V2.5/README.md index 3c77d200..6b37293b 100644 --- a/contrib/models/MiMo-V2.5/README.md +++ b/contrib/models/MiMo-V2.5/README.md @@ -36,7 +36,7 @@ Key features: ## Prerequisites -- **Instance**: trn2.48xlarge (32 NeuronCores, logical_nc_config=2 → 64 logical cores) +- **Instance**: trn2.48xlarge (128 NeuronCores, logical_nc_config=2 → 64 logical cores) - **Neuron SDK**: 2.29 (Python 3.12, PyTorch 2.9) - **Venv**: `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (ships with the DLAMI; has NxDI, vllm-neuron, and `huggingface_hub`/`s5cmd`). - **Disk**: ~900 GB free under `/opt/dlami/nvme` (HF FP8 checkpoint ~295 GB, Neuron-FP8 preprocessed output ~310 GB, and `save_sharded_checkpoint=true` writes another ~300 GB of per-rank sharded weights per compiled config). The DLAMI creates a 6.9 TB RAID0 at `/dev/md0` across the instance-store NVMes but does **not** add it to `/etc/fstab`, so it is not mounted automatically after a reboot. Before running any of the steps below, remount it if needed: @@ -312,6 +312,32 @@ python3 -m vllm.entrypoints.openai.api_server \ See `perf_test/bench_mimo_v2_5.sh` for the full benchmark recipe at BS=32 and BS=128. +### Testing the vLLM server + +Once `/v1/models` returns 200 (first-compile takes ~30 min; subsequent starts ~3 min), hit `/v1/chat/completions`. MiMo-V2.5's chat template expects the `<|im_start|>...<|im_end|>` ChatML format — vLLM applies it automatically when you use the chat endpoint, so just send a standard messages array: + +```bash +MODEL=/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8 + +# 1. Short sanity — should return a one-line MiMo self-introduction. +curl -s http://localhost:8000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d "{\"model\":\"$MODEL\", + \"messages\":[{\"role\":\"user\",\"content\":\"Hello! Introduce yourself in one sentence.\"}], + \"max_tokens\":64}" | python3 -m json.tool + +# 2. Long output — check for repetition collapse / gibberish on 500+ tokens. +curl -s http://localhost:8000/v1/chat/completions \ + -H 'Content-Type: application/json' \ + -d "{\"model\":\"$MODEL\", + \"messages\":[{\"role\":\"user\",\"content\":\"Explain the B-tree data structure in detail, including how insertions and deletions preserve balance.\"}], + \"max_tokens\":800}" | python3 -c "import sys,json; r=json.load(sys.stdin); print(r['choices'][0]['message']['content'])" +``` + +If you see a coherent MiMo introduction and a multi-paragraph technical explanation, the FP8 path is working end-to-end. Output collapse ("helpful helpful helpful ...") on either prompt indicates a broken FP8 recipe — re-check that `moe_tp_degree=1`, `moe_ep_degree=64`, `batch_size>=32`, and that the server is pointed at the Neuron-FP8 preprocessed directory (not the raw HF one). + +**Note on sampling determinism**: `on_device_sampling_config.do_sample=true` is the recommended setting; request-level `temperature` is ignored (sampling params are baked into the NEFF at compile time). + ### vllm-neuron patch summary The patch is applied to vllm-neuron 0.5.0 and: From 2149480ddb37be75630036e247b4561bc20cfbb2 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 11:43:53 +0800 Subject: [PATCH 07/27] 0_setup: resolve CONTRIB_SRC before cd-ing into vllm-neuron The CONTRIB_SRC lookup used $(cd "$(dirname "$0")/.." && pwd), which only works when $0 is an absolute path or dirname resolves from the current working directory. But by the time CONTRIB_SRC was computed, the script had already cd'd into $HOME/vllm-neuron, so a relative $0 like "contrib/models/MiMo-V2.5/perf_test/0_setup.sh" could not find the parent directory and the script failed with: cd: contrib/models/MiMo-V2.5/perf_test/..: No such file or directory Resolve SCRIPT_DIR, PATCH_FILE, and CONTRIB_SRC at the top of the script (before any cd), and reuse SCRIPT_DIR. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/perf_test/0_setup.sh | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/contrib/models/MiMo-V2.5/perf_test/0_setup.sh b/contrib/models/MiMo-V2.5/perf_test/0_setup.sh index 7dfec5ca..5344ded6 100755 --- a/contrib/models/MiMo-V2.5/perf_test/0_setup.sh +++ b/contrib/models/MiMo-V2.5/perf_test/0_setup.sh @@ -16,7 +16,11 @@ echo "==========================================" source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate -PATCH_FILE="$(cd "$(dirname "$0")" && pwd)/vllm-neuron-patch.patch" +# Resolve repo-relative paths up front — we cd into $HOME/vllm-neuron below, +# after which $0's relative form would no longer resolve. +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PATCH_FILE="$SCRIPT_DIR/vllm-neuron-patch.patch" +CONTRIB_SRC="$(cd "$SCRIPT_DIR/.." && pwd)/src" echo "" echo "[1/2] Installing vllm-neuron (release-0.5.0) with the contrib registration patch..." @@ -51,8 +55,6 @@ else echo " Download complete: $(du -sh $MIMO_PATH | cut -f1)" fi -CONTRIB_SRC="$(cd "$(dirname "$0")/.." && pwd)/src" - echo "" echo "Next, preprocess the FP8 checkpoint for Neuron (~15 min, ~15 GB peak RAM):" echo " python $CONTRIB_SRC/conversion_script/preprocess_mimo_v2_5_fp8.py \\" From 60c17a9259d3c5ba320896295d4db4e6c64004ac Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 11:45:24 +0800 Subject: [PATCH 08/27] README: fix maintainer name (Henan Wang, not Wan) Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/contrib/models/MiMo-V2.5/README.md b/contrib/models/MiMo-V2.5/README.md index 6b37293b..d8d6c09d 100644 --- a/contrib/models/MiMo-V2.5/README.md +++ b/contrib/models/MiMo-V2.5/README.md @@ -385,6 +385,6 @@ pytest contrib/models/MiMo-V2.5/test/integration/test_model.py -v ## Maintainer -Henan Wan (whn09) +Henan Wang (whn09) **Last Updated:** 2026-04-28 From 89e04cde0d6adb6d6bc81234496128f2297d1c4e Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 12:26:54 +0800 Subject: [PATCH 09/27] bench scripts: source vllm venv, not the nxd-training one bench_mimo_v2_5.sh and run_bench_single.sh sourced /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference, which doesn't have vllm installed. The rest of the Quick Start (preprocess, smoke, 0_setup.sh) already uses /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16, so align these two. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh | 2 +- contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh index 6606cb5d..67e14727 100755 --- a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh +++ b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh @@ -20,7 +20,7 @@ set -e # this model on Trn2 — use the BF16 checkpoint with the old bench recipe # (`moe_tp_degree=64, moe_ep_degree=1, batch_size=1`). -source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8}" # The NxDI contrib MiMo-V2.5 modeling code is registered into vLLM / diff --git a/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh index 54455731..5679b41d 100755 --- a/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh @@ -25,7 +25,7 @@ set -e -source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-BF16}" PORT="${PORT:-8000}" From 4e788f22e8477d1659ad0f1d4848dab2108d21b0 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 12:34:44 +0800 Subject: [PATCH 10/27] vllm: reuse MiMoV2Flash arch name to satisfy vLLM's builtin validator vLLM 0.16 validates ModelConfig.architectures against its builtin supported-archs list before plugins get a chance to register new classes. That list already contains MiMoV2FlashForCausalLM and MiMoV2ProForCausalLM (Xiaomi upstream PRs), but not the new V2.5 arch name MiMoV2ForCausalLM, so serving V2.5 via vLLM would fail the pydantic check at APIServer startup. Since the V2.5 and Flash NxDI modeling code are the same (modeling_mimo_v2.NeuronMiMoV2ForCausalLM), reuse the Flash arch name to piggyback on the existing vLLM support instead of trying to register a brand new arch from a plugin: - preprocess rewrites `architectures: ["MiMoV2ForCausalLM"]` in the copied config.json to `["MiMoV2FlashForCausalLM"]`. auto_map still points at the V2.5 configuration_mimo_v2 / modeling_mimo_v2 modules, so trust_remote_code loads V2.5 classes as expected. - vllm-neuron-patch.patch is replaced with the Flash-branch patch verbatim (registers mimov2flash in MODEL_TYPES and registers MiMoV2FlashForCausalLM in vllm's ModelRegistry via the worker loader hook). Exactly the same payload as Flash uses. - bench_mimo_v2_5.sh aliases NXDI_CONTRIB_MIMO_V2_FLASH_SRC to the V2.5 src so the Flash-keyed registration hook picks up our V2.5 modeling code. No new __init__.py surgery, no architecture spoofing at runtime; just one config.json rewrite during preprocess and one env var alias at serve time. Co-Authored-By: Claude Opus 4.7 --- .../MiMo-V2.5/perf_test/bench_mimo_v2_5.sh | 8 ++++++++ .../MiMo-V2.5/perf_test/vllm-neuron-patch.patch | 16 ++++++++-------- .../preprocess_mimo_v2_5_fp8.py | 17 +++++++++++++++++ 3 files changed, 33 insertions(+), 8 deletions(-) diff --git a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh index 67e14727..ebdc67cb 100755 --- a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh +++ b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh @@ -28,6 +28,14 @@ MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8}" # Default to this contrib package's own src/ relative to the script. : "${NXDI_CONTRIB_MIMO_V2_5_SRC:=$(cd "$(dirname "$0")/.." && pwd)/src}" export NXDI_CONTRIB_MIMO_V2_5_SRC +# Preprocess writes architectures=["MiMoV2FlashForCausalLM"] into the +# checkpoint's config.json to survive vLLM 0.16's builtin arch validator +# (which knows MiMoV2FlashForCausalLM but not MiMoV2ForCausalLM). Alias +# the V2.5 src as the Flash src so vllm-neuron's _register_contrib_models +# hook (triggered on Flash arch names) registers our V2.5 modeling class +# under MODEL_TYPES['mimov2flash']. The modeling module (modeling_mimo_v2) +# and NeuronMiMoV2ForCausalLM class are shared between Flash and V2.5. +export NXDI_CONTRIB_MIMO_V2_FLASH_SRC="$NXDI_CONTRIB_MIMO_V2_5_SRC" # First-time MiMo-V2.5 FP8 compile takes 30-60 minutes; extend vLLM's ready # timeout and the compiler's environment variables for FP8 numerics. diff --git a/contrib/models/MiMo-V2.5/perf_test/vllm-neuron-patch.patch b/contrib/models/MiMo-V2.5/perf_test/vllm-neuron-patch.patch index 67a88579..4a84a558 100644 --- a/contrib/models/MiMo-V2.5/perf_test/vllm-neuron-patch.patch +++ b/contrib/models/MiMo-V2.5/perf_test/vllm-neuron-patch.patch @@ -11,7 +11,7 @@ index d2099eb..0c162e4 100644 + """Monkey-patch ``AutoConfig.from_pretrained`` to default ``trust_remote_code=True``. + + NxDI's ``hf_adapter.load_config`` calls ``AutoConfig.from_pretrained(path)`` -+ without ``trust_remote_code``. Contrib models like MiMo-V2.5 that ++ without ``trust_remote_code``. Contrib models like MiMo-V2-Flash that + ship a ``configuration_*.py`` with the checkpoint require custom code + execution, so the default behaviour crashes with ``ValueError: The + repository ... contains custom code which must be executed``. @@ -39,7 +39,7 @@ index d2099eb..0c162e4 100644 + """Lazy-register NxDI contrib models on each process that calls the loader. + + Driven by env vars: -+ NXDI_CONTRIB_MIMO_V2_5_SRC -> path to contrib MiMo-V2.5 src/ ++ NXDI_CONTRIB_MIMO_V2_FLASH_SRC -> path to contrib MiMo-V2-Flash src/ + NXDI_CONTRIB_MINIMAX_M2_SRC -> path to contrib MiniMax-M2 src/ + + Registers the contrib model class into NxDI's MODEL_TYPES and, where @@ -55,26 +55,26 @@ index d2099eb..0c162e4 100644 + + _patch_autoconfig_trust_remote_code() + -+ mimo_src = _os.environ.get("NXDI_CONTRIB_MIMO_V2_5_SRC") -+ if mimo_src and _os.path.isdir(mimo_src) and "mimov2" not in MODEL_TYPES: ++ mimo_src = _os.environ.get("NXDI_CONTRIB_MIMO_V2_FLASH_SRC") ++ if mimo_src and _os.path.isdir(mimo_src) and "mimov2flash" not in MODEL_TYPES: + if mimo_src not in _sys.path: + _sys.path.insert(0, mimo_src) + try: + from modeling_mimo_v2 import NeuronMiMoV2ForCausalLM + MODEL_TYPES.setdefault( -+ "mimov2", {"causal-lm": NeuronMiMoV2ForCausalLM} ++ "mimov2flash", {"causal-lm": NeuronMiMoV2ForCausalLM} + ) + try: + from vllm.model_executor.models.registry import ModelRegistry -+ if "MiMoV2ForCausalLM" not in ModelRegistry.get_supported_archs(): ++ if "MiMoV2FlashForCausalLM" not in ModelRegistry.get_supported_archs(): + ModelRegistry.register_model( -+ "MiMoV2ForCausalLM", NeuronMiMoV2ForCausalLM ++ "MiMoV2FlashForCausalLM", NeuronMiMoV2ForCausalLM + ) + except ImportError: + pass + except Exception as e: + _w.warn( -+ f"Failed to register MiMo-V2.5 contrib model: {e}", ++ f"Failed to register MiMo-V2-Flash contrib model: {e}", + category=UserWarning, + ) + diff --git a/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py index 80838239..1530e93f 100644 --- a/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py +++ b/contrib/models/MiMo-V2.5/src/conversion_script/preprocess_mimo_v2_5_fp8.py @@ -671,6 +671,23 @@ def process_flash_checkpoint(hf_model_path: str, save_path: str, tp_degree: int) if os.path.isfile(src): shutil.copy(src, os.path.join(save_path, name)) + # --- Rewrite architectures in the copied config.json so vLLM's + # pydantic ModelConfig validator accepts the checkpoint. vLLM's + # builtin supported-archs list (Xiaomi's upstream PR) contains + # MiMoV2FlashForCausalLM but not the V2.5 arch `MiMoV2ForCausalLM`, + # and the vllm-neuron plugin registers contribs too late to patch + # that list. The NxDI side loads via auto_map + trust_remote_code, + # so the arch name only has to survive the vLLM pydantic check. + # auto_map still points at the V2.5 modeling/configuration modules. + cfg_path = os.path.join(save_path, "config.json") + if os.path.isfile(cfg_path): + with open(cfg_path) as _f: + _cfg = json.load(_f) + if _cfg.get("architectures") == ["MiMoV2ForCausalLM"]: + _cfg["architectures"] = ["MiMoV2FlashForCausalLM"] + with open(cfg_path, "w") as _f: + json.dump(_cfg, _f, indent=2) + print(f"\nPreprocess complete. total_size={total_size/1e9:.2f} GB", flush=True) print(f" tensors written: {len(weight_map_out)}", flush=True) print(f" output dir: {save_path}", flush=True) From d724a94d9509c63ecf6d3444dca97edb5dfb41ff Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 13:05:46 +0800 Subject: [PATCH 11/27] Consolidate env var docs + point scratch dirs at persistent storage Centralize the four env vars used across smoke / bench / manual vLLM launches: - 0_setup.sh: clearer "Next steps" output that prints all four exports (required + optional) with explanations. Replaces the two-line hint the previous version ended with. - bench_mimo_v2_5.sh: adds defaults for NEURON_COMPILED_ARTIFACTS (/opt/dlami/nvme/compiled/mimo_v2_5_bs32_moetp1_ep64_fp8_vllm, same layout as other contrib models on this instance) and BASE_COMPILE_WORK_DIR (/opt/dlami/nvme/tmp/nxd_model/, so NxDI's HLO/NEFF staging survives the nightly Trn2 reboot and parallel compiles can't clobber each other). - README: new "Environment variables" subsection under Quick Start tabulating required vs optional vars, defaults, and why each matters. Without NEURON_COMPILED_ARTIFACTS set, vllm-neuron falls back to /neuron-compiled-artifacts//, which buries the output inside the checkpoint dir and isn't what we want when iterating. Without BASE_COMPILE_WORK_DIR set, NxDI's /tmp/nxd_model/ default gets wiped by the reboot mid-compile. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/README.md | 23 ++++++++ contrib/models/MiMo-V2.5/perf_test/0_setup.sh | 53 ++++++++++++++++--- .../MiMo-V2.5/perf_test/bench_mimo_v2_5.sh | 12 +++++ 3 files changed, 80 insertions(+), 8 deletions(-) diff --git a/contrib/models/MiMo-V2.5/README.md b/contrib/models/MiMo-V2.5/README.md index d8d6c09d..12af0025 100644 --- a/contrib/models/MiMo-V2.5/README.md +++ b/contrib/models/MiMo-V2.5/README.md @@ -87,6 +87,29 @@ The bench script runs two configurations (BS=32 and BS=128, both `moe_tp_degree=1 / moe_ep_degree=64`) and logs results under `/opt/dlami/nvme/logs/bench_results/mimo_v2_5/`. +### Environment variables + +`0_setup.sh` prints these at the end; setting them explicitly makes the +smoke / bench / manual-launch paths all behave the same. All of them have +sensible defaults in the scripts — export them only if you want to +override or if you plan to launch vLLM outside of `bench_mimo_v2_5.sh`. + +**Required (at least for manual `vllm api_server` launches):** + +| Variable | Purpose | +|---|---| +| `NXDI_CONTRIB_MIMO_V2_5_SRC` | Path to `contrib/models/MiMo-V2.5/src/`. `vllm-neuron`'s registration hook reads it to plug `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` table. | +| `NXDI_CONTRIB_MIMO_V2_FLASH_SRC` | Alias of `NXDI_CONTRIB_MIMO_V2_5_SRC` — same value. vLLM's builtin arch validator only knows `MiMoV2FlashForCausalLM`, so preprocess rewrites the checkpoint's `architectures` to that name and we re-use the Flash registration key (`mimov2flash`) in vllm-neuron's lookup table. | +| `MIMO_V2_5_PATH` | Preprocessed Neuron-FP8 checkpoint dir (the `--save_path` output from preprocess). | + +**Optional (recommended):** + +| Variable | Default | Purpose | +|---|---|---| +| `NEURON_COMPILED_ARTIFACTS` | `/opt/dlami/nvme/compiled/mimo_v2_5_bs32_moetp1_ep64_fp8_vllm` | Where vLLM writes the NEFF + per-rank sharded weights. Default points at a persistent path under `/opt/dlami/nvme/compiled/` so multiple configs don't collide and runs after the nightly reboot can reuse the sharded weights. vLLM's fallback is `/neuron-compiled-artifacts//` which buries output inside the checkpoint dir. | +| `BASE_COMPILE_WORK_DIR` | `/opt/dlami/nvme/tmp/nxd_model/` | NxDI's HLO / NEFF staging workdir. Default is `/tmp/nxd_model/`, which is wiped by the nightly Trn2 reboot and can silently corrupt parallel compiles that share a basename; the pinned value lives on persistent storage and is unique per config. | +| `VLLM_ENGINE_READY_TIMEOUT_S` | `7200` | First-time compile of V2.5's 256-expert MoE is ~30 min dominated by `shard_checkpoint`, well past vLLM's default. | + For a quick `curl` sanity check while the server is up: ```bash diff --git a/contrib/models/MiMo-V2.5/perf_test/0_setup.sh b/contrib/models/MiMo-V2.5/perf_test/0_setup.sh index 5344ded6..68c3752e 100755 --- a/contrib/models/MiMo-V2.5/perf_test/0_setup.sh +++ b/contrib/models/MiMo-V2.5/perf_test/0_setup.sh @@ -55,13 +55,50 @@ else echo " Download complete: $(du -sh $MIMO_PATH | cut -f1)" fi +NEURON_FP8_PATH="${MIMO_PATH}-Neuron-FP8" +COMPILED_PATH="/opt/dlami/nvme/compiled/mimo_v2_5_bs32_moetp1_ep64_fp8_vllm" + +echo "" +echo "========================================================================" +echo "Next steps" +echo "========================================================================" +echo "" +echo "1. Preprocess the FP8 checkpoint for Neuron (~16 min, ~15 GB peak RAM):" +echo "" +echo " python $CONTRIB_SRC/conversion_script/preprocess_mimo_v2_5_fp8.py \\" +echo " --hf_model_path $MIMO_PATH \\" +echo " --save_path $NEURON_FP8_PATH \\" +echo " --tp_degree 64" +echo "" +echo "2. Export the environment variables used by the smoke / bench scripts:" +echo "" +echo " # --- Required ---" +echo " # Contrib package src (registers NeuronMiMoV2ForCausalLM with vllm-neuron)." +echo " export NXDI_CONTRIB_MIMO_V2_5_SRC=$CONTRIB_SRC" +echo " # vLLM's builtin arch validator only knows MiMoV2FlashForCausalLM, so the" +echo " # preprocess rewrites the checkpoint's config.json architectures to that" +echo " # name. Alias V2.5 src to the Flash env var so vllm-neuron's contrib hook" +echo " # registers mimov2flash -> our V2.5 NeuronMiMoV2ForCausalLM class." +echo " export NXDI_CONTRIB_MIMO_V2_FLASH_SRC=\"\$NXDI_CONTRIB_MIMO_V2_5_SRC\"" +echo " # Preprocessed Neuron-FP8 checkpoint." +echo " export MIMO_V2_5_PATH=$NEURON_FP8_PATH" +echo "" +echo " # --- Optional (recommended) ---" +echo " # vLLM compiles into /neuron-compiled-artifacts// by" +echo " # default. Pin it to a persistent shared location so multiple configs" +echo " # don't collide and you can reuse the NEFF / sharded weights across runs." +echo " export NEURON_COMPILED_ARTIFACTS=$COMPILED_PATH" +echo " # NxDI's HLO/NEFF staging workdir (.hlo_module.pb etc). Default is" +echo " # /tmp/nxd_model//; on Trn2 /tmp is wiped by the nightly" +echo " # reboot, and parallel compiles sharing the same basename silently" +echo " # overwrite each other's staged HLOs. Pin to a unique per-config" +echo " # directory under persistent storage." +echo " export BASE_COMPILE_WORK_DIR=/opt/dlami/nvme/tmp/nxd_model/\$(basename $COMPILED_PATH)" +echo " # First-time compile of V2.5's 256-expert MoE takes ~30 min (NEFF HLO +" +echo " # shard_checkpoint for 64 ranks). Extend vLLM's ready timeout." +echo " export VLLM_ENGINE_READY_TIMEOUT_S=7200" +echo "" +echo "3. Run the benchmark:" echo "" -echo "Next, preprocess the FP8 checkpoint for Neuron (~15 min, ~15 GB peak RAM):" -echo " python $CONTRIB_SRC/conversion_script/preprocess_mimo_v2_5_fp8.py \\" -echo " --hf_model_path $MIMO_PATH \\" -echo " --save_path ${MIMO_PATH}-Neuron-FP8 \\" -echo " --tp_degree 64" +echo " bash $SCRIPT_DIR/bench_mimo_v2_5.sh" echo "" -echo "Then before running the benchmark, export:" -echo " export MIMO_V2_5_PATH=${MIMO_PATH}-Neuron-FP8" -echo " export NXDI_CONTRIB_MIMO_V2_5_SRC=$CONTRIB_SRC" diff --git a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh index ebdc67cb..f793e240 100755 --- a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh +++ b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh @@ -37,6 +37,18 @@ export NXDI_CONTRIB_MIMO_V2_5_SRC # and NeuronMiMoV2ForCausalLM class are shared between Flash and V2.5. export NXDI_CONTRIB_MIMO_V2_FLASH_SRC="$NXDI_CONTRIB_MIMO_V2_5_SRC" +# vLLM writes the NEFF + sharded weights here. Default under +# /opt/dlami/nvme/compiled (persistent) so multiple configs don't clobber +# each other and the NEFF/shards are reusable after a reboot. +: "${NEURON_COMPILED_ARTIFACTS:=/opt/dlami/nvme/compiled/mimo_v2_5_bs32_moetp1_ep64_fp8_vllm}" +export NEURON_COMPILED_ARTIFACTS +# NxDI's HLO/NEFF staging directory (default /tmp/nxd_model//). Pin +# to persistent storage so it survives the nightly Trn2 reboot, and to a +# unique per-config subdir to stay safe under parallel compile invocations. +: "${BASE_COMPILE_WORK_DIR:=/opt/dlami/nvme/tmp/nxd_model/$(basename "$NEURON_COMPILED_ARTIFACTS")}" +export BASE_COMPILE_WORK_DIR +mkdir -p "$BASE_COMPILE_WORK_DIR" + # First-time MiMo-V2.5 FP8 compile takes 30-60 minutes; extend vLLM's ready # timeout and the compiler's environment variables for FP8 numerics. export VLLM_ENGINE_READY_TIMEOUT_S=7200 From 1469d1511de3a547fe7f68b59be5fae7ea478353 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 13:39:17 +0800 Subject: [PATCH 12/27] bench: drop the BS=128 config, keep only BS=32 BS=32 is the smallest batch size the FP8 path supports (num_experts/top_k = 32 requirement for EP>1 in the TKG graph), and it's already the target recipe for serving. Running BS=128 in the same bench script doubled compile time for no additional signal and produced a second NEFF + sharded-weights tree that we don't use. Also update the README description of the bench script. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/README.md | 6 +-- .../MiMo-V2.5/perf_test/bench_mimo_v2_5.sh | 45 ------------------- 2 files changed, 3 insertions(+), 48 deletions(-) diff --git a/contrib/models/MiMo-V2.5/README.md b/contrib/models/MiMo-V2.5/README.md index 12af0025..ab9f12e9 100644 --- a/contrib/models/MiMo-V2.5/README.md +++ b/contrib/models/MiMo-V2.5/README.md @@ -83,8 +83,8 @@ bash contrib/models/MiMo-V2.5/perf_test/0_setup.sh bash contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh ``` -The bench script runs two configurations (BS=32 and BS=128, both -`moe_tp_degree=1 / moe_ep_degree=64`) and logs results under +The bench script runs one configuration (BS=32, +`moe_tp_degree=1 / moe_ep_degree=64`) at three concurrency levels (1, 16, 32) and logs results under `/opt/dlami/nvme/logs/bench_results/mimo_v2_5/`. ### Environment variables @@ -333,7 +333,7 @@ python3 -m vllm.entrypoints.openai.api_server \ }' ``` -See `perf_test/bench_mimo_v2_5.sh` for the full benchmark recipe at BS=32 and BS=128. +See `perf_test/bench_mimo_v2_5.sh` for the full benchmark recipe at BS=32. ### Testing the vLLM server diff --git a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh index f793e240..407b5892 100755 --- a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh +++ b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh @@ -196,51 +196,6 @@ run_bench "$CONFIG_NAME" 16 128 run_bench "$CONFIG_NAME" 32 128 stop_server -############################################################################### -# Config 2: BS=128, TP=64 + moe_tp=1/moe_ep=64, CB + bucketing (throughput). -############################################################################### -CONFIG_NAME="bs128_tp64_moetp1_ep64" -echo "--- Config 2: BS=128, moe_tp=1/moe_ep=64, CB + bucketing ---" - -python3 -m vllm.entrypoints.openai.api_server \ - --model "$MODEL_PATH" \ - --tokenizer "$MODEL_PATH" \ - --tensor-parallel-size 64 \ - --max-model-len 1024 \ - --max-num-seqs 128 \ - --no-enable-chunked-prefill \ - --no-enable-prefix-caching \ - --port $PORT \ - --trust_remote_code \ - --additional-config '{ - "override_neuron_config": { - '"$COMMON_MIMO_CONFIG"', - "moe_tp_degree": 1, - "moe_ep_degree": 64, - "batch_size": 128, - "ctx_batch_size": 1, - "tkg_batch_size": 128, - "max_context_length": 1024, - "seq_len": 1024, - "is_continuous_batching": true, - "enable_bucketing": true, - "context_encoding_buckets": [1024], - "token_generation_buckets": [1024], - "async_mode": true, - "on_device_sampling_config": { - "do_sample": true, "temperature": 0.6, "top_k": 20, "top_p": 0.95 - } - } - }' & - -wait_for_server -sanity_check -run_bench "$CONFIG_NAME" 1 16 -run_bench "$CONFIG_NAME" 16 128 -run_bench "$CONFIG_NAME" 32 128 -run_bench "$CONFIG_NAME" 128 512 -stop_server - echo "==========================================" echo "MiMo-V2.5 FP8 benchmarks complete!" echo "Results saved to: $RESULTS_DIR" From 9b44f577a0e95864b7e1085dfd8fb42a4e544bb6 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 13:46:05 +0800 Subject: [PATCH 13/27] sanity_check + run_bench_single: default to Neuron-FP8 checkpoint Both scripts defaulted MODEL_PATH to the BF16 directory path, which is leftover from the Flash-era bench (Flash had a BF16 serving recipe at BS=1 alongside the FP8 recipe at BS=32). On V2.5 only FP8 is supported, so default to the -Neuron-FP8 directory instead. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh | 6 +++--- contrib/models/MiMo-V2.5/perf_test/sanity_check.sh | 2 +- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh index 5679b41d..5206ba57 100755 --- a/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh @@ -13,8 +13,8 @@ # # Environment knobs: # PORT vLLM server port (default 8000) -# MIMO_V2_5_PATH Path to the BF16 checkpoint (default -# /opt/dlami/nvme/models/MiMo-V2.5-BF16) +# MIMO_V2_5_PATH Path to the Neuron-FP8 checkpoint (default +# /opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8) # CONCURRENCY --max-concurrency (default 1) # NUM_PROMPTS --num-prompts (default 16) # INPUT_LEN --random-input-len (default 900) @@ -27,7 +27,7 @@ set -e source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate -MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-BF16}" +MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8}" PORT="${PORT:-8000}" CONCURRENCY="${CONCURRENCY:-1}" NUM_PROMPTS="${NUM_PROMPTS:-16}" diff --git a/contrib/models/MiMo-V2.5/perf_test/sanity_check.sh b/contrib/models/MiMo-V2.5/perf_test/sanity_check.sh index 8dca4c7c..3e43a279 100755 --- a/contrib/models/MiMo-V2.5/perf_test/sanity_check.sh +++ b/contrib/models/MiMo-V2.5/perf_test/sanity_check.sh @@ -11,7 +11,7 @@ set -e -MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-BF16}" +MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8}" PORT="${PORT:-8000}" PROMPT="${PROMPT:-What is 1+1? Answer briefly.}" MAX_TOKENS="${MAX_TOKENS:-64}" From 2e5348b55fdece75c177a57f9da255566a8a0692 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 13:59:51 +0800 Subject: [PATCH 14/27] perf_test: split out start_vllm_server.sh; make bench_mimo_v2_5 a thin wrapper Previously bench_mimo_v2_5.sh inlined server launch + sanity + three bench runs + teardown (and repeated env-var setup + the full additional-config JSON twice, once per Config). sanity_check.sh and run_bench_single.sh existed as standalone tools but there was no matching "just start the server" script, so the only way to get a running server was to invoke the bench driver. Users who wanted to keep a server up to iterate on prompt or concurrency choices had to either copy-paste bench's launch block or kill bench mid-run. Extract start_vllm_server.sh as the single place that: - sources the vllm venv - exports NXDI_CONTRIB_MIMO_V2_5_SRC, NXDI_CONTRIB_MIMO_V2_FLASH_SRC, NEURON_COMPILED_ARTIFACTS, BASE_COMPILE_WORK_DIR (with defaults) - execs `python3 -m vllm.entrypoints.openai.api_server` with the recipe bench_mimo_v2_5.sh is now a thin orchestrator: backgrounds start_vllm_server.sh, waits for readiness, invokes sanity_check.sh and run_bench_single.sh at c=1,16,32, tears down on exit. 205 lines -> 87. 0_setup.sh "Next steps" and the README now document both the one-shot path and the long-running-server + ad-hoc probe path. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/README.md | 21 ++ contrib/models/MiMo-V2.5/perf_test/0_setup.sh | 13 +- .../MiMo-V2.5/perf_test/bench_mimo_v2_5.sh | 206 ++++-------------- .../MiMo-V2.5/perf_test/start_vllm_server.sh | 104 +++++++++ 4 files changed, 181 insertions(+), 163 deletions(-) create mode 100755 contrib/models/MiMo-V2.5/perf_test/start_vllm_server.sh diff --git a/contrib/models/MiMo-V2.5/README.md b/contrib/models/MiMo-V2.5/README.md index ab9f12e9..5a3ff8ce 100644 --- a/contrib/models/MiMo-V2.5/README.md +++ b/contrib/models/MiMo-V2.5/README.md @@ -87,6 +87,27 @@ The bench script runs one configuration (BS=32, `moe_tp_degree=1 / moe_ep_degree=64`) at three concurrency levels (1, 16, 32) and logs results under `/opt/dlami/nvme/logs/bench_results/mimo_v2_5/`. +### Keeping a server up for ad-hoc testing + +`bench_mimo_v2_5.sh` is a one-shot wrapper (launch server → sanity → +3 bench runs → teardown). If you want a long-running server to iterate +against, use the three underlying scripts separately: + +```bash +# Terminal 1: launch the server in the foreground (Ctrl-C to stop). +bash contrib/models/MiMo-V2.5/perf_test/start_vllm_server.sh + +# Terminal 2: once "Application startup complete." prints, sanity-check: +bash contrib/models/MiMo-V2.5/perf_test/sanity_check.sh + +# Run a single bench pass with a chosen concurrency: +CONCURRENCY=16 NUM_PROMPTS=128 \ + bash contrib/models/MiMo-V2.5/perf_test/run_bench_single.sh +``` + +`bench_mimo_v2_5.sh` composes exactly these three pieces; use whichever +is more convenient. + ### Environment variables `0_setup.sh` prints these at the end; setting them explicitly makes the diff --git a/contrib/models/MiMo-V2.5/perf_test/0_setup.sh b/contrib/models/MiMo-V2.5/perf_test/0_setup.sh index 68c3752e..7c88c1af 100755 --- a/contrib/models/MiMo-V2.5/perf_test/0_setup.sh +++ b/contrib/models/MiMo-V2.5/perf_test/0_setup.sh @@ -98,7 +98,16 @@ echo " # First-time compile of V2.5's 256-expert MoE takes ~30 min (NEFF HLO echo " # shard_checkpoint for 64 ranks). Extend vLLM's ready timeout." echo " export VLLM_ENGINE_READY_TIMEOUT_S=7200" echo "" -echo "3. Run the benchmark:" +echo "3a. Run the one-shot benchmark (launches + benches + tears down):" echo "" -echo " bash $SCRIPT_DIR/bench_mimo_v2_5.sh" +echo " bash $SCRIPT_DIR/bench_mimo_v2_5.sh" +echo "" +echo "3b. ...OR keep a server up and probe it manually:" +echo "" +echo " # shell 1: server in foreground (Ctrl-C to stop)" +echo " bash $SCRIPT_DIR/start_vllm_server.sh" +echo "" +echo " # shell 2: once 'Application startup complete.' prints," +echo " bash $SCRIPT_DIR/sanity_check.sh" +echo " CONCURRENCY=16 NUM_PROMPTS=128 bash $SCRIPT_DIR/run_bench_single.sh" echo "" diff --git a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh index 407b5892..7f82cfcd 100755 --- a/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh +++ b/contrib/models/MiMo-V2.5/perf_test/bench_mimo_v2_5.sh @@ -1,97 +1,42 @@ #!/bin/bash set -e -# MiMo-V2.5 FP8 vLLM benchmark on Trn2. +# MiMo-V2.5 FP8 vLLM benchmark on Trn2. One-shot wrapper: +# launch server -> sanity check -> bench at c=1,16,32 -> stop server. # -# Requires a Neuron-FP8 preprocessed checkpoint (see -# `src/conversion_script/preprocess_mimo_v2_5_fp8.py`). The configs below -# all use moe_tp_degree=1 / moe_ep_degree=64 (experts sharded by expert -# parallelism only, no intra-expert TP split) because moe_tp_degree=64 collapses -# the per-rank FP8 blockwise scale to a singleton — per-rank expert -# intermediate is 32 rows, below the 128-row blockwise block, so -# NxDI's `_setup_for_scale` drops per-channel scale granularity. The resulting -# drift compounds across 47 MoE layers and gives repetition / output collapse. -# Using moe_ep_degree=64 keeps all of each expert's weight + scale on one rank -# (4 experts per rank), which preserves the blockwise scale intact. +# This script composes three building blocks in perf_test/: +# start_vllm_server.sh - server launch + env-var setup (backgrounded here) +# sanity_check.sh - one-shot curl against the running server +# run_bench_single.sh - one concurrency level of `vllm bench serve` # -# NxDI's TKG path refuses Expert Parallelism with BS < num_experts/top_k -# (256 / 8 = 32 for MiMo-V2.5), so the smallest working batch size here is 32. -# If you want BS=1 behaviour, the FP8 path is not currently supported on -# this model on Trn2 — use the BF16 checkpoint with the old bench recipe -# (`moe_tp_degree=64, moe_ep_degree=1, batch_size=1`). - -source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate - -MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8}" -# The NxDI contrib MiMo-V2.5 modeling code is registered into vLLM / -# NxDI lookup tables by vllm-neuron's register() hook using this env var. -# Default to this contrib package's own src/ relative to the script. -: "${NXDI_CONTRIB_MIMO_V2_5_SRC:=$(cd "$(dirname "$0")/.." && pwd)/src}" -export NXDI_CONTRIB_MIMO_V2_5_SRC -# Preprocess writes architectures=["MiMoV2FlashForCausalLM"] into the -# checkpoint's config.json to survive vLLM 0.16's builtin arch validator -# (which knows MiMoV2FlashForCausalLM but not MiMoV2ForCausalLM). Alias -# the V2.5 src as the Flash src so vllm-neuron's _register_contrib_models -# hook (triggered on Flash arch names) registers our V2.5 modeling class -# under MODEL_TYPES['mimov2flash']. The modeling module (modeling_mimo_v2) -# and NeuronMiMoV2ForCausalLM class are shared between Flash and V2.5. -export NXDI_CONTRIB_MIMO_V2_FLASH_SRC="$NXDI_CONTRIB_MIMO_V2_5_SRC" - -# vLLM writes the NEFF + sharded weights here. Default under -# /opt/dlami/nvme/compiled (persistent) so multiple configs don't clobber -# each other and the NEFF/shards are reusable after a reboot. -: "${NEURON_COMPILED_ARTIFACTS:=/opt/dlami/nvme/compiled/mimo_v2_5_bs32_moetp1_ep64_fp8_vllm}" -export NEURON_COMPILED_ARTIFACTS -# NxDI's HLO/NEFF staging directory (default /tmp/nxd_model//). Pin -# to persistent storage so it survives the nightly Trn2 reboot, and to a -# unique per-config subdir to stay safe under parallel compile invocations. -: "${BASE_COMPILE_WORK_DIR:=/opt/dlami/nvme/tmp/nxd_model/$(basename "$NEURON_COMPILED_ARTIFACTS")}" -export BASE_COMPILE_WORK_DIR -mkdir -p "$BASE_COMPILE_WORK_DIR" - -# First-time MiMo-V2.5 FP8 compile takes 30-60 minutes; extend vLLM's ready -# timeout and the compiler's environment variables for FP8 numerics. -export VLLM_ENGINE_READY_TIMEOUT_S=7200 +# Use those directly if you want to keep a long-running server and iterate +# on bench parameters from another shell. +# +# Server recipe: TP=64, moe_tp=1/moe_ep=64, BS=32, continuous batching. +# BS=32 is the smallest working batch size on the FP8 path (NxDI's TKG +# path refuses Expert Parallelism with BS < num_experts/top_k = 256/8 = 32). +# BS=1 single-stream latency demos are not currently supported on V2.5 FP8. + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PORT="${PORT:-8000}" +RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/logs/bench_results/mimo_v2_5}" +CONFIG_NAME="bs32_tp64_moetp1_ep64" -PORT=8000 -RESULTS_DIR="/opt/dlami/nvme/logs/bench_results/mimo_v2_5" mkdir -p "$RESULTS_DIR" -# Common neuron config shared across all MiMo-V2.5 FP8 configs. -# save_sharded_checkpoint=true persists per-rank sharded weights to -# /weights/tp{N}_sharded_checkpoint.safetensors during compile; -# load() then reads those directly (~30s) instead of re-sharding the entire -# checkpoint on every vllm-neuron startup (~10+ min). -COMMON_MIMO_CONFIG='"tp_degree": 64, - "logical_nc_config": 2, - "fused_qkv": false, - "sequence_parallel_enabled": false, - "glu_mlp": true, - "normalize_top_k_affinities": true, - "save_sharded_checkpoint": true, - "router_config": {"act_fn": "sigmoid", "dtype": "float32"}, - "quantized": true, - "quantized_checkpoints_path": "'"$MODEL_PATH"'", - "quantization_dtype": "f8e4m3", - "quantization_type": "blockwise_symmetric", - "quantization_block_axis": [1, 2], - "quantization_block_size": [128, 128], - "modules_to_not_convert": ["embed_tokens", "lm_head", "norm", "router", "o_proj"], - "blockwise_matmul_config": {"use_shard_on_block_dynamic_while": true, "block_sharding_strategy": "PING_PONG"}' - -# Helper: wait for vLLM server to be ready. First-time compilation of a -# 256-expert MoE model takes 30-90 minutes, so we poll for up to 2 hours. +# Wait for vLLM server to be ready. First-time compile of the 256-expert +# MoE model takes ~30 min and can stretch past 2 h under contention, so +# poll for up to 2 h. wait_for_server() { - echo " Waiting for vLLM server to be ready (up to 2h for first compile)..." + echo " Waiting for vLLM server on port $PORT (up to 2 h for first compile)..." local interval=10 - local max_attempts=720 # 720 * 10s = 7200s = 2h + local max_attempts=720 local start=$SECONDS for i in $(seq 1 $max_attempts); do - if curl -s http://localhost:$PORT/health > /dev/null 2>&1; then - echo " Server ready! (waited $((SECONDS - start))s)" + if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then + echo " Server ready after $((SECONDS - start))s." return 0 fi - # Show a progress blip every minute so the user knows we're alive if [ $((i % 6)) -eq 0 ]; then echo " ...still waiting ($((SECONDS - start))s elapsed)" fi @@ -101,103 +46,42 @@ wait_for_server() { return 1 } -# Helper: run benchmark -run_bench() { - local config_name=$1 - local concurrency=$2 - local num_prompts=$3 - - echo " Benchmark: concurrency=$concurrency, prompts=$num_prompts" - vllm bench serve \ - --backend vllm \ - --model "$MODEL_PATH" \ - --tokenizer "$MODEL_PATH" \ - --endpoint /v1/completions \ - --dataset-name random \ - --num-prompts "$num_prompts" \ - --random-input-len 900 \ - --random-output-len 90 \ - --random-range-ratio 0.03 \ - --max-concurrency "$concurrency" \ - 2>&1 | tee "$RESULTS_DIR/${config_name}_c${concurrency}.txt" - echo "" -} - -# Helper: stop server stop_server() { echo " Stopping vLLM server..." pkill -f "vllm.entrypoints.openai.api_server" 2>/dev/null || true sleep 5 } -# Helper: quick sanity check -sanity_check() { - echo " Running sanity check..." - curl -s http://localhost:$PORT/v1/chat/completions \ - -H 'Content-Type: application/json' \ - -d '{ - "messages": [{"role": "user", "content": "What is 1+1? Answer briefly."}], - "model": "'"$MODEL_PATH"'", - "max_tokens": 64, - "temperature": 0.0, - "stream": false - }' | python3 -c "import sys,json; r=json.load(sys.stdin); print(' Sanity:', r['choices'][0]['message']['content'][:100])" 2>/dev/null || echo " Sanity check: could not parse response" -} - echo "==========================================" echo "MiMo-V2.5 FP8 Performance Benchmark" echo "==========================================" -echo "Model: $MODEL_PATH" +echo "Port: $PORT" echo "Results: $RESULTS_DIR" echo "" -############################################################################### -# Config 1: BS=32, TP=64 + moe_tp=1/moe_ep=64, CB + bucketing (smallest BS -# that satisfies NxDI's Expert-Parallel BS >= num_experts/top_k requirement). -############################################################################### -CONFIG_NAME="bs32_tp64_moetp1_ep64" -echo "--- Config 1: BS=32, moe_tp=1/moe_ep=64, CB + bucketing ---" - -python3 -m vllm.entrypoints.openai.api_server \ - --model "$MODEL_PATH" \ - --tokenizer "$MODEL_PATH" \ - --tensor-parallel-size 64 \ - --max-model-len 1024 \ - --max-num-seqs 32 \ - --no-enable-chunked-prefill \ - --no-enable-prefix-caching \ - --port $PORT \ - --trust_remote_code \ - --additional-config '{ - "override_neuron_config": { - '"$COMMON_MIMO_CONFIG"', - "moe_tp_degree": 1, - "moe_ep_degree": 64, - "batch_size": 32, - "ctx_batch_size": 1, - "tkg_batch_size": 32, - "max_context_length": 1024, - "seq_len": 1024, - "is_continuous_batching": true, - "enable_bucketing": true, - "context_encoding_buckets": [1024], - "token_generation_buckets": [1024], - "async_mode": true, - "on_device_sampling_config": { - "do_sample": true, "temperature": 0.6, "top_k": 20, "top_p": 0.95 - } - } - }' & +# Start the server in the background. start_vllm_server.sh handles all the +# env vars (MODEL_PATH, NEURON_COMPILED_ARTIFACTS, BASE_COMPILE_WORK_DIR, +# contrib src registration, etc.) and execs `python3 -m vllm...`. +bash "$SCRIPT_DIR/start_vllm_server.sh" & +SERVER_PID=$! +trap stop_server EXIT wait_for_server -sanity_check -run_bench "$CONFIG_NAME" 1 16 -run_bench "$CONFIG_NAME" 16 128 -run_bench "$CONFIG_NAME" 32 128 -stop_server + +# One-shot sanity check (curl the chat endpoint). +PORT="$PORT" bash "$SCRIPT_DIR/sanity_check.sh" || true + +# Three concurrency levels. run_bench_single.sh reads knobs from the +# environment; see its header for all the options. +PORT="$PORT" RESULTS_DIR="$RESULTS_DIR" CONFIG_NAME="$CONFIG_NAME" \ + CONCURRENCY=1 NUM_PROMPTS=16 bash "$SCRIPT_DIR/run_bench_single.sh" +PORT="$PORT" RESULTS_DIR="$RESULTS_DIR" CONFIG_NAME="$CONFIG_NAME" \ + CONCURRENCY=16 NUM_PROMPTS=128 bash "$SCRIPT_DIR/run_bench_single.sh" +PORT="$PORT" RESULTS_DIR="$RESULTS_DIR" CONFIG_NAME="$CONFIG_NAME" \ + CONCURRENCY=32 NUM_PROMPTS=128 bash "$SCRIPT_DIR/run_bench_single.sh" echo "==========================================" -echo "MiMo-V2.5 FP8 benchmarks complete!" +echo "MiMo-V2.5 FP8 benchmark complete!" echo "Results saved to: $RESULTS_DIR" echo "==========================================" ls -la "$RESULTS_DIR" diff --git a/contrib/models/MiMo-V2.5/perf_test/start_vllm_server.sh b/contrib/models/MiMo-V2.5/perf_test/start_vllm_server.sh new file mode 100755 index 00000000..dc515ba0 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/start_vllm_server.sh @@ -0,0 +1,104 @@ +#!/bin/bash +# Start the MiMo-V2.5 FP8 vLLM OpenAI-compatible server in the foreground. +# +# The server stays up until you Ctrl-C it. Use sanity_check.sh and +# run_bench_single.sh in a separate shell to exercise / benchmark it. +# bench_mimo_v2_5.sh calls this script under the hood for its one-shot +# launch + bench + teardown flow. +# +# Recipe: TP=64, moe_tp=1/moe_ep=64, BS=32, continuous batching + bucketing. +# moe_tp=1/moe_ep=64 keeps each expert's weights and blockwise FP8 scales +# intact on a single rank (4 experts/rank), avoiding the per-rank scale +# collapse that comes from moe_tp=64 when intermediate=2048 is TP-sharded +# below the 128-row scale block boundary. +# +# NxDI's TKG path refuses Expert Parallelism with BS < num_experts/top_k +# (256 / 8 = 32), so BS=32 is the smallest working batch size on the FP8 +# path. BS=1 single-stream latency is not currently supported on V2.5 FP8. + +set -e + +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate + +MODEL_PATH="${MIMO_V2_5_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Neuron-FP8}" +PORT="${PORT:-8000}" + +# Contrib package src. vllm-neuron's registration hook reads these env vars +# to plug NeuronMiMoV2ForCausalLM into NxDI's MODEL_TYPES table. +: "${NXDI_CONTRIB_MIMO_V2_5_SRC:=$(cd "$(dirname "$0")/.." && pwd)/src}" +export NXDI_CONTRIB_MIMO_V2_5_SRC +# vLLM 0.16's builtin arch validator knows MiMoV2FlashForCausalLM but not +# MiMoV2ForCausalLM. Preprocess rewrites the checkpoint's config.json +# architectures to the Flash name, and we reuse the Flash registration +# key in vllm-neuron (MODEL_TYPES['mimov2flash']). The modeling module +# (modeling_mimo_v2) and class (NeuronMiMoV2ForCausalLM) are shared. +export NXDI_CONTRIB_MIMO_V2_FLASH_SRC="$NXDI_CONTRIB_MIMO_V2_5_SRC" + +# Persistent compile-artifact location (NEFF + per-rank sharded weights). +# Setting this overrides vLLM's fallback of /neuron-compiled-artifacts//. +: "${NEURON_COMPILED_ARTIFACTS:=/opt/dlami/nvme/compiled/mimo_v2_5_bs32_moetp1_ep64_fp8_vllm}" +export NEURON_COMPILED_ARTIFACTS +# NxDI HLO/NEFF staging directory, pinned to persistent storage so it +# survives the nightly Trn2 reboot and a unique per-config subdir. +: "${BASE_COMPILE_WORK_DIR:=/opt/dlami/nvme/tmp/nxd_model/$(basename "$NEURON_COMPILED_ARTIFACTS")}" +export BASE_COMPILE_WORK_DIR +mkdir -p "$BASE_COMPILE_WORK_DIR" + +# First-time compile of V2.5's 256-expert MoE takes ~30 min (HLO + shard). +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-7200}" + +echo "==========================================" +echo "Starting MiMo-V2.5 FP8 vLLM server" +echo "==========================================" +echo " Model path: $MODEL_PATH" +echo " Port: $PORT" +echo " Compiled artifacts: $NEURON_COMPILED_ARTIFACTS" +echo " Compile work dir: $BASE_COMPILE_WORK_DIR" +echo " NXDI_CONTRIB_MIMO_V2_5_SRC: $NXDI_CONTRIB_MIMO_V2_5_SRC" +echo "" + +exec python3 -m vllm.entrypoints.openai.api_server \ + --model "$MODEL_PATH" \ + --tokenizer "$MODEL_PATH" \ + --tensor-parallel-size 64 \ + --max-model-len 1024 \ + --max-num-seqs 32 \ + --no-enable-chunked-prefill \ + --no-enable-prefix-caching \ + --port "$PORT" \ + --trust_remote_code \ + --additional-config '{ + "override_neuron_config": { + "tp_degree": 64, + "logical_nc_config": 2, + "fused_qkv": false, + "sequence_parallel_enabled": false, + "glu_mlp": true, + "normalize_top_k_affinities": true, + "save_sharded_checkpoint": true, + "router_config": {"act_fn": "sigmoid", "dtype": "float32"}, + "quantized": true, + "quantized_checkpoints_path": "'"$MODEL_PATH"'", + "quantization_dtype": "f8e4m3", + "quantization_type": "blockwise_symmetric", + "quantization_block_axis": [1, 2], + "quantization_block_size": [128, 128], + "modules_to_not_convert": ["embed_tokens", "lm_head", "norm", "router", "o_proj"], + "blockwise_matmul_config": {"use_shard_on_block_dynamic_while": true, "block_sharding_strategy": "PING_PONG"}, + "moe_tp_degree": 1, + "moe_ep_degree": 64, + "batch_size": 32, + "ctx_batch_size": 1, + "tkg_batch_size": 32, + "max_context_length": 1024, + "seq_len": 1024, + "is_continuous_batching": true, + "enable_bucketing": true, + "context_encoding_buckets": [1024], + "token_generation_buckets": [1024], + "async_mode": true, + "on_device_sampling_config": { + "do_sample": true, "temperature": 0.6, "top_k": 20, "top_p": 0.95 + } + } + }' From 1ef5ef4f01ef77e823c2ef1979e64b4efd81caf4 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 14:33:57 +0800 Subject: [PATCH 15/27] README: record actual bench_mimo_v2_5 results (FP8 BS=32, c=1/16/32) Replace the "numbers pending" placeholder with the real vLLM serving numbers from trn2.48xlarge: output throughput, TPOT/TTFT medians and P99, plus a short analysis note explaining the 58 ms ITL floor (cost of one BS=32 TKG NEFF forward), the 576 tok/s peak at c=32, and why TPOT and TTFT degrade with concurrency under `enable_chunked_prefill=false`. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5/README.md | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/contrib/models/MiMo-V2.5/README.md b/contrib/models/MiMo-V2.5/README.md index 5a3ff8ce..88281a47 100644 --- a/contrib/models/MiMo-V2.5/README.md +++ b/contrib/models/MiMo-V2.5/README.md @@ -392,7 +392,20 @@ The patch is applied to vllm-neuron 0.5.0 and: ## Performance -> Benchmark numbers will be added once a stable bench run completes on the FP8 recipe. Preliminary single-stream sanity test produces fluent MiMo self-introduction output on the recipe below (`moe_tp=1, moe_ep=64, batch_size=32`). +### vLLM Serving (trn2.48xlarge, FP8, BS=32, TP=64 / moe_ep=64, CB + bucketing) + +Input/output: 900 / 90 tokens (random dataset). Recipe is the one `bench_mimo_v2_5.sh` drives; 16 prompts at c=1 and 128 prompts at c=16/c=32. + +| Concurrency | Output throughput (tok/s) | Total throughput (tok/s) | TPOT median (ms) | TTFT median (ms) | TTFT P99 (ms) | +|---|---|---|---|---|---| +| 1 | 15.88 | 174.16 | 58.28 | 485 | 485 | +| 16 | 113.92 | 1251.14 | 130.85 | 863 | 6371 | +| 32 | 147.39 | 1618.81 | 190.48 | 1798 | 13281 | + +Observations: +- **Median ITL stays at ~58 ms across all three concurrency levels** — that's the cost of one BS=32 TKG NEFF forward, which runs at fixed shape regardless of how many slots are actually occupied. +- **Peak output throughput at c=32 is 576 tok/s**, close to the theoretical `32 / 0.058 ≈ 552` ceiling. +- **TPOT and TTFT grow with concurrency** because `enable_chunked_prefill=false`: each new request's context-encoding pass (900 tokens) preempts TKG for a few hundred ms, and the higher the concurrency the more frequently that happens. > **Compile time:** the first MiMo-V2.5 compile on SDK 2.29 is ~30 minutes (TKG + CE HLO compilation, weight layout optimization, then `shard_checkpoint` for 64 ranks which dominates at ~27 minutes). Subsequent runs with the same `override_neuron_config` hit the neuronx-cc cache and the NEFF loads in ~1 minute. `save_sharded_checkpoint=true` persists per-rank FP8 shards under `/weights/`, letting future `load()` calls skip the `shard_checkpoint` pass entirely. From d47d9e80a1472dc8abb6da6ac74e4c1b224e7ef7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E7=8E=8B=E9=B9=A4=E7=94=B7?= Date: Tue, 21 Jul 2026 11:45:56 +0000 Subject: [PATCH 16/27] [contrib] MiMo-V2.5: fix decode sliding-window mask to use per-batch-slot position The SWA-layer decode mask used position_ids[0, 0] (batch slot 0) broadcast to the whole batch. That is only correct when every slot is at the same decode step (static/synchronous batch, e.g. the NxDI smoke path). Under vLLM continuous batching, requests occupy different decode positions, so a single broadcast window mis-clips the other slots KV. Fix: use position_ids[:, 0] and build a per-slot [bsz, 1, 1, kv_seq_len] mask. Verified equivalent to the old code for synchronous batches (numerically identical window) and correct per-slot for asynchronous ones. This is a real correctness fix but does NOT by itself resolve the issue-#31 garble: after this change, req#1 is still coherent and req#2+ are still garbled, because the dominant cause is stale KV in the full-attention layers (vLLM never resets the NxDI per-batch KV state that the smoke path resets each generate(), and MiMos full-attention decode masks prior KV only via position_ids). Tracking that separately. Co-Authored-By: Claude Opus 4.8 --- .../models/MiMo-V2.5/src/modeling_mimo_v2.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py index 947dfe2c..e0349404 100644 --- a/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py +++ b/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py @@ -691,13 +691,22 @@ def forward( else: prior_scores = prior_scores + attention_mask - # Apply sliding window mask for SWA layers + # Apply sliding window mask for SWA layers. + # NOTE: build the window per batch slot using each request's own + # decode position. The old code used position_ids[0, 0] (slot 0) + # broadcast to the whole batch, which is only correct when every + # slot is at the same decode step (static/synchronous batch, e.g. + # the smoke path). Under vLLM continuous batching, requests are at + # different positions, so a single broadcast window mis-clips the + # other slots' KV and produces garbled output on 2nd+ requests + # (the "first request ok, rest garbled" bug). Use per-slot positions. if self.is_sliding_window and self.sliding_window_size is not None and position_ids is not None: kv_seq_len = prior_scores.size(-1) - current_pos = position_ids[0, 0] - pos_indices = torch.arange(kv_seq_len, device=prior_scores.device) - sliding_mask = pos_indices >= (current_pos - self.sliding_window_size + 1) - sliding_mask = sliding_mask[None, None, None, :] + # position_ids: [bsz, q_len] -> per-slot current position [bsz, 1] + current_pos = position_ids[:, 0].reshape(-1, 1) # [bsz, 1] + pos_indices = torch.arange(kv_seq_len, device=prior_scores.device)[None, :] # [1, kv_seq_len] + sliding_mask = pos_indices >= (current_pos - self.sliding_window_size + 1) # [bsz, kv_seq_len] + sliding_mask = sliding_mask[:, None, None, :] # [bsz, 1, 1, kv_seq_len] prior_scores = prior_scores.masked_fill(~sliding_mask, float('-inf')) prior_scores = prior_scores.to(torch.float32) From b2488f63300b67e3e0094bb7b4f519b723536a4d Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 21 Jul 2026 21:44:28 +0800 Subject: [PATCH 17/27] Add H100 GPU baseline for MiMo-V2.5 (SGLang single-node) Mirrors the MiMo-V2.5-Pro perf_test/h100 layout, adapted for single-node 8xH100 (V2.5 FP8 ~295 GB fits on 640 GB). SGLang is the working baseline via --tp 8 --attention-context-parallel-size 2 (effective attn TP = 4, required by the TP=4-interleaved fused qkv); no DeepEP, no spec decode, for apples-to-apples comparison with the Trn2 numbers. vLLM cannot serve V2.5 FP8: both 0.25.1 and nightly crash in _shard_fp8_qkv_proj on the sliding-window layers (1856 vs 1792 scale-row mismatch, since V2.5 has 4 full-attn KV heads and no 8-GPU TP shape avoids the buggy g>1 path). run_vllm_h100.sh keeps the reference command and documents the bug. Measured (900/90, no spec decode): c=1 125.7 / c=16 792.7 / c=32 1570.8 out tok/s -- ~7-11x the Trn2 vLLM-Neuron baseline. Co-Authored-By: Claude Opus 4.8 --- .../models/MiMo-V2.5/perf_test/h100/README.md | 129 ++++++++++++++++++ .../MiMo-V2.5/perf_test/h100/bench_all.sh | 25 ++++ .../perf_test/h100/results/sglang_cp2_c1.txt | 43 ++++++ .../perf_test/h100/results/sglang_cp2_c16.txt | 43 ++++++ .../perf_test/h100/results/sglang_cp2_c32.txt | 43 ++++++ .../perf_test/h100/run_bench_single.sh | 78 +++++++++++ .../perf_test/h100/run_sglang_h100.sh | 114 ++++++++++++++++ .../MiMo-V2.5/perf_test/h100/run_vllm_h100.sh | 74 ++++++++++ 8 files changed, 549 insertions(+) create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/README.md create mode 100755 contrib/models/MiMo-V2.5/perf_test/h100/bench_all.sh create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_cp2_c1.txt create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_cp2_c16.txt create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_cp2_c32.txt create mode 100755 contrib/models/MiMo-V2.5/perf_test/h100/run_bench_single.sh create mode 100755 contrib/models/MiMo-V2.5/perf_test/h100/run_sglang_h100.sh create mode 100755 contrib/models/MiMo-V2.5/perf_test/h100/run_vllm_h100.sh diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/README.md b/contrib/models/MiMo-V2.5/perf_test/h100/README.md new file mode 100644 index 00000000..7be58f68 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/README.md @@ -0,0 +1,129 @@ +# H100 GPU baseline for MiMo-V2.5 + +Scripts to serve the **official HuggingFace OCP-FP8 checkpoint** +(`XiaomiMiMo/MiMo-V2.5`) on a **single** 8xH100-80GB node, for a cross-platform +throughput comparison against the Trn2 Neuron port. + +## Single node (unlike MiMo-V2.5-Pro) + +V2.5's FP8 weights are ~295 GB, which fits on one 8xH100 node (640 GB), so no +cross-node NCCL / EFA is needed (the Pro variant is ~963 GB and requires 2 +nodes). We still reuse the `sglang-efa` / `vllm-efa` images built for the Pro +work (strict supersets of the stock images); EFA just goes unused here. + +## TL;DR — use SGLang; vLLM is broken for this model + +| Backend | Status | +|---|---| +| **SGLang** (`run_sglang_h100.sh`) | ✅ **Works.** `--tp 8 --attention-context-parallel-size 2`, no DeepEP, no spec decode. | +| **vLLM** (`run_vllm_h100.sh`) | ❌ **Broken** (upstream bug, 0.25.1 **and** nightly). Kept only to document the reference command + reproduce the bug. | + +### Why vLLM fails + +vLLM crashes loading the fused `qkv_proj` on the **sliding-window** layers: + +``` +RuntimeError: The size of tensor a (1856) must match the size of tensor b (1792) + at non-singleton dimension 0 + (vllm/model_executor/models/mimo_v2.py::_shard_fp8_qkv_proj) +``` + +A SWA-layer fused-qkv group is `(64/8)*192 + 192 + 128 = 1856` rows = **14.5** +blocks of the 128-row FP8 scale, so it is **not 128-aligned**. vLLM's loader +slices the scale per KV-head group (`scale_rows_per_group = 116 // 8 = 14` -> +1792 rows) and multiplies it against the 1856-row weight group. Because +MiMo-V2.5 has only **4 full-attention KV heads** (vs 8 for the Pro), no 8-GPU TP +config avoids the buggy path: + +- `TP=8` — can't start: "TP size must evenly split the number of KV heads" (4). +- `TP=4` / `TP=2` / `DP=2 x TP=4` — each rank owns 2 SWA KV heads -> the `g>1` + de-interleave path -> the 1856-vs-1792 crash. + +(The Pro variant never hits this: with 8 KV heads and TP=8 each rank gets 1 KV +head and takes the trivial `kv_heads_per_rank == 1` fast path.) A shape-only +monkeypatch makes the server start but produces gibberish, so it is not a real +fix. This should be reported/fixed upstream in vLLM. + +### Why SGLang uses `--attention-context-parallel-size 2` + +SGLang also requires the *effective* attention TP size to be exactly 4 (the +qkv is TP=4-interleaved), so plain `--tp 8` is rejected. Two 8-GPU shapes give +effective attn TP = 4: + +- **DP=2 + `--enable-dp-attention`** (the model-card reference command): but + single-node it **deadlocks** the shared-MoE / lm-head collective when a + request occupies only one DP group (idle group never launches its matching + forward -> 300 s scheduler watchdog -> crash). The reference relies on + `--moe-a2a-backend deepep` to change that collective, but **DeepEP needs a + working nvshmem RDMA transport, which this plain P5 box's NICs (`rdmapXXs0`, + not `mlx5`) do not provide** — IBGDA init fails and the forward hangs (tested + `low_latency` + `normal` + `NVSHMEM_REMOTE_TRANSPORT=none` + `IBGDA=0`). DeepEP + *does* work on `p5en.48xlarge` (mlx5 EFA + a custom Mooncake/nvshmem image); + see `xiaomi_datalab/mimo_v25`. +- **`--attention-context-parallel-size 2`**: effective attn TP = tp/cp = 4, MoE + shards over TP=8, no DP idle-group problem, no DeepEP. **Works out of the box. + This is the default.** + +## Usage + +Download the checkpoint (~295 GB; uses the `/opt/pytorch` env on the P5 box): + +```bash +/opt/pytorch/bin/hf download XiaomiMiMo/MiMo-V2.5 \ + --local-dir /opt/dlami/nvme/models/MiMo-V2.5 --max-workers 16 +``` + +Serve + bench (SGLang, the working path): + +```bash +# Terminal 1: launch (first start ~4 min: weight load + DeepGEMM warmup + CUDA +# graph capture). Wait for "The server is fired up and ready to roll!". +bash run_sglang_h100.sh + +# Terminal 2: bench at c=1/16/32 (MiMo-V2.5 bs=32, so cap concurrency at 32). +# The bench client runs inside the vllm-efa container (has `vllm bench serve`). +bash bench_all.sh 30000 sglang_cp2 +``` + +`bench_all.sh` is a thin loop over `run_bench_single.sh` (c=1 x16 prompts, +c=16/c=32 x96 prompts), the same client and 900/90 token random dataset as the +Trn2 table. Per-run logs land in `results/`. + +For a single ad-hoc run: + +```bash +docker run --rm --network host -v /opt/dlami/nvme/models:/wk -v $PWD:/sc \ + --entrypoint bash sglang-efa:latest -c ' + export SERVED_MODEL_NAME=MiMo-V2.5 TOKENIZER_PATH=/wk/MiMo-V2.5 PORT=30000 \ + CONFIG_NAME=sglang_cp2 RESULTS_DIR=/wk/bench_results/mimo_v2_5_h100 + CONCURRENCY=32 NUM_PROMPTS=96 bash /sc/run_bench_single.sh' +``` + +## Performance (SGLang, single-node 8xH100, FP8, TP=8 + attn-CP=2) + +Input/output: 900 / 90 tokens (random dataset), no speculative decoding — same +recipe and shape as the Trn2 table in the main README, for direct comparison. +Raw logs in `results/sglang_cp2_c{1,16,32}.txt`. + +| Concurrency | Output throughput (tok/s) | Total throughput (tok/s) | TPOT median (ms) | TTFT median (ms) | TTFT P99 (ms) | +|---|---|---|---|---|---| +| 1 | 125.69 | 1378.59 | 6.93 | 95 | 157 | +| 16 | 792.68 | 8699.94 | 17.66 | 186 | 726 | +| 32 | 1570.83 | 17240.36 | 17.78 | 237 | 293 | + +### H100 (SGLang) vs Trn2 (vLLM-Neuron) — same 900/90 shape + +| Concurrency | H100 out tok/s | Trn2 out tok/s | H100 / Trn2 | +|---|---:|---:|---:| +| 1 | 125.69 | 15.88 | **7.9×** | +| 16 | 792.68 | 113.92 | **7.0×** | +| 32 | 1570.83 | 147.39 | **10.7×** | + +(Trn2 numbers from the main README's Performance section: BS=32, TP=64 / +moe_ep=64, CB + bucketing.) The H100 gap is largest at c=32 because SGLang's +median ITL stays ~13–18 ms across concurrency while the Trn2 BS=32 TKG NEFF runs +at a fixed ~58 ms/token regardless of occupancy. + +> **Note:** these are the *language-only, no-spec-decode* baselines for +> apples-to-apples comparison with Trn2. Turning on EAGLE (`SPEC=1`) would raise +> H100 decode throughput further; the model card's full recipe also enables it. diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/bench_all.sh b/contrib/models/MiMo-V2.5/perf_test/h100/bench_all.sh new file mode 100755 index 00000000..5bb433d2 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/bench_all.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# Bench a running MiMo-V2.5 H100 server at c=1/16/32 (bs=32 -> cap at 32). +# Runs the bench client inside the vllm-efa container (has `vllm bench serve`), +# reusing run_bench_single.sh. Prints the key metrics; full logs go to +# /opt/dlami/nvme/models/bench_results/mimo_v2_5_h100/_c.txt. +# +# Usage: bash bench_all.sh +# bash bench_all.sh 30000 sglang_cp2 # SGLang +# bash bench_all.sh 8000 vllm # vLLM (if/when it works) +set -e +PORT="${1:-30000}" +CONFIG_NAME="${2:-sglang_cp2}" +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" + +for C in 1 16 32; do + NP=96; [ "$C" = "1" ] && NP=16 + echo "===== $CONFIG_NAME c=$C np=$NP =====" + docker run --rm --network host \ + -v /opt/dlami/nvme/models:/wk -v "$SCRIPT_DIR":/sc \ + --entrypoint bash vllm-efa:latest -c " + export SERVED_MODEL_NAME=MiMo-V2.5 TOKENIZER_PATH=/wk/MiMo-V2.5 PORT=$PORT \ + RESULTS_DIR=/wk/bench_results/mimo_v2_5_h100 CONFIG_NAME=$CONFIG_NAME + CONCURRENCY=$C NUM_PROMPTS=$NP bash /sc/run_bench_single.sh" 2>&1 | \ + grep -E "Successful requests|Maximum request concurrency:|Benchmark duration|Total input tokens|Total generated tokens|Output token throughput|Total token throughput|Median TTFT|P99 TTFT|Mean TPOT|Median TPOT|Median ITL" +done diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_cp2_c1.txt b/contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_cp2_c1.txt new file mode 100644 index 00000000..0fd5c28b --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_cp2_c1.txt @@ -0,0 +1,43 @@ +INFO 07-21 12:53:28 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 12:53:28 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 12:53:29.997000 9 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 12:53:31 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 12:53:31 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +[transformers] You are using a model of type `mimo_v2` to instantiate a model of type ``. This may be expected if you are loading a checkpoint that shares a subset of the architecture (e.g., loading a `sam2_video` checkpoint into `Sam2Model`), but is otherwise not supported and can yield errors. Please verify that the checkpoint is compatible with the model you are instantiating. +[transformers] PreTrainedConfig got `key=rope_scaling` in kwargs but hasn't set it as attribute. For RoPE standardization you need to set `self.rope_parameters` in model's config. +Namespace(subparser='bench', bench_type='serve', dispatch_function=, trust_remote_code=False, seed=0, num_prompts=16, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=30000, endpoint='/v1/completions', header=None, max_concurrency=1, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-bba75904-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 12:53:32 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 1 + 0%| | 0/16 [00:00, trust_remote_code=False, seed=0, num_prompts=96, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=30000, endpoint='/v1/completions', header=None, max_concurrency=16, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-b1e6ada9-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 12:54:06 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 16 + 0%| | 0/96 [00:00, trust_remote_code=False, seed=0, num_prompts=96, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=30000, endpoint='/v1/completions', header=None, max_concurrency=32, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-ebb2ca59-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 12:54:40 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 32 + 0%| | 0/96 [00:00 cap at 32) +# NUM_PROMPTS --num-prompts (default 16) +# INPUT_LEN --random-input-len (default 900, matches Trn2 table) +# OUTPUT_LEN --random-output-len (default 90, matches Trn2 table) +# RANGE_RATIO --random-range-ratio (default 0.03) +# CONFIG_NAME used in the output filename (default vllm) +# RESULTS_DIR where to dump per-run log (default /opt/dlami/nvme/logs/bench_results/mimo_v2_5_h100) + +set -e + +NEURON_VENV="/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate" +[ -f "$NEURON_VENV" ] && source "$NEURON_VENV" + +MODEL_PATH="${MIMO_V2_5_PATH:-/models/MiMo-V2.5}" +SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-MiMo-V2.5}" +TOKENIZER_PATH="${TOKENIZER_PATH:-$MODEL_PATH}" +PORT="${PORT:-8000}" +CONCURRENCY="${CONCURRENCY:-1}" +NUM_PROMPTS="${NUM_PROMPTS:-16}" +INPUT_LEN="${INPUT_LEN:-900}" +OUTPUT_LEN="${OUTPUT_LEN:-90}" +RANGE_RATIO="${RANGE_RATIO:-0.03}" +CONFIG_NAME="${CONFIG_NAME:-vllm}" +RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/logs/bench_results/mimo_v2_5_h100}" + +mkdir -p "$RESULTS_DIR" + +echo "==========================================" +echo "MiMo-V2.5 H100 single-run benchmark" +echo "==========================================" +echo " Served model: $SERVED_MODEL_NAME" +echo " Tokenizer: $TOKENIZER_PATH" +echo " Port: $PORT" +echo " Config: $CONFIG_NAME" +echo " Concurrency: $CONCURRENCY Prompts: $NUM_PROMPTS" +echo " Input len: $INPUT_LEN Output len: $OUTPUT_LEN" +echo " Results: $RESULTS_DIR/${CONFIG_NAME}_c${CONCURRENCY}.txt" +echo "" + +if ! curl -sf "http://localhost:$PORT/health" > /dev/null; then + echo "ERROR: server is not responding on http://localhost:$PORT" + exit 1 +fi + +vllm bench serve \ + --backend vllm \ + --host localhost \ + --port "$PORT" \ + --model "$SERVED_MODEL_NAME" \ + --tokenizer "$TOKENIZER_PATH" \ + --endpoint /v1/completions \ + --dataset-name random \ + --num-prompts "$NUM_PROMPTS" \ + --random-input-len "$INPUT_LEN" \ + --random-output-len "$OUTPUT_LEN" \ + --random-range-ratio "$RANGE_RATIO" \ + --max-concurrency "$CONCURRENCY" \ + 2>&1 | tee "$RESULTS_DIR/${CONFIG_NAME}_c${CONCURRENCY}.txt" + +echo "" +echo "Saved to: $RESULTS_DIR/${CONFIG_NAME}_c${CONCURRENCY}.txt" diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/run_sglang_h100.sh b/contrib/models/MiMo-V2.5/perf_test/h100/run_sglang_h100.sh new file mode 100755 index 00000000..7137ff6b --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/run_sglang_h100.sh @@ -0,0 +1,114 @@ +#!/bin/bash +# Single-node 8xH100 SGLang server for MiMo-V2.5 (official HF OCP-FP8 checkpoint). +# This is the WORKING H100 baseline (vLLM cannot serve V2.5 FP8 -- see +# run_vllm_h100.sh). V2.5 FP8 (~295 GB) fits on one 8xH100 node (640 GB). +# +# Reuses sglang-efa:latest (superset of stock lmsysorg/sglang). +# +# ---- Why --attention-context-parallel-size 2 (not plain TP=8 or DP=2) ---- +# MiMoV2ForCausalLM's fused qkv_proj is TP=4-interleaved, so SGLang requires the +# *effective* attention TP size to be exactly 4 (plain --tp 8 is rejected with +# "requires effective attention TP size 4"). Two 8-GPU shapes give effective +# attn TP = 4: +# - DP=2 + --enable-dp-attention (the model-card reference): but single-node it +# deadlocks the shared-MoE / lm-head collective when a request occupies only +# one DP group -- the idle group never launches its matching forward, so the +# collective hangs (300s scheduler watchdog -> crash). The reference relies on +# --moe-a2a-backend deepep to change that collective, but DeepEP needs a +# working nvshmem RDMA transport which this box's NICs (rdmapXXs0, not mlx5) +# don't provide -- IBGDA init fails and the forward pass hangs (tested +# low_latency + normal + REMOTE_TRANSPORT=none + IBGDA=0). DeepEP does work on +# p5en.48xlarge (mlx5 EFA + a custom Mooncake/nvshmem image); see +# xiaomi_datalab/mimo_v25. +# - ATTN_CP=2 (attention context parallel): effective attn TP = tp/cp = 4, MoE +# shards over TP=8, no DP idle-group problem, no DeepEP. Works out of the box +# on plain P5. THIS IS THE DEFAULT. +# +# ---- Speculative decoding (MTP/EAGLE) ---- +# The checkpoint ships model_mtp.safetensors, and the reference command enables +# EAGLE. We LEAVE IT OFF (the Trn2 port has no spec decoding, so this is the +# apples-to-apples baseline). SPEC=1 turns on the reference EAGLE flags. +# +# Usage: +# bash run_sglang_h100.sh # working default (attn-cp=2, no spec) +# SPEC=1 bash run_sglang_h100.sh # + EAGLE speculative decoding +# DP=2 DEEPEP=1 bash run_sglang_h100.sh # model-card DP-attention + DeepEP +# # (only works where nvshmem/DeepEP does) +set -e + +MODEL_DIR="${MODEL_DIR:-/opt/dlami/nvme/models/MiMo-V2.5}" +PORT="${PORT:-30000}" +TP="${TP:-8}" +DP="${DP:-1}" # 1 => use ATTN_CP; >1 => DP-attention +ATTN_CP="${ATTN_CP:-2}" # effective attn TP = TP/ATTN_CP = 4 +MEM_FRAC="${MEM_FRAC:-0.9}" # weights ~295GB; rest for KV cache +IMAGE="${IMAGE:-sglang-efa:latest}" +CACHE_DIR="${CACHE_DIR:-/opt/dlami/nvme/sglang_cache}" +CTR_MODEL="/models/MiMo-V2.5" +mkdir -p "$CACHE_DIR" + +if [ ! -f "$MODEL_DIR/config.json" ]; then + echo "ERROR: model not found at $MODEL_DIR" >&2 + exit 1 +fi + +# Parallelism shape: DP-attention (reference) if DP>1, else attention-CP. +DP_ARGS=() +if [ "$DP" -gt 1 ]; then + DP_ARGS=( --dp "$DP" --enable-dp-attention --enable-dp-lm-head ) + echo " DP-attention: ON (dp=$DP, effective attn TP = $((TP/DP)))" +else + DP_ARGS=( --attention-context-parallel-size "$ATTN_CP" ) + echo " Attention context parallel: $ATTN_CP (effective attn TP = $((TP/ATTN_CP)))" +fi + +# DeepEP MoE all-to-all (needs a working nvshmem RDMA transport; see header). +A2A_ARGS=() +DEEPEP_ENV=() +if [ "${DEEPEP:-0}" = "1" ]; then + A2A_ARGS=( --moe-a2a-backend deepep --deepep-mode "${DEEPEP_MODE:-auto}" ) + DEEPEP_ENV=( -e SGLANG_DEEPEP_NUM_MAX_DISPATCH_TOKENS_PER_RANK=256 ) + echo " MoE a2a backend: DeepEP ON (mode=${DEEPEP_MODE:-auto})" +fi + +# EAGLE speculative decoding (off by default for parity with Trn2). +SPEC_ARGS=() +if [ "${SPEC:-0}" = "1" ]; then + SPEC_ARGS=( + --speculative-algorithm EAGLE + --speculative-num-steps 3 + --speculative-eagle-topk 1 + --speculative-num-draft-tokens 4 + --enable-multi-layer-eagle + ) + echo " Speculative decoding: EAGLE ON" +fi + +echo "==========================================" +echo "MiMo-V2.5 SGLang (single-node 8xH100, Docker)" +echo " TP=$TP DP=$DP Port: $PORT mem-frac: $MEM_FRAC" +echo "==========================================" + +exec docker run --rm --gpus all \ + --network host --privileged --ipc=host --shm-size=32g \ + -v "${MODEL_DIR}:${CTR_MODEL}:ro" \ + -v "${CACHE_DIR}:/root/.cache" \ + -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + "${DEEPEP_ENV[@]}" \ + "$IMAGE" \ + python3 -m sglang.launch_server \ + --model-path "$CTR_MODEL" \ + --served-model-name MiMo-V2.5 \ + --trust-remote-code \ + --tp "$TP" \ + "${DP_ARGS[@]}" \ + --moe-dense-tp-size 1 \ + --mem-fraction-static "$MEM_FRAC" \ + --max-running-requests 128 \ + --chunked-prefill-size 16384 \ + "${A2A_ARGS[@]}" \ + "${SPEC_ARGS[@]}" \ + --reasoning-parser mimo \ + --tool-call-parser mimo \ + --host 0.0.0.0 \ + --port "$PORT" diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/run_vllm_h100.sh b/contrib/models/MiMo-V2.5/perf_test/h100/run_vllm_h100.sh new file mode 100755 index 00000000..3e40715e --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/run_vllm_h100.sh @@ -0,0 +1,74 @@ +#!/bin/bash +# Single-node 8xH100 vLLM server for MiMo-V2.5 (official HF OCP-FP8 checkpoint). +# +# ============================ KNOWN BROKEN ============================ +# vLLM CANNOT currently serve MiMo-V2.5 FP8. Both vllm 0.25.1 and nightly +# crash while loading the fused qkv_proj on the sliding-window layers: +# +# RuntimeError: The size of tensor a (1856) must match the size of +# tensor b (1792) at non-singleton dimension 0 +# (vllm/model_executor/models/mimo_v2.py::_shard_fp8_qkv_proj) +# +# Root cause (upstream bug): a SWA layer's fused-qkv group is 1856 rows = +# 14.5 blocks of the 128-row FP8 scale, so it is NOT 128-aligned. The loader +# slices the scale per KV-head group (scale_rows_per_group = 116 // 8 = 14 -> +# 1792 rows) and multiplies it against the 1856-row weight group -> shape +# mismatch. This fires at EVERY 8-GPU shape: +# - TP=8: can't even start, "TP size must evenly split the 4 KV heads". +# - TP=4 / TP=2 (incl. DP=2 x TP=4): each rank gets 2 SWA KV heads -> the +# g>1 de-interleave path -> the 1856-vs-1792 crash. +# There is no vLLM TP config that both (a) splits the 4 full-attn KV heads and +# (b) gives 1 SWA KV head per rank, so the buggy path is unavoidable. +# +# apply_mimo_fp8_patch.py makes the server START (whole-tensor dequant), but the +# output is gibberish -- the re-quantization path is not numerically correct -- +# so it is NOT a usable fix, only a diagnostic. Use SGLang (run_sglang_h100.sh) +# for the H100 baseline. This script is kept to document the reference command +# and reproduce the bug. +# ===================================================================== +# +# The command below is the reference vLLM command from the model card +# (TP=8 + expert parallel, mimo parsers), which is what you WOULD run once the +# upstream loader is fixed. +set -e + +MODEL_DIR="${MODEL_DIR:-/opt/dlami/nvme/models/MiMo-V2.5}" +PORT="${PORT:-8000}" +TP="${TP:-8}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-32}" # MiMo-V2.5 bs=32 +IMAGE="${IMAGE:-vllm-efa:latest}" +CACHE_DIR="${CACHE_DIR:-/opt/dlami/nvme/vllm_cache}" +CTR_MODEL="/models/MiMo-V2.5" +mkdir -p "$CACHE_DIR" + +if [ ! -f "$MODEL_DIR/config.json" ]; then + echo "ERROR: model not found at $MODEL_DIR" >&2 + exit 1 +fi + +echo "==========================================" +echo "MiMo-V2.5 vLLM (single-node TP=$TP + EP, Docker) -- EXPECTED TO FAIL" +echo " Port: $PORT max-model-len: $MAX_MODEL_LEN max-num-seqs: $MAX_NUM_SEQS" +echo "==========================================" + +exec docker run --rm --gpus all \ + --network host --privileged --ipc=host --shm-size=32g \ + -v "${MODEL_DIR}:${CTR_MODEL}:ro" \ + -v "${CACHE_DIR}:/root/.cache" \ + -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + "$IMAGE" "$CTR_MODEL" \ + --served-model-name MiMo-V2.5 \ + --trust-remote-code \ + --generation-config vllm \ + --enable-expert-parallel \ + --tensor-parallel-size "$TP" \ + --max-model-len "$MAX_MODEL_LEN" \ + --max-num-seqs "$MAX_NUM_SEQS" \ + --no-enable-prefix-caching \ + --no-enable-chunked-prefill \ + --host 0.0.0.0 \ + --port "$PORT" \ + --tool-call-parser mimo \ + --enable-auto-tool-choice \ + --reasoning-parser mimo From 8a64b5f8c6f68e8d5fcd26f0a6afa83ad2b386d7 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 21 Jul 2026 22:47:00 +0800 Subject: [PATCH 18/27] Fix batched-decode garble: per-row causal upper bound on prior KV MiMo hand-rolls the token-generation attention as a decomposed prior(cached) + active(current) computation, reading the entire n_positions KV cache buffer as K_prior. The full-attention layers had no per-row upper bound and relied solely on the externally-supplied attention_mask to exclude invalid cache positions. Under vLLM batched continuous-batching decode, that mask can be width-collapsed to torch.max(position_ids) (model_base._infer_attention_mask), so any batch row whose decode position is below the batch-global max attends to stale KV physically left by a different request that previously occupied the reused cache slot/row. The result is deterministic garble from decode token 2 onward (byte-identical across identical prompts), triggered whenever more than one sequence is active in the batch. A single-active-sequence server happened to be coherent, which masked the bug. Add an explicit per-row causal bound (pos_indices < current_pos) on the prior scores for BOTH full-attention and sliding-window layers, so correctness no longer depends on the external mask width. The current token is scored separately via active_scores, and the KV cache is updated after all decoder layers run, so the strict-less-than bound introduces no double counting. Verified on V2.5 FP8 (vLLM, async_mode=false): single, 4x concurrent distinct, 3x concurrent identical, and post-concurrent sequential requests all decode coherently, where every one previously garbled. Co-Authored-By: Claude Opus 4.8 --- .../models/MiMo-V2.5/src/modeling_mimo_v2.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py index e0349404..3e13adc9 100644 --- a/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py +++ b/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py @@ -691,6 +691,36 @@ def forward( else: prior_scores = prior_scores + attention_mask + # Per-row causal UPPER bound on the prior (cached) KV. + # + # K_prior spans the whole n_positions cache buffer. Only positions + # strictly BEFORE this row's current decode position hold valid + # prior KV for THIS request; positions >= current_pos are either + # future slots or stale KV physically left by a *different* request + # that previously occupied this batch row / cache slot. MiMo hand- + # rolls the decode attention (it does not use the base flash-decode + # kernel, which would bound reads by per-row computed_context_lens), + # so without an explicit per-row bound the full-attention layers + # rely entirely on the externally-supplied attention_mask. Under + # vLLM batched decode that mask can be width-collapsed to + # torch.max(position_ids) (model_base._infer_attention_mask), so a + # row whose position < the batch-global max attends to stale KV and + # its output deterministically garbles from decode token 2 onward + # (byte-identical across identical prompts), keyed on having >1 + # active row. Bounding per-row here makes correctness independent of + # the external mask for BOTH full-attn and SWA layers. The active + # (current) token is scored separately via active_scores below, so + # the bound is strict-less-than current_pos. + if position_ids is not None: + kv_seq_len = prior_scores.size(-1) + current_pos = position_ids[:, 0].reshape(-1, 1) # [bsz, 1] + pos_indices = torch.arange( + kv_seq_len, device=prior_scores.device + )[None, :] # [1, kv_seq_len] + causal_mask = (pos_indices < current_pos) # [bsz, kv_seq_len] + causal_mask = causal_mask[:, None, None, :] # [bsz, 1, 1, kv_seq_len] + prior_scores = prior_scores.masked_fill(~causal_mask, float("-inf")) + # Apply sliding window mask for SWA layers. # NOTE: build the window per batch slot using each request's own # decode position. The old code used position_ids[0, 0] (slot 0) From 4e69aee260490b5a6f2352a0b315b1be3905815e Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 22 Jul 2026 06:33:20 +0800 Subject: [PATCH 19/27] h100/MiMo-V2.5: use cookbook DP=2 as SGLang default (no DeepEP needed) Correction to the initial H100 baseline: the SGLang cookbook command (--tp 8 --dp 2 --enable-dp-attention --enable-dp-lm-head --mm-enable-dp-encoder --mem-fraction-static 0.65) runs fine single-node with NO DeepEP. The earlier DP-attention hang was from using only --enable-dp-attention at mem-fraction 0.9 and omitting --enable-dp-lm-head / --mm-enable-dp-encoder; the full cookbook command works. DP=2 is now the default; attention-context-parallel (DP=1 ATTN_CP=2) is kept as a documented, slightly-faster-at-low-concurrency variant. Both verified to produce correct output. Added DP=2 bench results (sglang_dp2_c{1,16,32}) and updated the perf tables. DeepEP is optional (DEEPEP=1) and only for hosts with mlx5 EFA. Co-Authored-By: Claude Opus 4.8 --- .../models/MiMo-V2.5/perf_test/h100/README.md | 76 +++++++++++-------- .../MiMo-V2.5/perf_test/h100/bench_all.sh | 2 +- .../perf_test/h100/results/sglang_dp2_c1.txt | 43 +++++++++++ .../perf_test/h100/results/sglang_dp2_c16.txt | 43 +++++++++++ .../perf_test/h100/results/sglang_dp2_c32.txt | 43 +++++++++++ .../perf_test/h100/run_sglang_h100.sh | 67 ++++++++-------- 6 files changed, 209 insertions(+), 65 deletions(-) create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_dp2_c1.txt create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_dp2_c16.txt create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_dp2_c32.txt diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/README.md b/contrib/models/MiMo-V2.5/perf_test/h100/README.md index 7be58f68..8969e191 100644 --- a/contrib/models/MiMo-V2.5/perf_test/h100/README.md +++ b/contrib/models/MiMo-V2.5/perf_test/h100/README.md @@ -15,7 +15,7 @@ work (strict supersets of the stock images); EFA just goes unused here. | Backend | Status | |---|---| -| **SGLang** (`run_sglang_h100.sh`) | ✅ **Works.** `--tp 8 --attention-context-parallel-size 2`, no DeepEP, no spec decode. | +| **SGLang** (`run_sglang_h100.sh`) | ✅ **Works.** Cookbook `--tp 8 --dp 2 --enable-dp-attention` (no DeepEP, no spec decode). | | **vLLM** (`run_vllm_h100.sh`) | ❌ **Broken** (upstream bug, 0.25.1 **and** nightly). Kept only to document the reference command + reproduce the bug. | ### Why vLLM fails @@ -44,25 +44,30 @@ head and takes the trivial `kv_heads_per_rank == 1` fast path.) A shape-only monkeypatch makes the server start but produces gibberish, so it is not a real fix. This should be reported/fixed upstream in vLLM. -### Why SGLang uses `--attention-context-parallel-size 2` - -SGLang also requires the *effective* attention TP size to be exactly 4 (the -qkv is TP=4-interleaved), so plain `--tp 8` is rejected. Two 8-GPU shapes give -effective attn TP = 4: - -- **DP=2 + `--enable-dp-attention`** (the model-card reference command): but - single-node it **deadlocks** the shared-MoE / lm-head collective when a - request occupies only one DP group (idle group never launches its matching - forward -> 300 s scheduler watchdog -> crash). The reference relies on - `--moe-a2a-backend deepep` to change that collective, but **DeepEP needs a - working nvshmem RDMA transport, which this plain P5 box's NICs (`rdmapXXs0`, - not `mlx5`) do not provide** — IBGDA init fails and the forward hangs (tested - `low_latency` + `normal` + `NVSHMEM_REMOTE_TRANSPORT=none` + `IBGDA=0`). DeepEP - *does* work on `p5en.48xlarge` (mlx5 EFA + a custom Mooncake/nvshmem image); - see `xiaomi_datalab/mimo_v25`. -- **`--attention-context-parallel-size 2`**: effective attn TP = tp/cp = 4, MoE - shards over TP=8, no DP idle-group problem, no DeepEP. **Works out of the box. - This is the default.** +### SGLang parallelism: DP=2 (cookbook, default) — no DeepEP needed + +SGLang requires the *effective* attention TP size to be exactly 4 (the qkv is +TP=4-interleaved), so plain `--tp 8` is rejected. Two 8-GPU shapes give +effective attn TP = 4, and **both work single-node with no DeepEP**: + +- **DP=2 + DP-attention (default)** — the SGLang cookbook / model-card command: + `--tp 8 --dp 2 --enable-dp-attention --enable-dp-lm-head --mm-enable-dp-encoder + --mem-fraction-static 0.65`. Effective attn TP = tp/dp = 4. DeepEP is **not** + required — the three DP flags together (dp-attention + dp-lm-head + + dp-encoder) at mem-fraction 0.65 make the shared-MoE / lm-head collective work. + (An earlier attempt with only `--enable-dp-attention` and mem-fraction 0.9 hung + the collective; adding the other two DP flags + 0.65 fixes it — the full + cookbook command is what to use.) +- **`--attention-context-parallel-size 2`** (`DP=1 ATTN_CP=2`): effective attn + TP = tp/cp = 4, MoE shards over TP=8, no DP flags. Splits attention along the + sequence so a single request uses all 8 GPUs, measured **slightly faster at + low concurrency** (see table). Kept as a documented alternative. + +> **On DeepEP**: the cookbook *also* shows a DeepEP variant +> (`--moe-a2a-backend deepep`). DeepEP is not needed here, and on this plain P5 +> box it doesn't init anyway (its nvshmem RDMA/IBGDA transport wants `mlx5` NICs; +> this box has `rdmapXXs0` — set `DEEPEP=1` only on hosts where it works, e.g. +> `p5en.48xlarge` with mlx5 EFA, cf. `xiaomi_datalab/mimo_v25`). ## Usage @@ -82,7 +87,7 @@ bash run_sglang_h100.sh # Terminal 2: bench at c=1/16/32 (MiMo-V2.5 bs=32, so cap concurrency at 32). # The bench client runs inside the vllm-efa container (has `vllm bench serve`). -bash bench_all.sh 30000 sglang_cp2 +bash bench_all.sh 30000 sglang_dp2 ``` `bench_all.sh` is a thin loop over `run_bench_single.sh` (c=1 x16 prompts, @@ -95,15 +100,26 @@ For a single ad-hoc run: docker run --rm --network host -v /opt/dlami/nvme/models:/wk -v $PWD:/sc \ --entrypoint bash sglang-efa:latest -c ' export SERVED_MODEL_NAME=MiMo-V2.5 TOKENIZER_PATH=/wk/MiMo-V2.5 PORT=30000 \ - CONFIG_NAME=sglang_cp2 RESULTS_DIR=/wk/bench_results/mimo_v2_5_h100 + CONFIG_NAME=sglang_dp2 RESULTS_DIR=/wk/bench_results/mimo_v2_5_h100 CONCURRENCY=32 NUM_PROMPTS=96 bash /sc/run_bench_single.sh' ``` -## Performance (SGLang, single-node 8xH100, FP8, TP=8 + attn-CP=2) +## Performance (SGLang, single-node 8xH100, FP8) Input/output: 900 / 90 tokens (random dataset), no speculative decoding — same recipe and shape as the Trn2 table in the main README, for direct comparison. -Raw logs in `results/sglang_cp2_c{1,16,32}.txt`. + +**DP=2 (cookbook default)** — raw logs `results/sglang_dp2_c{1,16,32}.txt`: + +| Concurrency | Output throughput (tok/s) | Total throughput (tok/s) | TPOT median (ms) | TTFT median (ms) | TTFT P99 (ms) | +|---|---|---|---|---|---| +| 1 | 91.10 | 999.23 | 9.78 | 112 | 199 | +| 16 | 702.18 | 7706.69 | 20.33 | 210 | 735 | +| 32 | 1420.97 | 15595.60 | 18.63 | 299 | 440 | + +**attn-CP=2 alternative** (`DP=1 ATTN_CP=2`) — raw logs +`results/sglang_cp2_c{1,16,32}.txt`. Slightly faster (a single request uses all +8 GPUs instead of 4 per DP group): | Concurrency | Output throughput (tok/s) | Total throughput (tok/s) | TPOT median (ms) | TTFT median (ms) | TTFT P99 (ms) | |---|---|---|---|---|---| @@ -113,15 +129,15 @@ Raw logs in `results/sglang_cp2_c{1,16,32}.txt`. ### H100 (SGLang) vs Trn2 (vLLM-Neuron) — same 900/90 shape -| Concurrency | H100 out tok/s | Trn2 out tok/s | H100 / Trn2 | -|---|---:|---:|---:| -| 1 | 125.69 | 15.88 | **7.9×** | -| 16 | 792.68 | 113.92 | **7.0×** | -| 32 | 1570.83 | 147.39 | **10.7×** | +| Concurrency | H100 DP=2 out tok/s | H100 CP=2 out tok/s | Trn2 out tok/s | DP=2 / Trn2 | +|---|---:|---:|---:|---:| +| 1 | 91.10 | 125.69 | 15.88 | **5.7×** | +| 16 | 702.18 | 792.68 | 113.92 | **6.2×** | +| 32 | 1420.97 | 1570.83 | 147.39 | **9.6×** | (Trn2 numbers from the main README's Performance section: BS=32, TP=64 / moe_ep=64, CB + bucketing.) The H100 gap is largest at c=32 because SGLang's -median ITL stays ~13–18 ms across concurrency while the Trn2 BS=32 TKG NEFF runs +median ITL stays ~15–20 ms across concurrency while the Trn2 BS=32 TKG NEFF runs at a fixed ~58 ms/token regardless of occupancy. > **Note:** these are the *language-only, no-spec-decode* baselines for diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/bench_all.sh b/contrib/models/MiMo-V2.5/perf_test/h100/bench_all.sh index 5bb433d2..00f5c280 100755 --- a/contrib/models/MiMo-V2.5/perf_test/h100/bench_all.sh +++ b/contrib/models/MiMo-V2.5/perf_test/h100/bench_all.sh @@ -9,7 +9,7 @@ # bash bench_all.sh 8000 vllm # vLLM (if/when it works) set -e PORT="${1:-30000}" -CONFIG_NAME="${2:-sglang_cp2}" +CONFIG_NAME="${2:-sglang_dp2}" SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" for C in 1 16 32; do diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_dp2_c1.txt b/contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_dp2_c1.txt new file mode 100644 index 00000000..38e72c73 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/sglang_dp2_c1.txt @@ -0,0 +1,43 @@ +INFO 07-21 22:26:20 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 22:26:20 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 22:26:21.168000 9 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 22:26:22 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 22:26:22 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +[transformers] You are using a model of type `mimo_v2` to instantiate a model of type ``. This may be expected if you are loading a checkpoint that shares a subset of the architecture (e.g., loading a `sam2_video` checkpoint into `Sam2Model`), but is otherwise not supported and can yield errors. Please verify that the checkpoint is compatible with the model you are instantiating. +[transformers] PreTrainedConfig got `key=rope_scaling` in kwargs but hasn't set it as attribute. For RoPE standardization you need to set `self.rope_parameters` in model's config. +Namespace(subparser='bench', bench_type='serve', dispatch_function=, trust_remote_code=False, seed=0, num_prompts=16, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=30000, endpoint='/v1/completions', header=None, max_concurrency=1, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-9933cbf4-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 22:26:23 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 1 + 0%| | 0/16 [00:00, trust_remote_code=False, seed=0, num_prompts=96, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=30000, endpoint='/v1/completions', header=None, max_concurrency=16, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-70c8961d-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 22:27:02 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 16 + 0%| | 0/96 [00:00, trust_remote_code=False, seed=0, num_prompts=96, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=30000, endpoint='/v1/completions', header=None, max_concurrency=32, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-5c7be356-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 22:27:37 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 32 + 0%| | 0/96 [00:00 crash). The reference relies on -# --moe-a2a-backend deepep to change that collective, but DeepEP needs a -# working nvshmem RDMA transport which this box's NICs (rdmapXXs0, not mlx5) -# don't provide -- IBGDA init fails and the forward pass hangs (tested -# low_latency + normal + REMOTE_TRANSPORT=none + IBGDA=0). DeepEP does work on -# p5en.48xlarge (mlx5 EFA + a custom Mooncake/nvshmem image); see -# xiaomi_datalab/mimo_v25. -# - ATTN_CP=2 (attention context parallel): effective attn TP = tp/cp = 4, MoE -# shards over TP=8, no DP idle-group problem, no DeepEP. Works out of the box -# on plain P5. THIS IS THE DEFAULT. +# ---- Default = the model-card / SGLang cookbook command (DP-attention) ---- +# https://docs.sglang.io/cookbook/autoregressive/Xiaomi/MiMo-V2.5 +# --tp 8 --dp 2 --enable-dp-attention --enable-dp-lm-head --mm-enable-dp-encoder +# --mem-fraction-static 0.65 --chunked-prefill-size 16384 +# DP=2 gives effective attention TP = tp/dp = 4, which MiMoV2ForCausalLM requires +# (its fused qkv_proj is TP=4-interleaved; plain --tp 8 is rejected). This runs +# out of the box single-node -- NO DeepEP needed. The three DP flags +# (--enable-dp-attention + --enable-dp-lm-head + --mm-enable-dp-encoder) together +# with mem-fraction 0.65 are what make the shared-MoE / lm-head collective work; +# an earlier attempt with only --enable-dp-attention and mem-fraction 0.9 hung. # # ---- Speculative decoding (MTP/EAGLE) ---- -# The checkpoint ships model_mtp.safetensors, and the reference command enables -# EAGLE. We LEAVE IT OFF (the Trn2 port has no spec decoding, so this is the -# apples-to-apples baseline). SPEC=1 turns on the reference EAGLE flags. +# The checkpoint ships model_mtp.safetensors and the cookbook enables EAGLE. We +# LEAVE IT OFF by default (the Trn2 port has no spec decoding, so this is the +# apples-to-apples baseline). SPEC=1 turns on the cookbook EAGLE flags. +# +# ---- Alternative: attention context parallel (ATTN_CP) ---- +# Set DP=1 ATTN_CP=2 for --attention-context-parallel-size 2 instead of DP. Also +# gives effective attn TP=4, splits attention along the sequence so a single +# request uses all 8 GPUs (measured slightly faster at low concurrency), and +# needs no DP flags. Kept as a documented alternative. # # Usage: -# bash run_sglang_h100.sh # working default (attn-cp=2, no spec) +# bash run_sglang_h100.sh # cookbook DP=2 baseline (no spec) # SPEC=1 bash run_sglang_h100.sh # + EAGLE speculative decoding -# DP=2 DEEPEP=1 bash run_sglang_h100.sh # model-card DP-attention + DeepEP -# # (only works where nvshmem/DeepEP does) +# DP=1 ATTN_CP=2 bash run_sglang_h100.sh # attention-context-parallel variant set -e MODEL_DIR="${MODEL_DIR:-/opt/dlami/nvme/models/MiMo-V2.5}" PORT="${PORT:-30000}" TP="${TP:-8}" -DP="${DP:-1}" # 1 => use ATTN_CP; >1 => DP-attention -ATTN_CP="${ATTN_CP:-2}" # effective attn TP = TP/ATTN_CP = 4 -MEM_FRAC="${MEM_FRAC:-0.9}" # weights ~295GB; rest for KV cache +DP="${DP:-2}" # cookbook default; DP=1 => use ATTN_CP +ATTN_CP="${ATTN_CP:-2}" # only used when DP=1 +MEM_FRAC="${MEM_FRAC:-0.65}" # cookbook value IMAGE="${IMAGE:-sglang-efa:latest}" CACHE_DIR="${CACHE_DIR:-/opt/dlami/nvme/sglang_cache}" CTR_MODEL="/models/MiMo-V2.5" @@ -52,17 +49,19 @@ if [ ! -f "$MODEL_DIR/config.json" ]; then exit 1 fi -# Parallelism shape: DP-attention (reference) if DP>1, else attention-CP. -DP_ARGS=() +# Parallelism shape: cookbook DP-attention if DP>1, else attention-CP. +PAR_ARGS=() if [ "$DP" -gt 1 ]; then - DP_ARGS=( --dp "$DP" --enable-dp-attention --enable-dp-lm-head ) - echo " DP-attention: ON (dp=$DP, effective attn TP = $((TP/DP)))" + PAR_ARGS=( --dp "$DP" --enable-dp-attention --enable-dp-lm-head --mm-enable-dp-encoder ) + echo " DP-attention: dp=$DP (effective attn TP = $((TP/DP)))" else - DP_ARGS=( --attention-context-parallel-size "$ATTN_CP" ) + PAR_ARGS=( --attention-context-parallel-size "$ATTN_CP" ) echo " Attention context parallel: $ATTN_CP (effective attn TP = $((TP/ATTN_CP)))" fi -# DeepEP MoE all-to-all (needs a working nvshmem RDMA transport; see header). +# DeepEP MoE all-to-all: NOT needed single-node (and its nvshmem RDMA transport +# fails on this box's non-mlx5 NICs). Off by default; DEEPEP=1 to force it where +# it works (e.g. p5en.48xlarge with mlx5 EFA). A2A_ARGS=() DEEPEP_ENV=() if [ "${DEEPEP:-0}" = "1" ]; then @@ -101,7 +100,7 @@ exec docker run --rm --gpus all \ --served-model-name MiMo-V2.5 \ --trust-remote-code \ --tp "$TP" \ - "${DP_ARGS[@]}" \ + "${PAR_ARGS[@]}" \ --moe-dense-tp-size 1 \ --mem-fraction-static "$MEM_FRAC" \ --max-running-requests 128 \ From 6ea3156274700aee04f77f5b884170d98a8f9cbf Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 22 Jul 2026 06:51:37 +0800 Subject: [PATCH 20/27] h100/MiMo-V2.5: vLLM works with PR #42270 (verified) + add vLLM benchmarks Retested vLLM: the fused-qkv FP8 crash is fixed by vLLM PR #42270 ("MiMo V2: Pro fused-QKV FP8 loader + fix SWA wrong-data on V2.5 base"), which is still open upstream. Vendored the PR's two model files under h100/pr42270/; run_vllm_h100.sh now cp's them into the nightly image at container start and runs the reference command --tp 8 --enable-expert-parallel. Verified correct output (short answers + 500-token B-tree explanation, no gibberish). Added vLLM bench results (vllm_pr42270_c{1,16,32}) and a 3-backend perf table. vLLM has the lowest TPOT (5.4ms) but with chunked-prefill off its c=1/16 tail TTFT and throughput suffer from prefill preemption; it catches SGLang at c=32. SGLang remains the out-of-the-box path; both now documented as working. Upstream refs: PR #42270, issue #42803 (root bug). Co-Authored-By: Claude Opus 4.8 --- .../models/MiMo-V2.5/perf_test/h100/README.md | 84 +- .../perf_test/h100/pr42270/README.md | 26 + .../perf_test/h100/pr42270/mimo_v2.py | 1290 +++++++++++++++++ .../perf_test/h100/pr42270/mimo_v2_mtp.py | 435 ++++++ .../h100/results/vllm_pr42270_c1.txt | 43 + .../h100/results/vllm_pr42270_c16.txt | 43 + .../h100/results/vllm_pr42270_c32.txt | 43 + .../perf_test/h100/run_sglang_h100.sh | 2 +- .../MiMo-V2.5/perf_test/h100/run_vllm_h100.sh | 90 +- 9 files changed, 1980 insertions(+), 76 deletions(-) create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/pr42270/README.md create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/pr42270/mimo_v2.py create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/pr42270/mimo_v2_mtp.py create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_pr42270_c1.txt create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_pr42270_c16.txt create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_pr42270_c32.txt diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/README.md b/contrib/models/MiMo-V2.5/perf_test/h100/README.md index 8969e191..d241e870 100644 --- a/contrib/models/MiMo-V2.5/perf_test/h100/README.md +++ b/contrib/models/MiMo-V2.5/perf_test/h100/README.md @@ -11,16 +11,17 @@ cross-node NCCL / EFA is needed (the Pro variant is ~963 GB and requires 2 nodes). We still reuse the `sglang-efa` / `vllm-efa` images built for the Pro work (strict supersets of the stock images); EFA just goes unused here. -## TL;DR — use SGLang; vLLM is broken for this model +## TL;DR — both backends work; vLLM needs PR #42270 | Backend | Status | |---|---| -| **SGLang** (`run_sglang_h100.sh`) | ✅ **Works.** Cookbook `--tp 8 --dp 2 --enable-dp-attention` (no DeepEP, no spec decode). | -| **vLLM** (`run_vllm_h100.sh`) | ❌ **Broken** (upstream bug, 0.25.1 **and** nightly). Kept only to document the reference command + reproduce the bug. | +| **SGLang** (`run_sglang_h100.sh`) | ✅ **Works out of the box.** Cookbook `--tp 8 --dp 2 --enable-dp-attention` (no DeepEP, no spec decode). | +| **vLLM** (`run_vllm_h100.sh`) | ✅ **Works with vLLM PR #42270** (not yet in any released image; the script bind-mounts the PR's two model files into the nightly image). Reference command `--tp 8 --enable-expert-parallel`. | -### Why vLLM fails +### vLLM needs PR #42270 (fused-qkv FP8 loader fix) -vLLM crashes loading the fused `qkv_proj` on the **sliding-window** layers: +Stock vLLM (0.25.1 **and** nightly) crashes loading the fused `qkv_proj` on the +**sliding-window** layers: ``` RuntimeError: The size of tensor a (1856) must match the size of tensor b (1792) @@ -29,20 +30,26 @@ RuntimeError: The size of tensor a (1856) must match the size of tensor b (1792) ``` A SWA-layer fused-qkv group is `(64/8)*192 + 192 + 128 = 1856` rows = **14.5** -blocks of the 128-row FP8 scale, so it is **not 128-aligned**. vLLM's loader +blocks of the 128-row FP8 scale, so it is **not 128-aligned**. The stock loader slices the scale per KV-head group (`scale_rows_per_group = 116 // 8 = 14` -> 1792 rows) and multiplies it against the 1856-row weight group. Because MiMo-V2.5 has only **4 full-attention KV heads** (vs 8 for the Pro), no 8-GPU TP -config avoids the buggy path: - -- `TP=8` — can't start: "TP size must evenly split the number of KV heads" (4). -- `TP=4` / `TP=2` / `DP=2 x TP=4` — each rank owns 2 SWA KV heads -> the `g>1` - de-interleave path -> the 1856-vs-1792 crash. - -(The Pro variant never hits this: with 8 KV heads and TP=8 each rank gets 1 KV -head and takes the trivial `kv_heads_per_rank == 1` fast path.) A shape-only -monkeypatch makes the server start but produces gibberish, so it is not a real -fix. This should be reported/fixed upstream in vLLM. +config avoids the buggy path (`TP=8` can't split 4 KV heads; `TP=4`/`TP=2`/ +`DP=2 x TP=4` put 2 SWA KV heads per rank -> the `g>1` crash). The Pro variant +never hits this (8 KV heads + TP=8 -> 1 head/rank -> trivial fast path). + +**Fix**: vLLM [PR #42270](https://github.com/vllm-project/vllm/pull/42270) +("MiMo V2: Pro fused-QKV FP8 loader + fix SWA wrong-data on V2.5 base") replaces +the loader with one that "dequantizes and requantizes local shards for TP +configurations that cut through 128-row FP8 scale blocks". As of 2026-07 it is +still **open** (merge conflicts / awaiting review), so not in any release. This +directory vendors the PR's two model files under `pr42270/`, and +`run_vllm_h100.sh` `cp`s them into the nightly image at container start. +**Verified 2026-07-21**: with the PR applied, `--tp 8 --enable-expert-parallel` +loads and produces correct output (coherent short answers + a full 500-token +B-tree explanation, no gibberish). Related upstream: issue +[#42803](https://github.com/vllm-project/vllm/issues/42803) (root bug report), +which the PR closes. ### SGLang parallelism: DP=2 (cookbook, default) — no DeepEP needed @@ -104,12 +111,13 @@ docker run --rm --network host -v /opt/dlami/nvme/models:/wk -v $PWD:/sc \ CONCURRENCY=32 NUM_PROMPTS=96 bash /sc/run_bench_single.sh' ``` -## Performance (SGLang, single-node 8xH100, FP8) +## Performance (single-node 8xH100, FP8) Input/output: 900 / 90 tokens (random dataset), no speculative decoding — same recipe and shape as the Trn2 table in the main README, for direct comparison. +Raw logs in `results/`. -**DP=2 (cookbook default)** — raw logs `results/sglang_dp2_c{1,16,32}.txt`: +**SGLang, DP=2 (cookbook default)** — `results/sglang_dp2_c{1,16,32}.txt`: | Concurrency | Output throughput (tok/s) | Total throughput (tok/s) | TPOT median (ms) | TTFT median (ms) | TTFT P99 (ms) | |---|---|---|---|---|---| @@ -117,9 +125,8 @@ recipe and shape as the Trn2 table in the main README, for direct comparison. | 16 | 702.18 | 7706.69 | 20.33 | 210 | 735 | | 32 | 1420.97 | 15595.60 | 18.63 | 299 | 440 | -**attn-CP=2 alternative** (`DP=1 ATTN_CP=2`) — raw logs -`results/sglang_cp2_c{1,16,32}.txt`. Slightly faster (a single request uses all -8 GPUs instead of 4 per DP group): +**SGLang, attn-CP=2** (`DP=1 ATTN_CP=2`) — `results/sglang_cp2_c{1,16,32}.txt`. +Slightly faster (a single request uses all 8 GPUs instead of 4 per DP group): | Concurrency | Output throughput (tok/s) | Total throughput (tok/s) | TPOT median (ms) | TTFT median (ms) | TTFT P99 (ms) | |---|---|---|---|---|---| @@ -127,19 +134,34 @@ recipe and shape as the Trn2 table in the main README, for direct comparison. | 16 | 792.68 | 8699.94 | 17.66 | 186 | 726 | | 32 | 1570.83 | 17240.36 | 17.78 | 237 | 293 | -### H100 (SGLang) vs Trn2 (vLLM-Neuron) — same 900/90 shape +**vLLM + PR #42270, TP=8 + EP** — `results/vllm_pr42270_c{1,16,32}.txt`: + +| Concurrency | Output throughput (tok/s) | Total throughput (tok/s) | TPOT median (ms) | TTFT median (ms) | TTFT P99 (ms) | +|---|---|---|---|---|---| +| 1 | 48.26 | 529.33 | 5.42 | 84 | 12469 | +| 16 | 214.49 | 2354.09 | 16.74 | 224 | 16225 | +| 32 | 1459.66 | 16020.21 | 18.31 | 273 | 715 | + +vLLM's median TPOT is actually the lowest of the three (5.4 ms at c=1), but with +`--no-enable-chunked-prefill` (matching Trn2) its c=1/c=16 output throughput and +tail TTFT suffer badly — each new request's 900-token prefill preempts decode +(P99 TTFT 12–16 s at c≤16). By c=32 the pipeline stays saturated and vLLM catches +SGLang (1460 vs 1421 out tok/s). SGLang's DP-attention absorbs the prefill +interleaving far better at low/mid concurrency. + +### H100 vs Trn2 (vLLM-Neuron) — same 900/90 shape, out tok/s -| Concurrency | H100 DP=2 out tok/s | H100 CP=2 out tok/s | Trn2 out tok/s | DP=2 / Trn2 | -|---|---:|---:|---:|---:| -| 1 | 91.10 | 125.69 | 15.88 | **5.7×** | -| 16 | 702.18 | 792.68 | 113.92 | **6.2×** | -| 32 | 1420.97 | 1570.83 | 147.39 | **9.6×** | +| Concurrency | SGLang DP=2 | SGLang CP=2 | vLLM+PR#42270 | Trn2 | best / Trn2 | +|---|---:|---:|---:|---:|---:| +| 1 | 91.10 | 125.69 | 48.26 | 15.88 | **7.9×** | +| 16 | 702.18 | 792.68 | 214.49 | 113.92 | **7.0×** | +| 32 | 1420.97 | 1570.83 | 1459.66 | 147.39 | **10.7×** | (Trn2 numbers from the main README's Performance section: BS=32, TP=64 / -moe_ep=64, CB + bucketing.) The H100 gap is largest at c=32 because SGLang's -median ITL stays ~15–20 ms across concurrency while the Trn2 BS=32 TKG NEFF runs +moe_ep=64, CB + bucketing.) The H100 gap is largest at c=32 because the GPU +median ITL stays ~13–20 ms across concurrency while the Trn2 BS=32 TKG NEFF runs at a fixed ~58 ms/token regardless of occupancy. > **Note:** these are the *language-only, no-spec-decode* baselines for -> apples-to-apples comparison with Trn2. Turning on EAGLE (`SPEC=1`) would raise -> H100 decode throughput further; the model card's full recipe also enables it. +> apples-to-apples comparison with Trn2. Turning on EAGLE (`SPEC=1` on SGLang) +> would raise decode throughput further; the model card's full recipe enables it. diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/pr42270/README.md b/contrib/models/MiMo-V2.5/perf_test/h100/pr42270/README.md new file mode 100644 index 00000000..99d878a7 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/pr42270/README.md @@ -0,0 +1,26 @@ +# vLLM PR #42270 model files (MiMo-V2.5 FP8 fused-qkv loader fix) + +`mimo_v2.py` and `mimo_v2_mtp.py` from vLLM PR +[#42270](https://github.com/vllm-project/vllm/pull/42270) +("[models] MiMo V2: Pro fused-QKV FP8 loader + fix SWA wrong-data on V2.5 base"), +branch `amd-satre:mimo-v25-pro-fp8-qkv-loader`, fetched 2026-07-21. + +## Why these are here + +Stock vLLM (0.25.1 **and** nightly) crashes loading MiMo-V2.5 FP8 in +`_shard_fp8_qkv_proj` (`RuntimeError: size of tensor a (1856) must match b +(1792)`): a sliding-window layer's fused-qkv group is 1856 rows = 14.5 128-row +FP8 scale blocks, so per-KV-head scale slicing misaligns. This PR replaces the +loader with `_mimo_v2_copy_paired_qkv_fp8`, which "dequantizes and requantizes +local shards for TP configurations that cut through 128-row FP8 scale blocks, +such as TP4". As of 2026-07 the PR is still **open** (merge conflicts / awaiting +review), so it is not in any released image. + +## Verified (2026-07-21, single-node 8xH100, nightly + these files) + +Overwriting nightly's two files with these makes `--tensor-parallel-size 8 +--enable-expert-parallel` (the reference command) **load and generate correct +output** — coherent short answers and a full 500-token B-tree explanation with +no gibberish/collapse. `run_vllm_h100.sh` bind-mounts this dir and `cp`s the +files in at container start. Delete this dir and revert run_vllm_h100.sh once +the PR merges into the image. diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/pr42270/mimo_v2.py b/contrib/models/MiMo-V2.5/perf_test/h100/pr42270/mimo_v2.py new file mode 100644 index 00000000..f5414811 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/pr42270/mimo_v2.py @@ -0,0 +1,1290 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterable +from itertools import islice + +import torch +from torch import nn + +from vllm.compilation.decorators import support_torch_compile +from vllm.config import ( + CacheConfig, + VllmConfig, + get_current_vllm_config, + str_dtype_to_torch_dtype, +) +from vllm.distributed import ( + divide, + get_ep_group, + get_pp_group, + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_gather, +) +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import SiluAndMul +from vllm.model_executor.layers.attention import Attention +from vllm.model_executor.layers.fused_moe import ( + FusedMoE, + fused_moe_make_expert_params_mapping, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ( + MergedColumnParallelLinear, + QKVParallelLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.quantization.utils.quant_utils import ( + GroupShape, + scaled_dequantize, + scaled_quantize, +) +from vllm.model_executor.layers.rotary_embedding import get_rope +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import ( + default_weight_loader, + maybe_remap_kv_scale_name, +) +from vllm.model_executor.models.utils import sequence_parallel_chunk +from vllm.sequence import IntermediateTensors +from vllm.utils.math_utils import cdiv, round_up +from vllm.v1.attention.backend import AttentionType +from vllm.v1.attention.backends.registry import AttentionBackendEnum + +from .interfaces import MixtureOfExperts, SupportsPP +from .utils import ( + AutoWeightsLoader, + PPMissingLayer, + extract_layer_index, + is_pp_missing_parameter, + make_empty_intermediate_tensors_factory, + make_layers, + maybe_prefix, +) + +logger = init_logger(__name__) + + +def _mimo_v2_qkv_pair_key(name: str) -> tuple[str, str] | None: + if name.endswith(".qkv_proj.weight"): + return name[: -len(".weight")], "weight" + if name.endswith(".qkv_proj.weight_scale_inv"): + return name[: -len(".weight_scale_inv")], "scale" + return None + + +def _mimo_v2_kv_shard_range( + total_num_kv_heads: int, + head_size: int, + tp_rank: int, + tp_size: int, +) -> tuple[int, int]: + if tp_size >= total_num_kv_heads: + num_replicas = divide(tp_size, total_num_kv_heads) + kv_head_idx = tp_rank // num_replicas + return kv_head_idx * head_size, head_size + + num_heads = divide(total_num_kv_heads, tp_size) + shard_size = num_heads * head_size + return tp_rank * shard_size, shard_size + + +def _mimo_v2_qkv_dims( + config, + layer_idx: int | None, +) -> tuple[int, int, int, int, int, int, int]: + is_swa = ( + layer_idx is not None + and hasattr(config, "hybrid_layer_pattern") + and layer_idx < len(config.hybrid_layer_pattern) + and config.hybrid_layer_pattern[layer_idx] == 1 + ) + if is_swa: + head_dim = getattr(config, "swa_head_dim", config.head_dim) + v_head_dim = getattr(config, "swa_v_head_dim", config.v_head_dim) + num_heads = getattr( + config, "swa_num_attention_heads", config.num_attention_heads + ) + num_kv_heads = getattr( + config, "swa_num_key_value_heads", config.num_key_value_heads + ) + else: + head_dim = config.head_dim + v_head_dim = getattr(config, "v_head_dim", head_dim) + num_heads = config.num_attention_heads + num_kv_heads = config.num_key_value_heads + + q_rows = num_heads * head_dim + k_rows = num_kv_heads * head_dim + v_rows = num_kv_heads * v_head_dim + return q_rows, k_rows, v_rows, head_dim, v_head_dim, num_heads, num_kv_heads + + +def _mimo_v2_dequant_block_shard( + weight: torch.Tensor, + scale: torch.Tensor, + row_start: int, + row_count: int, + block_n: int, + block_k: int, + device: torch.device, +) -> torch.Tensor: + block_start = row_start // block_n + block_end = cdiv(row_start + row_count, block_n) + expanded_start = block_start * block_n + expanded_rows = (block_end - block_start) * block_n + available_rows = min(expanded_rows, weight.shape[0] - expanded_start) + expanded_weight = weight.narrow(0, expanded_start, available_rows).to(device) + if available_rows != expanded_rows: + padded = torch.empty( + (expanded_rows, weight.shape[1]), + dtype=expanded_weight.dtype, + device=device, + ) + padded.zero_() + padded[:available_rows] = expanded_weight + expanded_weight = padded + expanded_scale = scale.narrow(0, block_start, block_end - block_start).to(device) + dequant_weight = scaled_dequantize( + expanded_weight, + expanded_scale, + group_shape=GroupShape(block_n, block_k), + out_dtype=torch.float32, + ) + local_start = row_start - expanded_start + return dequant_weight.narrow(0, local_start, row_count) + + +def _mimo_v2_quantize_block_weight( + weight: torch.Tensor, + block_n: int, + block_k: int, + quant_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + rows, cols = weight.shape + padded_rows = round_up(rows, block_n) + padded_cols = round_up(cols, block_k) + if padded_rows != rows or padded_cols != cols: + padded = torch.empty( + (padded_rows, padded_cols), + dtype=weight.dtype, + device=weight.device, + ) + padded.zero_() + padded[:rows, :cols] = weight + weight = padded + + qweight, scale = scaled_quantize( + weight, + GroupShape(block_n, block_k), + quant_dtype=quant_dtype, + compute_dtype=torch.float32, + ) + return qweight[:rows, :cols].contiguous(), scale + + +def _mimo_v2_copy_presharded_qkv_bf16( + *, + config, + weight_name: str, + weight_param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + tp_rank: int, + tp_size: int, +) -> bool: + """Handle a pre-sharded fused-QKV BF16 weight at any local TP factor. + + Returns True if the layout was pre-sharded and handled here; False if the + layout isn't pre-sharded and the caller should fall back to the regular + per-rank weight_loader. + + The source ships the fused QKV as the row-wise concatenation of + ``num_kv_heads`` independently-sharded TP chunks, each laid out as + ``[Q_chunk, K_chunk, V_chunk]``. vLLM's ``QKVParallelLinear`` weight + loader assumes a canonical ``[Q_global, K_global, V_global]`` layout, so + handing it the pre-sharded tensor produces structurally wrong slices. + Reassemble per-rank Q / K / V from the relevant ckpt chunks instead. + """ + layer_idx = extract_layer_index(weight_name) + ( + q_rows, + k_rows, + v_rows, + head_dim, + v_head_dim, + num_heads, + num_kv_heads, + ) = _mimo_v2_qkv_dims(config, layer_idx) + + ckpt_tp = num_kv_heads + q_heads_per_ckpt_rank = divide(num_heads, ckpt_tp) + ckpt_q_rows = q_heads_per_ckpt_rank * head_dim + ckpt_k_rows = head_dim + ckpt_v_rows = v_head_dim + ckpt_chunk_rows = ckpt_q_rows + ckpt_k_rows + ckpt_v_rows + + if loaded_weight.shape[0] != ckpt_tp * ckpt_chunk_rows: + return False + + if tp_size == ckpt_tp: + local_weight = loaded_weight.narrow( + 0, tp_rank * ckpt_chunk_rows, ckpt_chunk_rows + ) + else: + + def chunk_slice(ckpt_rank: int, row_start: int, row_count: int) -> torch.Tensor: + base = ckpt_rank * ckpt_chunk_rows + row_start + return loaded_weight.narrow(0, base, row_count) + + q_heads_per_rank = divide(num_heads, tp_size) + q_head_start = tp_rank * q_heads_per_rank + q_head_end = q_head_start + q_heads_per_rank + q_parts: list[torch.Tensor] = [] + next_q_head = q_head_start + while next_q_head < q_head_end: + ckpt_rank = next_q_head // q_heads_per_ckpt_rank + ckpt_head_start = ckpt_rank * q_heads_per_ckpt_rank + part_head_end = min(q_head_end, ckpt_head_start + q_heads_per_ckpt_rank) + part_rows = (part_head_end - next_q_head) * head_dim + part_start = (next_q_head - ckpt_head_start) * head_dim + q_parts.append(chunk_slice(ckpt_rank, part_start, part_rows)) + next_q_head = part_head_end + + if tp_size >= num_kv_heads: + num_replicas = divide(tp_size, num_kv_heads) + kv_head_start = tp_rank // num_replicas + kv_head_count = 1 + else: + kv_head_count = divide(num_kv_heads, tp_size) + kv_head_start = tp_rank * kv_head_count + kv_head_end = kv_head_start + kv_head_count + k_parts = [ + chunk_slice(ckpt_rank, ckpt_q_rows, ckpt_k_rows) + for ckpt_rank in range(kv_head_start, kv_head_end) + ] + v_parts = [ + chunk_slice(ckpt_rank, ckpt_q_rows + ckpt_k_rows, ckpt_v_rows) + for ckpt_rank in range(kv_head_start, kv_head_end) + ] + local_weight = torch.cat([*q_parts, *k_parts, *v_parts], dim=0).contiguous() + + if tuple(local_weight.shape) != tuple(weight_param.shape): + raise ValueError( + f"{weight_name} local shard has shape {tuple(local_weight.shape)}, " + f"expected {tuple(weight_param.shape)}." + ) + weight_param.data.copy_(local_weight.to(weight_param.device)) + return True + + +def _mimo_v2_copy_paired_qkv_fp8( + *, + config, + weight_name: str, + scale_name: str, + weight_param: torch.nn.Parameter, + scale_param: torch.nn.Parameter, + loaded_weight: torch.Tensor, + loaded_scale: torch.Tensor, + tp_rank: int, + tp_size: int, + block_size: list[int], +) -> None: + block_n, block_k = block_size + layer_idx = extract_layer_index(weight_name) + ( + q_rows, + k_rows, + v_rows, + head_dim, + v_head_dim, + num_heads, + num_kv_heads, + ) = _mimo_v2_qkv_dims(config, layer_idx) + expected_rows = q_rows + k_rows + v_rows + if loaded_weight.shape[0] != expected_rows: + raise ValueError( + f"{weight_name} has {loaded_weight.shape[0]} rows, expected " + f"{expected_rows} from q={q_rows}, k={k_rows}, v={v_rows}." + ) + + # Detect ckpt_tp from scale shape; weight shape alone is degenerate. + ckpt_tp = None + q_heads_per_ckpt_rank = None + ckpt_q_rows = ckpt_k_rows = 0 + ckpt_chunk_rows = ckpt_chunk_scale_rows = 0 + for candidate in range(min(num_heads, num_kv_heads), 0, -1): + if num_heads % candidate != 0 or num_kv_heads % candidate != 0: + continue + cand_q_heads = num_heads // candidate + cand_kv_heads = num_kv_heads // candidate + cand_q_rows = cand_q_heads * head_dim + cand_k_rows = cand_kv_heads * head_dim + cand_v_rows = cand_kv_heads * v_head_dim + cand_chunk_rows = cand_q_rows + cand_k_rows + cand_v_rows + cand_chunk_scale_rows = cdiv(cand_chunk_rows, block_n) + if ( + loaded_weight.shape[0] == candidate * cand_chunk_rows + and loaded_scale.shape[0] == candidate * cand_chunk_scale_rows + ): + ckpt_tp = candidate + q_heads_per_ckpt_rank = cand_q_heads + ckpt_q_rows = cand_q_rows + ckpt_k_rows = cand_k_rows + ckpt_chunk_rows = cand_chunk_rows + ckpt_chunk_scale_rows = cand_chunk_scale_rows + break + if ckpt_tp is not None: + logger.info_once( + "Detected MiMo-V2 TP%d pre-sharded fused-QKV FP8 checkpoint layout.", + ckpt_tp, + ) + if tp_size == ckpt_tp: + local_weight = loaded_weight.narrow( + 0, tp_rank * ckpt_chunk_rows, ckpt_chunk_rows + ) + local_scale = loaded_scale.narrow( + 0, tp_rank * ckpt_chunk_scale_rows, ckpt_chunk_scale_rows + ) + else: + device = weight_param.device + + def dequant_ckpt_shard( + ckpt_rank: int, + row_start: int, + row_count: int, + ) -> torch.Tensor: + chunk_weight = loaded_weight.narrow( + 0, ckpt_rank * ckpt_chunk_rows, ckpt_chunk_rows + ) + chunk_scale = loaded_scale.narrow( + 0, ckpt_rank * ckpt_chunk_scale_rows, ckpt_chunk_scale_rows + ) + return _mimo_v2_dequant_block_shard( + chunk_weight, + chunk_scale, + row_start, + row_count, + block_n, + block_k, + device, + ) + + q_heads_per_rank = divide(num_heads, tp_size) + q_head_start = tp_rank * q_heads_per_rank + q_head_end = q_head_start + q_heads_per_rank + q_parts: list[torch.Tensor] = [] + next_q_head = q_head_start + while next_q_head < q_head_end: + ckpt_rank = next_q_head // q_heads_per_ckpt_rank + ckpt_head_start = ckpt_rank * q_heads_per_ckpt_rank + part_head_end = min(q_head_end, ckpt_head_start + q_heads_per_ckpt_rank) + part_rows = (part_head_end - next_q_head) * head_dim + part_start = (next_q_head - ckpt_head_start) * head_dim + q_parts.append(dequant_ckpt_shard(ckpt_rank, part_start, part_rows)) + next_q_head = part_head_end + + if tp_size >= num_kv_heads: + num_replicas = divide(tp_size, num_kv_heads) + kv_head_start = tp_rank // num_replicas + kv_head_count = 1 + else: + kv_head_count = divide(num_kv_heads, tp_size) + kv_head_start = tp_rank * kv_head_count + kv_head_end = kv_head_start + kv_head_count + + kv_heads_per_ckpt_rank = divide(num_kv_heads, ckpt_tp) + + def _kv_parts( + intra_chunk_offset: int, per_head_rows: int + ) -> list[torch.Tensor]: + parts: list[torch.Tensor] = [] + next_kv_head = kv_head_start + while next_kv_head < kv_head_end: + ckpt_rank = next_kv_head // kv_heads_per_ckpt_rank + ckpt_head_start = ckpt_rank * kv_heads_per_ckpt_rank + part_head_end = min( + kv_head_end, ckpt_head_start + kv_heads_per_ckpt_rank + ) + part_rows = (part_head_end - next_kv_head) * per_head_rows + part_start = ( + intra_chunk_offset + + (next_kv_head - ckpt_head_start) * per_head_rows + ) + parts.append(dequant_ckpt_shard(ckpt_rank, part_start, part_rows)) + next_kv_head = part_head_end + return parts + + k_parts = _kv_parts(ckpt_q_rows, head_dim) + v_parts = _kv_parts(ckpt_q_rows + ckpt_k_rows, v_head_dim) + local_dense = torch.cat([*q_parts, *k_parts, *v_parts], dim=0) + local_weight, local_scale = _mimo_v2_quantize_block_weight( + local_dense, + block_n, + block_k, + weight_param.dtype, + ) + + if tuple(local_weight.shape) != tuple(weight_param.shape): + raise ValueError( + f"{weight_name} local shard has shape " + f"{tuple(local_weight.shape)}, expected " + f"{tuple(weight_param.shape)}." + ) + if tuple(local_scale.shape) != tuple(scale_param.shape): + raise ValueError( + f"{scale_name} local shard has shape " + f"{tuple(local_scale.shape)}, expected " + f"{tuple(scale_param.shape)}." + ) + + weight_param.data.copy_(local_weight.to(weight_param.device)) + scale_param.data.copy_(local_scale.to(scale_param.device)) + return + + q_scale_rows = cdiv(q_rows, block_n) + k_scale_rows = cdiv(k_rows, block_n) + v_scale_rows = cdiv(v_rows, block_n) + expected_scale_rows = q_scale_rows + k_scale_rows + v_scale_rows + if loaded_scale.shape[0] < expected_scale_rows: + raise ValueError( + f"{scale_name} has {loaded_scale.shape[0]} rows, expected at " + f"least {expected_scale_rows}." + ) + extra_scale_rows = loaded_scale.shape[0] - expected_scale_rows + if extra_scale_rows: + logger.info_once( + "Dropping %d extra MiMo-V2 fused-QKV FP8 scale rows from %s. " + "The checkpoint has q/k/v scale rows %d/%d/%d plus extras, " + "while q/k/v weight rows imply %d/%d/%d.", + extra_scale_rows, + scale_name, + q_scale_rows, + k_scale_rows, + v_scale_rows, + q_rows, + k_rows, + v_rows, + ) + + q_weight = loaded_weight.narrow(0, 0, q_rows) + k_weight = loaded_weight.narrow(0, q_rows, k_rows) + v_weight = loaded_weight.narrow(0, q_rows + k_rows, v_rows) + q_scale = loaded_scale.narrow(0, 0, q_scale_rows) + k_scale = loaded_scale.narrow(0, q_scale_rows, k_scale_rows) + v_scale = loaded_scale.narrow(0, q_scale_rows + k_scale_rows, v_scale_rows) + + q_shard_rows = divide(q_rows, tp_size) + q_start = tp_rank * q_shard_rows + k_start, k_shard_rows = _mimo_v2_kv_shard_range( + config.num_key_value_heads, head_dim, tp_rank, tp_size + ) + v_start, v_shard_rows = _mimo_v2_kv_shard_range( + config.num_key_value_heads, v_head_dim, tp_rank, tp_size + ) + + direct_copy = all( + value % block_n == 0 + for value in ( + q_start, + q_shard_rows, + k_start, + k_shard_rows, + v_start, + v_shard_rows, + ) + ) + + if direct_copy: + local_weight = torch.cat( + [ + q_weight.narrow(0, q_start, q_shard_rows), + k_weight.narrow(0, k_start, k_shard_rows), + v_weight.narrow(0, v_start, v_shard_rows), + ], + dim=0, + ) + local_scale = torch.cat( + [ + q_scale.narrow(0, q_start // block_n, q_shard_rows // block_n), + k_scale.narrow(0, k_start // block_n, k_shard_rows // block_n), + v_scale.narrow(0, v_start // block_n, v_shard_rows // block_n), + ], + dim=0, + ) + else: + device = weight_param.device + local_dense = torch.cat( + [ + _mimo_v2_dequant_block_shard( + q_weight, q_scale, q_start, q_shard_rows, block_n, block_k, device + ), + _mimo_v2_dequant_block_shard( + k_weight, k_scale, k_start, k_shard_rows, block_n, block_k, device + ), + _mimo_v2_dequant_block_shard( + v_weight, v_scale, v_start, v_shard_rows, block_n, block_k, device + ), + ], + dim=0, + ) + local_weight, local_scale = _mimo_v2_quantize_block_weight( + local_dense, + block_n, + block_k, + weight_param.dtype, + ) + + if tuple(local_weight.shape) != tuple(weight_param.shape): + raise ValueError( + f"{weight_name} local shard has shape {tuple(local_weight.shape)}, " + f"expected {tuple(weight_param.shape)}." + ) + if tuple(local_scale.shape) != tuple(scale_param.shape): + raise ValueError( + f"{scale_name} local shard has shape {tuple(local_scale.shape)}, " + f"expected {tuple(scale_param.shape)}." + ) + + weight_param.data.copy_(local_weight.to(weight_param.device)) + scale_param.data.copy_(local_scale.to(scale_param.device)) + + +class MiMoV2MLP(nn.Module): + def __init__( + self, + hidden_size: int, + intermediate_size: int, + hidden_act: str, + quant_config: QuantizationConfig | None = None, + reduce_results: bool = True, + prefix: str = "", + ) -> None: + super().__init__() + self.gate_up_proj = MergedColumnParallelLinear( + hidden_size, + [intermediate_size] * 2, + bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", + ) + self.down_proj = RowParallelLinear( + intermediate_size, + hidden_size, + bias=False, + quant_config=quant_config, + reduce_results=reduce_results, + prefix=f"{prefix}.down_proj", + ) + if hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {hidden_act}. Only silu is supported for now." + ) + self.act_fn = SiluAndMul() + + def forward(self, x): + gate_up, _ = self.gate_up_proj(x) + x = self.act_fn(gate_up) + x, _ = self.down_proj(x) + return x + + +class MiMoV2MoE(nn.Module): + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + is_nextn: bool = False, + ): + super().__init__() + + config = vllm_config.model_config.hf_text_config + parallel_config = vllm_config.parallel_config + quant_config = vllm_config.quant_config + + self.tp_size = get_tensor_model_parallel_world_size() + + self.ep_group = get_ep_group().device_group + self.ep_rank = get_ep_group().rank_in_group + self.ep_size = self.ep_group.size() + self.n_routed_experts = config.n_routed_experts + + self.is_sequence_parallel = parallel_config.use_sequence_parallel_moe + + if self.tp_size > config.n_routed_experts: + raise ValueError( + f"Tensor parallel size {self.tp_size} is greater than " + f"the number of experts {config.n_routed_experts}." + ) + + if config.hidden_act != "silu": + raise ValueError( + f"Unsupported activation: {config.hidden_act}. " + "Only silu is supported for now." + ) + + vllm_config = get_current_vllm_config() + eplb_config = vllm_config.parallel_config.eplb_config + self.enable_eplb = parallel_config.enable_eplb + + self.n_logical_experts = self.n_routed_experts + self.n_redundant_experts = eplb_config.num_redundant_experts + self.n_physical_experts = self.n_logical_experts + self.n_redundant_experts + self.n_local_physical_experts = self.n_physical_experts // self.ep_size + + self.physical_expert_start = self.ep_rank * self.n_local_physical_experts + self.physical_expert_end = ( + self.physical_expert_start + self.n_local_physical_experts + ) + + dtype = getattr(config, "moe_router_dtype", "float32") + self.gate_dtype = str_dtype_to_torch_dtype(dtype) + self.gate = nn.Linear( + config.hidden_size, + config.n_routed_experts, + bias=False, + dtype=self.gate_dtype, + ) + self.gate.e_score_correction_bias = nn.Parameter( + torch.empty(config.n_routed_experts, dtype=self.gate_dtype) + ) + + self.experts = FusedMoE( + num_experts=self.n_routed_experts, + top_k=config.num_experts_per_tok, + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + renormalize=config.norm_topk_prob, + quant_config=quant_config, + prefix=f"{prefix}.experts", + e_score_correction_bias=self.gate.e_score_correction_bias, + enable_eplb=self.enable_eplb, + num_redundant_experts=self.n_redundant_experts, + is_sequence_parallel=self.is_sequence_parallel, + use_grouped_topk=True, + num_expert_group=config.n_group, + topk_group=config.topk_group, + scoring_func="sigmoid", + router_logits_dtype=self.gate_dtype, + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + assert hidden_states.dim() <= 2, "MiMoV2MoE only supports 1D or 2D inputs" + is_input_1d = hidden_states.dim() == 1 + num_tokens, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + + if self.is_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + + if self.gate_dtype is not None: + gate_input = hidden_states.to(self.gate_dtype) + else: + gate_input = hidden_states + router_logits = self.gate(gate_input) + final_hidden_states = self.experts( + hidden_states=hidden_states, router_logits=router_logits + ) + + if self.is_sequence_parallel: + final_hidden_states = tensor_model_parallel_all_gather( + final_hidden_states, 0 + ) + final_hidden_states = final_hidden_states[:num_tokens] + + return final_hidden_states.squeeze(0) if is_input_1d else final_hidden_states + + +class MiMoV2Attention(nn.Module): + def __init__( + self, + hidden_size: int, + num_heads: int, + num_kv_heads: int, + head_dim: int, + v_head_dim: int | None = None, + v_scale: float | None = None, + sliding_window_size: int = -1, + attention_bias: bool = False, + add_swa_attention_sink_bias: bool = False, + layer_id: int = 0, + rope_theta: float = 1000000, + max_position_embeddings: int = 32768, + cache_config: CacheConfig | None = None, + quant_config: QuantizationConfig | None = None, + partial_rotary_factor: float = 1.0, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = hidden_size + self.layer_id = layer_id + tp_size = get_tensor_model_parallel_world_size() + + self.total_num_heads = num_heads + self.num_heads = self.total_num_heads // tp_size + + self.total_num_kv_heads = num_kv_heads + self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size) + + self.head_dim = head_dim + + self.v_head_dim = v_head_dim if v_head_dim is not None else head_dim + + self.q_size = self.num_heads * self.head_dim + self.k_size = self.num_kv_heads * self.head_dim + self.v_size = self.num_kv_heads * self.v_head_dim + + self.v_scale = v_scale + self.scaling = self.head_dim**-0.5 + self.rope_theta = rope_theta + self.max_position_embeddings = max_position_embeddings + + self.qkv_proj = QKVParallelLinear( + hidden_size, + self.head_dim, + self.total_num_heads, + self.total_num_kv_heads, + bias=attention_bias, + quant_config=quant_config, + prefix=f"{prefix}.qkv_proj", + v_head_size=self.v_head_dim, + ) + + self.o_proj = RowParallelLinear( + self.total_num_heads * self.v_head_dim, + hidden_size, + bias=False, + quant_config=quant_config if "mtp.layers" not in prefix else None, + reduce_results=True, + prefix=f"{prefix}.o_proj", + ) + + self.rotary_emb = get_rope( + head_size=self.head_dim, + max_position=max_position_embeddings, + rope_parameters={ + "rope_type": "default", + "rope_theta": rope_theta, + "partial_rotary_factor": partial_rotary_factor, + }, + ) + + self.attention_sink_bias = ( + torch.nn.Parameter(torch.empty(self.num_heads), requires_grad=False) + if add_swa_attention_sink_bias + else None + ) + + sliding_window = sliding_window_size if sliding_window_size > -1 else None + + # Use DiffKV backend when V has a different head dim than K. + # Auto-pick FA-DiffKV when FA3/4 is usable on this device, else fall + # back to TRITON_ATTN_DIFFKV. Users can force a choice via + # `--attention-backend `. + if self.v_head_dim != self.head_dim: + requested = get_current_vllm_config().attention_config.backend + if requested is not None and requested.name.endswith("_DIFFKV"): + backend_enum = requested + else: + fa_backend = AttentionBackendEnum.FLASH_ATTN_DIFFKV.get_class() + if fa_backend.is_supported_on_current_device( + head_size=self.head_dim, + head_size_v=self.v_head_dim, + has_sinks=self.attention_sink_bias is not None, + ): + backend_enum = AttentionBackendEnum.FLASH_ATTN_DIFFKV + else: + backend_enum = AttentionBackendEnum.TRITON_ATTN_DIFFKV + attn_backend = backend_enum.get_class() + attn_backend.set_head_size_v(self.v_head_dim) + logger.info_once("Using %s for attention.", attn_backend.get_name()) + else: + attn_backend = None + + self.attn = Attention( + self.num_heads, + self.head_dim, + self.scaling, + num_kv_heads=self.num_kv_heads, + cache_config=cache_config, + quant_config=quant_config, + per_layer_sliding_window=sliding_window, + attn_type=AttentionType.DECODER, + prefix=f"{prefix}.attn", + sinks=self.attention_sink_bias, + attn_backend=attn_backend, + head_size_v=self.v_head_dim, + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + ) -> torch.Tensor: + qkv, _ = self.qkv_proj(hidden_states) + q, k, v = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) + q, k = self.rotary_emb(positions, q, k) + + # Apply v_scale before attention + if self.v_scale is not None: + v = v * self.v_scale + + attn_output = self.attn(q, k, v) + + output, _ = self.o_proj(attn_output) + return output + + +class MiMoV2FlashDecoderLayer(nn.Module): + def __init__(self, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + config = vllm_config.model_config.hf_text_config + quant_config = vllm_config.quant_config + layer_id = extract_layer_index(prefix) + + self.hidden_size = config.hidden_size + self.config = config + self.layer_id = layer_id + + rope_theta = getattr(config, "rope_theta", 1000000) + max_position_embeddings = getattr(config, "max_position_embeddings", 32768) + + v_scale = getattr(config, "attention_value_scale", None) + + if self.is_compressed_softmax_layer(): + self.self_attn = MiMoV2Attention( + hidden_size=self.hidden_size, + num_heads=config.swa_num_attention_heads, + num_kv_heads=config.swa_num_key_value_heads, + head_dim=config.swa_head_dim, + v_head_dim=getattr(config, "swa_v_head_dim", None), + v_scale=v_scale, + sliding_window_size=config.sliding_window_size, + attention_bias=config.attention_bias, + add_swa_attention_sink_bias=getattr( + config, "add_swa_attention_sink_bias", False + ), + layer_id=layer_id, + rope_theta=getattr(config, "swa_rope_theta", rope_theta), + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), + prefix=f"{prefix}.self_attn", + ) + else: + self.self_attn = MiMoV2Attention( + hidden_size=self.hidden_size, + num_heads=config.num_attention_heads, + num_kv_heads=config.num_key_value_heads, + head_dim=config.head_dim, + v_head_dim=getattr(config, "v_head_dim", None), + v_scale=v_scale, + sliding_window_size=-1, # normal attention + attention_bias=config.attention_bias, + layer_id=layer_id, + rope_theta=rope_theta, + max_position_embeddings=max_position_embeddings, + quant_config=quant_config, + partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), + prefix=f"{prefix}.self_attn", + ) + + self.is_layer_sparse = self.is_moe_layer(layer_id) + if self.is_layer_sparse: + self.mlp = MiMoV2MoE( + vllm_config=vllm_config, + prefix=f"{prefix}.mlp", + ) + else: + self.mlp = MiMoV2MLP( + hidden_size=self.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + self.post_attention_layernorm = RMSNorm( + config.hidden_size, eps=config.layernorm_epsilon + ) + + def forward( + self, + positions: torch.Tensor, + hidden_states: torch.Tensor, + residual: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if residual is None: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + else: + hidden_states, residual = self.input_layernorm(hidden_states, residual) + + hidden_states = self.self_attn( + positions=positions, + hidden_states=hidden_states, + ) + + hidden_states, residual = self.post_attention_layernorm(hidden_states, residual) + hidden_states = self.mlp(hidden_states) + return hidden_states, residual + + def is_moe_layer(self, layer_idx: int) -> bool: + return ( + hasattr(self.config, "moe_layer_freq") + and layer_idx >= 0 + and not isinstance(self.config.moe_layer_freq, int) + and self.config.moe_layer_freq[layer_idx] + ) + + def is_compressed_softmax_layer(self) -> bool: + return self.config.hybrid_layer_pattern[self.layer_id] == 1 + + +@support_torch_compile +class MiMoV2Model(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + config = vllm_config.model_config.hf_config.get_text_config() + quant_config = vllm_config.quant_config + eplb_config = vllm_config.parallel_config.eplb_config + + self.config = config + self.quant_config = quant_config + self.vocab_size = config.vocab_size + self.num_redundant_experts = eplb_config.num_redundant_experts + + if get_pp_group().is_first_rank or ( + config.tie_word_embeddings and get_pp_group().is_last_rank + ): + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=f"{prefix}.embed_tokens", + ) + else: + self.embed_tokens = PPMissingLayer() + + self.start_layer, self.end_layer, self.layers = make_layers( + config.num_hidden_layers, + lambda prefix: MiMoV2FlashDecoderLayer( + vllm_config=vllm_config, + prefix=prefix, + ), + prefix=f"{prefix}.layers", + ) + + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( + ["hidden_states", "residual"], config.hidden_size + ) + if get_pp_group().is_last_rank: + self.norm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + else: + self.norm = PPMissingLayer() + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.embed_input_ids(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + + for idx, layer in enumerate( + islice(self.layers, self.start_layer, self.end_layer) + ): + hidden_states, residual = layer(positions, hidden_states, residual) + + if not get_pp_group().is_last_rank: + return IntermediateTensors( + {"hidden_states": hidden_states, "residual": residual} + ) + + hidden_states, _ = self.norm(hidden_states, residual) + + return hidden_states + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + # Params for weights, fp8 weight scales, fp8 activation scales + # (param_name, weight_name, expert_id, shard_id) + return fused_moe_make_expert_params_mapping( + self, + ckpt_gate_proj_name="gate_proj", + ckpt_down_proj_name="down_proj", + ckpt_up_proj_name="up_proj", + num_experts=self.config.n_routed_experts, + num_redundant_experts=self.num_redundant_experts, + ) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + ] + + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + expert_params_mapping = self.get_expert_mapping() + qkv_buffers: dict[str, dict[str, torch.Tensor]] = {} + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + if "rotary_emb.cos_cached" in name or "rotary_emb.sin_cached" in name: + continue + if "mtp" in name: + continue + + expert_matched = False + for param_name, weight_name, expert_id, shard_id in expert_params_mapping: + if weight_name not in name: + continue + + name_rewritten = name.replace(weight_name, param_name) + + if is_pp_missing_parameter(name_rewritten, self): + continue + + if ( + name_rewritten.endswith(".bias") or name_rewritten.endswith("_bias") + ) and name_rewritten not in params_dict: + continue + + if name_rewritten not in params_dict: + continue + + param = params_dict[name_rewritten] + weight_loader = param.weight_loader + + weight_loader( + param, + loaded_weight, + name_rewritten, + shard_id=shard_id, + expert_id=expert_id, + ) + loaded_params.add(name_rewritten) + expert_matched = True + break + + if expert_matched: + continue + qkv_pair = _mimo_v2_qkv_pair_key(name) + if qkv_pair is not None: + qkv_base_name, qkv_kind = qkv_pair + weight_name = f"{qkv_base_name}.weight" + scale_name = f"{qkv_base_name}.weight_scale_inv" + has_paired_qkv = ( + weight_name in params_dict + and scale_name in params_dict + and getattr(self.quant_config, "weight_block_size", None) + is not None + ) + if not has_paired_qkv: + if name in params_dict: + param = params_dict[name] + # Detect pre-sharded BF16 layout first. This applies + # when the source's fused QKV is the row-wise concat + # of per-ckpt-TP chunks (e.g. after Quark dequantizes + # the source FP8 to BF16 in --file2file_quantization). + # Vllm's QKVParallelLinear loader assumes canonical + # [Q_global,K,V] and would incorrectly slice the pre-sharded + # tensor; reassemble per-rank Q/K/V from ckpt chunks. + handled = ( + qkv_kind == "weight" + and _mimo_v2_copy_presharded_qkv_bf16( + config=self.config, + weight_name=name, + weight_param=param, + loaded_weight=loaded_weight, + tp_rank=tp_rank, + tp_size=tp_size, + ) + ) + if not handled: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + continue + + qkv_buffer = qkv_buffers.setdefault(qkv_base_name, {}) + qkv_buffer[qkv_kind] = loaded_weight + if "weight" in qkv_buffer and "scale" in qkv_buffer: + _mimo_v2_copy_paired_qkv_fp8( + config=self.config, + weight_name=weight_name, + scale_name=scale_name, + weight_param=params_dict[weight_name], + scale_param=params_dict[scale_name], + loaded_weight=qkv_buffer["weight"], + loaded_scale=qkv_buffer["scale"], + tp_rank=tp_rank, + tp_size=tp_size, + block_size=self.quant_config.weight_block_size, + ) + loaded_params.add(weight_name) + loaded_params.add(scale_name) + del qkv_buffers[qkv_base_name] + continue + stacked_matched = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name_rewritten = name.replace(weight_name, param_name) + + if ( + name_rewritten.endswith(".bias") + and name_rewritten not in params_dict + ): + continue + + if is_pp_missing_parameter(name_rewritten, self): + continue + + if name_rewritten not in params_dict: + continue + + param = params_dict[name_rewritten] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name_rewritten) + + stacked_matched = True + break + + if stacked_matched: + continue + + if name.endswith(".bias") and name not in params_dict: + continue + + orig_name = name + mapped_name = maybe_remap_kv_scale_name(name, params_dict) + name = mapped_name if mapped_name is not None else orig_name + + if name not in params_dict: + continue + + param = params_dict[name] + + if "attention_sink_bias" in name: + total_heads = loaded_weight.shape[0] + heads_per_rank = total_heads // tp_size + head_start = tp_rank * heads_per_rank + narrow_weight = loaded_weight.narrow(0, head_start, heads_per_rank) + + param.data.copy_(narrow_weight) + loaded_params.add(name) + else: + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + if qkv_buffers: + missing = ", ".join(sorted(qkv_buffers)) + raise RuntimeError( + "Missing fused-QKV FP8 weight/scale pair for MiMo-V2 " + f"checkpoint tensors: {missing}" + ) + + return loaded_params + + +class MiMoV2FlashForCausalLM(nn.Module, SupportsPP, MixtureOfExperts): + packed_modules_mapping = { + "qkv_proj": ["q_proj", "k_proj", "v_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + + self.config = config + self.quant_config = quant_config + self.model = MiMoV2Model( + vllm_config=vllm_config, + prefix=maybe_prefix(prefix, "model"), + ) + + if get_pp_group().is_last_rank: + self.lm_head = ParallelLMHead( + config.vocab_size, + config.hidden_size, + quant_config=quant_config, + prefix=maybe_prefix(prefix, "lm_head"), + ) + else: + self.lm_head = PPMissingLayer() + + self.logits_processor = LogitsProcessor(config.vocab_size) + + self.make_empty_intermediate_tensors = ( + self.model.make_empty_intermediate_tensors + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + ) -> torch.Tensor | IntermediateTensors: + hidden_states = self.model( + input_ids, positions, intermediate_tensors, inputs_embeds + ) + return hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: + logits = self.logits_processor(self.lm_head, hidden_states) + return logits + + def get_expert_mapping(self) -> list[tuple[str, str, int, str]]: + return self.model.get_expert_mapping() + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader(self) + return loader.load_weights(weights) + + +class MiMoV2ForCausalLM(MiMoV2FlashForCausalLM): + packed_modules_mapping = { + "qkv_proj": ["qkv_proj"], + "gate_up_proj": ["gate_proj", "up_proj"], + } diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/pr42270/mimo_v2_mtp.py b/contrib/models/MiMo-V2.5/perf_test/h100/pr42270/mimo_v2_mtp.py new file mode 100644 index 00000000..2a2bba0a --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/pr42270/mimo_v2_mtp.py @@ -0,0 +1,435 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Inference-only MiMo-V2 MTP (Multi-Token Prediction) draft model. + +Supports both MiMo-V2-Pro and MiMo-V2-Flash checkpoints. + +Checkpoint weight layout (model.mtp.layers.{idx}.*): + enorm - RMSNorm for token embeddings + hnorm - RMSNorm for previous hidden states + eh_proj - ReplicatedLinear(hidden*2 -> hidden) + input_layernorm - pre-attention RMSNorm + self_attn.* - attention weights; format differs by variant: + Pro: fused qkv_proj [Q;K;V] concatenated + Flash: separate q_proj, k_proj, v_proj + pre_mlp_layernorm - post-attention / pre-MLP RMSNorm + mlp.* - dense MLP (gate_proj / up_proj / down_proj) + final_layernorm - norm applied before logit computation +""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn +from transformers import PretrainedConfig + +from vllm.config import VllmConfig +from vllm.distributed import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.sequence import IntermediateTensors + +from .interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + _require_is_multimodal, +) +from .mimo_v2 import ( + MiMoV2Attention, + MiMoV2MLP, + _mimo_v2_copy_paired_qkv_fp8, + _mimo_v2_copy_presharded_qkv_bf16, + _mimo_v2_qkv_pair_key, +) +from .utils import _merge_multimodal_embeddings, maybe_prefix + +# MiMo-V2 checkpoints contain multiple MTP layers, but vLLM currently supports +# only the first layer +_MIMO_V2_PRO_NUM_MTP_LAYERS = 1 +_MIMO_V2_FLASH_NUM_MTP_LAYERS = 1 + + +class MiMoV2MTPLayer(nn.Module): + """Single MTP predictor layer for MiMo-V2 (Pro and Flash). + + Mirrors the single-layer MiMo-V2 nextn reference implementation. + """ + + def __init__( + self, + config: PretrainedConfig, + prefix: str, + quant_config: QuantizationConfig | None = None, + ) -> None: + super().__init__() + + # Predictor head components + self.enorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + self.hnorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + self.eh_proj = ReplicatedLinear( + config.hidden_size * 2, config.hidden_size, bias=False + ) + + # MTP uses the SWA attention configuration + # implementation. + swa_rope_theta = getattr( + config, + "swa_rope_theta", + getattr(config, "rope_theta", 1000000), + ) + sliding_window_size = getattr(config, "sliding_window_size", -1) + + self.input_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + self.self_attn = MiMoV2Attention( + hidden_size=config.hidden_size, + num_heads=config.swa_num_attention_heads, + num_kv_heads=config.swa_num_key_value_heads, + head_dim=config.swa_head_dim, + v_head_dim=getattr(config, "swa_v_head_dim", None), + v_scale=getattr(config, "attention_value_scale", None), + sliding_window_size=sliding_window_size, + attention_bias=config.attention_bias, + add_swa_attention_sink_bias=getattr( + config, "add_swa_attention_sink_bias", False + ), + layer_id=0, + rope_theta=swa_rope_theta, + max_position_embeddings=getattr(config, "max_position_embeddings", 32768), + quant_config=quant_config, + partial_rotary_factor=getattr(config, "partial_rotary_factor", 1.0), + prefix=f"{prefix}.self_attn", + ) + self.pre_mlp_layernorm = RMSNorm( + config.hidden_size, eps=config.layernorm_epsilon + ) + self.mlp = MiMoV2MLP( + hidden_size=config.hidden_size, + intermediate_size=config.intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + self.final_layernorm = RMSNorm(config.hidden_size, eps=config.layernorm_epsilon) + + def forward( + self, + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + ) -> torch.Tensor: + # Combine token embedding and previous hidden state + h, _ = self.eh_proj( + torch.cat( + [self.enorm(inputs_embeds), self.hnorm(previous_hidden_states)], dim=-1 + ) + ) + + # Transformer block with fused residual norms + residual = h + h = self.input_layernorm(h) + h = self.self_attn(positions=positions, hidden_states=h) + h, residual = self.pre_mlp_layernorm(h, residual) + h = self.mlp(h) + h = h + residual + + return self.final_layernorm(h) + + +class _MiMoV2MTPLayers(nn.Module): + """Thin wrapper so parameter paths match checkpoint: model.mtp.layers.*""" + + def __init__( + self, + config: PretrainedConfig, + num_mtp_layers: int, + quant_config: QuantizationConfig | None, + prefix: str, + ) -> None: + super().__init__() + self.layers = nn.ModuleDict( + { + str(i): MiMoV2MTPLayer( + config=config, + prefix=f"{prefix}.{i}", + quant_config=quant_config, + ) + for i in range(num_mtp_layers) + } + ) + + +class MiMoV2MultiTokenPredictor(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + + config = vllm_config.model_config.hf_config + spec_cfg = vllm_config.speculative_config + assert spec_cfg is not None + num_mtp_layers = 1 + + self.num_mtp_layers = num_mtp_layers + + self.embed_tokens = VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + + self.mtp = _MiMoV2MTPLayers( + config=config, + num_mtp_layers=num_mtp_layers, + quant_config=vllm_config.quant_config, + prefix=maybe_prefix(prefix, "mtp.layers"), + ) + + self.logits_processor = LogitsProcessor(config.vocab_size) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.embed_tokens(input_ids) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + if inputs_embeds is None: + inputs_embeds = self.embed_input_ids(input_ids) + current_step_idx = spec_step_idx % self.num_mtp_layers + return self.mtp.layers[str(current_step_idx)]( + inputs_embeds, positions, previous_hidden_states + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + lm_head: ParallelLMHead, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.logits_processor(lm_head, hidden_states) + + +class MiMoV2MTP(nn.Module): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + super().__init__() + self.config = vllm_config.model_config.hf_config + self.quant_config = vllm_config.quant_config + self.model = MiMoV2MultiTokenPredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model") + ) + self.lm_head = ParallelLMHead( + self.config.vocab_size, + self.config.hidden_size, + prefix=maybe_prefix(prefix, "lm_head"), + ) + + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) + + def forward( + self, + input_ids: torch.Tensor | None, + positions: torch.Tensor, + hidden_states: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + spec_step_idx: int = 0, + ) -> torch.Tensor: + return self.model( + input_ids, positions, hidden_states, inputs_embeds, spec_step_idx + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + spec_step_idx: int = 0, + ) -> torch.Tensor | None: + return self.model.compute_logits(hidden_states, self.lm_head, spec_step_idx) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + tp_rank = get_tensor_model_parallel_rank() + tp_size = get_tensor_model_parallel_world_size() + + stacked_params_mapping = [ + ("gate_up_proj", "gate_proj", 0), + ("gate_up_proj", "up_proj", 1), + # Flash format: separate projections → fused qkv_proj + ("qkv_proj", "q_proj", "q"), + ("qkv_proj", "k_proj", "k"), + ("qkv_proj", "v_proj", "v"), + ] + + params_dict = dict(self.named_parameters()) + loaded_params: set[str] = set() + qkv_buffers: dict[str, dict[str, torch.Tensor]] = {} + + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + + # Only load MTP-related weights, shared embeddings, and lm_head + if ( + "model.mtp" not in name + and "model.embed_tokens" not in name + and not name.startswith("lm_head") + ): + continue + + # Pro fused-QKV checkpoint: the FP8 weight is stored pre-sharded + # for TP=8 with per-chunk block_n=128 scales, and chunk_rows + # (3392 for MiMo-V2.5-Pro) is not a multiple of block_n. A naive + # chunk(tp_size) at the block-misaligned boundaries produces wrong + # slices. Route through the same paired loader the main model + # uses, which pairs weight + weight_scale_inv via the buffer below + # and either does a direct copy (when local TP == ckpt TP) or + # per-chunk dequant→requant for arbitrary local TP. + qkv_pair = _mimo_v2_qkv_pair_key(name) + if qkv_pair is not None: + qkv_base_name, qkv_kind = qkv_pair + weight_name = f"{qkv_base_name}.weight" + scale_name = f"{qkv_base_name}.weight_scale_inv" + has_paired_qkv = ( + weight_name in params_dict + and scale_name in params_dict + and getattr(self.quant_config, "weight_block_size", None) + is not None + ) + if not has_paired_qkv: + if name in params_dict: + param = params_dict[name] + # Detect pre-sharded BF16 layout first. This applies + # whenever the source's fused QKV was originally a + # row-wise concat of per-ckpt-TP chunks (e.g. after + # Quark dequantizes the source FP8 to BF16 in + # --file2file_quantization). Vllm's QKVParallelLinear + # loader assumes a canonical [Q_global,K,V] layout + # and would incorrectly slice the pre-sharded tensor. + handled = ( + qkv_kind == "weight" + and _mimo_v2_copy_presharded_qkv_bf16( + config=self.config, + weight_name=name, + weight_param=param, + loaded_weight=loaded_weight, + tp_rank=tp_rank, + tp_size=tp_size, + ) + ) + if not handled: + weight_loader = getattr( + param, "weight_loader", default_weight_loader + ) + weight_loader(param, loaded_weight) + loaded_params.add(name) + continue + + qkv_buffer = qkv_buffers.setdefault(qkv_base_name, {}) + qkv_buffer[qkv_kind] = loaded_weight + if "weight" in qkv_buffer and "scale" in qkv_buffer: + _mimo_v2_copy_paired_qkv_fp8( + config=self.config, + weight_name=weight_name, + scale_name=scale_name, + weight_param=params_dict[weight_name], + scale_param=params_dict[scale_name], + loaded_weight=qkv_buffer["weight"], + loaded_scale=qkv_buffer["scale"], + tp_rank=tp_rank, + tp_size=tp_size, + block_size=self.quant_config.weight_block_size, + ) + loaded_params.add(weight_name) + loaded_params.add(scale_name) + del qkv_buffers[qkv_base_name] + continue + + # gate_proj/up_proj → gate_up_proj stacking (both formats); + # Flash: q_proj/k_proj/v_proj → qkv_proj merging. + stacked_matched = False + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name_rewritten = name.replace(weight_name, param_name) + if ( + name_rewritten.endswith(".bias") + and name_rewritten not in params_dict + ): + continue + if name_rewritten not in params_dict: + continue + param = params_dict[name_rewritten] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, shard_id) + loaded_params.add(name_rewritten) + stacked_matched = True + break + + if stacked_matched: + continue + + if name.endswith(".bias") and name not in params_dict: + continue + if name not in params_dict: + continue + + param = params_dict[name] + # attention_sink_bias is head-parallel; slice by tp + if "attention_sink_bias" in name: + total_heads = loaded_weight.shape[0] + heads_per_rank = total_heads // tp_size + loaded_weight = loaded_weight.narrow( + 0, tp_rank * heads_per_rank, heads_per_rank + ) + + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + if qkv_buffers: + missing = ", ".join(sorted(qkv_buffers)) + raise RuntimeError( + "Missing fused-QKV FP8 weight/scale pair for MiMo-V2 MTP " + f"checkpoint tensors: {missing}" + ) + + return loaded_params + + +class MiMoV2OmniMTP(MiMoV2MTP, SupportsMultiModal): + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + ) -> torch.Tensor: + inputs_embeds = self._embed_text_input_ids( + input_ids, + self.model.embed_input_ids, + is_multimodal=is_multimodal, + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + is_multimodal = _require_is_multimodal(is_multimodal) + + inputs_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + return inputs_embeds diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_pr42270_c1.txt b/contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_pr42270_c1.txt new file mode 100644 index 00000000..1eea6da9 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_pr42270_c1.txt @@ -0,0 +1,43 @@ +INFO 07-21 22:43:20 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 22:43:20 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 22:43:21.338000 9 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 22:43:22 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 22:43:22 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +[transformers] You are using a model of type `mimo_v2` to instantiate a model of type ``. This may be expected if you are loading a checkpoint that shares a subset of the architecture (e.g., loading a `sam2_video` checkpoint into `Sam2Model`), but is otherwise not supported and can yield errors. Please verify that the checkpoint is compatible with the model you are instantiating. +[transformers] PreTrainedConfig got `key=rope_scaling` in kwargs but hasn't set it as attribute. For RoPE standardization you need to set `self.rope_parameters` in model's config. +Namespace(subparser='bench', bench_type='serve', dispatch_function=, trust_remote_code=False, seed=0, num_prompts=16, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=8000, endpoint='/v1/completions', header=None, max_concurrency=1, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-c17cf181-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 22:43:23 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 1 + 0%| | 0/16 [00:00, trust_remote_code=False, seed=0, num_prompts=96, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=8000, endpoint='/v1/completions', header=None, max_concurrency=16, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-a142506d-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 22:44:15 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 16 + 0%| | 0/96 [00:00, trust_remote_code=False, seed=0, num_prompts=96, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=8000, endpoint='/v1/completions', header=None, max_concurrency=32, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-1ec5eac6-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 22:45:18 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 32 + 0%| | 0/96 [00:00 -# 1792 rows) and multiplies it against the 1856-row weight group -> shape -# mismatch. This fires at EVERY 8-GPU shape: -# - TP=8: can't even start, "TP size must evenly split the 4 KV heads". -# - TP=4 / TP=2 (incl. DP=2 x TP=4): each rank gets 2 SWA KV heads -> the -# g>1 de-interleave path -> the 1856-vs-1792 crash. -# There is no vLLM TP config that both (a) splits the 4 full-attn KV heads and -# (b) gives 1 SWA KV head per rank, so the buggy path is unavoidable. -# -# apply_mimo_fp8_patch.py makes the server START (whole-tensor dequant), but the -# output is gibberish -- the re-quantization path is not numerically correct -- -# so it is NOT a usable fix, only a diagnostic. Use SGLang (run_sglang_h100.sh) -# for the H100 baseline. This script is kept to document the reference command -# and reproduce the bug. -# ===================================================================== -# -# The command below is the reference vLLM command from the model card -# (TP=8 + expert parallel, mimo parsers), which is what you WOULD run once the -# upstream loader is fixed. +# Usage: +# bash run_vllm_h100.sh set -e MODEL_DIR="${MODEL_DIR:-/opt/dlami/nvme/models/MiMo-V2.5}" @@ -37,8 +27,9 @@ PORT="${PORT:-8000}" TP="${TP:-8}" MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}" MAX_NUM_SEQS="${MAX_NUM_SEQS:-32}" # MiMo-V2.5 bs=32 -IMAGE="${IMAGE:-vllm-efa:latest}" +IMAGE="${IMAGE:-vllm/vllm-openai:nightly}" # PR #42270 not in a release yet CACHE_DIR="${CACHE_DIR:-/opt/dlami/nvme/vllm_cache}" +PATCH_DIR="${PATCH_DIR:-$(cd "$(dirname "$0")/pr42270" && pwd)}" CTR_MODEL="/models/MiMo-V2.5" mkdir -p "$CACHE_DIR" @@ -46,29 +37,40 @@ if [ ! -f "$MODEL_DIR/config.json" ]; then echo "ERROR: model not found at $MODEL_DIR" >&2 exit 1 fi +if [ ! -f "$PATCH_DIR/mimo_v2.py" ]; then + echo "ERROR: PR #42270 files not found at $PATCH_DIR" >&2 + exit 1 +fi echo "==========================================" -echo "MiMo-V2.5 vLLM (single-node TP=$TP + EP, Docker) -- EXPECTED TO FAIL" +echo "MiMo-V2.5 vLLM (single-node TP=$TP + EP, PR #42270, Docker)" echo " Port: $PORT max-model-len: $MAX_MODEL_LEN max-num-seqs: $MAX_NUM_SEQS" echo "==========================================" +M=/usr/local/lib/python3.12/dist-packages/vllm/model_executor/models exec docker run --rm --gpus all \ --network host --privileged --ipc=host --shm-size=32g \ -v "${MODEL_DIR}:${CTR_MODEL}:ro" \ -v "${CACHE_DIR}:/root/.cache" \ + -v "${PATCH_DIR}:/pr:ro" \ -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ - "$IMAGE" "$CTR_MODEL" \ - --served-model-name MiMo-V2.5 \ - --trust-remote-code \ - --generation-config vllm \ - --enable-expert-parallel \ - --tensor-parallel-size "$TP" \ - --max-model-len "$MAX_MODEL_LEN" \ - --max-num-seqs "$MAX_NUM_SEQS" \ - --no-enable-prefix-caching \ - --no-enable-chunked-prefill \ - --host 0.0.0.0 \ - --port "$PORT" \ - --tool-call-parser mimo \ - --enable-auto-tool-choice \ - --reasoning-parser mimo + --entrypoint bash \ + "$IMAGE" -c " + cp /pr/mimo_v2.py $M/mimo_v2.py && + cp /pr/mimo_v2_mtp.py $M/mimo_v2_mtp.py && + echo '[run_vllm_h100] applied PR #42270 model files' && + exec vllm serve '$CTR_MODEL' \ + --served-model-name MiMo-V2.5 \ + --trust-remote-code \ + --generation-config vllm \ + --enable-expert-parallel \ + --tensor-parallel-size $TP \ + --max-model-len $MAX_MODEL_LEN \ + --max-num-seqs $MAX_NUM_SEQS \ + --no-enable-prefix-caching \ + --no-enable-chunked-prefill \ + --host 0.0.0.0 \ + --port $PORT \ + --tool-call-parser mimo \ + --enable-auto-tool-choice \ + --reasoning-parser mimo" From d87a41accda8ceba6fd1424ac45950e813331ea7 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 22 Jul 2026 07:17:46 +0800 Subject: [PATCH 21/27] h100/MiMo-V2.5: enable chunked prefill for vLLM (default) + add its bench Chunked prefill works fine with PR #42270 (an earlier "it crashes" reading was a test-script bug, not real). It's the one optimization Neuron can't match, so it belongs ON for vLLM's real capability: c=16 output throughput jumps 214 -> 924 tok/s (~4.3x) and P99 TTFT drops 16s -> 0.4s vs chunked-off, by interleaving the 900-token prefills with decode instead of preempting. run_vllm_h100.sh now enables it by default (CHUNKED=0 restores exact Trn2 parity). Prefix caching and spec decode stay OFF. Added vllm_chunked_c{1,16,32} results and refreshed the perf tables: with chunked prefill, vLLM wins c=16 (924) among the GPU configs; SGLang-CP2 wins c=1/c=32. Best GPU config is 7.9-10.7x the Trn2 baseline. Co-Authored-By: Claude Opus 4.8 --- .../models/MiMo-V2.5/perf_test/h100/README.md | 45 ++++++++++--------- .../h100/results/vllm_chunked_c1.txt | 43 ++++++++++++++++++ .../h100/results/vllm_chunked_c16.txt | 43 ++++++++++++++++++ .../h100/results/vllm_chunked_c32.txt | 43 ++++++++++++++++++ .../MiMo-V2.5/perf_test/h100/run_vllm_h100.sh | 23 ++++++++-- 5 files changed, 172 insertions(+), 25 deletions(-) create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_chunked_c1.txt create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_chunked_c16.txt create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_chunked_c32.txt diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/README.md b/contrib/models/MiMo-V2.5/perf_test/h100/README.md index d241e870..66b04ab3 100644 --- a/contrib/models/MiMo-V2.5/perf_test/h100/README.md +++ b/contrib/models/MiMo-V2.5/perf_test/h100/README.md @@ -134,34 +134,37 @@ Slightly faster (a single request uses all 8 GPUs instead of 4 per DP group): | 16 | 792.68 | 8699.94 | 17.66 | 186 | 726 | | 32 | 1570.83 | 17240.36 | 17.78 | 237 | 293 | -**vLLM + PR #42270, TP=8 + EP** — `results/vllm_pr42270_c{1,16,32}.txt`: +**vLLM + PR #42270, TP=8 + EP, chunked prefill ON (default)** — +`results/vllm_chunked_c{1,16,32}.txt`: | Concurrency | Output throughput (tok/s) | Total throughput (tok/s) | TPOT median (ms) | TTFT median (ms) | TTFT P99 (ms) | |---|---|---|---|---|---| -| 1 | 48.26 | 529.33 | 5.42 | 84 | 12469 | -| 16 | 214.49 | 2354.09 | 16.74 | 224 | 16225 | -| 32 | 1459.66 | 16020.21 | 18.31 | 273 | 715 | +| 1 | 109.39 | 1199.76 | 5.41 | 83 | 3584 | +| 16 | 923.54 | 10136.18 | 15.35 | 216 | 412 | +| 32 | 1484.95 | 16297.79 | 18.20 | 255 | 697 | -vLLM's median TPOT is actually the lowest of the three (5.4 ms at c=1), but with -`--no-enable-chunked-prefill` (matching Trn2) its c=1/c=16 output throughput and -tail TTFT suffer badly — each new request's 900-token prefill preempts decode -(P99 TTFT 12–16 s at c≤16). By c=32 the pipeline stays saturated and vLLM catches -SGLang (1460 vs 1421 out tok/s). SGLang's DP-attention absorbs the prefill -interleaving far better at low/mid concurrency. +vLLM has the **lowest TPOT** of all configs (5.4 ms at c=1) and, with chunked +prefill, the **highest c=16 throughput** (924 tok/s). Chunked prefill is what +makes the difference: with it OFF (`CHUNKED=0`, the exact Trn2-parity config, +`results/vllm_pr42270_*`) each new request's 900-token prefill preempts decode +and c=16 collapses to 214 tok/s with 16 s P99 TTFT. Turning it on interleaves +prefill with decode and recovers ~4.3x at c=16. ### H100 vs Trn2 (vLLM-Neuron) — same 900/90 shape, out tok/s -| Concurrency | SGLang DP=2 | SGLang CP=2 | vLLM+PR#42270 | Trn2 | best / Trn2 | +| Concurrency | SGLang DP=2 | SGLang CP=2 | vLLM (chunked) | Trn2 | best / Trn2 | |---|---:|---:|---:|---:|---:| -| 1 | 91.10 | 125.69 | 48.26 | 15.88 | **7.9×** | -| 16 | 702.18 | 792.68 | 214.49 | 113.92 | **7.0×** | -| 32 | 1420.97 | 1570.83 | 1459.66 | 147.39 | **10.7×** | +| 1 | 91.10 | 125.69 | 109.39 | 15.88 | **7.9×** | +| 16 | 702.18 | 792.68 | 923.54 | 113.92 | **8.1×** | +| 32 | 1420.97 | 1570.83 | 1484.95 | 147.39 | **10.7×** | (Trn2 numbers from the main README's Performance section: BS=32, TP=64 / -moe_ep=64, CB + bucketing.) The H100 gap is largest at c=32 because the GPU -median ITL stays ~13–20 ms across concurrency while the Trn2 BS=32 TKG NEFF runs -at a fixed ~58 ms/token regardless of occupancy. - -> **Note:** these are the *language-only, no-spec-decode* baselines for -> apples-to-apples comparison with Trn2. Turning on EAGLE (`SPEC=1` on SGLang) -> would raise decode throughput further; the model card's full recipe enables it. +moe_ep=64, CB + bucketing.) All three GPU configs land within ~2x of each other; +vLLM-chunked wins c=16, SGLang-CP2 wins c=1 and c=32. The H100/Trn2 gap is +largest at c=32 because the GPU median ITL stays ~13–20 ms across concurrency +while the Trn2 BS=32 TKG NEFF runs at a fixed ~58 ms/token regardless of occupancy. + +> **Note:** these are the *language-only, no-spec-decode* baselines. All GPU runs +> keep prefix caching + speculative decode OFF; vLLM keeps chunked prefill ON +> (the one optimization Neuron can't match — Trn2 uses fixed-shape CTE/TKG NEFFs). +> Turning on EAGLE (`SPEC=1` on SGLang) would raise decode throughput further. diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_chunked_c1.txt b/contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_chunked_c1.txt new file mode 100644 index 00000000..a472be0e --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/vllm_chunked_c1.txt @@ -0,0 +1,43 @@ +INFO 07-21 23:06:02 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 23:06:02 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 23:06:03.214000 9 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 23:06:04 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 23:06:04 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +[transformers] You are using a model of type `mimo_v2` to instantiate a model of type ``. This may be expected if you are loading a checkpoint that shares a subset of the architecture (e.g., loading a `sam2_video` checkpoint into `Sam2Model`), but is otherwise not supported and can yield errors. Please verify that the checkpoint is compatible with the model you are instantiating. +[transformers] PreTrainedConfig got `key=rope_scaling` in kwargs but hasn't set it as attribute. For RoPE standardization you need to set `self.rope_parameters` in model's config. +Namespace(subparser='bench', bench_type='serve', dispatch_function=, trust_remote_code=False, seed=0, num_prompts=16, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=8000, endpoint='/v1/completions', header=None, max_concurrency=1, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-d08c46c8-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 23:06:05 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 1 + 0%| | 0/16 [00:00, trust_remote_code=False, seed=0, num_prompts=96, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=8000, endpoint='/v1/completions', header=None, max_concurrency=16, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-aa005b1b-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 23:06:40 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 16 + 0%| | 0/96 [00:00, trust_remote_code=False, seed=0, num_prompts=96, dataset_name='random', no_stream=False, dataset_path=None, no_oversample=False, skip_chat_template=False, enable_multimodal_chat=False, disable_shuffle=False, custom_output_len=256, custom_ensure_client_side_data=False, spec_bench_output_len=256, spec_bench_category=None, sonnet_input_len=550, sonnet_output_len=150, sonnet_prefix_len=200, sharegpt_output_len=None, timed_trace_chunk_hash_size=16, timed_trace_sec_multiplier=1, timed_trace_label_timestamp='timestamp', timed_trace_label_input_length='input_length', timed_trace_label_output_length='output_length', timed_trace_label_hash_ids='hash_ids', blazedit_min_distance=0.0, blazedit_max_distance=1.0, asr_max_audio_len_sec=inf, asr_min_audio_len_sec=0.0, random_input_len=900, random_output_len=90, random_range_ratio='0.03', random_prefix_len=0, random_batch_size=1, no_reranker=False, random_mm_base_items_per_request=1, random_mm_num_mm_items_range_ratio=0.0, random_mm_limit_mm_per_prompt={'image': 255, 'video': 1}, random_mm_bucket_config={(256, 256, 1): 0.5, (720, 1280, 1): 0.5, (720, 1280, 16): 0.0}, hf_subset=None, hf_split=None, hf_name=None, hf_output_len=None, bfcl_categories=None, prefix_repetition_prefix_len=256, prefix_repetition_suffix_len=256, prefix_repetition_num_prefixes=10, prefix_repetition_output_len=128, speed_bench_dataset_subset='qualitative', speed_bench_output_len=4096, speed_bench_category=None, label=None, backend='vllm', base_url=None, host='localhost', port=8000, endpoint='/v1/completions', header=None, max_concurrency=32, model='MiMo-V2.5', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5', tokenizer_mode='auto', use_beam_search=False, logprobs=None, request_rate=inf, burstiness=1.0, disable_tqdm=False, num_warmups=0, profile=False, save_result=False, save_detailed=False, append_result=False, metadata=None, result_dir=None, result_filename=None, ignore_eos=False, self_timed=None, percentile_metrics=None, metric_percentiles='99', goodput=None, request_id_prefix='bench-59504eff-', top_p=None, top_k=None, min_p=None, temperature=None, frequency_penalty=None, presence_penalty=None, repetition_penalty=None, served_model_name=None, lora_modules=None, lora_assignment='random', ramp_up_strategy=None, ramp_up_start_rps=None, ramp_up_end_rps=None, ready_check_timeout_sec=0, chat_template_kwargs=None, extra_body=None, skip_tokenizer_init=False, insecure=False, plot_timeline=False, timeline_itl_thresholds='25,50', plot_dataset_stats=False) +INFO 07-21 23:07:11 [utils.py:90] Sampling input_len from [873, 927] and output_len from [87, 93] +WARNING: vllm bench serve no longer sets temperature==0 (greedy) in requests by default. The default will be determined on the server side and can be model/API specific. For the old behavior, include --temperature=0. +Starting initial single prompt test run... +Skipping endpoint ready check. +Starting main benchmark run... +Traffic request rate: inf +Burstiness factor: 1.0 (Poisson process) +Maximum request concurrency: 32 + 0%| | 0/96 [00:00 Date: Wed, 22 Jul 2026 08:03:38 +0800 Subject: [PATCH 22/27] h100/MiMo-V2.5: add 32K long-context concurrency exploration (SGLang) How many concurrent 32K-input requests fit on one 8xH100 node? Added run_sglang_h100_32k.sh (DP=2, context-length 34816, mem-frac 0.9) and bench_32k.sh (concurrency sweep, 32K in / 128 out). Findings (results/longctx_32k/): client concurrency 1..128 all complete with ZERO failures, sustained by SGLang's queue. Server-side resident requests peak at ~20-27 (KV pool max_total_num_tokens=678149 / 32768 ~= 20.7 full-32K reqs); higher concurrency queues (#queue-req peaked ~55 at c=128). Throughput saturates by c=16-32 at ~46K tok/s (prefill-bound); beyond that only TTFT grows (11s median at c=32 -> 65s at c=128). 32K fits because hybrid attention caps 39/48 layers' KV to the 128 window (full-token usage ~0.39, swa ~0.08 at 32K). Left the c=32 row as the intended Trn2 cross-platform comparison point (to fill). Co-Authored-By: Claude Opus 4.8 --- .../models/MiMo-V2.5/perf_test/h100/README.md | 43 +++++++++++ .../MiMo-V2.5/perf_test/h100/bench_32k.sh | 29 ++++++++ .../results/longctx_32k/longctx_32k_sweep.txt | 72 +++++++++++++++++++ .../perf_test/h100/run_sglang_h100_32k.sh | 59 +++++++++++++++ 4 files changed, 203 insertions(+) create mode 100755 contrib/models/MiMo-V2.5/perf_test/h100/bench_32k.sh create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/longctx_32k/longctx_32k_sweep.txt create mode 100755 contrib/models/MiMo-V2.5/perf_test/h100/run_sglang_h100_32k.sh diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/README.md b/contrib/models/MiMo-V2.5/perf_test/h100/README.md index 66b04ab3..6f2574ee 100644 --- a/contrib/models/MiMo-V2.5/perf_test/h100/README.md +++ b/contrib/models/MiMo-V2.5/perf_test/h100/README.md @@ -168,3 +168,46 @@ while the Trn2 BS=32 TKG NEFF runs at a fixed ~58 ms/token regardless of occupan > keep prefix caching + speculative decode OFF; vLLM keeps chunked prefill ON > (the one optimization Neuron can't match — Trn2 uses fixed-shape CTE/TKG NEFFs). > Turning on EAGLE (`SPEC=1` on SGLang) would raise decode throughput further. + +## Long-context: 32K input, how far does concurrency go? (SGLang) + +How many concurrent 32K-input requests can one 8xH100 node sustain? Server: +`run_sglang_h100_32k.sh` (cookbook DP=2, `--context-length 34816 +--mem-fraction-static 0.9`). Sweep: `bench_32k.sh 1 8 16 32 64 128` (32K input / +128 output, `2*c` prompts per level). Raw logs in `results/longctx_32k/`. + +| Client concurrency (`--max-concurrency`) | Succeeded | Total throughput (tok/s) | TTFT median (s) | TTFT P99 (s) | TPOT median (ms) | +|---|---|---|---|---|---| +| 1 | 2/2 ✅ | 11597 | 1.6 | 2.1 | 9.9 | +| 8 | 16/16 ✅ | 39038 | 3.6 | 5.3 | 24 | +| 16 | 32/32 ✅ | 43433 | 5.9 | 10.3 | 48 | +| 32 | 64/64 ✅ | 46423 | 11.0 | 20.6 | 92 | +| 64 | 128/128 ✅ | 46227 | 20.4 | 42.9 | 173 | +| 128 | 256/256 ✅ | 46253 | 64.9 | 87.7 | 196 | + +**Answer: up to c=128 completes with zero failures** — but that's the client's +in-flight count, absorbed by SGLang's queue. Two different "concurrency" numbers: + +- **Client concurrency** (`Maximum request concurrency` in the bench output) = the + value passed to `--max-concurrency`; the client keeps that many requests + in-flight. Everything from 1 to 128 runs clean. +- **Server-side resident** (`#running-req` in the server log) peaks at **~20–27**, + regardless of how high the client pushes; the rest sit in the queue + (`#queue-req` peaked at ~55 at c=128). This ceiling is the KV pool: + `max_total_num_tokens = 678149` / 32768 ≈ **20.7 fully-resident 32K requests**. + +**Why 32K fits at all**: hybrid attention. Only 9 of 48 layers are full attention; +the 39 sliding-window layers (window 128) have window-capped KV. At 32K the server +logs **full-token usage ≈ 0.39, SWA-token usage ≈ 0.08** — the long prompt's KV is +dominated by just the 9 full layers, so KV per request (~780 MB) is far below a +same-size dense model. + +**Practical guidance**: +- **Max throughput** saturates by **c=16–32** (~46K tok/s total, prefill-bound); + higher concurrency only grows the queue, not throughput. +- **Low TTFT** (no queueing): keep client concurrency **≤ ~20** (≈ resident cap). +- **c=64/128 works** but is pure queueing — throughput flat, TTFT climbs to the + minute range. Bump `--mem-fraction-static` toward 0.92+ to raise the resident cap. + +> **Trn2 comparison (pending)**: the c=32 row is the intended cross-platform point. +> Trn2 32K-input / c=32 numbers to be filled in once measured. diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/bench_32k.sh b/contrib/models/MiMo-V2.5/perf_test/h100/bench_32k.sh new file mode 100755 index 00000000..0ef85904 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/bench_32k.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# 32K-input concurrency sweep against a running long-context server +# (run_sglang_h100_32k.sh on :30000). Output 128 tok isolates the +# prefill + KV-capacity limit. Sends 2x concurrency prompts per level so the +# client keeps `--max-concurrency` requests in flight throughout. +# +# Usage: +# bash bench_32k.sh 1 8 16 32 64 128 # full sweep +# bash bench_32k.sh 32 # single level (e.g. to match Trn2) +set -e +PORT="${PORT:-30000}" +INPUT_LEN="${INPUT_LEN:-32768}" +OUTPUT_LEN="${OUTPUT_LEN:-128}" +RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/models/bench_results/mimo_v2_5_h100/longctx_32k}" +mkdir -p "$RESULTS_DIR" + +for C in "$@"; do + NP=$((C * 2)) + echo "===== ${INPUT_LEN} in / ${OUTPUT_LEN} out c=$C np=$NP =====" + docker run --rm --network host -v /opt/dlami/nvme/models:/wk --entrypoint bash \ + vllm-efa:latest -c " + vllm bench serve --backend vllm --host localhost --port $PORT \ + --model MiMo-V2.5 --tokenizer /wk/MiMo-V2.5 --endpoint /v1/completions \ + --dataset-name random --num-prompts $NP \ + --random-input-len $INPUT_LEN --random-output-len $OUTPUT_LEN \ + --random-range-ratio 0.02 --max-concurrency $C" 2>&1 | \ + tee "$RESULTS_DIR/c${C}.txt" | \ + grep -E "Successful requests|Failed|Maximum request concurrency:|Benchmark duration|Total input tokens|Output token throughput|Total token throughput|Median TTFT|P99 TTFT|Median TPOT|Median E2E|Mean E2E" +done diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/longctx_32k/longctx_32k_sweep.txt b/contrib/models/MiMo-V2.5/perf_test/h100/results/longctx_32k/longctx_32k_sweep.txt new file mode 100644 index 00000000..26b98e2d --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/longctx_32k/longctx_32k_sweep.txt @@ -0,0 +1,72 @@ +===== 32K in / 128 out c=1 np=2 ===== +Maximum request concurrency: 1 +Successful requests: 2 +Failed requests: 0 +Maximum request concurrency: 1 +Benchmark duration (s): 5.73 +Total input tokens: 66176 +Output token throughput (tok/s): 44.34 +Total token throughput (tok/s): 11597.34 +Median TTFT (ms): 1621.26 +P99 TTFT (ms): 2092.67 +Median TPOT (ms): 9.86 +===== 32K in / 128 out c=8 np=16 ===== +Maximum request concurrency: 8 +Successful requests: 16 +Failed requests: 0 +Maximum request concurrency: 8 +Benchmark duration (s): 13.48 +Total input tokens: 524380 +Output token throughput (tok/s): 152.24 +Total token throughput (tok/s): 39038.47 +Median TTFT (ms): 3635.72 +P99 TTFT (ms): 5271.67 +Median TPOT (ms): 24.01 +===== 32K in / 128 out c=16 np=32 ===== +Maximum request concurrency: 16 +Successful requests: 32 +Failed requests: 0 +Maximum request concurrency: 16 +Benchmark duration (s): 24.26 +Total input tokens: 1049711 +Output token throughput (tok/s): 168.65 +Total token throughput (tok/s): 43433.19 +Median TTFT (ms): 5894.95 +P99 TTFT (ms): 10309.44 +Median TPOT (ms): 48.07 +===== 32K in / 128 out c=32 np=64 ===== +Maximum request concurrency: 32 +Successful requests: 64 +Failed requests: 0 +Maximum request concurrency: 32 +Benchmark duration (s): 45.37 +Total input tokens: 2097979 +Output token throughput (tok/s): 180.72 +Total token throughput (tok/s): 46422.67 +Median TTFT (ms): 10967.79 +P99 TTFT (ms): 20627.90 +Median TPOT (ms): 91.52 +===== 32K in / 128 out c=64 np=128 ===== +Maximum request concurrency: 64 +Successful requests: 128 +Failed requests: 0 +Maximum request concurrency: 64 +Benchmark duration (s): 91.13 +Total input tokens: 4196114 +Output token throughput (tok/s): 180.09 +Total token throughput (tok/s): 46227.46 +Median TTFT (ms): 20360.99 +P99 TTFT (ms): 42893.79 +Median TPOT (ms): 173.11 +===== 32K in / 128 out c=128 np=256 ===== +Maximum request concurrency: 128 +Successful requests: 256 +Failed requests: 0 +Maximum request concurrency: 128 +Benchmark duration (s): 182.23 +Total input tokens: 8395750 +Output token throughput (tok/s): 180.00 +Total token throughput (tok/s): 46253.30 +Median TTFT (ms): 64895.67 +P99 TTFT (ms): 87692.25 +Median TPOT (ms): 196.05 diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/run_sglang_h100_32k.sh b/contrib/models/MiMo-V2.5/perf_test/h100/run_sglang_h100_32k.sh new file mode 100755 index 00000000..13173e57 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/run_sglang_h100_32k.sh @@ -0,0 +1,59 @@ +#!/bin/bash +# Long-context SGLang server for MiMo-V2.5 on 8xH100: 32K input capacity. +# Same cookbook DP=2 shape as run_sglang_h100.sh, but with a raised context +# length and mem-fraction to maximize the KV pool for long prompts. +# +# Why this fits so well: MiMo-V2.5 is HYBRID attention -- only 9 of 48 layers are +# full attention (4 KV heads); the other 39 are sliding-window (window=128), whose +# KV is capped by the window. So a 32K prompt's KV is dominated by just the 9 full +# layers. At 32K the server logs "full token usage ~0.39, swa token usage ~0.08". +# +# Observed KV pool (mem-fraction 0.9): max_total_num_tokens = 678149 +# -> ~20 fully-resident 32K requests (678149 / 32768). Higher client concurrency +# still completes via SGLang's queue/preemption, just with growing TTFT. +# +# Usage: +# bash run_sglang_h100_32k.sh +# MEM=0.92 MML=40960 bash run_sglang_h100_32k.sh # push the KV pool higher +set -e + +MODEL_DIR="${MODEL_DIR:-/opt/dlami/nvme/models/MiMo-V2.5}" +PORT="${PORT:-30000}" +MEM="${MEM:-0.9}" # weights ~37GB/GPU; rest -> KV +MML="${MML:-34816}" # 32768 in + 2048 out headroom +MAX_RUN="${MAX_RUN:-256}" # scheduler cap (not KV-resident cap) +IMAGE="${IMAGE:-sglang-efa:latest}" +CACHE_DIR="${CACHE_DIR:-/opt/dlami/nvme/sglang_cache}" +CTR_MODEL="/models/MiMo-V2.5" +mkdir -p "$CACHE_DIR" + +if [ ! -f "$MODEL_DIR/config.json" ]; then + echo "ERROR: model not found at $MODEL_DIR" >&2 + exit 1 +fi + +echo "==========================================" +echo "MiMo-V2.5 SGLang 32K-context (single-node 8xH100, Docker)" +echo " context-length: $MML mem-frac: $MEM max-running-requests: $MAX_RUN" +echo "==========================================" + +exec docker run --rm --gpus all \ + --network host --privileged --ipc=host --shm-size=32g \ + -v "${MODEL_DIR}:${CTR_MODEL}:ro" \ + -v "${CACHE_DIR}:/root/.cache" \ + -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + "$IMAGE" \ + python3 -m sglang.launch_server \ + --model-path "$CTR_MODEL" \ + --served-model-name MiMo-V2.5 \ + --trust-remote-code \ + --tp 8 --dp 2 --enable-dp-attention --enable-dp-lm-head --mm-enable-dp-encoder \ + --moe-dense-tp-size 1 \ + --mem-fraction-static "$MEM" \ + --context-length "$MML" \ + --chunked-prefill-size 16384 \ + --max-running-requests "$MAX_RUN" \ + --reasoning-parser mimo \ + --tool-call-parser mimo \ + --host 0.0.0.0 \ + --port "$PORT" From 163a7c6804da4f78df04163e42de5eb3702f3ac2 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 23 Jul 2026 02:02:37 +0800 Subject: [PATCH 23/27] Add data-parallel + context-parallel attention for long context (512->16K) Enables MiMo-V2.5 to serve long context (validated at seq_len=16384, 32x the prior 512) on Trn2 by combining three mechanisms, since MiMo's asymmetric head dims (Q/K=192, V=128) force a hand-rolled attention that bypassed NxDI's stock data-parallel / context-parallel machinery. 1. Data-parallel attention (attention_dp_degree). MiMo previously hard-wired all attention projections to the full TP=64 group, so tp_degree(64) > kv_heads(4/8) forced CONVERT_TO_MHA, replicating K/V up to 64 heads (an 8-16x KV cache bloat that capped seq_len at 512). We now bind the projections to the reduced data-parallel-attention TP group (get_data_parallel_attention_tp_group, size tp_degree//attention_dp_degree) so attn TP == kv_heads with ZERO replication. Decode splits the batch across DP groups and all-gathers the output before the residual/MoE. All local head counts, the KV-cache head count fed to DataParallelKVCacheManager, the CONVERT_TO_MHA gate, and the checkpoint-conversion replication gate now key off attn_tp_degree instead of the full tp_degree. attention_dp_degree=1 (default) reduces to the original behavior (backward compatible). The attention-sink bias is indexed by the rank WITHIN the DP-attention group. 2. Skip the base chunked-attention validator. The HF config carries attention_chunk_size (== sliding_window), which trips _validate_chunked_attention_support (hard-requires attention_dp_degree==1), but MiMo implements sliding-window attention entirely in its own modeling code and never uses the framework chunked-attention path, so the validator is a no-op here and is overridden. 3. Cache the FULL-S K/V under context parallelism (bugfix). In the CP prefill path the KV cache must hold ALL sequence positions; previously the local S/cp_degree chunk was saved BEFORE the all-gather, so decode attended to only ~1/cp_degree of the context, producing grammatical-but-wrong, repetitive, drifting output (e.g. "1 2 3...9" -> "1 1 1", "author of Hamlet" -> "William Smith"). We now gather K/V to full-S first, then cache, matching attention_base. This restored correct output at 16K. Validated config (seq 16384): tp_degree=64, attention_dp_degree=16, cp_degree=16, moe_ep_degree=64. Post-fix outputs are factually correct ("author of Hamlet is William Shakespeare", "1 2 3...9 10 11 12 13 14"). DP=1 cannot serve 16K (KV cache alone would need ~1.2TB, exceeding total HBM). Co-Authored-By: Claude Opus 4.8 --- .../models/MiMo-V2.5/src/modeling_mimo_v2.py | 186 +++++++++++++++--- 1 file changed, 162 insertions(+), 24 deletions(-) diff --git a/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py index 3e13adc9..6b00cbd3 100644 --- a/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py +++ b/contrib/models/MiMo-V2.5/src/modeling_mimo_v2.py @@ -196,6 +196,27 @@ def __init__(self, *args, **kwargs): # Parse hybrid layer pattern self._parse_hybrid_pattern() + def _validate_chunked_attention_support(self): + """Skip the base class's chunked-attention validator for MiMo. + + MiMo implements sliding-window attention ENTIRELY in its own modeling + code (the hand-rolled per-row sliding_mask, driven by + sliding_window_size), and never uses the base class's framework-level + "chunked attention" path. The HF config carries attention_chunk_size + (== sliding_window = 128), which InferenceConfig absorbs as an + attribute and which trips the base _validate_chunked_attention_support: + that validator hard-requires attention_dp_degree == 1 (among others), + which would forbid the data-parallel attention we use to eliminate the + CONVERT_TO_MHA KV replication. Since the chunked-attention code path is + inapplicable to MiMo, this validator is a no-op here. + + Note: the base calls validate_config() (hence this) at the END of + InferenceConfig.__init__, i.e. during MiMoV2InferenceConfig's + super().__init__() — so overriding the method (rather than deleting the + attribute afterward) is what actually takes effect in time. + """ + return + def _parse_hybrid_pattern(self): """Parse hybrid layer pattern to determine attention types.""" if hasattr(self, 'hybrid_layer_pattern') and self.hybrid_layer_pattern: @@ -359,7 +380,10 @@ def __init__( config.num_key_value_heads, getattr(config, 'swa_num_key_value_heads', config.num_key_value_heads) ) - self.local_cache_kv_heads = max(1, self.cache_num_kv_heads // tp_degree) + # Divide by the attention TP degree (== full tp_degree when DP=1), + # matching DataParallelKVCacheManager._get_num_kv_heads_per_rank + # which sizes the per-rank cache as num_kv_heads // (tp // dp). + self.local_cache_kv_heads = max(1, self.cache_num_kv_heads // self.attn_tp_degree) def init_gqa_properties(self): """Override base class to prevent creating incompatible QKV projections. @@ -383,8 +407,23 @@ def _init_projections(self, config: MiMoV2InferenceConfig): dtype = config.neuron_config.torch_dtype tp_degree = config.neuron_config.tp_degree - # Check if we need GQA CONVERT_TO_MHA (when tp_degree > num_kv_heads) - self.use_gqa_convert_to_mha = tp_degree > self.attn_num_kv_heads + # Data-parallel attention: when attention_dp_degree > 1, attention runs + # on a reduced TP group of size tp_degree // attention_dp_degree, while + # MoE keeps the full TP/EP. This lets attn TP == num_kv_heads so no + # CONVERT_TO_MHA replication is needed (killing the 8-16x KV cache + # bloat). We run BOTH prefill (CTE) and decode (TKG) on this same + # reduced group (unlike the base class which uses full-TP CTE + a + # DP-TP decode + rank reordering); this avoids the base's + # validate_tp_prefill_to_dp_decode assertion, which the SWA layers + # (8 kv heads) would otherwise fail. See the attn_dp design memo. + self.attention_dp_degree = getattr( + config.neuron_config, "attention_dp_degree", 1 + ) or 1 + self.attn_tp_degree = tp_degree // self.attention_dp_degree + + # Check if we need GQA CONVERT_TO_MHA. With DP attention the relevant + # comparison is the *attention* TP degree, not the full TP degree. + self.use_gqa_convert_to_mha = self.attn_tp_degree > self.attn_num_kv_heads # Store source heads for preshard_hook self._src_num_kv_heads = self.attn_num_kv_heads @@ -405,7 +444,17 @@ def _init_projections(self, config: MiMoV2InferenceConfig): o_hidden_size = self.attn_num_heads * self.attn_v_head_dim if parallel_state.model_parallel_is_initialized(): - tp_group = parallel_state.get_tensor_model_parallel_group() + if self.attention_dp_degree > 1: + # Bind attention projections to the reduced DP-attention TP + # group (size = attn_tp_degree). init is idempotent/global. + from neuronx_distributed_inference.modules.attention.attention_process_groups import ( + init_data_parallel_attention_process_groups, + get_data_parallel_attention_tp_group, + ) + init_data_parallel_attention_process_groups(config) + tp_group = get_data_parallel_attention_tp_group() + else: + tp_group = parallel_state.get_tensor_model_parallel_group() # Q projection self.q_proj = ColumnParallelLinear( @@ -449,13 +498,14 @@ def _init_projections(self, config: MiMoV2InferenceConfig): sequence_dimension=1 if self.sequence_parallel_enabled else None, ) - # Calculate local dimensions after TP split - self.local_num_heads = self.attn_num_heads // tp_degree + # Calculate local dimensions after the *attention* TP split + # (attn_tp_degree = tp_degree // attention_dp_degree). + self.local_num_heads = self.attn_num_heads // self.attn_tp_degree if self.use_gqa_convert_to_mha: # With CONVERT_TO_MHA, local KV heads = local Q heads self.local_num_kv_heads = self.local_num_heads else: - self.local_num_kv_heads = max(1, self.attn_num_kv_heads // tp_degree) + self.local_num_kv_heads = max(1, self.attn_num_kv_heads // self.attn_tp_degree) else: self.q_proj = nn.Linear(config.hidden_size, q_hidden_size, bias=config.attention_bias) self.k_proj = nn.Linear(config.hidden_size, k_hidden_size, bias=config.attention_bias) @@ -552,11 +602,47 @@ def forward( hidden_states, dim=1, rank=cp_rank, num_partitions=self.cp_degree ) - bsz, q_len, _ = hidden_states.size() - # Determine if this is token generation (past_key_value is not None) is_token_gen = past_key_value is not None + # Data-parallel attention (decode only): split the batch across the + # attention DP groups so each group processes batch // dp rows against + # its own slice of the KV cache (DataParallelKVCacheManager sizes the + # cache to kv_cache_batch_size = tkg_batch // dp). Prefill (CTE) runs on + # every DP group redundantly at ctx_batch_size=1 (only the owning + # group's KV write persists); we accept the redundant prefill compute + # to keep the CTE/TKG head layout identical (see attn_dp design memo). + is_data_parallel_attn = ( + is_token_gen + and self.attention_dp_degree > 1 + and parallel_state.model_parallel_is_initialized() + ) + if is_data_parallel_attn: + from neuronx_distributed_inference.modules.attention.attention_process_groups import ( + get_data_parallel_attention_dp_group, + ) + from neuronx_distributed_inference.utils.distributed import get_dp_rank + + dp_rank = get_dp_rank( + self.rank_util.get_rank(), + self.attn_tp_degree, + self.attention_dp_degree, + self.neuron_config.switch_cc, + ) + hidden_states = split_along_dim( + hidden_states, dim=0, rank=dp_rank, num_partitions=self.attention_dp_degree + ) + if attention_mask is not None: + attention_mask = split_along_dim( + attention_mask, dim=0, rank=dp_rank, num_partitions=self.attention_dp_degree + ) + if position_ids is not None: + position_ids = split_along_dim( + position_ids, dim=0, rank=dp_rank, num_partitions=self.attention_dp_degree + ) + + bsz, q_len, _ = hidden_states.size() + # Project Q, K, V query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) @@ -617,9 +703,19 @@ def forward( q_len = q_len // self.cp_degree # K/V stay at full S for attention computation else: - # Q/K/V have S/CP. Save local KV for cache, then all-gather K/V. - key_states_for_cache = key_states - value_states_for_cache = value_states + # Q/K/V have S/CP. All-gather K/V to full S, then save the + # FULL-S K/V for the cache. + # + # BUGFIX: the cache must hold ALL sequence positions, not just + # this CP rank's local S/CP chunk. Previously key/value_states_for_cache + # were assigned the local chunk BEFORE the gather, so the + # DataParallelKVCacheManager wrote only ~1/cp_degree of the + # sequence's KV. Decode then attended to a wildly incomplete + # context (only the owning CP-group chunk), producing + # grammatical-but-wrong / repetitive / drifting output that + # degraded with generation length. The base class gathers first + # then caches the full-S tensor (attention_base.py:548-563); we + # now match that. key_states = gather_from_tensor_model_parallel_region_with_dim( key_states, gather_dim=2, process_group=get_context_parallel_attention_cp_group(), @@ -628,6 +724,8 @@ def forward( value_states, gather_dim=2, process_group=get_context_parallel_attention_cp_group(), ) + key_states_for_cache = key_states + value_states_for_cache = value_states # Q stays at S/CP else: # Store key/value states BEFORE GQA repeat for KV cache @@ -752,8 +850,14 @@ def forward( # This must be applied to token generation as well! use_sink = self._use_sink_bias and self.attention_sink_bias is not None if use_sink: - tp_rank = parallel_state.get_tensor_model_parallel_rank() if parallel_state.model_parallel_is_initialized() else 0 - local_sink = self.attention_sink_bias[tp_rank * self.local_num_heads:(tp_rank + 1) * self.local_num_heads] + global_rank = parallel_state.get_tensor_model_parallel_rank() if parallel_state.model_parallel_is_initialized() else 0 + # Under DP attention the Q heads are sharded across the reduced + # TP group (attn_tp_degree), and every DP group shards the 64 + # heads identically. The sink bias (stored per global head) must + # be indexed by the rank WITHIN the attention TP group, not the + # global rank, else ranks >= attn_tp_degree index out of bounds. + attn_tp_rank = global_rank % self.attn_tp_degree + local_sink = self.attention_sink_bias[attn_tp_rank * self.local_num_heads:(attn_tp_rank + 1) * self.local_num_heads] sink_bias = local_sink.reshape(1, -1, 1, 1).expand(bsz, -1, q_len, 1) all_scores = torch.cat([all_scores, sink_bias], dim=-1) @@ -813,9 +917,12 @@ def forward( # This adds an extra "sink" column to attention weights use_sink = self._use_sink_bias and self.attention_sink_bias is not None if use_sink: - # Get local portion of sink bias for this TP rank - tp_rank = parallel_state.get_tensor_model_parallel_rank() if parallel_state.model_parallel_is_initialized() else 0 - local_sink = self.attention_sink_bias[tp_rank * self.local_num_heads:(tp_rank + 1) * self.local_num_heads] + # Get local portion of sink bias for this rank. Index by the + # rank WITHIN the attention TP group (see decode path above); + # equals the global rank when attention_dp_degree == 1. + global_rank = parallel_state.get_tensor_model_parallel_rank() if parallel_state.model_parallel_is_initialized() else 0 + attn_tp_rank = global_rank % self.attn_tp_degree + local_sink = self.attention_sink_bias[attn_tp_rank * self.local_num_heads:(attn_tp_rank + 1) * self.local_num_heads] # Reshape and expand: [local_num_heads] -> [bsz, local_num_heads, q_len, 1] sink_bias = local_sink.reshape(1, -1, 1, 1).expand(bsz, -1, q_len, 1) attn_weights = torch.cat([attn_weights, sink_bias], dim=-1) @@ -851,8 +958,19 @@ def forward( attn_output = self.o_proj(attn_output) - # Prepare KV cache output - return as tuple for KV cache manager - # Return LOCAL key/value states for cache (each CP rank stores its portion) + # Data-parallel attention: gather the per-DP-group batch shards back to + # the full batch (dim 0) before the residual add and the MoE, which run + # at full TP/EP and expect the full batch. Mirrors the base class's + # post-attention DP all-gather. + if is_data_parallel_attn: + attn_output = gather_from_tensor_model_parallel_region_with_dim( + attn_output, gather_dim=0, + process_group=get_data_parallel_attention_dp_group(), + ) + + # Prepare KV cache output - return as tuple for KV cache manager. + # Under DP attention these are the DP-local (batch // dp) shards, which + # DataParallelKVCacheManager writes into this group's cache slice. new_key_value = (key_states_for_cache, value_states_for_cache) return attn_output, new_key_value, cos_cache, sin_cache @@ -1000,13 +1118,23 @@ def setup_attr_for_model(self, config: MiMoV2InferenceConfig): self.hidden_size = config.hidden_size self.num_attention_heads = config.num_attention_heads - # Check if we need GQA CONVERT_TO_MHA mode - # When tp_degree > num_kv_heads, we replicate K/V to match num_attention_heads + # Attention TP degree under data-parallel attention (== tp_degree when + # attention_dp_degree == 1). Must match NeuronMiMoV2Attention so the KV + # cache shape agrees with the per-layer projections. + attention_dp_degree = getattr( + config.neuron_config, "attention_dp_degree", 1 + ) or 1 + attn_tp_degree = self.tp_degree // attention_dp_degree + + # Check if we need GQA CONVERT_TO_MHA mode. Under DP attention the + # decisive comparison is attn_tp_degree, not the full tp_degree: if + # attn_tp_degree <= num_kv_heads the KV heads shard cleanly with no + # replication, so MHA conversion (and its 8-16x cache bloat) is off. min_kv_heads = min( config.num_key_value_heads, getattr(config, 'swa_num_key_value_heads', config.num_key_value_heads) ) - self.use_gqa_convert_to_mha = self.tp_degree > min_kv_heads + self.use_gqa_convert_to_mha = attn_tp_degree > min_kv_heads if self.use_gqa_convert_to_mha: # With CONVERT_TO_MHA, KV cache stores num_attention_heads (same as Q) @@ -1124,9 +1252,19 @@ def convert_mimo_v2_hf_to_neuron_state_dict( full_num_kv_heads = config.num_key_value_heads # 4 swa_num_kv_heads = config.swa_num_key_value_heads # 8 + # Under data-parallel attention the K/V weights must be replicated only if + # the *attention* TP degree (tp_degree // attention_dp_degree) exceeds the + # kv-head count. With attention_dp_degree chosen so attn_tp <= kv_heads, + # both gates are False and the checkpoint loads at its native head width + # (no replication) — this is what removes the KV cache bloat. + attention_dp_degree = getattr( + config.neuron_config, "attention_dp_degree", 1 + ) or 1 + attn_tp_degree = tp_degree // attention_dp_degree + # Check if we need to replicate K/V weights - full_use_convert_to_mha = tp_degree > full_num_kv_heads - swa_use_convert_to_mha = tp_degree > swa_num_kv_heads + full_use_convert_to_mha = attn_tp_degree > full_num_kv_heads + swa_use_convert_to_mha = attn_tp_degree > swa_num_kv_heads for layer_idx in range(config.num_hidden_layers): # Add rank utility for attention From 94ce6b24231391f056b3604637e9735ec47b9ca8 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 23 Jul 2026 07:24:36 +0800 Subject: [PATCH 24/27] h100/MiMo-V2.5: add 15K-input / 1K-output bench (Trn2 16K-context point) 15000 in + 1000 out = 16000, matching the Trn2 port's max compiled context, so these are the H100 reference numbers for that long-context serving point. Single-node 8xH100, c=1/16/32, both backends, 0 failures. SGLang (DP=2, mem-frac 0.9, context 17408) wins output throughput at every concurrency (94.6/867.8/1255.2 out tok/s) and has ~10x lower c=1 TTFT (699ms vs 6.6s). vLLM (TP=8+EP+PR#42270, chunked prefill, max-model-len 17408) has the lowest TPOT (5.5ms @ c=1, once decoding) but its single-stream TTFT is dominated by the 15K chunked prefill; it converges with SGLang by c=32 (1181.7 out tok/s). Ran the two backends in parallel across P5-1 (SGLang) and P5-2 (vLLM). Raw logs in results/ctx15k/. Co-Authored-By: Claude Opus 4.8 --- .../perf_test/h100/results/ctx15k/README.md | 46 ++++++++++++++++++ .../h100/results/ctx15k/sglang_15k.log | 47 +++++++++++++++++++ .../h100/results/ctx15k/vllm_15k.log | 43 +++++++++++++++++ 3 files changed, 136 insertions(+) create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/README.md create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/sglang_15k.log create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/vllm_15k.log diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/README.md b/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/README.md new file mode 100644 index 00000000..ec622bd9 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/README.md @@ -0,0 +1,46 @@ +# MiMo-V2.5 single-node H100 @ 15K input / 1K output + +Matches the Trn2 port's **maximum compiled context (16K)** — 15000 input + 1000 +output = 16000 — so these numbers are the H100 reference for the long-context +point Trn2 can actually serve. Concurrency 1 / 16 / 32. + +Bench (per user template): `vllm bench serve --dataset-name random +--random-input-len 15000 --random-output-len 1000 --random-range-ratio 0.03 +--num-prompts 2*C --max-concurrency C`. Both servers single-node 8xH100, FP8, +no speculative decode. Raw logs: `sglang_15k.log`, `vllm_15k.log`. + +- **SGLang**: `run_sglang_h100.sh` shape (TP=8, DP=2 + DP-attention), but + `--mem-fraction-static 0.9 --context-length 17408` for the 16K KV pool + (max_total_num_tokens=678837). +- **vLLM**: `run_vllm_h100.sh` (TP=8 + EP + PR #42270, chunked prefill ON), + `--max-model-len 17408`. + +## Results (15K in / 1K out) + +| Concurrency | Backend | Output tput (tok/s) | Total tput (tok/s) | TTFT median (ms) | TTFT P99 (ms) | TPOT median (ms) | +|---|---|---|---|---|---|---| +| 1 | SGLang | **94.6** | **1544** | 699 | 898 | 9.9 | +| 1 | vLLM | 82.7 | 1350 | 6587 | 7487 | **5.5** | +| 16 | SGLang | **867.8** | **13912** | 1336 | 4501 | 16.7 | +| 16 | vLLM | 712.5 | 11423 | 1132 | 10021 | 16.3 | +| 32 | SGLang | **1255.2**| **20080** | 1473 | 8828 | 22.1 | +| 32 | vLLM | 1181.7 | 18904 | 2054 | 9957 | 24.1 | + +All runs: 0 failures (2 / 32 / 64 prompts at c=1 / 16 / 32). + +## Observations + +- **SGLang wins output throughput at every concurrency** (1.1–1.2x vLLM), and its + c=1 TTFT is ~10x lower (699 ms vs 6.6 s). SGLang chunks the 15K prefill more + aggressively into the decode loop, so first token lands fast even single-stream. +- **vLLM has the lowest TPOT** (5.5 ms at c=1) — once decoding, it's the fastest + per-token — but its c=1 TTFT is dominated by the 15K prefill (chunked-prefill + interleaves it over many steps before the first output token). +- Throughput scales cleanly 1→32 for both (SGLang 94→1255, vLLM 83→1182 out + tok/s); the crossover where they converge is c=32. +- At 15K/1K the workload is prefill-heavy (15:1 in:out); total throughput + (~20K tok/s SGLang @ c=32) is dominated by prefill, which is why it's ~13x the + output-token rate. + +> These are the H100 reference numbers for a direct comparison against the Trn2 +> 16K-compiled port at 15K/1K. Language-only, no prefix caching, no spec decode. diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/sglang_15k.log b/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/sglang_15k.log new file mode 100644 index 00000000..ddb79bad --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/sglang_15k.log @@ -0,0 +1,47 @@ +mkdir: cannot create directory '/opt/dlami/nvme/models/bench_results/mimo_v2_5_h100/ctx15k_sglang': Permission denied +===== sglang c=1 ISL=15000 OSL=1000 np=2 ===== +tee: /opt/dlami/nvme/models/bench_results/mimo_v2_5_h100/ctx15k_sglang/c1.txt: No such file or directory +WARNING 07-22 23:09:55 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 1 +Successful requests: 2 +Failed requests: 0 +Maximum request concurrency: 1 +Benchmark duration (s): 21.00 +Total input tokens: 30439 +Output token throughput (tok/s): 94.61 +Total token throughput (tok/s): 1544.02 +Median TTFT (ms): 698.86 +P99 TTFT (ms): 898.13 +Median TPOT (ms): 9.88 +--- c=1 done --- +===== sglang c=16 ISL=15000 OSL=1000 np=32 ===== +tee: /opt/dlami/nvme/models/bench_results/mimo_v2_5_h100/ctx15k_sglang/c16.txt: No such file or directory +WARNING 07-22 23:10:38 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 16 +Successful requests: 32 +Failed requests: 0 +Maximum request concurrency: 16 +Benchmark duration (s): 36.86 +Total input tokens: 480778 +Output token throughput (tok/s): 867.75 +Total token throughput (tok/s): 13911.65 +Median TTFT (ms): 1336.45 +P99 TTFT (ms): 4500.89 +Median TPOT (ms): 16.66 +--- c=16 done --- +===== sglang c=32 ISL=15000 OSL=1000 np=64 ===== +tee: /opt/dlami/nvme/models/bench_results/mimo_v2_5_h100/ctx15k_sglang/c32.txt: No such file or directory +WARNING 07-22 23:11:39 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 32 +Successful requests: 64 +Failed requests: 0 +Maximum request concurrency: 32 +Benchmark duration (s): 51.03 +Total input tokens: 960566 +Output token throughput (tok/s): 1255.16 +Total token throughput (tok/s): 20080.10 +Median TTFT (ms): 1473.42 +P99 TTFT (ms): 8828.01 +Median TPOT (ms): 22.10 +--- c=32 done --- +ALL_DONE diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/vllm_15k.log b/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/vllm_15k.log new file mode 100644 index 00000000..57733bc8 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/vllm_15k.log @@ -0,0 +1,43 @@ +===== vllm c=1 ISL=15000 OSL=1000 np=2 ===== +WARNING 07-22 23:19:57 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 1 +Successful requests: 2 +Failed requests: 0 +Maximum request concurrency: 1 +Benchmark duration (s): 24.02 +Total input tokens: 30439 +Output token throughput (tok/s): 82.72 +Total token throughput (tok/s): 1349.97 +Median TTFT (ms): 6587.39 +P99 TTFT (ms): 7486.82 +Median TPOT (ms): 5.46 +--- c=1 done --- +===== vllm c=16 ISL=15000 OSL=1000 np=32 ===== +WARNING 07-22 23:20:43 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 16 +Successful requests: 32 +Failed requests: 0 +Maximum request concurrency: 16 +Benchmark duration (s): 44.89 +Total input tokens: 480778 +Output token throughput (tok/s): 712.51 +Total token throughput (tok/s): 11422.87 +Median TTFT (ms): 1131.82 +P99 TTFT (ms): 10021.23 +Median TPOT (ms): 16.33 +--- c=16 done --- +===== vllm c=32 ISL=15000 OSL=1000 np=64 ===== +WARNING 07-22 23:21:52 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 32 +Successful requests: 64 +Failed requests: 0 +Maximum request concurrency: 32 +Benchmark duration (s): 54.20 +Total input tokens: 960566 +Output token throughput (tok/s): 1181.65 +Total token throughput (tok/s): 18904.03 +Median TTFT (ms): 2053.64 +P99 TTFT (ms): 9957.17 +Median TPOT (ms): 24.12 +--- c=32 done --- +ALL_DONE From b078f519537ea89f0c4b6e7d19433f10344c6627 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 23 Jul 2026 07:52:07 +0800 Subject: [PATCH 25/27] docs: document long-context (16K) support and benchmarks Add a "Long Context" section to the MiMo-V2.5 README covering the data-parallel + context-parallel attention recipe that lifts seq_len from ~512 to 16384 (32x): - the three override_neuron_config settings (attention_dp_degree=16, cp_degree=16, sequence_parallel_enabled=false) and why each is needed; - the cp_degree == attention_dp_degree constraint (cp != dp trips the base KV manager's CP->DP head-remap assertion that MiMo's 8 SWA KV heads violate); - long-context throughput at 15K/1K input/output (c=1/16/32); - prefill-only cost at 16K input (single 16K prefill ~6 s; TTFT queues ~linearly with concurrency since chunked_prefill is off); - a note on why 32K does not yet fit (single-core HBM at cp=dp=16, the cp!=dp assertion, and a compiler-backend failure at cp=dp=32). Also update the CONVERT_TO_MHA implementation note: the replication gate now keys off the attention TP degree (tp_degree // attention_dp_degree), so attention_dp_degree=16 disables replication while =1 preserves prior behavior. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5/README.md | 66 +++++++++++++++++++++++++++++- 1 file changed, 65 insertions(+), 1 deletion(-) diff --git a/contrib/models/MiMo-V2.5/README.md b/contrib/models/MiMo-V2.5/README.md index 88281a47..8c8e22d9 100644 --- a/contrib/models/MiMo-V2.5/README.md +++ b/contrib/models/MiMo-V2.5/README.md @@ -409,6 +409,70 @@ Observations: > **Compile time:** the first MiMo-V2.5 compile on SDK 2.29 is ~30 minutes (TKG + CE HLO compilation, weight layout optimization, then `shard_checkpoint` for 64 ranks which dominates at ~27 minutes). Subsequent runs with the same `override_neuron_config` hit the neuronx-cc cache and the NEFF loads in ~1 minute. `save_sharded_checkpoint=true` persists per-rank FP8 shards under `/weights/`, letting future `load()` calls skip the `shard_checkpoint` pass entirely. +## Long Context (up to 16K, data-parallel + context-parallel attention) + +The default recipe (`attention_dp_degree=1`) forces `CONVERT_TO_MHA` because +`tp_degree=64 > num_kv_heads (4 full / 8 SWA)`, replicating K/V up to 64 heads. +That 8-16x KV-cache bloat caps `seq_len` at ~512 on 64 cores. To serve long +context, combine three settings in `override_neuron_config`: + +| Setting | Purpose | +|---|---| +| `attention_dp_degree=16` | Attention runs at TP = 64/16 = 4 = num_kv_heads (zero KV replication, no CONVERT_TO_MHA). Decode splits the batch across 16 DP groups; MoE stays EP=64. | +| `cp_degree=16` | Context parallelism splits the prefill sequence 16 ways so the context-encoding compute graph stays under the compiler's ~10M-instruction limit (a single unsplit 16K CTE graph is ~10.1M and fails). | +| `sequence_parallel_enabled=false` | MiMo's hand-rolled CP path (asymmetric head_dim 192/128 precludes the base fused-QKV CP kernel) uses the SP=false branch. | + +**Constraint: `cp_degree == attention_dp_degree`.** With `cp != dp` the base KV +manager engages a CP→DP head-remap (`get_kv_head_indices_context_parallel_dp_decode`) +that asserts `tp_degree/dp_degree >= num_kv_heads`; MiMo's 8 SWA KV heads +violate it (64/16 = 4 < 8). Keeping `cp == dp` skips that path. + +Validated at **seq_len=16384** (`tp_degree=64, attention_dp_degree=16, +cp_degree=16, moe_ep_degree=64`, BS=32). Outputs are factually correct +(e.g. "The author of Hamlet is William Shakespeare"). DP=1 cannot serve 16K — +its KV cache alone would need ~1.2 TB, exceeding total HBM. + +### Long-context throughput (trn2.48xlarge, seq_len=16384, FP8, cp=dp=16, BS=32) + +Input/output 15000 / 1000 tokens (random dataset): + +| Concurrency | Output throughput (tok/s) | TTFT median (ms) | TTFT P99 (ms) | TPOT median (ms) | +|---|---|---|---|---| +| 1 | 4.49 | 6051 | 6059 | 217 | +| 16 | 48.49 | 17985 | 93566 | 275 | +| 32 | 74.01 | 23941 | 186927 | 352 | + +Observations: +- **Throughput scales with concurrency** (4.5 -> 48.5 -> 74.0 tok/s) as decode + batching fills up. +- **TTFT dominates at high concurrency**: a 15K-token prefill is expensive even + split 16 ways via CP, and with `chunked_prefill=false` those prefills queue — + P99 TTFT reaches ~187 s at c=32. Long context on a single node trades TTFT for + throughput; keep concurrency low for latency-sensitive use. + +### Prefill-only cost (input 16000 / output 10, isolates TTFT = prefill time) + +| Concurrency | TTFT median (ms) | TTFT mean (ms) | TTFT P99 (ms) | +|---|---|---|---| +| 1 | 6055 | 6055 | 6071 | +| 16 | 47955 | 46900 | 93495 | +| 32 | 89942 | 90806 | 186959 | + +- A single 16K prefill costs **~6.06 s** (cp=16, 48 layers of 256-expert MoE). + This matches the c=1 TTFT of the 15K/1K run, confirming TTFT there is + essentially the prefill. +- TTFT grows ~linearly with concurrency because `chunked_prefill=false` runs one + context-encoding pass at a time — 16 requests queue to ~48 s, 32 to ~90 s + median. The prefill queue, not decode, is the long-context bottleneck; enabling + chunked prefill would help but is not currently wired with DP attention. + +> **32K note:** seq_len=32768 does not currently fit. `cp=dp=16` compiles but a +> single core's HBM overflows (KV tensors ~12.7 GB + 32K activation scratchpad +> ~9.8 GB > 24 GB); `cp=32,dp=16` hits the `cp != dp` head-remap assertion; and +> `cp=dp=32` triggers a compiler-backend (walrus) failure. 32K needs further +> work (lower batch while satisfying the MoE EP floor, or adopting the base +> class's CP machinery). + ## Compatibility Matrix | Instance | Neuron SDK 2.29+ (PyTorch 2.9) | 2.21 and earlier | @@ -426,7 +490,7 @@ pytest contrib/models/MiMo-V2.5/test/integration/test_model.py -v ## Key Implementation Notes 1. **Hybrid Attention**: `hybrid_layer_pattern` list determines full vs sliding window per layer; the modeling code constructs one `NeuronMiMoV2Attention` per layer with the correct `is_sliding_window` flag and rope_theta. -2. **CONVERT_TO_MHA**: When `tp_degree > num_kv_heads` (64 > 4 full / 64 > 8 SWA), K/V are replicated to `num_attention_heads` (64) during state-dict conversion; this applies to both `.weight` and the per-row `.scale` on the FP8 path. +2. **CONVERT_TO_MHA**: When the *attention* TP degree `tp_degree // attention_dp_degree > num_kv_heads` (default: 64 > 4 full / 64 > 8 SWA), K/V are replicated to `num_attention_heads` (64) during state-dict conversion; this applies to both `.weight` and the per-row `.scale` on the FP8 path. Setting `attention_dp_degree=16` makes attn TP = 4 = num_kv_heads, disabling replication (see "Long Context" above) — the gate keys off `attn_tp_degree`, so `attention_dp_degree=1` preserves the original full-replication behavior. 3. **Attention Sink Bias**: Learnable per-head bias added as an extra "sink" column to attention scores in sliding window layers (not added in full-attention layers). Per-rank slicing of the bias happens inside `forward()` based on `parallel_state.get_tensor_model_parallel_rank()`. 4. **Fused qkv split in preprocess**: V2.5's HF checkpoint stores `self_attn.qkv_proj.weight` as 4 interleaved Q/K/V groups (see "Checkpoint Preparation" above). The preprocess script must slice these groups — naïve `[Q|K|V]` concat slicing produces garbage outputs. 5. **weight_map rebuild**: V2.5's `model.safetensors.index.json` references legacy `model_N-00001-of-00002.safetensors` filenames that do not match the actual `model_pp0_epN_shardM.safetensors` objects on disk. `LazyWeightMap` scans the on-disk shards at startup and rebuilds `weight_map` directly from each file's manifest; the inconsistent index is ignored. From 0063b0074cedb3a192e458017e37caaa7ddb74a0 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 23 Jul 2026 07:55:56 +0800 Subject: [PATCH 26/27] h100/MiMo-V2.5: add prefill benchmark (16K in / 10 out) + root-cause vLLM gap MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Isolates prefill (16000 in, 10 out) at c=1/16/32, both backends, single-node 8xH100 run in parallel on P5-1/P5-2. 0 failures. Prefill throughput (tok/s): SGLang 20539/51958/52984, vLLM 5875/29913/39605 for c=1/16/32. SGLang is 3.5x at c=1, saturating ~53K by c=16; vLLM still climbing at c=32. Investigated the gap — it's NOT a config mismatch: both chunk prefill at 8192. vLLM must use the FLASH_ATTN_DIFFKV backend for V2.5's asymmetric head dims (Q/K=192, V=128), the young path PR #42270 just enabled. That kernel is pinned to the 8192 chunk — raising --max-num-batched-tokens to 16384 drops it back to plain FlashAttention, which crashes with "TP must evenly split KV heads" (4 full-attn KV heads, TP=8). So it's a DiffKV kernel-maturity gap vs SGLang's fa3, not tunable away. SGLang is the better prefill-heavy long-context choice today. Co-Authored-By: Claude Opus 4.8 --- .../h100/results/prefill_16k/README.md | 63 +++++++++++++++++++ .../results/prefill_16k/sglang_prefill.log | 43 +++++++++++++ .../h100/results/prefill_16k/vllm_prefill.log | 43 +++++++++++++ 3 files changed, 149 insertions(+) create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/README.md create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/sglang_prefill.log create mode 100644 contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/vllm_prefill.log diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/README.md b/contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/README.md new file mode 100644 index 00000000..d85ee59e --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/README.md @@ -0,0 +1,63 @@ +# MiMo-V2.5 single-node H100 — prefill throughput (16K input / 10 output) + +Isolates **prefill** performance: 16000 input, only 10 output tokens, so the run +is dominated by context encoding and total throughput ≈ prefill token rate. +Concurrency 1 / 16 / 32. Same servers as the 15K/1K bench (see `../ctx15k/`). + +Bench: `vllm bench serve --dataset-name random --random-input-len 16000 +--random-output-len 10 --random-range-ratio 0.03 --num-prompts 2*C +--max-concurrency C`. Single-node 8xH100, FP8, no spec decode. Ran the two +backends in parallel on P5-1 (SGLang) and P5-2 (vLLM). Raw logs: +`sglang_prefill.log`, `vllm_prefill.log`. + +## Results (16K in / 10 out ≈ pure prefill) + +| Concurrency | Backend | Prefill throughput (tok/s) | Batch wall (s) | TTFT median (ms) | TTFT P99 (ms) | +|---|---|---|---|---|---| +| 1 | **SGLang** | **20539** | 1.6 | 709 | 913 | +| 1 | vLLM | 5875 | 5.5 | 2716 | 4980 | +| 16 | **SGLang** | **51958** | 9.9 | 2806 | 4810 | +| 16 | vLLM | 29913 | 17.2 | 3917 | 8520 | +| 32 | **SGLang** | **52984** | 19.4 | 5221 | 9444 | +| 32 | vLLM | 39605 | 25.9 | 8954 | 11796 | + +All runs: 0 failures. + +## Observations + +- **SGLang prefill is markedly faster than vLLM here** — 3.5x at c=1 (20.5K vs + 5.9K tok/s), 1.7x at c=16, 1.3x at c=32. The gap is widest single-stream and + narrows as concurrency fills the pipeline. +- **SGLang prefill throughput saturates by c=16** (~52K tok/s; c=32 barely higher + at 53K) — the H100 prefill compute ceiling for this model. vLLM is still + climbing at c=32 (30K → 40K), i.e. it needs more concurrency to hit peak. + +## Why vLLM is slower — it's the DiffKV kernel, not a config gap + +Investigated the 3.5x c=1 gap. It is **not** an unaligned setting: + +- **Both backends chunk prefill at the same 8192** (`chunked_prefill_size=8192` + on SGLang, forced by DP-attention; `max_num_batched_tokens=8192` vLLM default). + So the chunk size is identical. +- **vLLM must use the `FLASH_ATTN_DIFFKV` attention backend** for MiMo-V2.5 (log: + `Using FLASH_ATTN_DIFFKV for attention` / `Diff-KV with sinks: upgrading + FlashAttention 3 -> 4`). This is the kernel that handles V2.5's asymmetric head + dims (Q/K=192, V=128) + attention-sink bias; it's the new path PR #42270 just + made loadable. +- **That DiffKV kernel is pinned to the 8192 chunk and cannot be enlarged.** + Setting `--max-num-batched-tokens 16384` (to encode the 16K prompt in one pass) + makes vLLM fall back to the plain FlashAttention path, which then crashes at + init: `AssertionError: TP size must evenly split the number of KV heads` (V2.5 + has 4 full-attn KV heads, TP=8 can't split them — the exact reason plain vLLM + needs DiffKV in the first place). So the prefill chunk can't be tuned up. + +**Conclusion:** vLLM's lower prefill throughput is the efficiency of its DiffKV +kernel (still a young code path — only just enabled by PR #42270) versus SGLang's +mature `fa3` backend, at the same 8192 chunk. It is a kernel-maturity gap for +asymmetric-head-dim / hybrid-attention models, not a misconfiguration. SGLang is +the better choice for prefill-heavy long-context serving of MiMo-V2.5 today. + +- The 16K prefill compute ceiling (~53K tok/s, SGLang) is why the 15K/1K total + throughput in `../ctx15k/` tops out around 20K tok/s once the 1K decode mixes in. + +> H100 reference for the prefill-bound regime at Trn2's 16K compiled context. diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/sglang_prefill.log b/contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/sglang_prefill.log new file mode 100644 index 00000000..ea9d8b64 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/sglang_prefill.log @@ -0,0 +1,43 @@ +===== sglang c=1 ISL=16000 OSL=10 np=2 ===== +WARNING 07-22 23:32:28 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 1 +Successful requests: 2 +Failed requests: 0 +Maximum request concurrency: 1 +Benchmark duration (s): 1.58 +Total input tokens: 32469 +Output token throughput (tok/s): 12.01 +Total token throughput (tok/s): 20538.97 +Median TTFT (ms): 709.39 +P99 TTFT (ms): 913.47 +Median TPOT (ms): 9.57 +--- c=1 done --- +===== sglang c=16 ISL=16000 OSL=10 np=32 ===== +WARNING 07-22 23:32:51 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 16 +Successful requests: 32 +Failed requests: 0 +Maximum request concurrency: 16 +Benchmark duration (s): 9.88 +Total input tokens: 512831 +Output token throughput (tok/s): 32.60 +Total token throughput (tok/s): 51958.05 +Median TTFT (ms): 2806.42 +P99 TTFT (ms): 4809.99 +Median TPOT (ms): 225.31 +--- c=16 done --- +===== sglang c=32 ISL=16000 OSL=10 np=64 ===== +WARNING 07-22 23:33:26 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 32 +Successful requests: 64 +Failed requests: 0 +Maximum request concurrency: 32 +Benchmark duration (s): 19.35 +Total input tokens: 1024606 +Output token throughput (tok/s): 33.28 +Total token throughput (tok/s): 52984.49 +Median TTFT (ms): 5220.65 +P99 TTFT (ms): 9443.69 +Median TPOT (ms): 492.45 +--- c=32 done --- +ALL_DONE diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/vllm_prefill.log b/contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/vllm_prefill.log new file mode 100644 index 00000000..335df0b6 --- /dev/null +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/prefill_16k/vllm_prefill.log @@ -0,0 +1,43 @@ +===== vllm c=1 ISL=16000 OSL=10 np=2 ===== +WARNING 07-22 23:32:30 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 1 +Successful requests: 2 +Failed requests: 0 +Maximum request concurrency: 1 +Benchmark duration (s): 5.53 +Total input tokens: 32469 +Output token throughput (tok/s): 3.44 +Total token throughput (tok/s): 5874.85 +Median TTFT (ms): 2715.70 +P99 TTFT (ms): 4980.15 +Median TPOT (ms): 5.66 +--- c=1 done --- +===== vllm c=16 ISL=16000 OSL=10 np=32 ===== +WARNING 07-22 23:32:57 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 16 +Successful requests: 32 +Failed requests: 0 +Maximum request concurrency: 16 +Benchmark duration (s): 17.15 +Total input tokens: 512831 +Output token throughput (tok/s): 18.77 +Total token throughput (tok/s): 29913.16 +Median TTFT (ms): 3916.67 +P99 TTFT (ms): 8519.61 +Median TPOT (ms): 170.13 +--- c=16 done --- +===== vllm c=32 ISL=16000 OSL=10 np=64 ===== +WARNING 07-22 23:33:38 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 32 +Successful requests: 64 +Failed requests: 0 +Maximum request concurrency: 32 +Benchmark duration (s): 25.89 +Total input tokens: 1024606 +Output token throughput (tok/s): 24.88 +Total token throughput (tok/s): 39605.02 +Median TTFT (ms): 8954.28 +P99 TTFT (ms): 11796.27 +Median TPOT (ms): 170.01 +--- c=32 done --- +ALL_DONE From a07e2cb8db547f5a4bf478c50e61df6594291ccc Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 23 Jul 2026 07:58:32 +0800 Subject: [PATCH 27/27] h100/MiMo-V2.5: add H100-vs-Trn2 cross-platform table to 15K/1K results Now that the Trn2 16K long-context numbers landed in the main README, add a side-by-side comparison in the ctx15k results: at 15K/1K the H100 lead widens to ~17-21x throughput (vs ~7-11x at short 900/90) and 9-16x lower TTFT. Both are prefill-bound; Trn2 more so (cp=dp=16, chunked_prefill off -> 16K prefills queue, P99 TTFT ~187s at c=32; single 16K prefill ~6.06s vs ~0.7s on H100 SGLang). Co-Authored-By: Claude Opus 4.8 --- .../perf_test/h100/results/ctx15k/README.md | 29 +++++++++++++++++++ 1 file changed, 29 insertions(+) diff --git a/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/README.md b/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/README.md index ec622bd9..e2cab903 100644 --- a/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/README.md +++ b/contrib/models/MiMo-V2.5/perf_test/h100/results/ctx15k/README.md @@ -44,3 +44,32 @@ All runs: 0 failures (2 / 32 / 64 prompts at c=1 / 16 / 32). > These are the H100 reference numbers for a direct comparison against the Trn2 > 16K-compiled port at 15K/1K. Language-only, no prefix caching, no spec decode. + +## Cross-platform: H100 vs Trn2 (15K in / 1K out) + +Trn2 numbers from the main README's "Long Context" section (trn2.48xlarge, +seq_len=16384, `tp=64, attention_dp=16, cp=16, moe_ep=64`, BS=32). + +**Output throughput (tok/s):** + +| Concurrency | Trn2 | H100 SGLang | H100 vLLM | SGLang / Trn2 | +|---|---:|---:|---:|---:| +| 1 | 4.49 | 94.6 | 82.7 | **21×** | +| 16 | 48.49 | 867.8 | 712.5 | **18×** | +| 32 | 74.01 | 1255.2 | 1181.7 | **17×** | + +**TTFT median (ms, lower = faster first token):** + +| Concurrency | Trn2 | H100 SGLang | H100 vLLM | Trn2 / SGLang | +|---|---:|---:|---:|---:| +| 1 | 6051 | 699 | 6587 | 8.7× slower | +| 16 | 17985 | 1336 | 1132 | 13× slower | +| 32 | 23941 | 1473 | 2054 | 16× slower | + +At 16K context the H100/Trn2 gap widens to **~17–21× throughput** (vs ~7–11× at +the short 900/90 shape) — long context is where Trn2 is hit hardest. Both +platforms are prefill-bound here, but Trn2 more severely: it must serve 16K via +`cp=dp=16` with `chunked_prefill=false`, so 16K prefills run one-at-a-time and +queue (main-README Trn2 P99 TTFT reaches ~187 s at c=32). A single 16K prefill is +~6.06 s on Trn2 vs ~0.7 s on H100 SGLang. Trn2 also cannot currently fit 32K +(HBM overflow), while single-node H100 can (see `../../` 32K exploration).