From 82d7eade1d9bb358e98982c70ea635a38e69851d Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 05:52:07 +0800 Subject: [PATCH 01/47] [contrib] Add MiMo-V2.5-Pro initial port (copied from MiMo-V2-Pro, renamed) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Bootstrap contrib entry for XiaomiMiMo/MiMo-V2.5-Pro on Trn2 via NxDI. Same starting point as the MiMo-V2-Pro port: src/modeling_mimo_v2.py (was modeling_mimo_v2_pro.py) src/conversion_script/preprocess_mimo_v2_fp8.py perf_test/{smoke_compile,smoke_generate,bench}_mimo_v2.{py,sh} Rename-only changes in this commit: MiMoV2Pro* identifiers -> MiMoV2* (classes, configs, modules) mimo_v2_pro paths -> mimo_v2 / mimo_v25_pro (compile dirs) HF repo XiaomiMiMo/MiMo-V2-Pro -> XiaomiMiMo/MiMo-V2.5-Pro README architecture table updated to V2.5-Pro config (70 layers, 6144 hidden, 128 heads, 384 experts, etc.) README disk footprint updated to match V2.5-Pro actual size (~962GB HF) Not yet adapted to V2.5-specific differences — these still need work: - attention_chunk_size=128 (new in V2.5, not handled in V2-Pro code) - MoE group-limited noaux_tc (n_group, topk_group) — V2.5 config sets 1,1 so it degenerates to plain noaux_tc; the Pro monkey-patch already matches - FP8 recipe verification on V2.5 weights (V2-Pro workarounds may or may not apply: mean-subtract router bias, split_qkv_fused interleaved layout, blockwise scale stride fix) Subsequent commits will adapt each of the above after validation on Trn2. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 364 ++++ .../models/MiMo-V2.5-Pro/perf_test/0_setup.sh | 61 + .../MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh | 228 +++ .../perf_test/run_bench_single.sh | 76 + .../MiMo-V2.5-Pro/perf_test/sanity_check.sh | 59 + .../perf_test/smoke_compile_mimo_v2.py | 195 ++ .../perf_test/smoke_generate_mimo_v2.py | 231 +++ .../perf_test/vllm-neuron-patch.patch | 107 ++ contrib/models/MiMo-V2.5-Pro/src/__init__.py | 0 .../preprocess_mimo_v2_fp8.py | 641 +++++++ .../MiMo-V2.5-Pro/src/modeling_mimo_v2.py | 1676 +++++++++++++++++ contrib/models/MiMo-V2.5-Pro/test/__init__.py | 0 .../test/integration/__init__.py | 0 .../test/integration/test_model.py | 53 + .../MiMo-V2.5-Pro/test/unit/__init__.py | 0 15 files changed, 3691 insertions(+) create mode 100644 contrib/models/MiMo-V2.5-Pro/README.md create mode 100755 contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh create mode 100755 contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh create mode 100755 contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh create mode 100755 contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh create mode 100755 contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py create mode 100755 contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch create mode 100644 contrib/models/MiMo-V2.5-Pro/src/__init__.py create mode 100644 contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py create mode 100644 contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py create mode 100644 contrib/models/MiMo-V2.5-Pro/test/__init__.py create mode 100644 contrib/models/MiMo-V2.5-Pro/test/integration/__init__.py create mode 100644 contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py create mode 100644 contrib/models/MiMo-V2.5-Pro/test/unit/__init__.py diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md new file mode 100644 index 00000000..38623380 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -0,0 +1,364 @@ +# Contrib Model: MiMo-V2.5-Pro + +NeuronX Distributed Inference implementation of [XiaomiMiMo/MiMo-V2.5-Pro](https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro). + +## Model Information + +- **HuggingFace ID:** `XiaomiMiMo/MiMo-V2.5-Pro` +- **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 | 6144 | +| Layers | 70 | +| Attention Heads | 128 Q | +| KV Heads (full & sliding window) | 8 | +| Q/K Head Dim | 192 | +| V Head Dim | 128 | +| Experts | 384 routed (top-8 routing), no shared expert | +| Expert Intermediate | 2048 | +| Dense MLP Intermediate (layer 0) | 16,384 | +| Vocab Size | 152,576 | +| RoPE | Partial (33.4% → 64 of 192 dims), theta=10M (full) / 10K (SWA) | +| Sliding Window | 128 | +| Max Position | 1,048,576 (1M) | +| Attention Projection | `fused_qkv` (single `qkv_proj.weight`) | + +Key features: +- **Hybrid Attention**: 10 full attention layers (0, 7, 15, 23, 31, 39, 47, 55, 62, 69) + 60 sliding window layers, per `hybrid_layer_pattern` +- **Asymmetric Head Dims**: Q/K use head_dim=192, V uses v_head_dim=128 +- **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**: `sigmoid(logits) + e_score_correction_bias` is used to pick top-8 experts; unbiased `sigmoid(logits)` becomes the affinity weights. `n_group=1, topk_group=1` degenerates group-limited routing to plain noaux_tc. +- **attention_value_scale = 0.612**: HF reference multiplies `value_states` by this before `softmax(QK^T) × V` (NOT applied post-attention); the NxDI port 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**: ~3 TB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~962 GB, the Neuron-FP8 preprocessed output is ~1 TB, and `save_sharded_checkpoint=true` writes another ~300-1000 GB per compiled config (varies with recipe)). + +## 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-Pro 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-Pro # 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-Pro \ + --local-dir /opt/dlami/nvme/models/MiMo-V2.5-Pro + +# 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-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py \ + --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro \ + --save_path /opt/dlami/nvme/models/MiMo-V2.5-Pro-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-Pro/perf_test/smoke_compile_mimo_v2.py # compile +python contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py # 20-token generate + +# 5. Install vllm-neuron with the contrib registration patch +bash contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh + +# 6. Start vLLM serving MiMo-V2.5-Pro FP8 (first compile ~60 min; subsequent ~3 min) +bash contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh +``` + +The bench script runs two configurations (BS=32 and BS=128, both +`moe_tp_degree=X / moe_ep_degree=Y (see bench script)`) and logs results under +`/tmp/bench_results/mimo_v25_pro/`. + +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-Pro-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_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 ~1 TB across 70 per-layer safetensors shards. + +```bash +python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py \ + --hf_model_path /path/to/MiMo-V2.5-Pro \ + --save_path /path/to/MiMo-V2.5-Pro-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-Pro/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-Pro-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 Flash'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.py` — compile + load (STAGE=instantiate|compile|load|all, DRY_RUN, SKIP_WARMUP) +- `perf_test/smoke_generate_mimo_v2.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 Flash (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 Flash'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-Pro FP8. + +### batch_size >= 32 + +NxDI's TKG (token generation) path refuses Expert Parallelism when `batch_size < num_experts / top_k`. For Flash 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-Pro 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 Flash +# weights (BF16 by default; set MIMO_V2_FLASH_PATH to override). +bash contrib/models/MiMo-V2.5-Pro/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_FLASH_SRC` is set, registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under the key `mimo_v2` **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_FLASH_SRC=/path/to/neuronx-distributed-inference/contrib/models/MiMo-V2.5-Pro/src +export MIMO_V2_FLASH_PATH=/path/to/MiMo-V2.5-Pro-Neuron-FP8 +# First-time compile of Flash's 256-expert MoE takes 30-60 minutes. +export VLLM_ENGINE_READY_TIMEOUT_S=7200 +# Optional: isolate compile cache per config so parallel Flash/Pro/etc. compiles +# don't race on the default /var/tmp/neuron-compile-cache lock files. +export NEURON_COMPILED_ARTIFACTS=/path/to/compiled/mimo_v25_pro_bs32_moetp1_ep64_fp8 + +python3 -m vllm.entrypoints.openai.api_server \ + --model "$MIMO_V2_FLASH_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-Pro-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.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 Flash's model loader (reusing the Qwen2-family loader path, which Flash'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 Flash 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-Pro/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-Pro](https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro) — HF FP8 source checkpoint + +## Maintainer + +Henan Wan (whn09) + +**Last Updated:** 2026-04-25 diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh new file mode 100755 index 00000000..5d43c529 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh @@ -0,0 +1,61 @@ +#!/bin/bash +# Setup for MiMo-V2.5-Pro vLLM benchmarking on Trn2. +# +# This clones upstream vllm-project/vllm-neuron at release-0.5.0 (which already +# has the mimov2flash -> mimo_v2 model_type rewrite), then applies +# vllm-neuron-patch.patch to add a runtime registration hook so the contrib +# NeuronMiMoV2ForCausalLM is plugged into both NxDI's MODEL_TYPES and vLLM's +# ModelRegistry at vllm-neuron plugin init time. +set -e + +echo "==========================================" +echo "Setup: vllm-neuron + MiMo-V2.5-Pro weights" +echo "==========================================" + +source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/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 . +pip install s5cmd + +python3 -c "import vllm_neuron; print('vllm-neuron installed:', vllm_neuron.__file__)" + +echo "" +echo "[2/2] Downloading MiMo-V2.5-Pro BF16 weights..." + +MIMO_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-BF16}" +if [ -d "$MIMO_PATH" ] && [ "$(ls "$MIMO_PATH"/*.safetensors 2>/dev/null | wc -l)" -gt 0 ]; then + echo " MiMo weights already exist at $MIMO_PATH, skipping download" +else + echo " Downloading BF16 weights from your S3 bucket (edit the URI if needed)..." + mkdir -p "$MIMO_PATH" + s5cmd cp "s3://datalab/xiaomi/models/MiMo-V2.5-Pro-BF16/**" "$MIMO_PATH/" + echo " Download complete: $(du -sh $MIMO_PATH | cut -f1)" +fi + +# Figure out where this contrib package's src/ lives so the registration hook +# can add it to sys.path inside vllm-neuron. +CONTRIB_SRC="$(cd "$(dirname "$0")/.." && pwd)/src" + +echo "" +echo "Setup complete. Before running the benchmark, export:" +echo " export MIMO_V2_FLASH_PATH=$MIMO_PATH" +echo " export NXDI_CONTRIB_MIMO_V2_FLASH_SRC=$CONTRIB_SRC" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh new file mode 100755 index 00000000..69c7d417 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh @@ -0,0 +1,228 @@ +#!/bin/bash +set -e + +# MiMo-V2.5-Pro FP8 vLLM benchmark on Trn2. +# +# Requires a Neuron-FP8 preprocessed checkpoint (see +# `src/conversion_script/preprocess_mimo_v2_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 Flash), 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_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" +# The NxDI contrib MiMo-V2.5-Pro 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_FLASH_SRC:=$(cd "$(dirname "$0")/.." && pwd)/src}" +export NXDI_CONTRIB_MIMO_V2_FLASH_SRC + +# First-time Flash 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_v25_pro" +mkdir -p "$RESULTS_DIR" + +# Common neuron config shared across all MiMo-V2.5-Pro 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-Pro 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-Pro FP8 benchmarks complete!" +echo "Results saved to: $RESULTS_DIR" +echo "==========================================" +ls -la "$RESULTS_DIR" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh new file mode 100755 index 00000000..45729cd2 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/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.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.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_FLASH_PATH Path to the BF16 checkpoint (default +# /opt/dlami/nvme/models/MiMo-V2.5-Pro-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_v25_pro) + +set -e + +source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate + +MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-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_v25_pro}" + +mkdir -p "$RESULTS_DIR" + +echo "==========================================" +echo "MiMo-V2.5-Pro 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.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-Pro/perf_test/sanity_check.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh new file mode 100755 index 00000000..9cef80ab --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/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-Pro +# 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_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-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.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-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py new file mode 100755 index 00000000..b67a95f8 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -0,0 +1,195 @@ +#!/usr/bin/env python3 +"""Minimal compile+load smoke test for MiMo-V2.5-Pro 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 Flash 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_V25_PRO_MODEL_PATH", + "/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8", +) +COMPILED_PATH = os.environ.get( + "MIMO_V25_PRO_COMPILED_PATH", + "/opt/dlami/nvme/compiled/mimo_v25_pro_moetp16_ep4_bs48/", +) + +TP_DEGREE = int(os.environ.get("TP_DEGREE", "64")) +SEQ_LEN = int(os.environ.get("SEQ_LEN", "1024")) +# BS=48 is the minimum that avoids forward_selective_loading on decode: +# `BS * top_k / num_experts >= 1.0` → BS >= 384/8 = 48. At BS=1 the TKG +# path raises `NotImplementedError: Selective Loading with Expert parallelism`. +BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "48")) +CTX_BATCH_SIZE = int(os.environ.get("CTX_BATCH_SIZE", "1")) +# moe_tp=16 / moe_ep=4: after moetp1/ep64 (prefill broken, E_local=6) and +# moetp32/ep2 (HBM OOM by 28MB at load, E_local=192), try the middle ground. +# E_local=384/4=96. Also intermediate/moe_tp = 2048/16 = 128 which matches +# the FP8 scale block size exactly — avoids the stride_fix workaround path. +# world_size = 16*4 = 64 OK. +MOE_TP = int(os.environ.get("MOE_TP", "16")) +MOE_EP = int(os.environ.get("MOE_EP", "4")) + +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. Matches Flash + Kimi. + 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-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py new file mode 100755 index 00000000..f357aee5 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -0,0 +1,231 @@ +#!/usr/bin/env python3 +"""Minimal generate smoke test for MiMo-V2.5-Pro FP8 on Trn2. + +Assumes the compiled NEFF already exists at MIMO_V25_PRO_COMPILED_PATH +(from smoke_compile_mimo_v2.py). Rebuilds the same MoENeuronConfig / +Flash 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_V25_PRO_MODEL_PATH", + "/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8", +) +COMPILED_PATH = os.environ.get( + "MIMO_V25_PRO_COMPILED_PATH", + "/opt/dlami/nvme/compiled/mimo_v25_pro_tp64_moetp1_ep64_fp8/", +) + +# Must match smoke_compile_mimo_v2.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", "48")) # must match smoke_compile +CTX_BATCH_SIZE = int(os.environ.get("CTX_BATCH_SIZE", "1")) +MOE_TP = int(os.environ.get("MOE_TP", "16")) +MOE_EP = int(os.environ.get("MOE_EP", "4")) + +PROMPT = os.environ.get( + "MIMO_V25_PRO_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.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.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) + + # When CHAT_TEMPLATE=1, wrap the raw prompt in the checkpoint's chat + # template (system + user turns with <|im_start|>/<|im_end|> markers and + # trailing assistant cue). Matches how vllm /v1/chat/completions prepares + # inputs. Without this, the model free-continues the prompt as raw text + # instead of answering it. + use_chat_template = os.environ.get("CHAT_TEMPLATE", "0") == "1" + minimal_chat = os.environ.get("MINIMAL_CHAT", "0") == "1" + if minimal_chat: + # Skip the Pro default system prompt entirely; wrap prompt in bare + # <|im_start|>user ... <|im_end|><|im_start|>assistant\n framing. + templated = ( + f"<|im_start|>user\n{PROMPT}<|im_end|>" + f"<|im_start|>assistant\n" + ) + print(f"[gen] minimal-chat prompt ({len(templated)} chars, no system)") + inputs = tokenizer( + [templated] * BATCH_SIZE, + return_tensors="pt", + padding=True, + add_special_tokens=False, + ) + elif use_chat_template: + system = os.environ.get("CHAT_SYSTEM", "") + messages = [] + if system: + messages.append({"role": "system", "content": system}) + messages.append({"role": "user", "content": PROMPT}) + templated = tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True, + ) + print(f"[gen] chat-templated prompt ({len(templated)} chars)") + inputs = tokenizer([templated] * BATCH_SIZE, return_tensors="pt", padding=True) + else: + 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, + ) + + # DUMP_LOGITS=1 -> request scores so we can see top-k per step. + dump_logits = os.environ.get("DUMP_LOGITS", "0") == "1" + if dump_logits: + gen_config.output_scores = True + gen_config.return_dict_in_generate = True + + print(f"[gen] prompt: {PROMPT!r}") + print(f"[gen] input_ids.shape={tuple(inputs['input_ids'].shape)}") + t0 = time.time() + output = adapter.generate( + input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + generation_config=gen_config, + ) + dt = time.time() - t0 + + if dump_logits and hasattr(output, "sequences"): + output_ids = output.sequences + scores = output.scores # tuple of [bs, vocab] per step + else: + output_ids = output + scores = None + + 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}") + + if scores is not None: + import torch as _t + print("[gen] === top-5 per decode step (batch slot 0) ===") + for step, step_logits in enumerate(scores): + lp = _t.log_softmax(step_logits[0].float(), dim=-1) + top_lp, top_id = _t.topk(lp, 5) + parts = [] + for l, i in zip(top_lp.tolist(), top_id.tolist()): + tok = tokenizer.decode([i]).replace("\n", "\\n") + parts.append(f"({tok!r}:{i}:{l:.2f})") + chosen = new_tokens[step].item() + print(f" step {step:3d} chose id={chosen} top5={' '.join(parts)}") + + print("[gen] Done.") + + +if __name__ == "__main__": + try: + main() + except Exception: + traceback.print_exc() + sys.exit(1) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch b/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch new file mode 100644 index 00000000..d8a85b89 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/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-Pro 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_FLASH_SRC -> path to contrib MiMo-V2.5-Pro 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_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( ++ "mimov2flash", {"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-Pro 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-Pro/src/__init__.py b/contrib/models/MiMo-V2.5-Pro/src/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py new file mode 100644 index 00000000..56eabeaa --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py @@ -0,0 +1,641 @@ +""" +Preprocess MiMo-V2.5-Pro 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-Pro checkpoint layout: + - q_proj, k_proj, v_proj are FUSED into a single `qkv_proj` tensor per + layer (num_kv_heads interleaved groups, Pro-specific). We split into + three per-row-quantized projections via `split_qkv_fused()`. + - 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..69 are MoE with 384 experts each. + - Hybrid attention: 10 "full" layers (hybrid_layer_pattern[i] == 0) and + 60 "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 Pro/original-Flash): + - 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_fp8.py \\ + --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro \\ + --save_path /opt/dlami/nvme/models/MiMo-V2.5-Pro-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 + + def _open(self, filename: str): + 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 helpers +# --------------------------------------------------------------------------- + +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: int, + head_dim: int, + v_head_dim: int, +) -> Dict[str, Tuple[torch.Tensor, Optional[torch.Tensor]]]: + """Split Pro's pre-fused qkv_proj into q/k/v. MiMo-V2.5-Pro specific. + + HF layout — cross-validated against sglang on H200: + `qkv_proj.weight` is NOT `[all_Q | all_K | all_V]`. It is num_kv_heads + interleaved groups, each holding (heads_per_group Q heads, 1 K head, + 1 V head) packed contiguously: + + group g (g = 0 .. num_kv_heads-1): + rows [g*R : g*R + qg] = Q heads [g*hpg : (g+1)*hpg] + rows [g*R+qg : g*R + qg + kg] = K head g + rows [g*R+qg+kg : g*R + R] = V head g + where + hpg = num_q_heads / num_kv_heads (e.g. 128/8 = 16) + qg = hpg * head_dim (e.g. 16 * 192 = 3072) + kg = 1 * head_dim (e.g. 192) + vg = 1 * v_head_dim (e.g. 128) + R = qg + kg + vg (e.g. 3392) + + Scale: per-group 27 scale rows covering 27*128 = 3456 "padded" rows: + 24 rows for Q (24 * 128 = 3072 real Q rows) + 2 rows for K (1 full block + 1 half-real/half-phantom block) + 1 row for V (128 rows) + Total: 8 * 27 = 216 scale rows, 8 * 3392 = 27136 weight rows. + + The "phantom" 64 rows sit between each group's K tail and V start in + *scale block coordinates* only; in the physical weight tensor, group g's + V is immediately followed by group (g+1)'s Q. We recover the correct + dequant by padding each group up to 3456 rows before applying the scale, + then stripping the phantom rows. + """ + in_features = qkv_weight.shape[1] + hpg = num_q_heads // num_kv_heads + qg_rows = hpg * head_dim + kg_rows = 1 * head_dim + vg_rows = 1 * v_head_dim + real_rows_per_group = qg_rows + kg_rows + vg_rows + total_real_rows = num_kv_heads * real_rows_per_group + + BLOCK = 128 + q_scale_rows_per_group = qg_rows // BLOCK + 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]} != " + f"expected {total_real_rows} " + f"(num_kv_heads={num_kv_heads}, R={real_rows_per_group})" + ) + + if qkv_weight.dtype != torch.float8_e4m3fn or qkv_scale is None: + # BF16 path + w = qkv_weight.view(num_kv_heads, real_rows_per_group, in_features) + q_w = w[:, :qg_rows, :].reshape(num_kv_heads * qg_rows, in_features).contiguous() + k_w = w[:, qg_rows:qg_rows + kg_rows, :].reshape(num_kv_heads * kg_rows, in_features).contiguous() + v_w = w[:, qg_rows + kg_rows:, :].reshape(num_kv_heads * 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_kv_heads * 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})" + ) + + w = qkv_weight.to(torch.float32).view( + num_kv_heads, real_rows_per_group, in_features + ) + w_padded = torch.zeros( + num_kv_heads, padded_rows_per_group, in_features, dtype=torch.float32 + ) + w_padded[:, :real_rows_per_group, :] = w + + s = qkv_scale.to(torch.float32).view( + num_kv_heads, 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_kv_heads * qg_rows, in_features).contiguous() + k_deq = deq[:, qg_rows:qg_rows + kg_rows, :].reshape(num_kv_heads * kg_rows, in_features).contiguous() + v_deq = deq[:, qg_rows + kg_rows:, :].reshape(num_kv_heads * 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)} + + +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: Pro ships a pre-fused qkv_proj; Flash ships q/k/v split. + # Support both: detect qkv_proj.weight first. If present, split it using + # Pro's interleaved num_kv_heads-group layout (with phantom-row FP8 scale + # handling). Otherwise fall back to the Flash-style per-proj path. + qkv_w = lazy.get(f"{prefix}self_attn.qkv_proj.weight") + if qkv_w is not None: + qkv_s = lazy.get(f"{prefix}self_attn.qkv_proj.weight_scale_inv") + # Attention heads: use swa_* for SWA layers, else main. + if is_swa: + num_q = config.get("swa_num_attention_heads", config["num_attention_heads"]) + num_kv = config.get("swa_num_key_value_heads", config["num_key_value_heads"]) + hd = config.get("swa_head_dim", config.get("head_dim")) + vhd = config.get("swa_v_head_dim", config.get("v_head_dim", hd)) + else: + num_q = config["num_attention_heads"] + num_kv = config["num_key_value_heads"] + hd = config.get("head_dim") + vhd = config.get("v_head_dim", hd) + split = split_qkv_fused(qkv_w, qkv_s, num_q, num_kv, hd, vhd) + 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.{proj}.scale"] = s2 + else: + # Flash-style: q/k/v stored separately. + 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-Pro. + # 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 Flash 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: + # Pro-specific: HF bias has mean ~71 with per-expert std ~3e-4. NxDI + # casts router parameters to bf16 at load time, and bf16 step size at + # magnitude 71 is ~0.5 — which completely wipes out the per-expert + # std=3e-4 variation, collapsing all 384 experts to a single bias + # value (all 71.0) and reducing noaux_tc topk to plain sigmoid topk. + # Subtracting the mean first puts the bias at ~0, where bf16 step is + # 2.4e-4 (small enough to preserve the variation). topk is invariant + # to additive constants across all experts, so this is safe. + bias_f32 = router_bias.detach().float().clone() + bias_f32 = bias_f32 - bias_f32.mean() + out[f"{out_prefix}mlp.router.e_score_correction_bias"] = bias_f32 + + 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 Flash (experts ship in FP8); flag loudly. + raise NotImplementedError( + f"Layer {layer_idx} expert 0 gate_proj is BF16; Flash 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-Pro 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-Pro/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py new file mode 100644 index 00000000..85014672 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py @@ -0,0 +1,1676 @@ +# 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-Pro model from Xiaomi. +# Reference: https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro + +"""MiMo-V2.5-Pro 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-Pro. + + 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-Pro 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-Pro 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-Pro 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 Flash) 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-Pro 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-Pro 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 Flash) + # 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. + # Pro: full=config.rope_theta (10M), SWA=config.swa_rope_theta (10K). + # Cannot reuse cached cos/sin across layer types. + # + # 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-Pro.""" + + 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-Pro 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-Pro 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-Pro 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-Pro 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-Pro 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-Pro 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. + + Flash'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 Flash's noaux_tc routing matches HF reference. + + Flash'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): + """Flash 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-Pro.""" + 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-Pro/test/__init__.py b/contrib/models/MiMo-V2.5-Pro/test/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/contrib/models/MiMo-V2.5-Pro/test/integration/__init__.py b/contrib/models/MiMo-V2.5-Pro/test/integration/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py b/contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py new file mode 100644 index 00000000..bcbc368e --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py @@ -0,0 +1,53 @@ +#!/usr/bin/env python3 +"""Integration tests for MiMo-V2-Flash 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-Pro/test/unit/__init__.py b/contrib/models/MiMo-V2.5-Pro/test/unit/__init__.py new file mode 100644 index 00000000..e69de29b From e4a13c14fefe505edfa6c748a6a8abacff27e583 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 06:26:18 +0800 Subject: [PATCH 02/47] [contrib] MiMo-V2.5-Pro: apply review findings from HF reference diff MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit P0 fix: - MiMoV2InferenceConfig now stashes `dense_intermediate_size` from HF's `intermediate_size` BEFORE overwriting `self.intermediate_size` with the MoE value. MiMoV2MLP reads this explicit field instead of the brittle `config.intermediate_size * 8` fallback (which happened to equal 16384 for V2.5-Pro by coincidence). P3 — stale V2-Pro / Flash comments updated: - attention_value_scale comments: "(0.707 for Flash)" → "(0.612 for V2.5-Pro)" - convert_mimo_v2_hf_to_neuron_state_dict kv heads comments: V2.5-Pro has num_key_value_heads=8 (same as SWA), not 4 as in V2-Pro. - smoke_compile docstring reworded to drop "Flash BS=1 recipe" wording. - smoke_compile default recipe changed to moe_tp=1/moe_ep=64/BS=48 (per user request: first V2.5-Pro test uses this recipe because it compiles fastest; bug surface on V2-Pro under this recipe was FP8 precision loss in expert MLP weights, which may not reproduce on V2.5). - preprocess router bias comment: noted measured mean=70.906 std=2.4e-4 (identical pathology to V2-Pro, mean-subtract still required). No behavioral change to FP8 monkey-patches or qkv interleaved-group split logic — HF reference diff confirmed V2.5-Pro ships the same interleaved `[16Q|1K|1V]*8` FP8 qkv layout and the same noaux_tc routing (n_group=1, topk_group=1 degenerate to plain noaux_tc). Co-Authored-By: Claude Opus 4.7 --- .../perf_test/smoke_compile_mimo_v2.py | 30 +++++++++---------- .../perf_test/smoke_generate_mimo_v2.py | 6 ++-- .../preprocess_mimo_v2_fp8.py | 4 +-- .../MiMo-V2.5-Pro/src/modeling_mimo_v2.py | 22 ++++++++++---- 4 files changed, 36 insertions(+), 26 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index b67a95f8..ef5d0049 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -2,11 +2,11 @@ """Minimal compile+load smoke test for MiMo-V2.5-Pro 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 Flash 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. +checkpoint without paying vllm-neuron's startup cost. Compiles the model +(TP=64, configurable MoE TP/EP, blockwise FP8 for routed experts) to a +temp dir, then loads. For `moe_ep_degree > 1` the TKG path raises +`NotImplementedError: Selective Loading with Expert parallelism` unless +`batch_size * top_k / num_experts >= 1.0` → `batch_size >= 384 / 8 = 48`. STAGE controls how far we go: instantiate | compile | load | all (default: all) @@ -31,7 +31,7 @@ ) COMPILED_PATH = os.environ.get( "MIMO_V25_PRO_COMPILED_PATH", - "/opt/dlami/nvme/compiled/mimo_v25_pro_moetp16_ep4_bs48/", + "/opt/dlami/nvme/compiled/mimo_v25_pro_moetp1_ep64_bs48/", ) TP_DEGREE = int(os.environ.get("TP_DEGREE", "64")) @@ -41,13 +41,14 @@ # path raises `NotImplementedError: Selective Loading with Expert parallelism`. BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "48")) CTX_BATCH_SIZE = int(os.environ.get("CTX_BATCH_SIZE", "1")) -# moe_tp=16 / moe_ep=4: after moetp1/ep64 (prefill broken, E_local=6) and -# moetp32/ep2 (HBM OOM by 28MB at load, E_local=192), try the middle ground. -# E_local=384/4=96. Also intermediate/moe_tp = 2048/16 = 128 which matches -# the FP8 scale block size exactly — avoids the stride_fix workaround path. -# world_size = 16*4 = 64 OK. -MOE_TP = int(os.environ.get("MOE_TP", "16")) -MOE_EP = int(os.environ.get("MOE_EP", "4")) +# moe_tp=1 / moe_ep=64: first recipe to try on V2.5-Pro. Lowest compile time +# (no intra-expert TP split) and output quality should be comparable to +# Flash, which uses the same recipe. On V2-Pro this produced garbage +# prefill ("0.0.0.0:8080"), but we're re-testing on V2.5-Pro because the +# V2-Pro root cause ended up being FP8 expert-MLP precision loss, which +# V2.5 may or may not inherit. +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() @@ -96,8 +97,7 @@ def main(): # (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. + # controlled solely by moe_ep_degree. Keep ep_degree=1 unconditionally. neuron_config = MoENeuronConfig( tp_degree=TP_DEGREE, ep_degree=1, diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index f357aee5..b3318714 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -21,7 +21,7 @@ ) COMPILED_PATH = os.environ.get( "MIMO_V25_PRO_COMPILED_PATH", - "/opt/dlami/nvme/compiled/mimo_v25_pro_tp64_moetp1_ep64_fp8/", + "/opt/dlami/nvme/compiled/mimo_v25_pro_moetp1_ep64_bs48/", ) # Must match smoke_compile_mimo_v2.py exactly, else load() sees a @@ -30,8 +30,8 @@ SEQ_LEN = int(os.environ.get("SEQ_LEN", "1024")) BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "48")) # must match smoke_compile CTX_BATCH_SIZE = int(os.environ.get("CTX_BATCH_SIZE", "1")) -MOE_TP = int(os.environ.get("MOE_TP", "16")) -MOE_EP = int(os.environ.get("MOE_EP", "4")) +MOE_TP = int(os.environ.get("MOE_TP", "1")) +MOE_EP = int(os.environ.get("MOE_EP", "64")) PROMPT = os.environ.get( "MIMO_V25_PRO_PROMPT", diff --git a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py index 56eabeaa..b9882fbb 100644 --- a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py +++ b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py @@ -10,7 +10,7 @@ MiMo-V2.5-Pro checkpoint layout: - q_proj, k_proj, v_proj are FUSED into a single `qkv_proj` tensor per - layer (num_kv_heads interleaved groups, Pro-specific). We split into + layer (num_kv_heads interleaved groups, MiMo-V2.5-Pro-specific). We split into three per-row-quantized projections via `split_qkv_fused()`. - o_proj is BF16 (listed in quantization_config.ignored_layers); kept as BF16 on the Neuron side (RowParallelLinear, not QuantizedRowParallel). @@ -409,7 +409,7 @@ def process_layer( 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: - # Pro-specific: HF bias has mean ~71 with per-expert std ~3e-4. NxDI + # V2.5-Pro: HF bias has mean ~71 (same pathology as V2-Pro; measured mean=70.906, std=2.4e-4) with per-expert std ~3e-4. NxDI # casts router parameters to bf16 at load time, and bf16 step size at # magnitude 71 is ~0.5 — which completely wipes out the per-expert # std=3e-4 variation, collapsing all 384 experts to a single bias diff --git a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py index 85014672..230f627c 100644 --- a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py @@ -177,6 +177,14 @@ def __init__(self, *args, **kwargs): self.num_local_experts = self.n_routed_experts self.n_shared_experts = 0 # MiMo-V2.5-Pro has no shared experts + # Stash the HF config's `intermediate_size` (used by the dense MLP + # in layer 0) BEFORE we overwrite `self.intermediate_size` with the + # MoE value. `MiMoV2MLP` reads `dense_intermediate_size` and falls + # back to `config.intermediate_size * 8` if absent, which happens + # to equal 16384 for V2.5-Pro (2048 * 8) but is brittle if Xiaomi + # ever tweaks the ratio. + self.dense_intermediate_size = self.intermediate_size + # Set intermediate_size for MoE layers self.intermediate_size = self.moe_intermediate_size @@ -340,7 +348,7 @@ def __init__( # 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 Flash) right after the V + # by config.attention_value_scale (0.612 for MiMo-V2.5-Pro) 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)) @@ -562,7 +570,7 @@ def forward( 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 Flash) + # HF MiMoV2Attention scales V by attention_value_scale (0.612 for MiMo-V2.5-Pro) # 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. @@ -824,8 +832,10 @@ class MiMoV2MLP(nn.Module): 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) + # Use the dense intermediate size for non-MoE layers. + # `dense_intermediate_size` is stashed in MiMoV2InferenceConfig.__init__ + # before `self.intermediate_size` is overwritten with the MoE value. + self.intermediate_size = config.dense_intermediate_size dtype = config.neuron_config.torch_dtype @@ -1081,8 +1091,8 @@ def convert_mimo_v2_hf_to_neuron_state_dict( num_attention_heads = config.num_attention_heads # MiMo-V2.5-Pro 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 + full_num_kv_heads = config.num_key_value_heads # V2.5-Pro: 8 + swa_num_kv_heads = config.swa_num_key_value_heads # V2.5-Pro: 8 # Check if we need to replicate K/V weights full_use_convert_to_mha = tp_degree > full_num_kv_heads From 78b8e0a97f1b592b84d06c842ce97fe371d2396e Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 10:50:17 +0800 Subject: [PATCH 03/47] [contrib] MiMo-V2.5-Pro: document current WIP status in README MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Record what works and what doesn't on 2026-04-28: - Compile + load succeed on Trn2 (moe_tp=1/ep=64/BS=48 recipe). - Prefill produces coherent English but off-topic output ("100% of the time..." loop for a "explain transformer" prompt). Same signature as V2-Pro's earlier FP8 failures — per-expert weight distribution too narrow for FP8 e4m3 precision. - Note observed token IDs 15/16/4/315/279/882 look suspiciously small but are just " of/ the/ time" etc. — top-BPE English subwords. Greedy decode is correct, the logit distribution itself is wrong. - List recipes still to try (moe_tp=16/ep=4, moe_tp=32/ep=2 etc.) and NxDI constraints that rule out BS=1 when moe_ep>1. Points future debuggers at Jim Burtoft's Flash FP8 observation and his Kimi PR #131 SDK 2.28 recommendation. No code changes. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 16 ++++++++++++++++ 1 file changed, 16 insertions(+) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 38623380..6e44fd82 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -35,6 +35,22 @@ Key features: - **Sigmoid Router + noaux_tc**: `sigmoid(logits) + e_score_correction_bias` is used to pick top-8 experts; unbiased `sigmoid(logits)` becomes the affinity weights. `n_group=1, topk_group=1` degenerates group-limited routing to plain noaux_tc. - **attention_value_scale = 0.612**: HF reference multiplies `value_states` by this before `softmax(QK^T) × V` (NOT applied post-attention); the NxDI port matches. +## Status (work-in-progress) + +**This port compiles cleanly and loads on Trn2 but does not yet produce coherent output.** Current symptoms under the default recipe (`tp_degree=64, moe_tp_degree=1, moe_ep_degree=64, batch_size=48, seq_len=1024`) on 2026-04-28: + +- Prefill drifts from token 1: "Explain in one sentence what a transformer neural network is." → `"100% of the time, 100% of the time, ..."` (greedy decode, temperature=0). Decode speed 0.72 tok/s. Note the output is syntactically valid English ("100% of the time" = BPE tokens `15/16/4/315/279/882` = `"1"/"0"/"%"/" of"/" the"/" time"`, all high-frequency) — not a sampling bug: greedy argmax is correctly picking the model's top token, but the logit distribution itself is wrong (output unrelated to the prompt). Same signature as Jim Burtoft's "Flash FP8 → `erotici` repeat" symptom. +- Same failure pattern was observed on the MiMo-V2-Pro port under the same recipe (`"0.0.0.0:8080"` etc.). Root cause identified there was **FP8 expert-MLP precision loss**: Pro's expert weight std ≈ 0.0018 (10× smaller than Flash's ≈ 0.019), landing right at FP8 e4m3's subnormal threshold. V2.5-Pro inherits the same per-expert scale pathology (router bias mean ≈ 70.906 std ≈ 2.4e-4, verified via preprocess; mean-subtract workaround applied). +- Reference: Jim Burtoft observed similar prompt-dependent FP8 degradation on Flash (see PR notes) and recommends selective BF16 retention for precision-sensitive layers. + +Other recipes to try (none verified yet on V2.5-Pro): +- `moe_tp_degree=16, moe_ep_degree=4, BS=48` — balances E_local=96 vs HBM. +- `moe_tp_degree=32, moe_ep_degree=2, BS=1` — mirrors Jim's Kimi-K2 PR, but V2-Pro OOM'd by 28MB on load at BS=48; BS=1 hits `NotImplementedError: Selective Loading with Expert parallelism`. + +Known NxDI limits that constrain recipe choice: +- `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else NotImplementedError). With `num_experts=384, top_k=8` this forces `BS >= 48`. +- `n_routed_experts=384 = 2^7 * 3` → `384 / ep_degree` is never a power of 2 (6, 12, 24, 48, 96, 192, 384). Kimi PR #131 says NKI `_bwmm_shard_on_block_nki_call` on SDK 2.29 has "depressed logits with EP=2" and recommends SDK 2.28. SDK 2.28 venv is not currently installed on the target DLAMI. + ## Prerequisites - **Instance**: trn2.48xlarge (32 NeuronCores, logical_nc_config=2 → 64 logical cores) From e990f76004ca914aff30e326f21b2ef20540142e Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 11:08:08 +0800 Subject: [PATCH 04/47] [contrib] MiMo-V2.5-Pro: correct FP8 root-cause framing in Status MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Earlier wording said "Pro's expert weight std is too small for FP8 precision" in absolute terms. That's misleading — sglang on H100/H200 runs the exact same OCP FP8 checkpoint and produces correct output, because GPU cutlass/sglang paths dequantize FP8 to BF16 before the matmul. The actual issue appears to be Neuron's NKI blockwise FP8 compute kernel (_bwmm_shard_on_block_nki_call) running FP8 compute directly on subnormal-leaning tensors. Jim Burtoft's Kimi PR #131 names the Neuron SDK 2.29 blockwise kernel as producing "depressed logits with EP=2" and recommends SDK 2.28. Also noted: V2.5-Pro MoE expert weights are byte-identical to V2-Pro (measured layer 1 expert 0 gate_proj stats match to 6 decimals), so all V2-Pro FP8 workarounds remain required — not a new bug. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 6e44fd82..552354d9 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -40,8 +40,9 @@ Key features: **This port compiles cleanly and loads on Trn2 but does not yet produce coherent output.** Current symptoms under the default recipe (`tp_degree=64, moe_tp_degree=1, moe_ep_degree=64, batch_size=48, seq_len=1024`) on 2026-04-28: - Prefill drifts from token 1: "Explain in one sentence what a transformer neural network is." → `"100% of the time, 100% of the time, ..."` (greedy decode, temperature=0). Decode speed 0.72 tok/s. Note the output is syntactically valid English ("100% of the time" = BPE tokens `15/16/4/315/279/882` = `"1"/"0"/"%"/" of"/" the"/" time"`, all high-frequency) — not a sampling bug: greedy argmax is correctly picking the model's top token, but the logit distribution itself is wrong (output unrelated to the prompt). Same signature as Jim Burtoft's "Flash FP8 → `erotici` repeat" symptom. -- Same failure pattern was observed on the MiMo-V2-Pro port under the same recipe (`"0.0.0.0:8080"` etc.). Root cause identified there was **FP8 expert-MLP precision loss**: Pro's expert weight std ≈ 0.0018 (10× smaller than Flash's ≈ 0.019), landing right at FP8 e4m3's subnormal threshold. V2.5-Pro inherits the same per-expert scale pathology (router bias mean ≈ 70.906 std ≈ 2.4e-4, verified via preprocess; mean-subtract workaround applied). -- Reference: Jim Burtoft observed similar prompt-dependent FP8 degradation on Flash (see PR notes) and recommends selective BF16 retention for precision-sensitive layers. +- Same failure pattern was observed on the MiMo-V2-Pro port under the same recipe (`"0.0.0.0:8080"` etc.). Root cause identified there appears to be **Neuron's NKI blockwise FP8 compute kernel handling Pro's tight expert-weight distribution** (std ≈ 0.0018, ~10× smaller than Flash's ≈ 0.019). This is NOT an FP8-format problem per se — sglang on H100/H200 runs the exact same OCP FP8 checkpoint and produces correct output, because GPU paths dequantize FP8→BF16 before the matmul. Neuron NKI does FP8 compute directly and seems to lose precision on subnormal-leaning tensors. +- V2.5-Pro's MoE expert weights are byte-identical to V2-Pro (verified layer 1 expert 0 dequant stats match to 6 decimal places on 2026-04-28), so all V2-Pro workarounds remain required (router bias mean-subtract, qkv interleaved split, `_apply_2d_per_channel_fix`, `_apply_blockwise_scale_stride_fix`). +- Reference: Jim Burtoft observed similar prompt-dependent FP8 degradation on Flash and his Kimi PR #131 names "blockwise kernel padding produces depressed logits with EP=2 on SDK 2.29; SDK 2.28 recommended". Other recipes to try (none verified yet on V2.5-Pro): - `moe_tp_degree=16, moe_ep_degree=4, BS=48` — balances E_local=96 vs HBM. From f0d9c0b5f036c1e9a8d4442ec8f5585e9dca075b Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 11:59:29 +0800 Subject: [PATCH 05/47] [contrib] MiMo-V2.5-Pro: parallel preprocess + NVMe mount docs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Parallel preprocess wrapper: - preprocess_mimo_v2_parallel.py: multiprocessing Pool wrapper around preprocess_mimo_v2_fp8.process_layer. Each worker opens its own LazyWeightMap and processes one layer at a time. N_WORKERS default raised to 12 (user request: "越多越好"); 70 layers * ~25 GB peak/layer stays under ~300 GB RAM on a 2 TB trn2.48xl. - run_preprocess_parallel.sh: thin shell wrapper exposing HF_MODEL_PATH, SAVE_PATH, TP_DEGREE, N_WORKERS env vars. Defaults to the 2_9_nxd_inference venv (same one used by the serial preprocess). Wall-clock ~30 min serial → ~5-6 min at 12 workers on fresh cache. README: - Added "NVMe mount" subsection under Prerequisites. trn2.48xl DLAMI assembles four NVMe into RAID0 at /opt/dlami/nvme but does NOT remount automatically after a reboot. Document mdadm --assemble + mount /dev/md0 /opt/dlami/nvme before any path in the recipes resolves. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 19 ++ .../preprocess_mimo_v2_parallel.py | 179 ++++++++++++++++++ .../run_preprocess_parallel.sh | 34 ++++ 3 files changed, 232 insertions(+) create mode 100644 contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_parallel.py create mode 100755 contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 552354d9..1b35c567 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -59,6 +59,25 @@ Known NxDI limits that constrain recipe choice: - **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**: ~3 TB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~962 GB, the Neuron-FP8 preprocessed output is ~1 TB, and `save_sharded_checkpoint=true` writes another ~300-1000 GB per compiled config (varies with recipe)). +### NVMe mount + +The Trn2 DLAMI ships with four local NVMe SSDs that are assembled into a +RAID0 array at `/opt/dlami/nvme`. After a reboot the mount is **NOT** +reassembled automatically — you must re-mount manually before the paths +below resolve: + +```bash +lsblk # confirm you see nvme0n1..nvme3n1 devices +sudo mdadm --assemble /dev/md0 /dev/nvme[0-3]n1 2>/dev/null || true +sudo mount /dev/md0 /opt/dlami/nvme +df -h /opt/dlami/nvme # should show ~6.9 TB total +``` + +If `mdadm --assemble` says the array is already assembled, the mount +step alone is enough. If `/dev/md0` doesn't exist, the array was never +created on this instance — run `/opt/dlami/setup-nvme.sh` (or the +DLAMI's built-in helper; consult `ls /opt/dlami/*.sh`) before mounting. + ## 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-Pro FP8. First-time compile takes ~45-60 minutes; subsequent runs hit the neuronx-cc cache and start in a few minutes. diff --git a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_parallel.py b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_parallel.py new file mode 100644 index 00000000..23f05bcc --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_parallel.py @@ -0,0 +1,179 @@ +#!/usr/bin/env python3 +"""Parallel wrapper around preprocess_mimo_v2_fp8.process_layer. + +Each layer is independent: `process_layer(L, lazy, config, ...)` reads only +keys under `model.layers.{L}.*` from the HF shards and returns the Neuron +layer-shard dict. With 70 layers and per-MoE-layer cost ~60s serial, 4-8 +workers cuts wallclock from ~70 min to ~15-20 min (I/O + CPU FP8 math). + +Each worker opens its own LazyWeightMap so there's no shared safetensors +handle. Output dir is a CLI arg so it can write to a clean path without +touching the serial run's output. +""" +import argparse +import gc +import json +import multiprocessing as mp +import os +import shutil +import sys +import time + +# Resolve the sibling single-layer preprocess module. This file lives at +# .../MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_parallel.py, +# so the importable parent is two levels up. +_SRC_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) +if _SRC_DIR not in sys.path: + sys.path.insert(0, _SRC_DIR) +from conversion_script.preprocess_mimo_v2_fp8 import ( # noqa: E402 + LazyWeightMap, + process_layer, + save_shard, +) +from safetensors.torch import save_file # noqa: E402 + + +def _worker(task): + layer_idx, hf_model_path, save_path, config = task + hybrid = config.get( + "hybrid_layer_pattern", [0] * config["num_hidden_layers"] + ) + moe_freq = config.get("moe_layer_freq", [1] * config["num_hidden_layers"]) + is_dense = moe_freq[layer_idx] == 0 + is_swa = hybrid[layer_idx] == 1 + + with open( + os.path.join(hf_model_path, "model.safetensors.index.json") + ) as fh: + weight_map_in = json.load(fh)["weight_map"] + lazy = LazyWeightMap(hf_model_path, weight_map_in) + try: + t0 = time.time() + layer_sd = process_layer( + layer_idx, lazy, config, is_dense=is_dense, is_swa=is_swa + ) + filename = f"model_layer{layer_idx}.safetensors" + path = os.path.join(save_path, filename) + materialized = {} + total_bytes = 0 + for k, v in layer_sd.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) + keys = list(materialized.keys()) + del materialized, layer_sd + gc.collect() + elapsed = time.time() - t0 + finally: + lazy.close() + return layer_idx, is_dense, is_swa, keys, filename, total_bytes, elapsed + + +def main(): + p = argparse.ArgumentParser() + p.add_argument("--hf_model_path", required=True) + p.add_argument("--save_path", required=True) + p.add_argument("--tp_degree", type=int, default=64) + p.add_argument( + "--workers", + type=int, + default=int(os.environ.get("N_WORKERS", "12")), + ) + args = p.parse_args() + + os.makedirs(args.save_path, exist_ok=True) + with open(os.path.join(args.hf_model_path, "config.json")) as fh: + config = json.load(fh) + num_layers = config["num_hidden_layers"] + + print( + f"[par] {num_layers} layers x {args.workers} workers -> {args.save_path}", + flush=True, + ) + tasks = [ + (L, args.hf_model_path, args.save_path, config) for L in range(num_layers) + ] + weight_map_out = {} + t_start = time.time() + + ctx = mp.get_context("spawn") + with ctx.Pool(args.workers) as pool: + done = 0 + for li, is_dense, is_swa, keys, filename, total_bytes, elapsed in pool.imap_unordered( + _worker, tasks + ): + done += 1 + for k in keys: + weight_map_out[k] = filename + tag = "dense" if is_dense else "moe" + attn = "swa" if is_swa else "full" + print( + f" [{done:2d}/{num_layers}] layer {li:2d} [{tag:5s} {attn:4s}] " + f"{total_bytes/1e9:6.2f} GB in {elapsed:5.1f}s " + f"(wall {time.time()-t_start:5.1f}s)", + flush=True, + ) + + print( + f"[par] all {num_layers} layers done in {time.time()-t_start:.1f}s", + flush=True, + ) + + print("[par] processing embed_tokens / norm / lm_head ...", flush=True) + with open( + os.path.join(args.hf_model_path, "model.safetensors.index.json") + ) as fh: + weight_map_in = json.load(fh)["weight_map"] + lazy = LazyWeightMap(args.hf_model_path, weight_map_in) + extras = {} + try: + 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" WARN: missing {src}", flush=True) + if "lm_head.weight" not in extras and "embed_tokens.weight" in extras: + extras["lm_head.weight"] = extras["embed_tokens.weight"].detach().clone() + finally: + lazy.close() + save_shard(extras, args.save_path, "model_extras.safetensors", weight_map_out) + del extras + + total_size = 0 + for f in set(weight_map_out.values()): + total_size += os.path.getsize(os.path.join(args.save_path, f)) + index = { + "metadata": {"total_size": total_size}, + "weight_map": weight_map_out, + } + with open( + os.path.join(args.save_path, "model.safetensors.index.json"), "w" + ) as fh: + json.dump(index, fh, indent=2) + + for name in sorted(os.listdir(args.hf_model_path)): + if name.endswith(".safetensors"): + continue + if name == "model.safetensors.index.json": + continue + src = os.path.join(args.hf_model_path, name) + if os.path.isfile(src): + shutil.copy(src, os.path.join(args.save_path, name)) + + print( + f"\n[par] DONE. total_size={total_size/1e9:.2f} GB " + f"tensors={len(weight_map_out)} -> {args.save_path}", + flush=True, + ) + + +if __name__ == "__main__": + main() diff --git a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh new file mode 100755 index 00000000..5001ecd6 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh @@ -0,0 +1,34 @@ +#!/bin/bash +# Parallel wrapper around preprocess_mimo_v2_parallel.py. +# +# Each worker dequants one MoE layer at a time (peak ~25 GB per layer on +# V2.5-Pro's 6144 hidden / 384 experts / 2048 intermediate shape). 12 +# workers stay under ~300 GB CPU RAM on a 2 TB box while keeping the +# 192-core CPU busy. On a trn2.48xl that brings total wall time from +# ~30 min (serial) to ~5-6 min. +# +# Env: +# HF_MODEL_PATH raw HF checkpoint (default: /opt/dlami/nvme/models/MiMo-V2.5-Pro) +# SAVE_PATH output Neuron checkpoint (default: /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8) +# TP_DEGREE tensor-parallel degree used at compile time (default: 64) +# N_WORKERS concurrent layer workers (default: 12) +# VENV venv with torch + safetensors + contrib pkg on sys.path +set -e + +VENV=${VENV:-/opt/aws_neuronx_venv_pytorch_2_9_nxd_inference} +source "$VENV/bin/activate" + +HF_MODEL_PATH=${HF_MODEL_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro} +SAVE_PATH=${SAVE_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8} +TP_DEGREE=${TP_DEGREE:-64} +N_WORKERS=${N_WORKERS:-12} + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +SRC_DIR="$(cd "$SCRIPT_DIR/.." && pwd)" +export PYTHONPATH="$SRC_DIR:$PYTHONPATH" + +exec python3 "$SCRIPT_DIR/preprocess_mimo_v2_parallel.py" \ + --hf_model_path "$HF_MODEL_PATH" \ + --save_path "$SAVE_PATH" \ + --tp_degree "$TP_DEGREE" \ + --workers "$N_WORKERS" From b455a1ef850cf3a91a76fef24ae07d23181244cf Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 12:16:04 +0800 Subject: [PATCH 06/47] [contrib] MiMo-V2.5-Pro: set AWS Llama-405B FP8 env vars in smoke scripts MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit AWS trn2-llama3.1-405b-speculative FP8 tutorial ("Scenario 2, Step 2") requires XLA_HANDLE_SPECIAL_SCALAR=1 and UNSAFE_FP8FNCAST=1 for OCP-sourced FP8 checkpoints on Neuron. Setting them in both smoke_compile_mimo_v2.py and smoke_generate_mimo_v2.py via os.environ.setdefault (user-level env overrides still win). Note: our preprocess output has 0 bytes in the IEEE-NaN-adjacent range (byte exp=0b1111), verified on layers.1 attn q/k/v and MoE gate_up/down in /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8. So these flags are theoretically optional for our pipeline, but they match the exact surface of AWS's reference FP8 tutorial — cheap safety. Also corrected stale docstrings: smoke_compile now says the NxDI venv (pytorch_2_9_nxd_inference) is the target, not the vllm venv. Co-Authored-By: Claude Opus 4.7 --- .../perf_test/smoke_compile_mimo_v2.py | 16 ++++++++++++++-- .../perf_test/smoke_generate_mimo_v2.py | 6 +++++- 2 files changed, 19 insertions(+), 3 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index ef5d0049..c92946e4 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -16,8 +16,9 @@ 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). +Run under /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference (NxDI direct). +The `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` venv is only +needed for vllm serving (bench_mimo_v2.sh). """ import os @@ -25,6 +26,17 @@ import time import traceback +# AWS Llama-3.1-405B FP8 tutorial requires these two env vars to correctly +# handle OCP-derived FP8 checkpoints on Neuron: XLA_HANDLE_SPECIAL_SCALAR=1 +# opts in to XLA emitting the bit-reinterpretation path for fp8_e4m3fn scalars, +# and UNSAFE_FP8FNCAST=1 mirrors it for torch-side casts. Our preprocess output +# has 0 bytes in the IEEE-NaN range (verified 2026-04-28), so these flags are +# theoretically unnecessary, but setting them matches the AWS tutorial +# surface exactly. Source: trn2-llama3.1-405b-speculative-tutorial.html +# "Scenario 2, Step 2". +os.environ.setdefault("XLA_HANDLE_SPECIAL_SCALAR", "1") +os.environ.setdefault("UNSAFE_FP8FNCAST", "1") + MODEL_PATH = os.environ.get( "MIMO_V25_PRO_MODEL_PATH", "/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8", diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index b3318714..2eb7ab13 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -7,7 +7,7 @@ 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. +Run under /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference. """ import os @@ -15,6 +15,10 @@ import time import traceback +# Same env vars as smoke_compile — match AWS Llama-405B FP8 tutorial. +os.environ.setdefault("XLA_HANDLE_SPECIAL_SCALAR", "1") +os.environ.setdefault("UNSAFE_FP8FNCAST", "1") + MODEL_PATH = os.environ.get( "MIMO_V25_PRO_MODEL_PATH", "/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8", From 9d29bbf27864fa228ba67f4deb5e9e5ef26654e8 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 14:06:32 +0800 Subject: [PATCH 07/47] [contrib] MiMo-V2.5-Pro: fix tokenizer padding_side='left' for decoder-only LM HuggingFace tokenizer defaults to padding_side='right', which silently corrupts batched prefill on a causal LM: the last token of each slot becomes a pad token, and the logit used for generating the next token is predicting "what comes after the pad", not "what comes after the real prompt". Observed when running a 6-prompt probe at BS=48: prompts that nearly fill the 267-token batch dimension produced garbage output like "all spaces" (token 220) or random short-id BPE noise. Fix: explicitly set padding_side='left' after tokenizer load. Single-prompt smoke (all slots == same prompt, so no padding triggered) was not affected by this bug, but was producing wrong output for a different reason (the underlying FP8 expert-MLP precision issue). Co-Authored-By: Claude Opus 4.7 --- .../MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index 2eb7ab13..9f678a09 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -15,7 +15,9 @@ import time import traceback -# Same env vars as smoke_compile — match AWS Llama-405B FP8 tutorial. +# AWS Llama-3.1-405B FP8 tutorial env vars — these are RUNTIME flags that +# affect XLA's fp8 special-scalar handling. Safe to set at generate time; +# do NOT set at compile time (it changes HLO and busts the neuronx-cc cache). os.environ.setdefault("XLA_HANDLE_SPECIAL_SCALAR", "1") os.environ.setdefault("UNSAFE_FP8FNCAST", "1") @@ -134,6 +136,10 @@ def main(): print(f"[gen] Loaded in {time.time() - t0:.1f}s") tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) + # Decoder-only LM requires left-padding so the last token of each batch + # slot is the real prompt ending, not a pad token. Default HF tokenizer + # padding_side is 'right' which silently corrupts batched prefill. + tokenizer.padding_side = "left" adapter = HuggingFaceGenerationAdapter(model) # When CHAT_TEMPLATE=1, wrap the raw prompt in the checkpoint's chat From 2a4c9ff11ad2e248e93e40a977003dc8b80ae0a9 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 17:05:48 +0800 Subject: [PATCH 08/47] [contrib] MiMo-V2.5-Pro: skip FP8 quant on q/k/v_proj (attention path) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit V2.5-Pro's attention Q/K/V weights have dequantized abs_mean ~0.001-0.005, roughly 4x smaller than V2.5 (which works). Preprocess has been patched to rewrite the q/k/v_proj tensors in the preprocessed checkpoint as BF16 (matching how o_proj is already handled). Add q_proj/k_proj/v_proj to modules_to_not_convert so NxDI does not try to swap them to QuantizedColumnParallel at convert() time — they remain plain ColumnParallelLinear with BF16 weights. MoE expert weights (gate_up_proj, down_proj) stay FP8 blockwise; their weights saturate the full FP8 ±240 range so quantization is lossless there. Only the attention path goes BF16. Co-Authored-By: Claude Opus 4.7 --- .../perf_test/smoke_compile_mimo_v2.py | 19 ++++++++----------- .../perf_test/smoke_generate_mimo_v2.py | 2 +- 2 files changed, 9 insertions(+), 12 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index c92946e4..6b5fc2e7 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -26,16 +26,13 @@ import time import traceback -# AWS Llama-3.1-405B FP8 tutorial requires these two env vars to correctly -# handle OCP-derived FP8 checkpoints on Neuron: XLA_HANDLE_SPECIAL_SCALAR=1 -# opts in to XLA emitting the bit-reinterpretation path for fp8_e4m3fn scalars, -# and UNSAFE_FP8FNCAST=1 mirrors it for torch-side casts. Our preprocess output -# has 0 bytes in the IEEE-NaN range (verified 2026-04-28), so these flags are -# theoretically unnecessary, but setting them matches the AWS tutorial -# surface exactly. Source: trn2-llama3.1-405b-speculative-tutorial.html -# "Scenario 2, Step 2". -os.environ.setdefault("XLA_HANDLE_SPECIAL_SCALAR", "1") -os.environ.setdefault("UNSAFE_FP8FNCAST", "1") +# NOTE: AWS Llama-3.1-405B FP8 tutorial recommends XLA_HANDLE_SPECIAL_SCALAR=1 +# and UNSAFE_FP8FNCAST=1 for OCP-derived FP8 checkpoints. Setting them at +# compile time, however, appears to change the HLO that gets emitted (likely +# because XLA lowering of fp8_e4m3fn special scalars switches paths), which +# busts the neuronx-cc cache and forces a full recompile (~90 min). If you +# need these flags, set them ONLY at generate time and rely on the compiled +# NEFF's built-in `--experimental-unsafe-fp8e4m3fn-as-fp8e4m3` handling. MODEL_PATH = os.environ.get( "MIMO_V25_PRO_MODEL_PATH", @@ -150,7 +147,7 @@ def main(): "lm_head", "norm", "router", - "o_proj", + "o_proj", "q_proj", "k_proj", "v_proj", ], ) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index 9f678a09..762fb0d8 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -114,7 +114,7 @@ def main(): "lm_head", "norm", "router", - "o_proj", + "o_proj", "q_proj", "k_proj", "v_proj", ], ) From 2b8b5771c53e227b0e3bc7dcbeb3af3862fa7608 Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 28 Apr 2026 19:54:56 +0800 Subject: [PATCH 09/47] [contrib] MiMo-V2.5-Pro: try use_torch_block_wise + restore FP8 q/k/v Previous attempt (2a4c9ff) rewrote q/k/v_proj as BF16 to work around Pro's attention weight precision (q_proj abs_mean ~0.00124, 4x smaller than V2.5). Compile succeeded, but load failed with HBM OOM: the BF16 attention weights added ~2 GB per rank, pushing Tensors to 20.93 GB on a 24 GB Neuron HBM and leaving no room for collective DMA rings. Back off on the BF16-attn approach and try a different hypothesis: the NKI blockwise matmul kernel has accumulator precision issues on Pro's MoE expert weights (scale_mean ~5e-5 vs 2.5e-4 on V2.5). Switch blockwise_matmul_config from use_shard_on_block_dynamic_while to use_torch_block_wise=True, which uses a PyTorch fallback that dequantizes each block to BF16 before matmul. Slower but more precise in the accumulator. q/k/v_proj return to FP8 (back out of modules_to_not_convert) so the attention weights don't blow HBM. Co-Authored-By: Claude Opus 4.7 --- .../perf_test/smoke_compile_mimo_v2.py | 17 +++++++++++------ .../perf_test/smoke_generate_mimo_v2.py | 7 ++++--- 2 files changed, 15 insertions(+), 9 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index 6b5fc2e7..be7f72e7 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -124,12 +124,17 @@ def main(): 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. Matches Flash + Kimi. + # V2.5-Pro experiment: use PyTorch-fallback blockwise matmul instead + # of the NKI _call_shard_hidden_kernel path. Rationale: Pro's MoE + # expert weight scales run ~3-7x smaller than V2.5 (scale_mean + # ~5e-5 vs 2.5e-4); the NKI blockwise kernel's lower accumulator + # precision compounds across 70 layers into decode-time gibberish. + # PyTorch fallback dequantizes each block to BF16 before matmul, + # trading throughput for accumulator precision. blockwise_matmul_config={ - "use_shard_on_block_dynamic_while": True, - "block_sharding_strategy": "PING_PONG", + "use_torch_block_wise": True, + "use_shard_on_intermediate_dynamic_while": True, + "skip_dma_token": True, }, # Persist sharded FP8 weights to disk so subsequent load()s skip the # ~10-minute shard_checkpoint step (writes weights/tp{0..63}_*.safetensors @@ -147,7 +152,7 @@ def main(): "lm_head", "norm", "router", - "o_proj", "q_proj", "k_proj", "v_proj", + "o_proj", ], ) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index 762fb0d8..c3388e89 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -99,8 +99,9 @@ def main(): 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", + "use_torch_block_wise": True, + "use_shard_on_intermediate_dynamic_while": True, + "skip_dma_token": True, }, save_sharded_checkpoint=True, quantized=True, @@ -114,7 +115,7 @@ def main(): "lm_head", "norm", "router", - "o_proj", "q_proj", "k_proj", "v_proj", + "o_proj", ], ) From b3a8487f08437bed3858ad781e434d0cb9c57da0 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 10:30:07 +0800 Subject: [PATCH 10/47] [contrib] MiMo-V2.5-Pro: wire up vLLM serving and record FP8 perf Pro is now serveable via vllm-neuron 0.5.0 on Trn2 (TP=64, moe_ep=64, BS=48). Output quality under the FP8 recipe is still prompt-dependent (drift on most prompts, coherent on self-intro style), consistent with Pro's 4-7x smaller MoE FP8 scales compared to V2.5 and the V2-Pro symptom. Changes: - Revert blockwise_matmul_config back to use_shard_on_block_dynamic_while + PING_PONG (Flash/Kimi recipe). The use_torch_block_wise + BF16-attn experiments both OOM on load. - Fix bench_mimo_v2.sh / smoke configs from BS=32 (Flash) to BS=48 (Pro: 384/8=48), plus all accompanying text in the README. - vLLM patch now registers both MiMoV2FlashForCausalLM and MiMoV2ProForCausalLM in vLLM's ModelRegistry, overriding the built-in GPU stubs; patch works against vllm-neuron release-0.5.0. - Point sanity_check.sh, run_bench_single.sh, 0_setup.sh defaults at the Neuron-FP8 checkpoint (not BF16). - Record measured vLLM serving throughput at c=1/16/48 in the README Performance section (replaces stale BF16 numbers). - Rewrite the Status section: document the drift pattern with prompt examples, the recipes that were tried and failed (BF16-attn, torch blockwise), and the two-node BF16 experiment queued next. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 85 +++++++++---------- .../models/MiMo-V2.5-Pro/perf_test/0_setup.sh | 8 +- .../MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh | 19 +++-- .../perf_test/run_bench_single.sh | 6 +- .../MiMo-V2.5-Pro/perf_test/sanity_check.sh | 2 +- .../perf_test/smoke_compile_mimo_v2.py | 15 ++-- .../perf_test/smoke_generate_mimo_v2.py | 5 +- .../perf_test/vllm-neuron-patch.patch | 20 +++-- 8 files changed, 79 insertions(+), 81 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 1b35c567..7078e480 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -37,16 +37,21 @@ Key features: ## Status (work-in-progress) -**This port compiles cleanly and loads on Trn2 but does not yet produce coherent output.** Current symptoms under the default recipe (`tp_degree=64, moe_tp_degree=1, moe_ep_degree=64, batch_size=48, seq_len=1024`) on 2026-04-28: +**This port compiles cleanly and serves via vLLM on Trn2, but output quality is not production-ready under the default FP8 recipe.** Sanity checks on 2026-04-29 under `tp_degree=64, moe_tp_degree=1, moe_ep_degree=64, batch_size=48, seq_len=1024`: -- Prefill drifts from token 1: "Explain in one sentence what a transformer neural network is." → `"100% of the time, 100% of the time, ..."` (greedy decode, temperature=0). Decode speed 0.72 tok/s. Note the output is syntactically valid English ("100% of the time" = BPE tokens `15/16/4/315/279/882` = `"1"/"0"/"%"/" of"/" the"/" time"`, all high-frequency) — not a sampling bug: greedy argmax is correctly picking the model's top token, but the logit distribution itself is wrong (output unrelated to the prompt). Same signature as Jim Burtoft's "Flash FP8 → `erotici` repeat" symptom. -- Same failure pattern was observed on the MiMo-V2-Pro port under the same recipe (`"0.0.0.0:8080"` etc.). Root cause identified there appears to be **Neuron's NKI blockwise FP8 compute kernel handling Pro's tight expert-weight distribution** (std ≈ 0.0018, ~10× smaller than Flash's ≈ 0.019). This is NOT an FP8-format problem per se — sglang on H100/H200 runs the exact same OCP FP8 checkpoint and produces correct output, because GPU paths dequantize FP8→BF16 before the matmul. Neuron NKI does FP8 compute directly and seems to lose precision on subnormal-leaning tensors. -- V2.5-Pro's MoE expert weights are byte-identical to V2-Pro (verified layer 1 expert 0 dequant stats match to 6 decimal places on 2026-04-28), so all V2-Pro workarounds remain required (router bias mean-subtract, qkv interleaved split, `_apply_2d_per_channel_fix`, `_apply_blockwise_scale_stride_fix`). +- Prompt-dependent drift. Self-intro style prompts ("introduce yourself in one sentence") return coherent Chinese output; most other prompts collapse to repetition or unrelated text within a few tokens (e.g. `"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`). +- Same failure pattern was observed on the MiMo-V2-Pro port under the same recipe (`"0.0.0.0:8080"`, etc.). Root cause appears to be **Neuron's NKI blockwise FP8 compute kernel handling Pro's tight expert-weight distribution**: Pro's per-expert weight `abs_mean ≈ 0.00124` and blockwise `scale_mean ≈ 2.3e-5` are 4-7× smaller than V2.5 (256 experts) on the same recipe, and the NKI accumulator loses precision compounded across 70 layers. V2.5-Pro's MoE expert weights are byte-identical to V2-Pro (verified layer 1 expert 0 dequant stats match to 6 decimal places on 2026-04-28), so all V2-Pro workarounds remain required (router bias mean-subtract, qkv interleaved split, `_apply_2d_per_channel_fix`, `_apply_blockwise_scale_stride_fix`). +- GPU stacks (sglang on H100/H200) run the exact same OCP FP8 checkpoint correctly, because they dequantize FP8→BF16 before the matmul. Neuron NKI does FP8 compute directly and drifts on subnormal-leaning tensors. - Reference: Jim Burtoft observed similar prompt-dependent FP8 degradation on Flash and his Kimi PR #131 names "blockwise kernel padding produces depressed logits with EP=2 on SDK 2.29; SDK 2.28 recommended". -Other recipes to try (none verified yet on V2.5-Pro): -- `moe_tp_degree=16, moe_ep_degree=4, BS=48` — balances E_local=96 vs HBM. -- `moe_tp_degree=32, moe_ep_degree=2, BS=1` — mirrors Jim's Kimi-K2 PR, but V2-Pro OOM'd by 28MB on load at BS=48; BS=1 hits `NotImplementedError: Selective Loading with Expert parallelism`. +Recipes that were tried and did not resolve the drift (all on 2026-04-28/29): +- **q/k/v_proj dequant to BF16** (compile-time `modules_to_not_convert` includes `q_proj, k_proj, v_proj`): compiled but HBM OOM on load — BF16 attention weights pushed per-rank tensors to 20.93/24 GB, leaving no room for collective DMA rings. +- **`use_torch_block_wise=True`** (PyTorch-fallback blockwise matmul for higher accumulator precision): compile+shard succeeded after ~2 h, but `model.load()` crashed with `status=4 Allocation Failure` — the fallback path raises HBM demand. + +Next experiments queued (none verified yet): +- BF16 weights across two Trn2 instances with cross-node TP/PP (single-instance HBM cannot hold BF16 Pro). +- SDK 2.28 venv test once installed, per Kimi PR #131. +- Selective BF16 only on `gate_up_proj` (smallest MoE scales) while keeping `down_proj` FP8, if HBM allows. Known NxDI limits that constrain recipe choice: - `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else NotImplementedError). With `num_experts=384, top_k=8` this forces `BS >= 48`. @@ -114,8 +119,8 @@ bash contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh bash contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh ``` -The bench script runs two configurations (BS=32 and BS=128, both -`moe_tp_degree=X / moe_ep_degree=Y (see bench script)`) and logs results under +The bench script runs two configurations (BS=48 and BS=128, both +`moe_tp_degree=1 / moe_ep_degree=64`) and logs results under `/tmp/bench_results/mimo_v25_pro/`. For a quick `curl` sanity check while the server is up: @@ -182,10 +187,10 @@ neuron_config = MoENeuronConfig( 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, + batch_size=48, # must be >= num_experts / top_k = 384 / 8 = 48 + max_batch_size=48, ctx_batch_size=1, - tkg_batch_size=32, + tkg_batch_size=48, seq_len=1024, n_active_tokens=128, torch_dtype=torch.bfloat16, @@ -245,15 +250,15 @@ Both default to the recommended FP8 recipe (`moe_tp=1`, `moe_ep=64`). ### 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 Flash (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 Flash's 47 MoE layers and manifests as output collapse ("helpful helpful helpful ...") after roughly 30 decode tokens. +**Why**: at `moe_tp_degree=64` each rank owns 1/64 of the intermediate dim, which for MiMo-V2.5-Pro (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 Pro's 69 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. +`moe_tp_degree=1, moe_ep_degree=64` keeps each expert's weights and blockwise scales intact on a single rank (6 experts per rank for Pro's 384 experts), which preserves per-channel scale. On V2.5 (256 experts) this recipe yields coherent output; on V2.5-Pro it still exhibits prompt-dependent drift (see Status). -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-Pro FP8. +Intermediate ratios (`moe_tp=32/ep=2`, `moe_tp=16/ep=4`) have been empirically tested and still produce gibberish, so `moe_tp=1/moe_ep=64` is the only currently-usable moe_tp/ep combination. -### batch_size >= 32 +### batch_size >= 48 -NxDI's TKG (token generation) path refuses Expert Parallelism when `batch_size < num_experts / top_k`. For Flash 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. +NxDI's TKG (token generation) path refuses Expert Parallelism when `batch_size < num_experts / top_k`. For Pro that is 384 / 8 = 48, so the smallest working BS on the FP8 path is 48. BS=1 latency demos are not 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 @@ -272,7 +277,7 @@ MiMo-V2.5-Pro can be served via [vllm-neuron](https://github.com/aws-neuron/vllm bash contrib/models/MiMo-V2.5-Pro/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_FLASH_SRC` is set, registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under the key `mimo_v2` **and** registers the `MiMoV2ForCausalLM` architecture into vLLM's `ModelRegistry`. No upstream vLLM or NxDI source is modified. +The patch (`perf_test/vllm-neuron-patch.patch`) touches `vllm_neuron/worker/neuronx_distributed_model_loader.py`. It adds a `_register_contrib_models()` hook that, when `NXDI_CONTRIB_MIMO_V2_FLASH_SRC` is set, registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under keys `mimov2flash` **and** `mimov2pro`, **and** overrides vLLM's built-in `MiMoV2FlashForCausalLM` / `MiMoV2ProForCausalLM` (GPU-only stubs) in `ModelRegistry` with the Neuron wrapper so ModelConfig validation accepts either architecture. No upstream vLLM or NxDI source is modified. The checkpoint's `config.json` must set `architectures` to `["MiMoV2ProForCausalLM"]` (or `MiMoV2FlashForCausalLM` for V2.5); the preprocess script takes care of this. ### Serving (FP8, recommended) @@ -289,7 +294,7 @@ python3 -m vllm.entrypoints.openai.api_server \ --model "$MIMO_V2_FLASH_PATH" \ --tensor-parallel-size 64 \ --max-model-len 1024 \ - --max-num-seqs 32 \ + --max-num-seqs 48 \ --no-enable-chunked-prefill \ --no-enable-prefix-caching \ --trust_remote_code \ @@ -313,9 +318,9 @@ python3 -m vllm.entrypoints.openai.api_server \ "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, + "batch_size": 48, "ctx_batch_size": 1, - "tkg_batch_size": 32, + "tkg_batch_size": 48, "max_context_length": 1024, "seq_len": 1024, "is_continuous_batching": true, @@ -330,39 +335,33 @@ python3 -m vllm.entrypoints.openai.api_server \ }' ``` -See `perf_test/bench_mimo_v2.sh` for the full benchmark recipe at BS=32 and BS=128. +See `perf_test/bench_mimo_v2.sh` for the full benchmark recipe at BS=48 and BS=128. ### vllm-neuron patch summary The patch is applied to vllm-neuron 0.5.0 and: -- Maps the `MiMoV2ForCausalLM` architecture to Flash's model loader (reusing the Qwen2-family loader path, which Flash'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. +- Patches `AutoConfig.from_pretrained` to default `trust_remote_code=True` so NxDI's `hf_adapter.load_config` can load the `MiMoV2Config` custom code that ships with the checkpoint. +- Registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under `mimov2flash` and `mimov2pro` so the NxDI loader resolves either model_type to the contrib Neuron wrapper. +- Overrides vLLM's built-in `MiMoV2FlashForCausalLM` and `MiMoV2ProForCausalLM` GPU stubs in `ModelRegistry`, since vLLM's ModelConfig validator rejects any architecture not in its registry and the Neuron path never instantiates vLLM's stub class anyway. ## 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) +> The throughput numbers below are from a working vLLM server run on 2026-04-29 under the recommended FP8 recipe. Output quality under this recipe is **not production-usable** (see Status); the numbers show that the serving infrastructure runs end-to-end, not that the model answers correctly. -| Batch Size | Throughput (tok/s) | -|------------|-------------------| -| 1 | 29.92 | -| 8 | 215.94 | -| 32 | 649.14 | +### vLLM Serving (trn2.48xlarge, FP8, BS=48, TP=64, moe_tp=1/moe_ep=64, CB + bucketing) -### vLLM Serving (trn2.48xlarge, BF16, BS=32, TP=64/EP=64, CB) +Input/output: 900/90 tokens (`vllm bench serve --dataset-name random`), `on_device_sampling_config={do_sample:true, temperature:0.6, top_k:20, top_p:0.95}`. -Input/output: 900/90 tokens (random dataset) +| Concurrency | Total tok/s | Output tok/s | TTFT median (ms) | TTFT P99 (ms) | TPOT median (ms) | +|-------------|-------------|--------------|------------------|---------------|------------------| +| 1 | 47 | 4.3 | 1,392 | 1,393 | 220 | +| 16 | 391 | 35.6 | 2,361 | 17,394 | 422 | +| 48 | 606 | 55 | 7,322 | 54,413 | 752 | -| 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 | +Per-stream ITL median holds at ~220 ms across all concurrency levels; TPOT/TTFT growth at higher concurrency comes from continuous-batching queue pressure, not per-step compute. -> **Compile time:** the first Flash 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 Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. ## Compatibility Matrix @@ -385,7 +384,7 @@ pytest contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py -v 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 use `batch_size >= 48` (NxDI EP>1 requirement, `384 / 8 = 48`). - 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. @@ -397,4 +396,4 @@ pytest contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py -v Henan Wan (whn09) -**Last Updated:** 2026-04-25 +**Last Updated:** 2026-04-29 diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh index 5d43c529..bdcbe11c 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh @@ -39,15 +39,15 @@ pip install s5cmd python3 -c "import vllm_neuron; print('vllm-neuron installed:', vllm_neuron.__file__)" echo "" -echo "[2/2] Downloading MiMo-V2.5-Pro BF16 weights..." +echo "[2/2] Downloading MiMo-V2.5-Pro Neuron-FP8 weights..." -MIMO_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-BF16}" +MIMO_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" if [ -d "$MIMO_PATH" ] && [ "$(ls "$MIMO_PATH"/*.safetensors 2>/dev/null | wc -l)" -gt 0 ]; then echo " MiMo weights already exist at $MIMO_PATH, skipping download" else - echo " Downloading BF16 weights from your S3 bucket (edit the URI if needed)..." + echo " Downloading Neuron-FP8 weights from your S3 bucket (edit the URI if needed)..." mkdir -p "$MIMO_PATH" - s5cmd cp "s3://datalab/xiaomi/models/MiMo-V2.5-Pro-BF16/**" "$MIMO_PATH/" + s5cmd cp "s3://datalab/xiaomi/models/MiMo-V2.5-Pro-Neuron-FP8/**" "$MIMO_PATH/" echo " Download complete: $(du -sh $MIMO_PATH | cut -f1)" fi diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh index 69c7d417..45c38be1 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh @@ -15,7 +15,7 @@ set -e # (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 Flash), so the smallest working batch size here is 32. +# (384 / 8 = 48 for V2.5-Pro), so the smallest working batch size here is 48. # 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`). @@ -132,18 +132,19 @@ 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 1: BS=48, TP=64 + moe_tp=1/moe_ep=64, CB + bucketing (smallest BS +# that satisfies NxDI's Expert-Parallel BS >= num_experts/top_k requirement: +# 384 / 8 = 48). ############################################################################### -CONFIG_NAME="bs32_tp64_moetp1_ep64" -echo "--- Config 1: BS=32, moe_tp=1/moe_ep=64, CB + bucketing ---" +CONFIG_NAME="bs48_tp64_moetp1_ep64" +echo "--- Config 1: BS=48, 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 \ + --max-num-seqs 48 \ --no-enable-chunked-prefill \ --no-enable-prefix-caching \ --port $PORT \ @@ -153,9 +154,9 @@ python3 -m vllm.entrypoints.openai.api_server \ '"$COMMON_MIMO_CONFIG"', "moe_tp_degree": 1, "moe_ep_degree": 64, - "batch_size": 32, + "batch_size": 48, "ctx_batch_size": 1, - "tkg_batch_size": 32, + "tkg_batch_size": 48, "max_context_length": 1024, "seq_len": 1024, "is_continuous_batching": true, @@ -173,7 +174,7 @@ 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" 48 192 stop_server ############################################################################### diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh index 45729cd2..81f978bb 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh @@ -13,8 +13,8 @@ # # Environment knobs: # PORT vLLM server port (default 8000) -# MIMO_V2_FLASH_PATH Path to the BF16 checkpoint (default -# /opt/dlami/nvme/models/MiMo-V2.5-Pro-BF16) +# MIMO_V2_FLASH_PATH Path to the Neuron-FP8 checkpoint (default +# /opt/dlami/nvme/models/MiMo-V2.5-Pro-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_2_9_nxd_inference/bin/activate -MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-BF16}" +MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" PORT="${PORT:-8000}" CONCURRENCY="${CONCURRENCY:-1}" NUM_PROMPTS="${NUM_PROMPTS:-16}" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh index 9cef80ab..684211f4 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh @@ -11,7 +11,7 @@ set -e -MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-BF16}" +MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" PORT="${PORT:-8000}" PROMPT="${PROMPT:-What is 1+1? Answer briefly.}" MAX_TOKENS="${MAX_TOKENS:-64}" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index be7f72e7..a0933f7f 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -124,17 +124,12 @@ def main(): moe_tp_degree=MOE_TP, context_encoding_buckets=[SEQ_LEN], router_config={"act_fn": "sigmoid", "dtype": "float32"}, - # V2.5-Pro experiment: use PyTorch-fallback blockwise matmul instead - # of the NKI _call_shard_hidden_kernel path. Rationale: Pro's MoE - # expert weight scales run ~3-7x smaller than V2.5 (scale_mean - # ~5e-5 vs 2.5e-4); the NKI blockwise kernel's lower accumulator - # precision compounds across 70 layers into decode-time gibberish. - # PyTorch fallback dequantizes each block to BF16 before matmul, - # trading throughput for accumulator precision. + # 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. Matches Flash + Kimi. blockwise_matmul_config={ - "use_torch_block_wise": True, - "use_shard_on_intermediate_dynamic_while": True, - "skip_dma_token": True, + "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 diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index c3388e89..9f678a09 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -99,9 +99,8 @@ def main(): context_encoding_buckets=[SEQ_LEN], router_config={"act_fn": "sigmoid", "dtype": "float32"}, blockwise_matmul_config={ - "use_torch_block_wise": True, - "use_shard_on_intermediate_dynamic_while": True, - "skip_dma_token": True, + "use_shard_on_block_dynamic_while": True, + "block_sharding_strategy": "PING_PONG", }, save_sharded_checkpoint=True, quantized=True, diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch b/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch index d8a85b89..a29814d4 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch @@ -61,20 +61,24 @@ index d2099eb..0c162e4 100644 + _sys.path.insert(0, mimo_src) + try: + from modeling_mimo_v2 import NeuronMiMoV2ForCausalLM -+ MODEL_TYPES.setdefault( -+ "mimov2flash", {"causal-lm": NeuronMiMoV2ForCausalLM} -+ ) ++ # Register under both Flash and Pro model_type keys so the same ++ # NxDI wrapper serves MiMo-V2-Flash, MiMo-V2.5-Pro, and any sibling ++ # that inherits the same config. ++ for _mt in ("mimov2flash", "mimov2pro"): ++ MODEL_TYPES.setdefault(_mt, {"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 -+ ) ++ # Override vLLM's GPU-only MiMoV2* stubs with our Neuron ++ # wrapper so ModelConfig validation accepts the architecture ++ # regardless of whether the checkpoint calls itself Flash or ++ # Pro. ++ for _arch in ("MiMoV2FlashForCausalLM", "MiMoV2ProForCausalLM"): ++ ModelRegistry.register_model(_arch, NeuronMiMoV2ForCausalLM) + except ImportError: + pass + except Exception as e: + _w.warn( -+ f"Failed to register MiMo-V2.5-Pro contrib model: {e}", ++ f"Failed to register MiMo-V2 contrib model: {e}", + category=UserWarning, + ) + From bbb1e1fd187ed6c0a855c1f2e7d053e8750f7bc4 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 10:42:43 +0800 Subject: [PATCH 11/47] [contrib] MiMo-V2.5-Pro: fix physical NC count (128, not 32) trn2.48xlarge has 16 Trainium2 chips x 8 cores = 128 physical NeuronCores. logical_nc_config=2 halves that to 64 logical cores, which matches tp_degree=64. Previous Prerequisites line said "32 NeuronCores" which is wrong. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 7078e480..dcf2c30e 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -59,7 +59,7 @@ Known NxDI limits that constrain recipe choice: ## Prerequisites -- **Instance**: trn2.48xlarge (32 NeuronCores, logical_nc_config=2 → 64 logical cores) +- **Instance**: trn2.48xlarge (128 physical 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**: ~3 TB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~962 GB, the Neuron-FP8 preprocessed output is ~1 TB, and `save_sharded_checkpoint=true` writes another ~300-1000 GB per compiled config (varies with recipe)). From d0eb413f5da79b5d3cadd78a43f977bb68432442 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 11:42:15 +0800 Subject: [PATCH 12/47] [contrib] MiMo-V2.5-Pro: split vLLM launcher into start/bench/sanity trio Mirror the V2.5 structure so Pro has: - start_vllm_server.sh (new): foreground launcher baking in the full override_neuron_config, persistent NEURON_COMPILED_ARTIFACTS path, and all env-var plumbing. Stays up for ad-hoc curl/sanity. - bench_mimo_v2.sh: rewritten as a one-shot composer (start_vllm_server in background + wait + sanity + run_bench_single at c=1/16/48). Replaces the old inline-launch-with-full-JSON version (~110 lines shorter). - run_bench_single.sh: default CONFIG_NAME/RESULTS_DIR brought in line with bench_mimo_v2.sh and the V2.5 port. README: - Add "Keeping a server up for ad-hoc testing" section and an Environment variables table (NXDI_CONTRIB_MIMO_V2_FLASH_SRC, NEURON_COMPILED_ARTIFACTS, BASE_COMPILE_WORK_DIR, etc.). - Replace the ~60-line inline vllm api_server invocation with pointers to start_vllm_server.sh / bench_mimo_v2.sh; the README no longer duplicates the config that lives in the scripts. - Fix "downloads Flash weights" text in the 0_setup.sh blurb (now downloads Pro Neuron-FP8 weights). - Bench results dir default moved to /opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro/ to align with V2.5. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 121 +++++----- .../MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh | 228 ++++-------------- .../perf_test/run_bench_single.sh | 9 +- .../perf_test/start_vllm_server.sh | 100 ++++++++ 4 files changed, 207 insertions(+), 251 deletions(-) create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index dcf2c30e..b0f371e4 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -119,9 +119,52 @@ bash contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh bash contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh ``` -The bench script runs two configurations (BS=48 and BS=128, both -`moe_tp_degree=1 / moe_ep_degree=64`) and logs results under -`/tmp/bench_results/mimo_v25_pro/`. +The bench script runs one configuration (BS=48, +`moe_tp_degree=1 / moe_ep_degree=64`) at three concurrency levels (1, 16, 48) +and logs results under `/opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro/`. + +### Keeping a server up for ad-hoc testing + +`bench_mimo_v2.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-Pro/perf_test/start_vllm_server.sh + +# Terminal 2: once "Application startup complete." prints, sanity-check: +bash contrib/models/MiMo-V2.5-Pro/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-Pro/perf_test/run_bench_single.sh +``` + +`bench_mimo_v2.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 +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.sh`. + +**Required (at least for manual `vllm api_server` launches):** + +| Variable | Purpose | +|---|---| +| `NXDI_CONTRIB_MIMO_V2_FLASH_SRC` | Path to `contrib/models/MiMo-V2.5-Pro/src/`. `vllm-neuron`'s registration hook reads it to plug `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` table. The `_FLASH_` suffix is kept for backward compatibility with the shared registration hook that also serves V2-Flash and V2.5. | +| `MIMO_V2_FLASH_PATH` | Preprocessed Neuron-FP8 checkpoint dir (the `--save_path` output from preprocess). Same naming rationale as above. | + +**Optional (recommended):** + +| Variable | Default | Purpose | +|---|---|---| +| `NEURON_COMPILED_ARTIFACTS` | `/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_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 Pro's 384-expert MoE is ~60 min TKG + ~15 min CTE + ~30 min shard, well past vLLM's default. | For a quick `curl` sanity check while the server is up: @@ -133,11 +176,9 @@ curl -s http://localhost:8000/v1/chat/completions \ "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). +Output quality is currently prompt-dependent under the FP8 recipe (see +Status). A successful sanity check confirms the serving path works; it +does not yet confirm that all prompts produce coherent text. ## Checkpoint Preparation @@ -272,8 +313,8 @@ MiMo-V2.5-Pro can be served via [vllm-neuron](https://github.com/aws-neuron/vllm ```bash # The setup script clones vllm-project/vllm-neuron at release-0.5.0, applies -# the contrib registration patch, installs it editable, and downloads Flash -# weights (BF16 by default; set MIMO_V2_FLASH_PATH to override). +# the contrib registration patch, installs it editable, and downloads +# Pro Neuron-FP8 weights from S3 (set MIMO_V2_FLASH_PATH to override). bash contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh ``` @@ -281,61 +322,17 @@ The patch (`perf_test/vllm-neuron-patch.patch`) touches `vllm_neuron/worker/neur ### Serving (FP8, recommended) +Use `perf_test/start_vllm_server.sh` for a foreground launch (stays up until Ctrl-C), or `perf_test/bench_mimo_v2.sh` for the one-shot launch → sanity → bench → teardown flow. Both scripts bake in the full `override_neuron_config` (TP=64, moe_tp=1, moe_ep=64, BS=48, CB + bucketing, blockwise FP8 MoE with `PING_PONG`, on-device sampling), the required env vars, and the persistent compile-artifact path. See "Keeping a server up for ad-hoc testing" above for the three-terminal workflow. + ```bash -export NXDI_CONTRIB_MIMO_V2_FLASH_SRC=/path/to/neuronx-distributed-inference/contrib/models/MiMo-V2.5-Pro/src -export MIMO_V2_FLASH_PATH=/path/to/MiMo-V2.5-Pro-Neuron-FP8 -# First-time compile of Flash's 256-expert MoE takes 30-60 minutes. -export VLLM_ENGINE_READY_TIMEOUT_S=7200 -# Optional: isolate compile cache per config so parallel Flash/Pro/etc. compiles -# don't race on the default /var/tmp/neuron-compile-cache lock files. -export NEURON_COMPILED_ARTIFACTS=/path/to/compiled/mimo_v25_pro_bs32_moetp1_ep64_fp8 - -python3 -m vllm.entrypoints.openai.api_server \ - --model "$MIMO_V2_FLASH_PATH" \ - --tensor-parallel-size 64 \ - --max-model-len 1024 \ - --max-num-seqs 48 \ - --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-Pro-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": 48, - "ctx_batch_size": 1, - "tkg_batch_size": 48, - "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 - } - } - }' +# One-shot launch + bench + teardown (~2 h on cold cache, ~5 min on warm cache). +bash contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh + +# Or keep the server up for interactive work: +bash contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh ``` -See `perf_test/bench_mimo_v2.sh` for the full benchmark recipe at BS=48 and BS=128. +See "Environment variables" above for all the knobs (`NEURON_COMPILED_ARTIFACTS`, `BASE_COMPILE_WORK_DIR`, etc.) and their defaults. ### vllm-neuron patch summary diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh index 45c38be1..23c3f59d 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh @@ -1,77 +1,42 @@ #!/bin/bash set -e -# MiMo-V2.5-Pro FP8 vLLM benchmark on Trn2. +# MiMo-V2.5-Pro FP8 vLLM benchmark on Trn2. One-shot wrapper: +# launch server -> sanity check -> bench at c=1,16,48 -> stop server. # -# Requires a Neuron-FP8 preprocessed checkpoint (see -# `src/conversion_script/preprocess_mimo_v2_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 -# (384 / 8 = 48 for V2.5-Pro), so the smallest working batch size here is 48. -# 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_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" -# The NxDI contrib MiMo-V2.5-Pro 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_FLASH_SRC:=$(cd "$(dirname "$0")/.." && pwd)/src}" -export NXDI_CONTRIB_MIMO_V2_FLASH_SRC - -# First-time Flash 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=48, continuous batching. +# BS=48 is the smallest working batch size on the FP8 path (NxDI's TKG +# path refuses Expert Parallelism with BS < num_experts/top_k = 384/8 = 48). +# BS=1 single-stream latency demos are not currently supported on Pro FP8. + +SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)" +PORT="${PORT:-8000}" +RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro}" +CONFIG_NAME="bs48_tp64_moetp1_ep64" -PORT=8000 -RESULTS_DIR="/tmp/bench_results/mimo_v25_pro" mkdir -p "$RESULTS_DIR" -# Common neuron config shared across all MiMo-V2.5-Pro 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 384-expert +# MoE model takes ~90 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 @@ -81,149 +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-Pro FP8 Performance Benchmark" echo "==========================================" -echo "Model: $MODEL_PATH" +echo "Port: $PORT" echo "Results: $RESULTS_DIR" echo "" -############################################################################### -# Config 1: BS=48, TP=64 + moe_tp=1/moe_ep=64, CB + bucketing (smallest BS -# that satisfies NxDI's Expert-Parallel BS >= num_experts/top_k requirement: -# 384 / 8 = 48). -############################################################################### -CONFIG_NAME="bs48_tp64_moetp1_ep64" -echo "--- Config 1: BS=48, 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 48 \ - --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": 48, - "ctx_batch_size": 1, - "tkg_batch_size": 48, - "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" 48 192 -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 - } - } - }' & +# One-shot sanity check (curl the chat endpoint). +PORT="$PORT" bash "$SCRIPT_DIR/sanity_check.sh" || true -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 +# 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=48 NUM_PROMPTS=192 bash "$SCRIPT_DIR/run_bench_single.sh" echo "==========================================" -echo "MiMo-V2.5-Pro FP8 benchmarks complete!" +echo "MiMo-V2.5-Pro FP8 benchmark complete!" echo "Results saved to: $RESULTS_DIR" echo "==========================================" ls -la "$RESULTS_DIR" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh index 81f978bb..2e4e4a3e 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh @@ -20,8 +20,9 @@ # 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_v25_pro) +# CONFIG_NAME Used in the output filename (default bs48_tp64_moetp1_ep64) +# RESULTS_DIR Where to dump per-run log +# (default /opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro) set -e @@ -34,8 +35,8 @@ 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_v25_pro}" +CONFIG_NAME="${CONFIG_NAME:-bs48_tp64_moetp1_ep64}" +RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro}" mkdir -p "$RESULTS_DIR" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh new file mode 100644 index 00000000..6848b604 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh @@ -0,0 +1,100 @@ +#!/bin/bash +# Start the MiMo-V2.5-Pro 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.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=48, continuous batching + bucketing. +# moe_tp=1/moe_ep=64 keeps each expert's weights and blockwise FP8 scales +# intact on a single rank (6 experts/rank for Pro's 384 experts), 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 +# (384 / 8 = 48), so BS=48 is the smallest working batch size on the FP8 +# path. BS=1 single-stream latency is not currently supported on Pro FP8. + +set -e + +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate + +MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" +PORT="${PORT:-8000}" + +# Contrib package src. vllm-neuron's registration hook reads this env var +# to plug NeuronMiMoV2ForCausalLM into NxDI's MODEL_TYPES table. +: "${NXDI_CONTRIB_MIMO_V2_FLASH_SRC:=$(cd "$(dirname "$0")/.." && pwd)/src}" +export NXDI_CONTRIB_MIMO_V2_FLASH_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_pro_bs48_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 Pro's 384-expert MoE takes ~60 min TKG + ~15 min +# CTE + ~30 min shard; plan for 2 h. +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-7200}" + +echo "==========================================" +echo "Starting MiMo-V2.5-Pro 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_FLASH_SRC: $NXDI_CONTRIB_MIMO_V2_FLASH_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 48 \ + --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": 48, + "ctx_batch_size": 1, + "tkg_batch_size": 48, + "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 2a181c1e466182b5a8a7c69e3e38db4ae649180c Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 11:53:25 +0800 Subject: [PATCH 13/47] [contrib] MiMo-V2.5-Pro: swap sanity_check default prompt to self-intro "What is 1+1?" drifts to unrelated text under the current FP8 recipe. "Introduce yourself in one sentence." is a high-signal self-identifying prompt that still answers coherently (e.g. "I'm MiMo, developed by Xiaomi LLM Core Team.") and gives a sensible first-run demo. Also drop the explicit `temperature: 0.0` from the request body: vllm-neuron honours the compile-time on_device_sampling_config, not the request-side temperature, so sanity output is always sampled at T=0.6. Note this in a comment. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh index 684211f4..ed2f6b76 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh @@ -13,7 +13,10 @@ set -e MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" PORT="${PORT:-8000}" -PROMPT="${PROMPT:-What is 1+1? Answer briefly.}" +# "Introduce yourself" is a high-signal self-identification prompt that the +# FP8 path answers coherently even under current MoE drift (see README +# Status). Swap PROMPT=... if you want to probe other prompts. +PROMPT="${PROMPT:-Hello! Please introduce yourself in one sentence.}" MAX_TOKENS="${MAX_TOKENS:-64}" echo "Sanity check: POST /v1/chat/completions on port $PORT" @@ -29,6 +32,10 @@ if ! curl -sf "http://localhost:$PORT/health" > /dev/null; then exit 1 fi +# NOTE: request-side `temperature` is ignored by vllm-neuron on this model: +# on-device sampling_config (set at compile time in start_vllm_server.sh as +# do_sample=true, T=0.6, top_k=20, top_p=0.95) is baked into the NEFF and +# request params don't override it. Output will be stochastic. RESPONSE=$(curl -s "http://localhost:$PORT/v1/chat/completions" \ -H 'Content-Type: application/json' \ -d "$(cat < Date: Wed, 29 Apr 2026 13:58:15 +0800 Subject: [PATCH 14/47] [contrib] MiMo-V2.5-Pro: BF16-attn recipe restores coherent output Root cause of the FP8 drift is narrowed to the attention path, not the MoE experts. Pro's q/k/v weights have abs_mean ~0.00124, 4x smaller than V2.5 (256 experts), and the NKI blockwise FP8 accumulator loses enough precision at this magnitude to drift the logits across 70 layers. Dequantizing q/k/v to BF16 while keeping MoE experts FP8 restores coherent output on smoke_generate, e.g.: Okay, the user is asking for a simple self-introduction in one sentence, with no deeper or hidden needs apparent. As MiMo, based on Xiaomi's self-developed large model, I need to respond in a friendly, positive, and helpful way that aligns with providing assistance ... Changes: - Add src/conversion_script/repatch_qkv_bf16.py (promoted from /opt/dlami/nvme/scripts/), now argparse-driven. Reads HF fused qkv_proj + weight_scale_inv, dequants per kv-head group, writes BF16 q/k/v into the preprocessed Neuron-FP8 checkpoint in place, drops scale entries from the safetensors index. ~22 min runtime. - smoke_compile_mimo_v2.py / smoke_generate_mimo_v2.py: add q_proj/k_proj/v_proj to modules_to_not_convert, drop seq_len from 1024 to 256 (BF16 q/k/v adds ~2 GB per rank; seq_len=1024 OOMed on load last time), switch default COMPILED_PATH to the new BF16-attn directory name to avoid clobbering earlier artifacts. - README: rewrite Status to separate the all-FP8 result (drifted) from the BF16-attn result (coherent); document the required repatch step, the HBM / seq_len trade-off, and a warning that listing q/k/v in modules_to_not_convert without running repatch first produces nonsense (NxDI casts fp8 bytes to bf16 without applying the scale). Update Quick Start to include the repatch step. Flag that vLLM scripts still use the all-FP8 recipe and the bench numbers haven't been re-measured on BF16-attn. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 91 +++++++-- .../perf_test/smoke_compile_mimo_v2.py | 21 +- .../perf_test/smoke_generate_mimo_v2.py | 7 +- .../src/conversion_script/repatch_qkv_bf16.py | 190 ++++++++++++++++++ 4 files changed, 285 insertions(+), 24 deletions(-) create mode 100644 contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index b0f371e4..d7f0c223 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -37,25 +37,47 @@ Key features: ## Status (work-in-progress) -**This port compiles cleanly and serves via vLLM on Trn2, but output quality is not production-ready under the default FP8 recipe.** Sanity checks on 2026-04-29 under `tp_degree=64, moe_tp_degree=1, moe_ep_degree=64, batch_size=48, seq_len=1024`: - -- Prompt-dependent drift. Self-intro style prompts ("introduce yourself in one sentence") return coherent Chinese output; most other prompts collapse to repetition or unrelated text within a few tokens (e.g. `"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`). -- Same failure pattern was observed on the MiMo-V2-Pro port under the same recipe (`"0.0.0.0:8080"`, etc.). Root cause appears to be **Neuron's NKI blockwise FP8 compute kernel handling Pro's tight expert-weight distribution**: Pro's per-expert weight `abs_mean ≈ 0.00124` and blockwise `scale_mean ≈ 2.3e-5` are 4-7× smaller than V2.5 (256 experts) on the same recipe, and the NKI accumulator loses precision compounded across 70 layers. V2.5-Pro's MoE expert weights are byte-identical to V2-Pro (verified layer 1 expert 0 dequant stats match to 6 decimal places on 2026-04-28), so all V2-Pro workarounds remain required (router bias mean-subtract, qkv interleaved split, `_apply_2d_per_channel_fix`, `_apply_blockwise_scale_stride_fix`). -- GPU stacks (sglang on H100/H200) run the exact same OCP FP8 checkpoint correctly, because they dequantize FP8→BF16 before the matmul. Neuron NKI does FP8 compute directly and drifts on subnormal-leaning tensors. +**This port compiles cleanly and serves via vLLM on Trn2. Under the default all-FP8 recipe the output drifts on most prompts; a BF16-attn recipe (keep MoE FP8, dequant q/k/v_proj to BF16, compile at seq_len=256) restores coherent output and isolates the root cause.** Findings as of 2026-04-29: + +### All-FP8 recipe (`tp_degree=64, moe_tp=1, moe_ep=64, BS=48, seq_len=1024`) + +- **Prompt-dependent drift.** Self-intro prompts sometimes return coherent answers (sampling gets lucky on a strong self-identifying logit); most other prompts collapse to repetition or unrelated text within a few tokens (e.g. `"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`, chat continuations that drift into RLHF-style self-reflection with Chinese/Thai). +- **Not a sampling artifact.** `temperature` in the request is ignored — vllm-neuron's on-device sampling config (`do_sample=true, T=0.6, top_k=20, top_p=0.95`) is baked into the NEFF at compile time. Output is always stochastic, but the underlying logits are already drifted. +- Same failure pattern was observed on MiMo-V2-Pro under the same recipe (`"0.0.0.0:8080"`, etc.). V2.5-Pro's MoE expert weights are byte-identical to V2-Pro (verified layer 1 expert 0 dequant stats match to 6 decimal places on 2026-04-28). + +### BF16-attn recipe (`q_proj/k_proj/v_proj` dequant to BF16, MoE kept FP8, `seq_len=256`) + +- **Output is coherent.** On the same self-intro prompt, smoke_generate with a minimal chat template produces a well-formed reasoning trace that correctly identifies the model: + ``` + Okay, the user is asking for a simple self-introduction in one sentence, + with no deeper or hidden needs apparent. As MiMo, based on Xiaomi's self-developed + large model, I need to respond in a friendly, positive, and helpful way that aligns + with providing assistance ... + ``` +- **This narrows the root cause to attention-path FP8 precision**, not the MoE experts. Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2.5 (256 experts). The NKI blockwise FP8 accumulator on attention q/k/v at this magnitude loses enough precision to drift the logits across 70 layers; dequantizing q/k/v to BF16 before the matmul restores correct output. MoE experts (scales `≈ 2.3e-5`, also small) can stay FP8 under this recipe. +- **Cost: HBM headroom and seq_len.** BF16 q/k/v adds ~2 GB per rank. At `seq_len=1024` this OOMs on load (previous attempt failed allocating 41 MB for rdh/alltoall rings). `seq_len=256` frees enough full-attention softmax scratch to fit; longer context needs a different HBM plan (cross-instance TP/PP, or larger instance). +- GPU stacks (sglang on H100/H200) run the exact same OCP FP8 checkpoint correctly because they always dequantize FP8 → BF16 before the matmul. The issue is specific to Neuron's direct-FP8 compute path on subnormal-leaning tensors. - Reference: Jim Burtoft observed similar prompt-dependent FP8 degradation on Flash and his Kimi PR #131 names "blockwise kernel padding produces depressed logits with EP=2 on SDK 2.29; SDK 2.28 recommended". -Recipes that were tried and did not resolve the drift (all on 2026-04-28/29): -- **q/k/v_proj dequant to BF16** (compile-time `modules_to_not_convert` includes `q_proj, k_proj, v_proj`): compiled but HBM OOM on load — BF16 attention weights pushed per-rank tensors to 20.93/24 GB, leaving no room for collective DMA rings. -- **`use_torch_block_wise=True`** (PyTorch-fallback blockwise matmul for higher accumulator precision): compile+shard succeeded after ~2 h, but `model.load()` crashed with `status=4 Allocation Failure` — the fallback path raises HBM demand. +### Preprocess step required for BF16 attn + +`src/conversion_script/repatch_qkv_bf16.py` reads the original HF fused `qkv_proj` weight + scale, dequants per-group (`num_q_heads_per_kv_group` × `head_dim` rows per group) to BF16, and overwrites `q_proj/k_proj/v_proj` in the preprocessed Neuron-FP8 checkpoint in place. Runs in ~22 min. Simply adding `q_proj/k_proj/v_proj` to `modules_to_not_convert` without this preprocess is **not** equivalent — NxDI casts the raw fp8_e4m3fn bytes to bfloat16 without applying the scale, which produces nonsense weights. + +### Recipes that were tried and did not resolve the drift (all on 2026-04-28/29) + +- **`use_torch_block_wise=True`** (PyTorch-fallback blockwise matmul for higher accumulator precision): compile+shard succeeded after ~2 h, but `model.load()` crashed with `status=4 Allocation Failure` — the fallback path raises HBM demand even with MoE-only scope. -Next experiments queued (none verified yet): -- BF16 weights across two Trn2 instances with cross-node TP/PP (single-instance HBM cannot hold BF16 Pro). -- SDK 2.28 venv test once installed, per Kimi PR #131. -- Selective BF16 only on `gate_up_proj` (smallest MoE scales) while keeping `down_proj` FP8, if HBM allows. +### Next experiments queued -Known NxDI limits that constrain recipe choice: -- `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else NotImplementedError). With `num_experts=384, top_k=8` this forces `BS >= 48`. -- `n_routed_experts=384 = 2^7 * 3` → `384 / ep_degree` is never a power of 2 (6, 12, 24, 48, 96, 192, 384). Kimi PR #131 says NKI `_bwmm_shard_on_block_nki_call` on SDK 2.29 has "depressed logits with EP=2" and recommends SDK 2.28. SDK 2.28 venv is not currently installed on the target DLAMI. +- **BF16-attn + larger `seq_len`.** `seq_len=256` is tight; Pro's chat template with the default system prompt is already 260 tokens. Either shrink the system prompt, or try BF16-attn at `seq_len=512` on a less-full HBM plan (e.g. MoE-EP with fewer experts per rank, or drop batch size below 48 via different recipe trade-offs). +- **Cross-instance BF16** via pipeline/tensor parallelism on 2× Trn2 (single-instance HBM cannot hold full BF16 Pro). +- **Selective BF16 only on MoE `gate_up_proj`** (smallest expert scales) while keeping `down_proj` FP8 — another axis to probe if attn drift returns at longer contexts. +- **SDK 2.28 venv** test once installed, per Kimi PR #131. + +### Known NxDI limits that constrain recipe choice + +- `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else `NotImplementedError`). With `num_experts=384, top_k=8` this forces `BS >= 48`. +- `n_routed_experts=384 = 2^7 × 3` → `384 / ep_degree` is never a power of 2 (6, 12, 24, 48, 96, 192, 384). Kimi PR #131 says NKI `_bwmm_shard_on_block_nki_call` on SDK 2.29 has "depressed logits with EP=2" and recommends SDK 2.28. ## Prerequisites @@ -106,8 +128,14 @@ python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8 --save_path /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8 \ --tp_degree 64 +# 3b. Dequant attention q/k/v to BF16 in place (~22 min, required for +# coherent output under the current FP8 recipe; see Status). +python contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py \ + --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro \ + --neuron_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8 + # 4. (Optional) sanity-check the Neuron-FP8 checkpoint without vLLM -# ~45 min first compile; subsequent runs ~30s to load the pre-sharded NEFF. +# ~90 min first compile; subsequent runs ~60s to load the pre-sharded NEFF. source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate python contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py # compile python contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py # 20-token generate @@ -197,6 +225,20 @@ python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8 Peak RAM during preprocessing is ~24 GB; total runtime ~20 minutes on a trn2.48xlarge instance. +### Required follow-up: dequant q/k/v to BF16 (`repatch_qkv_bf16.py`) + +Under the default all-FP8 recipe, Pro's attention weights drift enough to produce gibberish output (see Status). The fix is to keep MoE experts FP8 but dequant q/k/v to BF16 before compile. `src/conversion_script/repatch_qkv_bf16.py` reads the HF fused `qkv_proj` + `weight_scale_inv`, dequants per-group (each group is one kv-head: `hpg` Q rows, `1×head_dim` K rows, `1×v_head_dim` V rows), and overwrites the q_proj/k_proj/v_proj entries in the preprocessed Neuron-FP8 checkpoint in place. The scale entries for q/k/v are dropped from the safetensors index. + +```bash +python contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py \ + --hf_model_path /path/to/MiMo-V2.5-Pro \ + --neuron_model_path /path/to/MiMo-V2.5-Pro-Neuron-FP8 +``` + +Runtime is ~22 minutes on a trn2.48xlarge, peak RAM ~20 GB. After running this, the compile-time `modules_to_not_convert` list in `smoke_compile_mimo_v2.py` / `start_vllm_server.sh` must include `q_proj`, `k_proj`, `v_proj` so NxDI keeps them as BF16 and routes them through a plain `ColumnParallelLinear` rather than the FP8 `QuantizedColumnParallel` path. + +**Do not skip this step.** Adding q_proj/k_proj/v_proj to `modules_to_not_convert` without running repatch first leaves the fp8 bytes in the checkpoint; NxDI then casts them to bf16 **without applying the blockwise scale**, producing nonsense weights and broken output. + ### 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. @@ -220,9 +262,13 @@ from modeling_mimo_v2 import NeuronMiMoV2ForCausalLM, MiMoV2InferenceConfig model_path = "/path/to/MiMo-V2.5-Pro-Neuron-FP8/" compiled_path = "/path/to/compiled/" -# Recommended FP8 recipe: +# Recommended recipe: BF16 attn + FP8 MoE. # moe_tp_degree = 1, moe_ep_degree = 64 -# See "FP8 Configuration Notes" below for why other moe_tp/ep ratios collapse. +# q_proj/k_proj/v_proj in modules_to_not_convert (BF16; requires +# repatch_qkv_bf16.py to have been run on the checkpoint first) +# seq_len = 256 (HBM-tight with BF16 attn; see Status) +# 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 @@ -232,7 +278,7 @@ neuron_config = MoENeuronConfig( max_batch_size=48, ctx_batch_size=1, tkg_batch_size=48, - seq_len=1024, + seq_len=256, # HBM is tight with BF16 attn; seq_len=1024 OOMs n_active_tokens=128, torch_dtype=torch.bfloat16, logical_nc_config=2, @@ -253,6 +299,7 @@ neuron_config = MoENeuronConfig( quantization_block_size=[128, 128], modules_to_not_convert=[ "embed_tokens", "lm_head", "norm", "router", "o_proj", + "q_proj", "k_proj", "v_proj", # BF16 attn — requires repatch ], on_device_sampling_config=OnDeviceSamplingConfig( do_sample=True, temperature=0.6, top_k=20, top_p=0.95, @@ -334,6 +381,8 @@ bash contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh See "Environment variables" above for all the knobs (`NEURON_COMPILED_ARTIFACTS`, `BASE_COMPILE_WORK_DIR`, etc.) and their defaults. +> **Note on output quality:** the shipped vLLM scripts use the **all-FP8 recipe** (`seq_len=1024`, attn+MoE both FP8). This currently produces prompt-dependent drift — see Status. The **BF16-attn recipe** that restores coherent output has so far only been validated end-to-end via `smoke_generate_mimo_v2.py` (direct NxDI, `seq_len=256`); porting it into `start_vllm_server.sh` requires also adding `q_proj/k_proj/v_proj` to the `modules_to_not_convert` list and dropping `seq_len` / `max_model_len` to 256, and the perf numbers below have not been re-measured in that configuration. + ### vllm-neuron patch summary The patch is applied to vllm-neuron 0.5.0 and: @@ -346,7 +395,7 @@ The patch is applied to vllm-neuron 0.5.0 and: > The throughput numbers below are from a working vLLM server run on 2026-04-29 under the recommended FP8 recipe. Output quality under this recipe is **not production-usable** (see Status); the numbers show that the serving infrastructure runs end-to-end, not that the model answers correctly. -### vLLM Serving (trn2.48xlarge, FP8, BS=48, TP=64, moe_tp=1/moe_ep=64, CB + bucketing) +### vLLM Serving (trn2.48xlarge, all-FP8 recipe, BS=48, TP=64, moe_tp=1/moe_ep=64, CB + bucketing, `seq_len=1024`) Input/output: 900/90 tokens (`vllm bench serve --dataset-name random`), `on_device_sampling_config={do_sample:true, temperature:0.6, top_k:20, top_p:0.95}`. @@ -358,6 +407,8 @@ Input/output: 900/90 tokens (`vllm bench serve --dataset-name random`), `on_devi Per-stream ITL median holds at ~220 ms across all concurrency levels; TPOT/TTFT growth at higher concurrency comes from continuous-batching queue pressure, not per-step compute. +> **Numbers are from the all-FP8 recipe, which produces drifted output (see Status).** The BF16-attn recipe that restores coherent output has not yet been re-benchmarked on vLLM; throughput should be comparable (only q/k/v go BF16, MoE stays FP8) but at `seq_len=256` instead of 1024, so TTFT/latency characteristics will differ. + > **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. ## Compatibility Matrix diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index a0933f7f..6475eea2 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -40,11 +40,15 @@ ) COMPILED_PATH = os.environ.get( "MIMO_V25_PRO_COMPILED_PATH", - "/opt/dlami/nvme/compiled/mimo_v25_pro_moetp1_ep64_bs48/", + "/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq256/", ) TP_DEGREE = int(os.environ.get("TP_DEGREE", "64")) -SEQ_LEN = int(os.environ.get("SEQ_LEN", "1024")) +# Drop seq_len to 256 to free ~200 MB of full-attention softmax scratch per +# rank. The previous BF16-attn attempt at seq_len=1024 OOM'd by 40 MB on load +# (failed to allocate 41943040 bytes for rdh/alltoall); seq_len=256 reclaims +# enough HBM to fit the extra BF16 q/k/v weights. +SEQ_LEN = int(os.environ.get("SEQ_LEN", "256")) # BS=48 is the minimum that avoids forward_selective_loading on decode: # `BS * top_k / num_experts >= 1.0` → BS >= 384/8 = 48. At BS=1 the TKG # path raises `NotImplementedError: Selective Loading with Expert parallelism`. @@ -142,12 +146,25 @@ def main(): quantization_type="blockwise_symmetric", quantization_block_axis=[1, 2], quantization_block_size=[128, 128], + # BF16 attention: keep q/k/v_proj in BF16 (not FP8). Pro's q/k/v + # abs_mean ~0.00124 is 4x smaller than V2.5 and the NKI blockwise + # FP8 accumulator drifts across 70 layers, producing gibberish + # output under the all-FP8 recipe. Dequantizing q/k/v to BF16 + # restores coherent output while keeping MoE experts FP8. + # Prerequisite: run src/conversion_script/repatch_qkv_bf16.py on + # the preprocessed Neuron-FP8 checkpoint first; simply listing + # q/k/v here without the repatch step leaves the fp8 bytes in the + # checkpoint and NxDI silently casts them to bf16 without applying + # the scale, which produces nonsense weights. modules_to_not_convert=[ "embed_tokens", "lm_head", "norm", "router", "o_proj", + "q_proj", + "k_proj", + "v_proj", ], ) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index 9f678a09..4c1a3bd2 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -27,13 +27,13 @@ ) COMPILED_PATH = os.environ.get( "MIMO_V25_PRO_COMPILED_PATH", - "/opt/dlami/nvme/compiled/mimo_v25_pro_moetp1_ep64_bs48/", + "/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq256/", ) # Must match smoke_compile_mimo_v2.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")) +SEQ_LEN = int(os.environ.get("SEQ_LEN", "256")) BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "48")) # must match smoke_compile CTX_BATCH_SIZE = int(os.environ.get("CTX_BATCH_SIZE", "1")) MOE_TP = int(os.environ.get("MOE_TP", "1")) @@ -115,6 +115,9 @@ def main(): "norm", "router", "o_proj", + "q_proj", + "k_proj", + "v_proj", ], ) diff --git a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py new file mode 100644 index 00000000..6ac7fb3a --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py @@ -0,0 +1,190 @@ +#!/usr/bin/env python3 +"""In-place patch: replace q/k/v FP8+scale with BF16 in a MiMo-V2.5-Pro +preprocessed Neuron checkpoint, for every decoder layer. Leaves MoE experts, +norms, embed, lm_head, o_proj untouched. + +Rationale: Pro's attention q/k/v weights have abs_mean ~0.001-0.005, roughly +4x smaller than V2.5. The NKI blockwise FP8 accumulator on the attention +path loses enough precision at this magnitude to drift the logits across +70 layers; dequantizing q/k/v to BF16 before the matmul restores coherent +output. MoE experts (also small-scale) can stay FP8. + +Note: simply adding q_proj/k_proj/v_proj to NxDI's `modules_to_not_convert` +at compile time is NOT equivalent — NxDI casts the raw fp8_e4m3fn bytes to +bfloat16 without applying the blockwise scale, which produces nonsense +weights. This script reads the HF fused qkv + scale, dequants per-group, +and writes BF16 weights back into the preprocessed Neuron checkpoint in +place. + +Run under /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16 (or any venv +with torch and safetensors). Takes ~22 min on a trn2.48xlarge for 70 +layers. +""" +import argparse +import glob +import json +import math +import os +import sys +import time + +import torch +from safetensors import safe_open +from safetensors.torch import save_file + + +def main(): + parser = argparse.ArgumentParser( + description="In-place dequant q/k/v from FP8+scale to BF16 in the " + "preprocessed Neuron checkpoint.", + ) + parser.add_argument( + "--hf_model_path", + required=True, + help="Path to the original HuggingFace MiMo-V2.5-Pro checkpoint " + "(fused qkv_proj + qkv_proj.weight_scale_inv).", + ) + parser.add_argument( + "--neuron_model_path", + required=True, + help="Path to the preprocessed Neuron-FP8 checkpoint. q/k/v entries " + "in its model_layer{N}.safetensors shards will be overwritten in " + "place with BF16 values; the scale entries are dropped from the " + "index.", + ) + args = parser.parse_args() + + # Import split_qkv_fused from the neighbouring preprocess script so the + # group layout math (hpg, qg_rows, kg_rows, vg_rows) stays in one place. + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from preprocess_mimo_v2_fp8 import LazyWeightMap # noqa: F401 + + hf_src = args.hf_model_path + neuron = args.neuron_model_path + cfg = json.load(open(os.path.join(hf_src, "config.json"))) + hp = cfg["hybrid_layer_pattern"] + num_hidden_layers = cfg["num_hidden_layers"] + + with open(os.path.join(hf_src, "model.safetensors.index.json")) as f: + hf_wm = json.load(f)["weight_map"] + lazy = LazyWeightMap(hf_src, hf_wm) + + print( + f"Patching q/k/v -> BF16 in {neuron}/model_layer{{0..{num_hidden_layers - 1}}}.safetensors", + flush=True, + ) + t0 = time.time() + + for li in range(num_hidden_layers): + layer_file = os.path.join(neuron, f"model_layer{li}.safetensors") + if not os.path.exists(layer_file): + print(f" layer {li}: FILE MISSING, skip", flush=True) + continue + with safe_open(layer_file, framework="pt") as fp: + layer_sd = {k: fp.get_tensor(k) for k in fp.keys()} + + is_swa = hp[li] == 1 + num_q = cfg["swa_num_attention_heads" if is_swa else "num_attention_heads"] + num_kv = cfg["swa_num_key_value_heads" if is_swa else "num_key_value_heads"] + hd = cfg["swa_head_dim" if is_swa else "head_dim"] + vhd = cfg["swa_v_head_dim" if is_swa else "v_head_dim"] + prefix = f"model.layers.{li}.self_attn" + qkv_w = lazy.get(f"{prefix}.qkv_proj.weight") + qkv_s = lazy.get(f"{prefix}.qkv_proj.weight_scale_inv") + + BLOCK = 128 + hpg = num_q // num_kv + qg_rows = hpg * hd + kg_rows = 1 * hd + vg_rows = 1 * vhd + R = qg_rows + kg_rows + vg_rows + in_features = qkv_w.shape[1] + q_blk = qg_rows // BLOCK + k_blk = (kg_rows + BLOCK - 1) // BLOCK + v_blk = (vg_rows + BLOCK - 1) // BLOCK + per = q_blk + k_blk + v_blk + padded = per * BLOCK + + w = qkv_w.to(torch.float32).view(num_kv, R, in_features) + w_padded = torch.zeros(num_kv, padded, in_features, dtype=torch.float32) + w_padded[:, :R, :] = w + s = qkv_s.to(torch.float32).view(num_kv, per, (in_features + BLOCK - 1) // BLOCK) + s_exp = s.repeat_interleave(BLOCK, dim=1).repeat_interleave(BLOCK, dim=2) + s_exp = s_exp[:, :padded, :in_features] + deq_padded = w_padded * s_exp + deq = deq_padded[:, :R, :] + + q_bf16 = ( + deq[:, :qg_rows, :] + .reshape(num_kv * qg_rows, in_features) + .contiguous() + .to(torch.bfloat16) + ) + k_bf16 = ( + deq[:, qg_rows : qg_rows + kg_rows, :] + .reshape(num_kv * kg_rows, in_features) + .contiguous() + .to(torch.bfloat16) + ) + v_bf16 = ( + deq[:, qg_rows + kg_rows :, :] + .reshape(num_kv * vg_rows, in_features) + .contiguous() + .to(torch.bfloat16) + ) + + for key in list(layer_sd): + if any( + key.endswith(f".{p}.weight") or key.endswith(f".{p}.scale") + for p in ("q_proj", "k_proj", "v_proj") + ): + del layer_sd[key] + + out_prefix = f"layers.{li}.self_attn" + layer_sd[f"{out_prefix}.q_proj.weight"] = q_bf16 + layer_sd[f"{out_prefix}.k_proj.weight"] = k_bf16 + layer_sd[f"{out_prefix}.v_proj.weight"] = v_bf16 + + save_file(layer_sd, layer_file) + dt = time.time() - t0 + if li % 5 == 0 or li == num_hidden_layers - 1: + print( + f" layer {li:2d} [{'swa' if is_swa else 'full'}]: " + f"q{list(q_bf16.shape)} k{list(k_bf16.shape)} v{list(v_bf16.shape)} " + f"elapsed={dt:.1f}s", + flush=True, + ) + + print("\nRewrite weight_map to reflect dtype change.", flush=True) + idx_path = os.path.join(neuron, "model.safetensors.index.json") + with open(idx_path) as f: + idx = json.load(f) + keys_to_drop = [ + k + for k in idx["weight_map"] + if any(k.endswith(f".{p}.scale") for p in ("q_proj", "k_proj", "v_proj")) + ] + for k in keys_to_drop: + idx["weight_map"].pop(k, None) + + total = 0 + for f_path in sorted(glob.glob(os.path.join(neuron, "*.safetensors"))): + with safe_open(f_path, framework="pt") as fp: + for k in fp.keys(): + t = fp.get_slice(k) + shape = list(t.get_shape()) + dt_bytes = {"F32": 4, "F16": 2, "BF16": 2, "F8_E4M3": 1}.get( + t.get_dtype(), 2 + ) + total += dt_bytes * max(1, int(math.prod(shape))) + idx["metadata"] = idx.get("metadata", {}) + idx["metadata"]["total_size"] = total + with open(idx_path, "w") as f: + json.dump(idx, f, indent=2) + print(f" dropped {len(keys_to_drop)} scale entries from index", flush=True) + print(f" total_size now {total / 1e9:.2f} GB", flush=True) + print(f"\nDone in {time.time() - t0:.1f}s", flush=True) + + +if __name__ == "__main__": + main() From cf60b9186f0ff681afc4ac4c9035a4acdd35f1dd Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 14:10:42 +0800 Subject: [PATCH 15/47] [contrib] MiMo-V2.5-Pro: fold BF16 attn into preprocess; drop repatch MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The separate repatch_qkv_bf16.py was a diagnostic workaround: preprocess FP8 first, discover drift, then retro-fit BF16. Now that BF16 attn is the confirmed recipe, fold the per-group dequant directly into the preprocess so there is one script, one output, no "forgot to run repatch" trap. Changes: - preprocess_mimo_v2_fp8.py::split_qkv_fused now returns BF16 per-proj tensors directly (Dict[str, Tensor] instead of Dict[str, Tuple[...]]). The FP8+blockwise path still unwinds the phantom-row padding, then dequants to BF16 in one go. BF16-source path collapses to the same reshape without requant. - Add _dequant_attn_to_bf16() for the Flash-style non-fused q/k/v fallback path; process_layer calls it so those projections also come out BF16. - No compile-time flag or branch for "all-FP8 attn" — that recipe is known broken for Pro (produces gibberish), preserving the branch only invites re-discovering the same trap. - Delete src/conversion_script/repatch_qkv_bf16.py. - README: drop the "Required follow-up: repatch" subsection, simplify the Status writeup (one recipe, one outcome), remove step 3b from Quick Start, clarify in "Preprocess emits BF16 q/k/v" that modules_to_not_convert still needs q/k/v so NxDI routes them through the non-quantized ColumnParallelLinear. - smoke_compile_mimo_v2.py: tighten the inline comment on q/k/v in modules_to_not_convert (no more "Prerequisite: run repatch"). Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 34 +--- .../perf_test/smoke_compile_mimo_v2.py | 12 +- .../preprocess_mimo_v2_fp8.py | 143 ++++++++----- .../src/conversion_script/repatch_qkv_bf16.py | 190 ------------------ 4 files changed, 106 insertions(+), 273 deletions(-) delete mode 100644 contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index d7f0c223..4788275c 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -59,9 +59,9 @@ Key features: - GPU stacks (sglang on H100/H200) run the exact same OCP FP8 checkpoint correctly because they always dequantize FP8 → BF16 before the matmul. The issue is specific to Neuron's direct-FP8 compute path on subnormal-leaning tensors. - Reference: Jim Burtoft observed similar prompt-dependent FP8 degradation on Flash and his Kimi PR #131 names "blockwise kernel padding produces depressed logits with EP=2 on SDK 2.29; SDK 2.28 recommended". -### Preprocess step required for BF16 attn +### Preprocess emits BF16 q/k/v -`src/conversion_script/repatch_qkv_bf16.py` reads the original HF fused `qkv_proj` weight + scale, dequants per-group (`num_q_heads_per_kv_group` × `head_dim` rows per group) to BF16, and overwrites `q_proj/k_proj/v_proj` in the preprocessed Neuron-FP8 checkpoint in place. Runs in ~22 min. Simply adding `q_proj/k_proj/v_proj` to `modules_to_not_convert` without this preprocess is **not** equivalent — NxDI casts the raw fp8_e4m3fn bytes to bfloat16 without applying the scale, which produces nonsense weights. +`src/conversion_script/preprocess_mimo_v2_fp8.py` now dequants q/k/v to BF16 directly (`split_qkv_fused` for the fused Pro layout, `_dequant_attn_to_bf16` for the Flash-style split layout). Output checkpoint has no `q_proj.scale` / `k_proj.scale` / `v_proj.scale` entries. Compile-time `modules_to_not_convert` must include `q_proj`, `k_proj`, `v_proj` so NxDI routes them through a plain `ColumnParallelLinear`; `smoke_compile_mimo_v2.py` and `start_vllm_server.sh` already do. ### Recipes that were tried and did not resolve the drift (all on 2026-04-28/29) @@ -121,19 +121,15 @@ git checkout contrib/MiMo-V2.5-Pro # the branch this README lives on huggingface-cli download XiaomiMiMo/MiMo-V2.5-Pro \ --local-dir /opt/dlami/nvme/models/MiMo-V2.5-Pro -# 3. Preprocess HF FP8 -> Neuron FP8 (~20 min, ~24 GB peak RAM) +# 3. Preprocess HF FP8 -> Neuron-FP8 (BF16 attn, FP8 MoE). ~20 min, ~24 GB +# peak RAM. The preprocess dequants q/k/v to BF16 in one pass — see +# "Checkpoint Preparation" below for why BF16 attn is the only recipe. source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py \ --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro \ --save_path /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8 \ --tp_degree 64 -# 3b. Dequant attention q/k/v to BF16 in place (~22 min, required for -# coherent output under the current FP8 recipe; see Status). -python contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py \ - --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro \ - --neuron_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8 - # 4. (Optional) sanity-check the Neuron-FP8 checkpoint without vLLM # ~90 min first compile; subsequent runs ~60s to load the pre-sharded NEFF. source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate @@ -225,19 +221,11 @@ python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8 Peak RAM during preprocessing is ~24 GB; total runtime ~20 minutes on a trn2.48xlarge instance. -### Required follow-up: dequant q/k/v to BF16 (`repatch_qkv_bf16.py`) - -Under the default all-FP8 recipe, Pro's attention weights drift enough to produce gibberish output (see Status). The fix is to keep MoE experts FP8 but dequant q/k/v to BF16 before compile. `src/conversion_script/repatch_qkv_bf16.py` reads the HF fused `qkv_proj` + `weight_scale_inv`, dequants per-group (each group is one kv-head: `hpg` Q rows, `1×head_dim` K rows, `1×v_head_dim` V rows), and overwrites the q_proj/k_proj/v_proj entries in the preprocessed Neuron-FP8 checkpoint in place. The scale entries for q/k/v are dropped from the safetensors index. - -```bash -python contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py \ - --hf_model_path /path/to/MiMo-V2.5-Pro \ - --neuron_model_path /path/to/MiMo-V2.5-Pro-Neuron-FP8 -``` +### Why q/k/v are BF16 in the preprocessed output -Runtime is ~22 minutes on a trn2.48xlarge, peak RAM ~20 GB. After running this, the compile-time `modules_to_not_convert` list in `smoke_compile_mimo_v2.py` / `start_vllm_server.sh` must include `q_proj`, `k_proj`, `v_proj` so NxDI keeps them as BF16 and routes them through a plain `ColumnParallelLinear` rather than the FP8 `QuantizedColumnParallel` path. +Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2.5 (256 experts). The NKI blockwise FP8 accumulator at this magnitude drifts the logits across 70 layers and produces gibberish output — `"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`, etc. Dequantizing q/k/v to BF16 while keeping MoE experts FP8 restores coherent output (verified on 2026-04-29 via `smoke_generate_mimo_v2.py`). -**Do not skip this step.** Adding q_proj/k_proj/v_proj to `modules_to_not_convert` without running repatch first leaves the fp8 bytes in the checkpoint; NxDI then casts them to bf16 **without applying the blockwise scale**, producing nonsense weights and broken output. +The preprocess handles this in a single pass: `split_qkv_fused()` unfuses Pro's `qkv_proj` into per-proj BF16 tensors directly, and the Flash-style per-proj fallback path dequants via `_dequant_attn_to_bf16()`. The checkpoint emitted by preprocess has no `q_proj.scale` / `k_proj.scale` / `v_proj.scale` entries. Compile-time `modules_to_not_convert` must therefore include `q_proj`, `k_proj`, `v_proj` so NxDI routes them through a plain `ColumnParallelLinear` rather than the FP8 `QuantizedColumnParallel` path — `smoke_compile_mimo_v2.py` already does this. ### Fallback: FP8 → BF16 @@ -264,8 +252,8 @@ compiled_path = "/path/to/compiled/" # Recommended recipe: BF16 attn + FP8 MoE. # moe_tp_degree = 1, moe_ep_degree = 64 -# q_proj/k_proj/v_proj in modules_to_not_convert (BF16; requires -# repatch_qkv_bf16.py to have been run on the checkpoint first) +# q_proj/k_proj/v_proj in modules_to_not_convert (BF16; preprocess +# emits BF16 for q/k/v, no separate step needed) # seq_len = 256 (HBM-tight with BF16 attn; see Status) # See "FP8 Configuration Notes" below for why other moe_tp/ep ratios # collapse. @@ -299,7 +287,7 @@ neuron_config = MoENeuronConfig( quantization_block_size=[128, 128], modules_to_not_convert=[ "embed_tokens", "lm_head", "norm", "router", "o_proj", - "q_proj", "k_proj", "v_proj", # BF16 attn — requires repatch + "q_proj", "k_proj", "v_proj", # BF16 attn — preprocess emits BF16 ], on_device_sampling_config=OnDeviceSamplingConfig( do_sample=True, temperature=0.6, top_k=20, top_p=0.95, diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index 6475eea2..73b8abbe 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -148,14 +148,10 @@ def main(): quantization_block_size=[128, 128], # BF16 attention: keep q/k/v_proj in BF16 (not FP8). Pro's q/k/v # abs_mean ~0.00124 is 4x smaller than V2.5 and the NKI blockwise - # FP8 accumulator drifts across 70 layers, producing gibberish - # output under the all-FP8 recipe. Dequantizing q/k/v to BF16 - # restores coherent output while keeping MoE experts FP8. - # Prerequisite: run src/conversion_script/repatch_qkv_bf16.py on - # the preprocessed Neuron-FP8 checkpoint first; simply listing - # q/k/v here without the repatch step leaves the fp8 bytes in the - # checkpoint and NxDI silently casts them to bf16 without applying - # the scale, which produces nonsense weights. + # FP8 accumulator drifts across 70 layers. Preprocess already + # emits BF16 for q/k/v (see preprocess_mimo_v2_fp8.py docstring); + # this list just tells NxDI to skip FP8 quantization and route + # through ColumnParallelLinear instead of QuantizedColumnParallel. modules_to_not_convert=[ "embed_tokens", "lm_head", diff --git a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py index b9882fbb..e63b5913 100644 --- a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py +++ b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py @@ -10,8 +10,11 @@ MiMo-V2.5-Pro checkpoint layout: - q_proj, k_proj, v_proj are FUSED into a single `qkv_proj` tensor per - layer (num_kv_heads interleaved groups, MiMo-V2.5-Pro-specific). We split into - three per-row-quantized projections via `split_qkv_fused()`. + layer (num_kv_heads interleaved groups, MiMo-V2.5-Pro-specific). We + split into three per-proj BF16 tensors via `split_qkv_fused()`. BF16 + (not FP8) is required: Pro's attention weights are small-magnitude + and the NKI blockwise FP8 accumulator drifts over 70 layers, producing + gibberish output. MoE experts can stay FP8. - 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 @@ -172,6 +175,33 @@ def _requantize_per_row(dequant: torch.Tensor) -> Tuple[torch.Tensor, torch.Tens return quantized, scales.to(torch.float32) +def _dequant_attn_to_bf16( + weight: torch.Tensor, scale: Optional[torch.Tensor] +) -> torch.Tensor: + """Dequantize an FP8 blockwise attention weight to BF16. + + Used by the Flash-style path where q/k/v ship as separate per-proj + tensors (not fused). The fused-qkv path handles dequant inside + split_qkv_fused because it also has to unwind the phantom-row padding. + """ + if weight.dtype != torch.float8_e4m3fn or scale is None: + return weight.to(torch.bfloat16) + + 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 + + wf = weight.float() + dequant = 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) + dequant[h0:h1, w0:w1] = wf[h0:h1, w0:w1] * scale[i, j].item() + return dequant.to(torch.bfloat16) + + def split_qkv_fused( qkv_weight: torch.Tensor, qkv_scale: Optional[torch.Tensor], @@ -179,10 +209,10 @@ def split_qkv_fused( num_kv_heads: int, head_dim: int, v_head_dim: int, -) -> Dict[str, Tuple[torch.Tensor, Optional[torch.Tensor]]]: - """Split Pro's pre-fused qkv_proj into q/k/v. MiMo-V2.5-Pro specific. +) -> Dict[str, torch.Tensor]: + """Split Pro's pre-fused qkv_proj into q/k/v (BF16 output). - HF layout — cross-validated against sglang on H200: + MiMo-V2.5-Pro specific. HF layout — cross-validated against sglang on H200: `qkv_proj.weight` is NOT `[all_Q | all_K | all_V]`. It is num_kv_heads interleaved groups, each holding (heads_per_group Q heads, 1 K head, 1 V head) packed contiguously: @@ -209,6 +239,13 @@ def split_qkv_fused( V is immediately followed by group (g+1)'s Q. We recover the correct dequant by padding each group up to 3456 rows before applying the scale, then stripping the phantom rows. + + Output dtype is always BF16 (no scale). Pro's q/k/v weights are + small-magnitude (abs_mean ~0.00124, 4x smaller than V2.5); the NKI + blockwise FP8 accumulator drifts at this scale and produces gibberish + output. Keeping q/k/v as BF16 while MoE experts stay FP8 is the only + configuration verified to produce coherent output, so this is the + single supported attention recipe. """ in_features = qkv_weight.shape[1] hpg = num_q_heads // num_kv_heads @@ -234,50 +271,53 @@ def split_qkv_fused( ) if qkv_weight.dtype != torch.float8_e4m3fn or qkv_scale is None: - # BF16 path + # BF16 source path (rare — most Pro checkpoints ship as FP8+scale). w = qkv_weight.view(num_kv_heads, real_rows_per_group, in_features) - q_w = w[:, :qg_rows, :].reshape(num_kv_heads * qg_rows, in_features).contiguous() - k_w = w[:, qg_rows:qg_rows + kg_rows, :].reshape(num_kv_heads * kg_rows, in_features).contiguous() - v_w = w[:, qg_rows + kg_rows:, :].reshape(num_kv_heads * 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_kv_heads * 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})" - ) - - w = qkv_weight.to(torch.float32).view( - num_kv_heads, real_rows_per_group, in_features - ) - w_padded = torch.zeros( - num_kv_heads, padded_rows_per_group, in_features, dtype=torch.float32 - ) - w_padded[:, :real_rows_per_group, :] = w + else: + # FP8 + blockwise scale path: dequant with phantom-row padding. + expected_scale_rows = num_kv_heads * 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})" + ) - s = qkv_scale.to(torch.float32).view( - num_kv_heads, 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] + wf = qkv_weight.to(torch.float32).view( + num_kv_heads, real_rows_per_group, in_features + ) + w_padded = torch.zeros( + num_kv_heads, padded_rows_per_group, in_features, dtype=torch.float32 + ) + w_padded[:, :real_rows_per_group, :] = wf - deq_padded = w_padded * s_exp - deq = deq_padded[:, :real_rows_per_group, :] + s = qkv_scale.to(torch.float32).view( + num_kv_heads, 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] - q_deq = deq[:, :qg_rows, :].reshape(num_kv_heads * qg_rows, in_features).contiguous() - k_deq = deq[:, qg_rows:qg_rows + kg_rows, :].reshape(num_kv_heads * kg_rows, in_features).contiguous() - v_deq = deq[:, qg_rows + kg_rows:, :].reshape(num_kv_heads * vg_rows, in_features).contiguous() + w = (w_padded * s_exp)[:, :real_rows_per_group, :] - 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) + q_bf16 = ( + w[:, :qg_rows, :] + .reshape(num_kv_heads * qg_rows, in_features) + .contiguous() + .to(torch.bfloat16) + ) + k_bf16 = ( + w[:, qg_rows:qg_rows + kg_rows, :] + .reshape(num_kv_heads * kg_rows, in_features) + .contiguous() + .to(torch.bfloat16) + ) + v_bf16 = ( + w[:, qg_rows + kg_rows:, :] + .reshape(num_kv_heads * vg_rows, in_features) + .contiguous() + .to(torch.bfloat16) + ) - return {"q_proj": (q_w2, q_s2), "k_proj": (k_w2, k_s2), "v_proj": (v_w2, v_s2)} + return {"q_proj": q_bf16, "k_proj": k_bf16, "v_proj": v_bf16} def _maybe_fp8_to_neuron_per_row( @@ -330,22 +370,21 @@ def process_layer( num_kv = config["num_key_value_heads"] hd = config.get("head_dim") vhd = config.get("v_head_dim", hd) + # split_qkv_fused returns BF16 weights only (no .scale); see its + # docstring for the rationale on why attn stays BF16 while MoE is FP8. split = split_qkv_fused(qkv_w, qkv_s, num_q, num_kv, hd, vhd) - 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.{proj}.scale"] = s2 + for proj, w_bf16 in split.items(): + out[f"{out_prefix}self_attn.{proj}.weight"] = w_bf16 else: - # Flash-style: q/k/v stored separately. + # Flash-style: q/k/v stored separately. Dequant to BF16 for the same + # reason as the fused path. 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 + w_bf16 = _dequant_attn_to_bf16(w, s) + out[f"{out_prefix}self_attn.{proj}.weight"] = w_bf16 # o_proj is listed in HF quantization_config.ignored_layers and ships as # BF16; on Neuron it binds to a plain RowParallelLinear (see diff --git a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py deleted file mode 100644 index 6ac7fb3a..00000000 --- a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/repatch_qkv_bf16.py +++ /dev/null @@ -1,190 +0,0 @@ -#!/usr/bin/env python3 -"""In-place patch: replace q/k/v FP8+scale with BF16 in a MiMo-V2.5-Pro -preprocessed Neuron checkpoint, for every decoder layer. Leaves MoE experts, -norms, embed, lm_head, o_proj untouched. - -Rationale: Pro's attention q/k/v weights have abs_mean ~0.001-0.005, roughly -4x smaller than V2.5. The NKI blockwise FP8 accumulator on the attention -path loses enough precision at this magnitude to drift the logits across -70 layers; dequantizing q/k/v to BF16 before the matmul restores coherent -output. MoE experts (also small-scale) can stay FP8. - -Note: simply adding q_proj/k_proj/v_proj to NxDI's `modules_to_not_convert` -at compile time is NOT equivalent — NxDI casts the raw fp8_e4m3fn bytes to -bfloat16 without applying the blockwise scale, which produces nonsense -weights. This script reads the HF fused qkv + scale, dequants per-group, -and writes BF16 weights back into the preprocessed Neuron checkpoint in -place. - -Run under /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16 (or any venv -with torch and safetensors). Takes ~22 min on a trn2.48xlarge for 70 -layers. -""" -import argparse -import glob -import json -import math -import os -import sys -import time - -import torch -from safetensors import safe_open -from safetensors.torch import save_file - - -def main(): - parser = argparse.ArgumentParser( - description="In-place dequant q/k/v from FP8+scale to BF16 in the " - "preprocessed Neuron checkpoint.", - ) - parser.add_argument( - "--hf_model_path", - required=True, - help="Path to the original HuggingFace MiMo-V2.5-Pro checkpoint " - "(fused qkv_proj + qkv_proj.weight_scale_inv).", - ) - parser.add_argument( - "--neuron_model_path", - required=True, - help="Path to the preprocessed Neuron-FP8 checkpoint. q/k/v entries " - "in its model_layer{N}.safetensors shards will be overwritten in " - "place with BF16 values; the scale entries are dropped from the " - "index.", - ) - args = parser.parse_args() - - # Import split_qkv_fused from the neighbouring preprocess script so the - # group layout math (hpg, qg_rows, kg_rows, vg_rows) stays in one place. - sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) - from preprocess_mimo_v2_fp8 import LazyWeightMap # noqa: F401 - - hf_src = args.hf_model_path - neuron = args.neuron_model_path - cfg = json.load(open(os.path.join(hf_src, "config.json"))) - hp = cfg["hybrid_layer_pattern"] - num_hidden_layers = cfg["num_hidden_layers"] - - with open(os.path.join(hf_src, "model.safetensors.index.json")) as f: - hf_wm = json.load(f)["weight_map"] - lazy = LazyWeightMap(hf_src, hf_wm) - - print( - f"Patching q/k/v -> BF16 in {neuron}/model_layer{{0..{num_hidden_layers - 1}}}.safetensors", - flush=True, - ) - t0 = time.time() - - for li in range(num_hidden_layers): - layer_file = os.path.join(neuron, f"model_layer{li}.safetensors") - if not os.path.exists(layer_file): - print(f" layer {li}: FILE MISSING, skip", flush=True) - continue - with safe_open(layer_file, framework="pt") as fp: - layer_sd = {k: fp.get_tensor(k) for k in fp.keys()} - - is_swa = hp[li] == 1 - num_q = cfg["swa_num_attention_heads" if is_swa else "num_attention_heads"] - num_kv = cfg["swa_num_key_value_heads" if is_swa else "num_key_value_heads"] - hd = cfg["swa_head_dim" if is_swa else "head_dim"] - vhd = cfg["swa_v_head_dim" if is_swa else "v_head_dim"] - prefix = f"model.layers.{li}.self_attn" - qkv_w = lazy.get(f"{prefix}.qkv_proj.weight") - qkv_s = lazy.get(f"{prefix}.qkv_proj.weight_scale_inv") - - BLOCK = 128 - hpg = num_q // num_kv - qg_rows = hpg * hd - kg_rows = 1 * hd - vg_rows = 1 * vhd - R = qg_rows + kg_rows + vg_rows - in_features = qkv_w.shape[1] - q_blk = qg_rows // BLOCK - k_blk = (kg_rows + BLOCK - 1) // BLOCK - v_blk = (vg_rows + BLOCK - 1) // BLOCK - per = q_blk + k_blk + v_blk - padded = per * BLOCK - - w = qkv_w.to(torch.float32).view(num_kv, R, in_features) - w_padded = torch.zeros(num_kv, padded, in_features, dtype=torch.float32) - w_padded[:, :R, :] = w - s = qkv_s.to(torch.float32).view(num_kv, per, (in_features + BLOCK - 1) // BLOCK) - s_exp = s.repeat_interleave(BLOCK, dim=1).repeat_interleave(BLOCK, dim=2) - s_exp = s_exp[:, :padded, :in_features] - deq_padded = w_padded * s_exp - deq = deq_padded[:, :R, :] - - q_bf16 = ( - deq[:, :qg_rows, :] - .reshape(num_kv * qg_rows, in_features) - .contiguous() - .to(torch.bfloat16) - ) - k_bf16 = ( - deq[:, qg_rows : qg_rows + kg_rows, :] - .reshape(num_kv * kg_rows, in_features) - .contiguous() - .to(torch.bfloat16) - ) - v_bf16 = ( - deq[:, qg_rows + kg_rows :, :] - .reshape(num_kv * vg_rows, in_features) - .contiguous() - .to(torch.bfloat16) - ) - - for key in list(layer_sd): - if any( - key.endswith(f".{p}.weight") or key.endswith(f".{p}.scale") - for p in ("q_proj", "k_proj", "v_proj") - ): - del layer_sd[key] - - out_prefix = f"layers.{li}.self_attn" - layer_sd[f"{out_prefix}.q_proj.weight"] = q_bf16 - layer_sd[f"{out_prefix}.k_proj.weight"] = k_bf16 - layer_sd[f"{out_prefix}.v_proj.weight"] = v_bf16 - - save_file(layer_sd, layer_file) - dt = time.time() - t0 - if li % 5 == 0 or li == num_hidden_layers - 1: - print( - f" layer {li:2d} [{'swa' if is_swa else 'full'}]: " - f"q{list(q_bf16.shape)} k{list(k_bf16.shape)} v{list(v_bf16.shape)} " - f"elapsed={dt:.1f}s", - flush=True, - ) - - print("\nRewrite weight_map to reflect dtype change.", flush=True) - idx_path = os.path.join(neuron, "model.safetensors.index.json") - with open(idx_path) as f: - idx = json.load(f) - keys_to_drop = [ - k - for k in idx["weight_map"] - if any(k.endswith(f".{p}.scale") for p in ("q_proj", "k_proj", "v_proj")) - ] - for k in keys_to_drop: - idx["weight_map"].pop(k, None) - - total = 0 - for f_path in sorted(glob.glob(os.path.join(neuron, "*.safetensors"))): - with safe_open(f_path, framework="pt") as fp: - for k in fp.keys(): - t = fp.get_slice(k) - shape = list(t.get_shape()) - dt_bytes = {"F32": 4, "F16": 2, "BF16": 2, "F8_E4M3": 1}.get( - t.get_dtype(), 2 - ) - total += dt_bytes * max(1, int(math.prod(shape))) - idx["metadata"] = idx.get("metadata", {}) - idx["metadata"]["total_size"] = total - with open(idx_path, "w") as f: - json.dump(idx, f, indent=2) - print(f" dropped {len(keys_to_drop)} scale entries from index", flush=True) - print(f" total_size now {total / 1e9:.2f} GB", flush=True) - print(f"\nDone in {time.time() - t0:.1f}s", flush=True) - - -if __name__ == "__main__": - main() From 62fea531bd500a3624aa4df02fa651a85b23cded Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 14:14:44 +0800 Subject: [PATCH 16/47] [contrib] MiMo-V2.5-Pro: standardize on pytorch_inference_vllm_0_16 venv Previously preprocess / smoke / vLLM-serving used three different venvs depending on which stage of the port we were in; both 2_9_nxd and inference_vllm_0_16 happen to have working NxDI + torch installs, so everything ran but the split was noise. Pick one and stick with it. pytorch_inference_vllm_0_16 is the right choice because: - 0_setup.sh installs vllm-neuron (editable) there, so vllm serving has no alternative. - NxDI direct calls from smoke_compile / smoke_generate also work there (nxdi is preinstalled by the DLAMI in both venvs). - Keeping one venv means no confusion about which python to invoke. Files updated: 0_setup.sh, run_bench_single.sh, smoke_compile_mimo_v2.py and smoke_generate_mimo_v2.py docstrings, run_preprocess_parallel.sh, README Prerequisites. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 4 ++-- contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh | 2 +- contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh | 2 +- .../models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py | 5 ++--- .../models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py | 2 +- .../src/conversion_script/run_preprocess_parallel.sh | 2 +- 6 files changed, 8 insertions(+), 9 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 4788275c..55a846f3 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -83,7 +83,7 @@ Key features: - **Instance**: trn2.48xlarge (128 physical 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. +- **Venv**: `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (used by preprocess, smoke, and vLLM serving alike; ships with the DLAMI and is where `0_setup.sh` installs the patched `vllm-neuron`). - **Disk**: ~3 TB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~962 GB, the Neuron-FP8 preprocessed output is ~1 TB, and `save_sharded_checkpoint=true` writes another ~300-1000 GB per compiled config (varies with recipe)). ### NVMe mount @@ -124,7 +124,7 @@ huggingface-cli download XiaomiMiMo/MiMo-V2.5-Pro \ # 3. Preprocess HF FP8 -> Neuron-FP8 (BF16 attn, FP8 MoE). ~20 min, ~24 GB # peak RAM. The preprocess dequants q/k/v to BF16 in one pass — see # "Checkpoint Preparation" below for why BF16 attn is the only recipe. -source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py \ --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro \ --save_path /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8 \ diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh index bdcbe11c..eb0ee005 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh @@ -12,7 +12,7 @@ echo "==========================================" echo "Setup: vllm-neuron + MiMo-V2.5-Pro weights" echo "==========================================" -source /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/bin/activate +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate PATCH_FILE="$(cd "$(dirname "$0")" && pwd)/vllm-neuron-patch.patch" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh index 2e4e4a3e..318ef477 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh @@ -26,7 +26,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_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" PORT="${PORT:-8000}" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index 73b8abbe..0f992dab 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -16,9 +16,8 @@ forward pass that allocates the shared scratchpad — useful when HBM is tight. -Run under /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference (NxDI direct). -The `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` venv is only -needed for vllm serving (bench_mimo_v2.sh). +Run under /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16 (same venv +as vllm serving; both NxDI direct and vllm-neuron are installed there). """ import os diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index 4c1a3bd2..1b76f4a4 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -7,7 +7,7 @@ single prompt via HuggingFaceGenerationAdapter. Purpose: sanity-check that the FP8 MoE + preprocessed scales actually produce coherent tokens. -Run under /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference. +Run under /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16. """ import os diff --git a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh index 5001ecd6..623946c0 100755 --- a/contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh +++ b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh @@ -15,7 +15,7 @@ # VENV venv with torch + safetensors + contrib pkg on sys.path set -e -VENV=${VENV:-/opt/aws_neuronx_venv_pytorch_2_9_nxd_inference} +VENV=${VENV:-/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16} source "$VENV/bin/activate" HF_MODEL_PATH=${HF_MODEL_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro} From e9ae0949a5c9de8305ab67fc07b8687ea8b77136 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 14:34:16 +0800 Subject: [PATCH 17/47] [contrib] MiMo-V2.5-Pro: trim Status, reframe perf numbers as historical Status had a 3-way split (all-FP8 vs BF16-attn vs preprocess emits BF16) that made sense during the diagnosis but doesn't once BF16-attn is the only shipping recipe. Collapse it into four focused subsections: * Why BF16 attn + FP8 MoE * Cost and constraints (HBM, seq_len=256, BS>=48, EP constraints) * Recipes tried that did not work (all-FP8, use_torch_block_wise) * Next experiments queued Performance: reframe the vLLM throughput table as a historical all-FP8 capture kept for infra validation and order-of-magnitude reference. The shipping recipe (BF16 attn + seq_len=256) hasn't been re-benchmarked yet; note the expected delta (only q/k/v change, MoE unchanged) so readers can project. vLLM Serving note: since the shipped start_vllm_server.sh still has seq_len=1024 and doesn't list q/k/v in modules_to_not_convert, spell out exactly what to change if the BF16-attn checkpoint OOMs on load. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 49 ++++++++++---------------- 1 file changed, 18 insertions(+), 31 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 55a846f3..9f77f799 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -37,48 +37,35 @@ Key features: ## Status (work-in-progress) -**This port compiles cleanly and serves via vLLM on Trn2. Under the default all-FP8 recipe the output drifts on most prompts; a BF16-attn recipe (keep MoE FP8, dequant q/k/v_proj to BF16, compile at seq_len=256) restores coherent output and isolates the root cause.** Findings as of 2026-04-29: +**This port compiles cleanly and serves via vLLM on Trn2. The shipping recipe is BF16 attention + FP8 MoE at `seq_len=256` — verified to produce coherent output end-to-end via `smoke_generate_mimo_v2.py`.** Last updated 2026-04-29. -### All-FP8 recipe (`tp_degree=64, moe_tp=1, moe_ep=64, BS=48, seq_len=1024`) +### Why BF16 attn + FP8 MoE -- **Prompt-dependent drift.** Self-intro prompts sometimes return coherent answers (sampling gets lucky on a strong self-identifying logit); most other prompts collapse to repetition or unrelated text within a few tokens (e.g. `"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`, chat continuations that drift into RLHF-style self-reflection with Chinese/Thai). -- **Not a sampling artifact.** `temperature` in the request is ignored — vllm-neuron's on-device sampling config (`do_sample=true, T=0.6, top_k=20, top_p=0.95`) is baked into the NEFF at compile time. Output is always stochastic, but the underlying logits are already drifted. -- Same failure pattern was observed on MiMo-V2-Pro under the same recipe (`"0.0.0.0:8080"`, etc.). V2.5-Pro's MoE expert weights are byte-identical to V2-Pro (verified layer 1 expert 0 dequant stats match to 6 decimal places on 2026-04-28). +Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2.5 (256 experts). Under an all-FP8 recipe, the NKI blockwise FP8 accumulator on attention q/k/v at this magnitude drifts the logits across 70 layers and produces prompt-dependent gibberish (`"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`, etc.). Dequantizing q/k/v to BF16 before the matmul restores coherent output. MoE experts (scales `≈ 2.3e-5`, similarly small) can stay FP8. -### BF16-attn recipe (`q_proj/k_proj/v_proj` dequant to BF16, MoE kept FP8, `seq_len=256`) +Verified end-to-end: `smoke_generate_mimo_v2.py` with a minimal chat template returns a well-formed reasoning trace that correctly identifies the model ("As MiMo, based on Xiaomi's self-developed large model..."). `preprocess_mimo_v2_fp8.py` emits BF16 q/k/v directly so no separate step is required. -- **Output is coherent.** On the same self-intro prompt, smoke_generate with a minimal chat template produces a well-formed reasoning trace that correctly identifies the model: - ``` - Okay, the user is asking for a simple self-introduction in one sentence, - with no deeper or hidden needs apparent. As MiMo, based on Xiaomi's self-developed - large model, I need to respond in a friendly, positive, and helpful way that aligns - with providing assistance ... - ``` -- **This narrows the root cause to attention-path FP8 precision**, not the MoE experts. Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2.5 (256 experts). The NKI blockwise FP8 accumulator on attention q/k/v at this magnitude loses enough precision to drift the logits across 70 layers; dequantizing q/k/v to BF16 before the matmul restores correct output. MoE experts (scales `≈ 2.3e-5`, also small) can stay FP8 under this recipe. -- **Cost: HBM headroom and seq_len.** BF16 q/k/v adds ~2 GB per rank. At `seq_len=1024` this OOMs on load (previous attempt failed allocating 41 MB for rdh/alltoall rings). `seq_len=256` frees enough full-attention softmax scratch to fit; longer context needs a different HBM plan (cross-instance TP/PP, or larger instance). -- GPU stacks (sglang on H100/H200) run the exact same OCP FP8 checkpoint correctly because they always dequantize FP8 → BF16 before the matmul. The issue is specific to Neuron's direct-FP8 compute path on subnormal-leaning tensors. -- Reference: Jim Burtoft observed similar prompt-dependent FP8 degradation on Flash and his Kimi PR #131 names "blockwise kernel padding produces depressed logits with EP=2 on SDK 2.29; SDK 2.28 recommended". +GPU stacks (sglang on H100/H200) run the same OCP FP8 checkpoint correctly because they always dequantize FP8 → BF16 before the matmul. The issue is specific to Neuron's direct-FP8 compute path on small-magnitude tensors. Kimi PR #131 observes similar FP8 degradation on Flash and recommends SDK 2.28. -### Preprocess emits BF16 q/k/v +### Cost and constraints -`src/conversion_script/preprocess_mimo_v2_fp8.py` now dequants q/k/v to BF16 directly (`split_qkv_fused` for the fused Pro layout, `_dequant_attn_to_bf16` for the Flash-style split layout). Output checkpoint has no `q_proj.scale` / `k_proj.scale` / `v_proj.scale` entries. Compile-time `modules_to_not_convert` must include `q_proj`, `k_proj`, `v_proj` so NxDI routes them through a plain `ColumnParallelLinear`; `smoke_compile_mimo_v2.py` and `start_vllm_server.sh` already do. +- **HBM headroom.** BF16 q/k/v adds ~2 GB per rank. `seq_len=1024` OOMs on load (the previous attempt failed allocating ~40 MB for rdh/alltoall rings after per-rank tensors already reached 20.9/24 GB). `seq_len=256` frees enough full-attention softmax scratch to fit. +- **Short context.** At `seq_len=256`, Pro's own chat template with the default system prompt is already 260 tokens. Longer context needs a different HBM plan (cross-instance TP/PP, or a larger instance). +- `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else `NotImplementedError`). With `num_experts=384, top_k=8` this forces `BS >= 48`. +- `n_routed_experts=384 = 2^7 × 3` → `384 / ep_degree` is never a power of 2 (6, 12, 24, 48, 96, 192, 384). Kimi PR #131 says NKI `_bwmm_shard_on_block_nki_call` on SDK 2.29 has "depressed logits with EP=2" and recommends SDK 2.28. -### Recipes that were tried and did not resolve the drift (all on 2026-04-28/29) +### Recipes tried that did not work -- **`use_torch_block_wise=True`** (PyTorch-fallback blockwise matmul for higher accumulator precision): compile+shard succeeded after ~2 h, but `model.load()` crashed with `status=4 Allocation Failure` — the fallback path raises HBM demand even with MoE-only scope. +- **All-FP8 attention (`modules_to_not_convert` without q/k/v).** Drifts as described above. Known broken; `preprocess_mimo_v2_fp8.py` no longer emits it. +- **`use_torch_block_wise=True`** (PyTorch-fallback blockwise matmul for higher accumulator precision): compile+shard succeeded after ~2 h, but `model.load()` crashed with `status=4 Allocation Failure` — the fallback path raises HBM demand even when scoped to MoE. ### Next experiments queued -- **BF16-attn + larger `seq_len`.** `seq_len=256` is tight; Pro's chat template with the default system prompt is already 260 tokens. Either shrink the system prompt, or try BF16-attn at `seq_len=512` on a less-full HBM plan (e.g. MoE-EP with fewer experts per rank, or drop batch size below 48 via different recipe trade-offs). +- **BF16-attn at `seq_len=512`** (needs a tighter HBM plan — smaller batch, different EP ratio, or shrinking the default system prompt). - **Cross-instance BF16** via pipeline/tensor parallelism on 2× Trn2 (single-instance HBM cannot hold full BF16 Pro). - **Selective BF16 only on MoE `gate_up_proj`** (smallest expert scales) while keeping `down_proj` FP8 — another axis to probe if attn drift returns at longer contexts. - **SDK 2.28 venv** test once installed, per Kimi PR #131. -### Known NxDI limits that constrain recipe choice - -- `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else `NotImplementedError`). With `num_experts=384, top_k=8` this forces `BS >= 48`. -- `n_routed_experts=384 = 2^7 × 3` → `384 / ep_degree` is never a power of 2 (6, 12, 24, 48, 96, 192, 384). Kimi PR #131 says NKI `_bwmm_shard_on_block_nki_call` on SDK 2.29 has "depressed logits with EP=2" and recommends SDK 2.28. - ## Prerequisites - **Instance**: trn2.48xlarge (128 physical NeuronCores, logical_nc_config=2 → 64 logical cores) @@ -369,7 +356,7 @@ bash contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh See "Environment variables" above for all the knobs (`NEURON_COMPILED_ARTIFACTS`, `BASE_COMPILE_WORK_DIR`, etc.) and their defaults. -> **Note on output quality:** the shipped vLLM scripts use the **all-FP8 recipe** (`seq_len=1024`, attn+MoE both FP8). This currently produces prompt-dependent drift — see Status. The **BF16-attn recipe** that restores coherent output has so far only been validated end-to-end via `smoke_generate_mimo_v2.py` (direct NxDI, `seq_len=256`); porting it into `start_vllm_server.sh` requires also adding `q_proj/k_proj/v_proj` to the `modules_to_not_convert` list and dropping `seq_len` / `max_model_len` to 256, and the perf numbers below have not been re-measured in that configuration. +> **Note on the shipped vLLM scripts:** the current `start_vllm_server.sh` still uses `seq_len=1024` and does not list `q_proj/k_proj/v_proj` in `modules_to_not_convert`. Coupled with a BF16-attn preprocessed checkpoint this runs correctly (NxDI just sees BF16 tensors where it expected FP8 and casts them as-is) but at a longer context than the BF16-attn recipe has been HBM-validated for. If you hit a `status=4 Allocation Failure` on load, drop `seq_len` / `max_model_len` / `context_encoding_buckets` / `token_generation_buckets` to 256 and add `q_proj/k_proj/v_proj` to `modules_to_not_convert` to match the smoke-verified configuration. The bench numbers below were taken on the older all-FP8 checkpoint and have not been re-measured since the preprocess switched to BF16 attn. ### vllm-neuron patch summary @@ -381,9 +368,9 @@ The patch is applied to vllm-neuron 0.5.0 and: ## Performance -> The throughput numbers below are from a working vLLM server run on 2026-04-29 under the recommended FP8 recipe. Output quality under this recipe is **not production-usable** (see Status); the numbers show that the serving infrastructure runs end-to-end, not that the model answers correctly. +> The throughput numbers below were captured on 2026-04-29 against a pre-BF16-attn checkpoint (all-FP8, `seq_len=1024`). They are historical — the shipping recipe is BF16 attn + FP8 MoE at `seq_len=256` and has not yet been re-benchmarked. The numbers are kept here for infra validation (continuous batching + bucketing + on-device sampling + vllm-neuron plugin all wire up end-to-end) and order-of-magnitude reference. -### vLLM Serving (trn2.48xlarge, all-FP8 recipe, BS=48, TP=64, moe_tp=1/moe_ep=64, CB + bucketing, `seq_len=1024`) +### vLLM Serving (trn2.48xlarge, historical all-FP8 run, BS=48, TP=64, moe_tp=1/moe_ep=64, CB + bucketing, `seq_len=1024`) Input/output: 900/90 tokens (`vllm bench serve --dataset-name random`), `on_device_sampling_config={do_sample:true, temperature:0.6, top_k:20, top_p:0.95}`. @@ -395,7 +382,7 @@ Input/output: 900/90 tokens (`vllm bench serve --dataset-name random`), `on_devi Per-stream ITL median holds at ~220 ms across all concurrency levels; TPOT/TTFT growth at higher concurrency comes from continuous-batching queue pressure, not per-step compute. -> **Numbers are from the all-FP8 recipe, which produces drifted output (see Status).** The BF16-attn recipe that restores coherent output has not yet been re-benchmarked on vLLM; throughput should be comparable (only q/k/v go BF16, MoE stays FP8) but at `seq_len=256` instead of 1024, so TTFT/latency characteristics will differ. +> Expected BF16-attn delta: only q/k/v go from FP8 to BF16 (MoE is unchanged), so steady-state throughput should be within a few percent. TTFT should drop proportionally with `seq_len` (256 vs 1024 prefill tokens). > **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. From d1ea946867385e6e4037a02892e16d7576707b35 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 14:36:31 +0800 Subject: [PATCH 18/47] [contrib] MiMo-V2.5-Pro: fix maintainer name typo --- contrib/models/MiMo-V2.5-Pro/README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 9f77f799..c28aa4cb 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -417,6 +417,6 @@ pytest contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py -v ## Maintainer -Henan Wan (whn09) +Henan Wang (whn09) **Last Updated:** 2026-04-29 From ad81f3f3c7ea3ac2211dca0a833525b16a3ed1a7 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 15:26:06 +0800 Subject: [PATCH 19/47] [contrib] MiMo-V2.5-Pro: align start_vllm_server.sh with BF16-attn ckpt The preprocess now emits BF16 q/k/v (no .scale entries), so vllm-neuron must route attention through the non-quantized ColumnParallelLinear. Three required changes: - Add q_proj/k_proj/v_proj to modules_to_not_convert. Without this, NxDI tries to load q_proj.scale and bails with "Cannot find layers.0.self_attn.q_proj.scale in state_dict". - Drop seq_len / max_model_len / context_encoding_buckets / token_generation_buckets from 1024 to 256. BF16 q/k/v adds ~2 GB per rank and seq_len=1024 OOMs on load; seq_len=256 is the smoke-verified upper bound. - Move NEURON_COMPILED_ARTIFACTS default to a new path (mimo_v2_5_pro_bs48_moetp1_ep64_bf16attn_seq256_vllm) so it doesn't collide with the old all-FP8 compile dir that's been S3-backed up. Note for longer context: seq_len is the single biggest HBM constraint on this recipe; raising it will require either a smaller batch, a different EP ratio, or cross-instance sharding (see README "Next experiments queued"). Co-Authored-By: Claude Opus 4.7 --- .../MiMo-V2.5-Pro/perf_test/start_vllm_server.sh | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh index 6848b604..04c1c5a3 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh @@ -31,7 +31,7 @@ export NXDI_CONTRIB_MIMO_V2_FLASH_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_pro_bs48_moetp1_ep64_fp8_vllm}" +: "${NEURON_COMPILED_ARTIFACTS:=/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_bf16attn_seq256_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. @@ -57,7 +57,7 @@ exec python3 -m vllm.entrypoints.openai.api_server \ --model "$MODEL_PATH" \ --tokenizer "$MODEL_PATH" \ --tensor-parallel-size 64 \ - --max-model-len 1024 \ + --max-model-len 256 \ --max-num-seqs 48 \ --no-enable-chunked-prefill \ --no-enable-prefix-caching \ @@ -79,19 +79,19 @@ exec python3 -m vllm.entrypoints.openai.api_server \ "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"], + "modules_to_not_convert": ["embed_tokens", "lm_head", "norm", "router", "o_proj", "q_proj", "k_proj", "v_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": 48, "ctx_batch_size": 1, "tkg_batch_size": 48, - "max_context_length": 1024, - "seq_len": 1024, + "max_context_length": 256, + "seq_len": 256, "is_continuous_batching": true, "enable_bucketing": true, - "context_encoding_buckets": [1024], - "token_generation_buckets": [1024], + "context_encoding_buckets": [256], + "token_generation_buckets": [256], "async_mode": true, "on_device_sampling_config": { "do_sample": true, "temperature": 0.6, "top_k": 20, "top_p": 0.95 From 5a297cba03fb930539d6ef628e79feecb1877a88 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 17:47:53 +0800 Subject: [PATCH 20/47] [contrib] MiMo-V2.5-Pro: fit sanity + bench within seq_len=256 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The start_vllm_server.sh now compiles with seq_len=256 (BF16-attn HBM constraint). Pro's default chat template prepends a ~240-token system prompt that by itself busts the bucket, and the old bench default (input 900, output 90) is also way over. sanity_check.sh: - Switch from /v1/chat/completions to /v1/completions with a hand-rolled <|im_start|>user... <|im_end|><|im_start|>assistant frame that tokenises to ~17 tokens. - Do the HTTP POST from python (bash heredoc mangles the \n inside the chat template, which used to make the model emit a garbage first token — UTF-8 replacement char "?" at the start of every reply). - Note in-comment that request-side temperature / top_k / top_p are ignored; the NEFF's on_device_sampling_config wins. run_bench_single.sh: - Default INPUT_LEN 900 -> 180, OUTPUT_LEN 90 -> 60 (180+60 = 240, fits under seq_len=256 with a small margin for random-range-ratio). - Comment explains the seq_len=256 constraint. bench_mimo_v2.sh is unchanged; it delegates length knobs to run_bench_single.sh. Co-Authored-By: Claude Opus 4.7 --- .../perf_test/run_bench_single.sh | 10 +- .../MiMo-V2.5-Pro/perf_test/sanity_check.sh | 94 +++++++++++-------- 2 files changed, 63 insertions(+), 41 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh index 318ef477..0ec5e90e 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh @@ -17,8 +17,8 @@ # /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8) # 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) +# INPUT_LEN --random-input-len (default 180; matches seq_len=256) +# OUTPUT_LEN --random-output-len (default 60; matches seq_len=256) # RANGE_RATIO --random-range-ratio (default 0.03) # CONFIG_NAME Used in the output filename (default bs48_tp64_moetp1_ep64) # RESULTS_DIR Where to dump per-run log @@ -32,9 +32,11 @@ MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP PORT="${PORT:-8000}" CONCURRENCY="${CONCURRENCY:-1}" NUM_PROMPTS="${NUM_PROMPTS:-16}" -INPUT_LEN="${INPUT_LEN:-900}" -OUTPUT_LEN="${OUTPUT_LEN:-90}" +INPUT_LEN="${INPUT_LEN:-180}" +OUTPUT_LEN="${OUTPUT_LEN:-60}" RANGE_RATIO="${RANGE_RATIO:-0.03}" +# seq_len=256 on the compiled server, so input+output must stay under 256. +# Default 180+60=240 leaves a small margin for random-range-ratio expansion. CONFIG_NAME="${CONFIG_NAME:-bs48_tp64_moetp1_ep64}" RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro}" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh index ed2f6b76..a0ed893c 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh @@ -1,25 +1,28 @@ #!/bin/bash # Quick sanity check against an already-running vLLM server. # -# Assumes vLLM is already listening on $PORT (default 8000) with MiMo-V2.5-Pro -# loaded. Sends a single chat completion and prints the model's reply. +# Posts a minimally-templated chat request to /v1/completions and prints the +# model's reply. We go through /v1/completions (not /v1/chat/completions) +# because Pro's default chat template prepends a ~240-token system prompt +# that by itself overflows the seq_len=256 compile-time bucket; building the +# im_start/im_end/assistant frame by hand keeps the prompt under ~30 tokens +# and fits cleanly. # # Usage: # bash sanity_check.sh # uses defaults # PORT=8001 bash sanity_check.sh # custom port -# PROMPT="..." bash sanity_check.sh # custom prompt +# PROMPT="..." bash sanity_check.sh # custom user content set -e MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" PORT="${PORT:-8000}" -# "Introduce yourself" is a high-signal self-identification prompt that the -# FP8 path answers coherently even under current MoE drift (see README -# Status). Swap PROMPT=... if you want to probe other prompts. +# "Introduce yourself" is the self-identification prompt that consistently +# lands in the model's MiMo-aware region. Swap PROMPT=... to probe others. PROMPT="${PROMPT:-Hello! Please introduce yourself in one sentence.}" -MAX_TOKENS="${MAX_TOKENS:-64}" +MAX_TOKENS="${MAX_TOKENS:-80}" -echo "Sanity check: POST /v1/chat/completions on port $PORT" +echo "Sanity check: POST /v1/completions on port $PORT" echo " Model: $MODEL_PATH" echo " Prompt: $PROMPT" echo " Max tokens: $MAX_TOKENS" @@ -28,38 +31,55 @@ 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.sh' or your own launcher first." + echo "Start it with 'bash start_vllm_server.sh' (or bench_mimo_v2.sh)" + echo "first and wait for 'Application startup complete.'" exit 1 fi -# NOTE: request-side `temperature` is ignored by vllm-neuron on this model: -# on-device sampling_config (set at compile time in start_vllm_server.sh as -# do_sample=true, T=0.6, top_k=20, top_p=0.95) is baked into the NEFF and -# request params don't override it. Output will be stochastic. -RESPONSE=$(curl -s "http://localhost:$PORT/v1/chat/completions" \ - -H 'Content-Type: application/json' \ - -d "$(cat </dev/null || echo "$RESPONSE" -echo "" +# NOTE: request-side `temperature` / `top_k` / `top_p` are ignored by +# vllm-neuron on this model: the on_device_sampling_config baked into the +# NEFF at compile time wins. Output is always stochastic; re-run to see +# variance. +# +# Build the chat framing in python so newlines and special tokens survive +# JSON encoding without shell escape pitfalls, then POST to /v1/completions. +python3 <user\n" + + user + + "<|im_end|>\n<|im_start|>assistant\n" +) +body = json.dumps({ + "model": model, + "prompt": prompt, + "max_tokens": int("$MAX_TOKENS"), + "stream": False, +}).encode() +req = urllib.request.Request( + "http://localhost:$PORT/v1/completions", + data=body, + headers={"Content-Type": "application/json"}, +) 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) + with urllib.request.urlopen(req, timeout=120) as r: + resp = json.load(r) +except urllib.error.HTTPError as e: + print("HTTP error:", e.code, e.read().decode(errors="replace")) + sys.exit(1) + +if "error" in resp: + print("Error from server:", json.dumps(resp["error"], indent=2)) + sys.exit(1) -echo "Model reply: $REPLY" +text = resp["choices"][0]["text"] +print("Response:") +print(text) +print() +print("Usage:", resp.get("usage", {})) +PYEOF From 965a947839bf626d22a17aef3bb2f23c23c91de6 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 17:53:26 +0800 Subject: [PATCH 21/47] =?UTF-8?q?[contrib]=20MiMo-V2.5-Pro:=20sanity=5Fche?= =?UTF-8?q?ck.sh=20=E2=80=94=20short=20system=20prompt,=20/v1/chat?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Switch back to /v1/chat/completions with an explicit short system message ("You are MiMo, a helpful assistant..."). apply_chat_template then uses our system turn instead of Pro's ~240-token default, and the prompt comes out to ~25 tokens — well under seq_len=256. This is simpler than the /v1/completions + manually-framed-chat route (no shell-escape \n landmines, native OpenAI API shape) and composes cleanly with other chat clients that assume /v1/chat. Override via SYSTEM=... / PROMPT=... / MAX_TOKENS=... env vars. Co-Authored-By: Claude Opus 4.7 --- .../MiMo-V2.5-Pro/perf_test/sanity_check.sh | 43 ++++++++++--------- 1 file changed, 23 insertions(+), 20 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh index a0ed893c..a80e85c1 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh @@ -1,29 +1,34 @@ #!/bin/bash # Quick sanity check against an already-running vLLM server. # -# Posts a minimally-templated chat request to /v1/completions and prints the -# model's reply. We go through /v1/completions (not /v1/chat/completions) -# because Pro's default chat template prepends a ~240-token system prompt -# that by itself overflows the seq_len=256 compile-time bucket; building the -# im_start/im_end/assistant frame by hand keeps the prompt under ~30 tokens -# and fits cleanly. +# Posts a chat request to /v1/completions and prints the reply. +# +# Pro's default chat template prepends a ~240-token system prompt that by +# itself overflows the seq_len=256 compile-time bucket, so we send an +# explicit short system message — apply_chat_template then uses ours +# instead of the default and the whole prompt fits in ~25 tokens. # # Usage: # bash sanity_check.sh # uses defaults # PORT=8001 bash sanity_check.sh # custom port # PROMPT="..." bash sanity_check.sh # custom user content +# SYSTEM="..." bash sanity_check.sh # custom system message set -e MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" PORT="${PORT:-8000}" +# Short system message (keeps total prompt ~25 tokens) — the checkpoint's +# default system prompt is ~240 tokens and would overflow seq_len=256. +SYSTEM="${SYSTEM:-You are MiMo, a helpful assistant developed by Xiaomi.}" # "Introduce yourself" is the self-identification prompt that consistently # lands in the model's MiMo-aware region. Swap PROMPT=... to probe others. PROMPT="${PROMPT:-Hello! Please introduce yourself in one sentence.}" MAX_TOKENS="${MAX_TOKENS:-80}" -echo "Sanity check: POST /v1/completions on port $PORT" +echo "Sanity check: POST /v1/chat/completions on port $PORT" echo " Model: $MODEL_PATH" +echo " System: $SYSTEM" echo " Prompt: $PROMPT" echo " Max tokens: $MAX_TOKENS" echo "" @@ -39,30 +44,28 @@ fi # NOTE: request-side `temperature` / `top_k` / `top_p` are ignored by # vllm-neuron on this model: the on_device_sampling_config baked into the # NEFF at compile time wins. Output is always stochastic; re-run to see -# variance. -# -# Build the chat framing in python so newlines and special tokens survive -# JSON encoding without shell escape pitfalls, then POST to /v1/completions. +# variance, or restart the server with `do_sample=false` in +# start_vllm_server.sh to force deterministic greedy decoding. python3 <user\n" - + user - + "<|im_end|>\n<|im_start|>assistant\n" -) body = json.dumps({ "model": model, - "prompt": prompt, + "messages": [ + {"role": "system", "content": system}, + {"role": "user", "content": user}, + ], "max_tokens": int("$MAX_TOKENS"), "stream": False, }).encode() req = urllib.request.Request( - "http://localhost:$PORT/v1/completions", + "http://localhost:$PORT/v1/chat/completions", data=body, headers={"Content-Type": "application/json"}, ) @@ -77,7 +80,7 @@ if "error" in resp: print("Error from server:", json.dumps(resp["error"], indent=2)) sys.exit(1) -text = resp["choices"][0]["text"] +text = resp["choices"][0]["message"]["content"] print("Response:") print(text) print() From 935510aca026c9d58706d79d60cf5b08a50f92c1 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 18:39:32 +0800 Subject: [PATCH 22/47] [contrib] MiMo-V2.5-Pro: default vLLM to smoke NEFF (workaround) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit vllm-neuron's own compile path — -O3, --enable-internal-neff-wrapper, on_device_sampling baked into the NEFF, continuous batching — produces garbled first-decode output on Pro: every reply starts with a UTF-8 replacement char and then coherent but completely off-topic text. V2.5 under the same vllm-neuron compile path works fine, so the trigger is Pro-specific (likely SWA + attention sink bias interacting with one of the compile / runtime options above, root cause not isolated). The NxDI-smoke compile path (-O1, no on-device sampler, static batch, produced by perf_test/smoke_compile_mimo_v2.py) does not hit the problem. vllm-neuron can load that NEFF at runtime and serves coherent chat completions with proper `` traces. As a workaround, default NEURON_COMPILED_ARTIFACTS to the smoke compile dir. Users can still override the env var to point at a vllm-neuron-compiled NEFF for testing. Co-Authored-By: Claude Opus 4.7 --- .../MiMo-V2.5-Pro/perf_test/start_vllm_server.sh | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh index 04c1c5a3..95bd8484 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh @@ -31,7 +31,15 @@ export NXDI_CONTRIB_MIMO_V2_FLASH_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_pro_bs48_moetp1_ep64_bf16attn_seq256_vllm}" +# Default points at the NxDI-smoke compile artifacts (produced by +# perf_test/smoke_compile_mimo_v2.py) instead of a vllm-neuron-compiled +# NEFF. Empirically, vllm-neuron's own compile path (-O3 + +# --enable-internal-neff-wrapper + on_device_sampling + CB-baked NEFF) +# produces garbled first-decode output on Pro (per-request `?` UTF-8 +# replacement chars followed by off-topic coherent text). The smoke +# compile (-O1, no on-device sampler, static batch) works correctly +# when vLLM loads it at runtime. Root cause TBD; not yet isolated. +: "${NEURON_COMPILED_ARTIFACTS:=/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq256}" 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. From 6da7188c5ffb161b4dc4b23585b02b671bce8e33 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Apr 2026 18:40:20 +0800 Subject: [PATCH 23/47] Revert "[contrib] MiMo-V2.5-Pro: default vLLM to smoke NEFF (workaround)" This reverts commit 935510aca026c9d58706d79d60cf5b08a50f92c1. --- .../MiMo-V2.5-Pro/perf_test/start_vllm_server.sh | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh index 95bd8484..04c1c5a3 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh @@ -31,15 +31,7 @@ export NXDI_CONTRIB_MIMO_V2_FLASH_SRC # Persistent compile-artifact location (NEFF + per-rank sharded weights). # Setting this overrides vLLM's fallback of # /neuron-compiled-artifacts//. -# Default points at the NxDI-smoke compile artifacts (produced by -# perf_test/smoke_compile_mimo_v2.py) instead of a vllm-neuron-compiled -# NEFF. Empirically, vllm-neuron's own compile path (-O3 + -# --enable-internal-neff-wrapper + on_device_sampling + CB-baked NEFF) -# produces garbled first-decode output on Pro (per-request `?` UTF-8 -# replacement chars followed by off-topic coherent text). The smoke -# compile (-O1, no on-device sampler, static batch) works correctly -# when vLLM loads it at runtime. Root cause TBD; not yet isolated. -: "${NEURON_COMPILED_ARTIFACTS:=/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq256}" +: "${NEURON_COMPILED_ARTIFACTS:=/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_bf16attn_seq256_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. From af27106b9a49e68a5a0026802e7e559721bb4a07 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 30 Apr 2026 10:39:37 +0800 Subject: [PATCH 24/47] [contrib] MiMo-V2.5-Pro: bump default seq_len 256 -> 512; document vLLM bug seq_len=512 under the BF16-attn recipe was verified end-to-end (compile + shard + load + 5x deterministic greedy generate) via smoke. HBM fits; seq_len=1024 still OOMs. Also documents the vllm-neuron "first request coherent, subsequent requests garbled" bug (tracked upstream at vllm-project/vllm-neuron#31). Every configuration knob we tried (all-FP8 attn, BF16 attn at 256 or 512, CB on/off, on-device sampling on/off, -O3 -> -O1) reproduced the same symptom on Pro but not on V2.5; the same NEFF serves 5 successive greedy generates byte-identically under smoke_generate_mimo_v2.py, so the bug is in vllm-neuron's runtime, not the NEFF. README changes: - Status opener now says the smoke path is verified and the vLLM serving path is blocked on issue #31. - Bump seq_len=256 references to seq_len=512 in HBM/constraints, Usage example, and the MoENeuronConfig code block. - Rewrite the vLLM "Note" callout to point at issue #31 as the single source of truth for the broken state, drop the obsolete "drop to 256" recovery hints. Script changes: - smoke_compile_mimo_v2.py: SEQ_LEN default 256 -> 512; COMPILED_PATH suffix seq256 -> seq512. Comment rewritten. - smoke_generate_mimo_v2.py: matching SEQ_LEN and COMPILED_PATH default changes so a bare `python smoke_generate_mimo_v2.py` picks up the seq_len=512 NEFF. Co-Authored-By: Claude Opus 4.7 --- contrib/models/MiMo-V2.5-Pro/README.md | 21 +++++++++++-------- .../perf_test/smoke_compile_mimo_v2.py | 11 +++++----- .../perf_test/smoke_generate_mimo_v2.py | 4 ++-- 3 files changed, 19 insertions(+), 17 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index c28aa4cb..1bdf8766 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -37,7 +37,9 @@ Key features: ## Status (work-in-progress) -**This port compiles cleanly and serves via vLLM on Trn2. The shipping recipe is BF16 attention + FP8 MoE at `seq_len=256` — verified to produce coherent output end-to-end via `smoke_generate_mimo_v2.py`.** Last updated 2026-04-29. +**This port compiles cleanly and is verified to produce coherent output end-to-end via the NxDI direct smoke path (`smoke_generate_mimo_v2.py`). The shipping recipe is BF16 attention + FP8 MoE at `seq_len=512` (largest seq_len that fits HBM).** + +**Known issue — vLLM serving is broken.** The first `/v1/chat/completions` request against `vllm-neuron` returns coherent output; every subsequent request returns garbled text. Same compiled NEFF serves 5 successive greedy generations byte-identically via the smoke path, so the bug is specifically in vllm-neuron's runtime / request-state handling. Tracking upstream at https://github.com/vllm-project/vllm-neuron/issues/31. Last updated 2026-04-30. ### Why BF16 attn + FP8 MoE @@ -49,8 +51,8 @@ GPU stacks (sglang on H100/H200) run the same OCP FP8 checkpoint correctly becau ### Cost and constraints -- **HBM headroom.** BF16 q/k/v adds ~2 GB per rank. `seq_len=1024` OOMs on load (the previous attempt failed allocating ~40 MB for rdh/alltoall rings after per-rank tensors already reached 20.9/24 GB). `seq_len=256` frees enough full-attention softmax scratch to fit. -- **Short context.** At `seq_len=256`, Pro's own chat template with the default system prompt is already 260 tokens. Longer context needs a different HBM plan (cross-instance TP/PP, or a larger instance). +- **HBM headroom.** BF16 q/k/v adds ~2 GB per rank. `seq_len=1024` OOMs on load (the previous attempt failed allocating ~40 MB for rdh/alltoall rings after per-rank tensors already reached 20.9/24 GB). `seq_len=512` is the largest value empirically verified to fit HBM at BS=48. +- **Short context.** Even at `seq_len=512`, Pro's full chat template with the default system prompt is ~260 tokens; that leaves ~250 tokens for user input + generation. Longer context needs a different HBM plan (cross-instance TP/PP, or a larger instance). - `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else `NotImplementedError`). With `num_experts=384, top_k=8` this forces `BS >= 48`. - `n_routed_experts=384 = 2^7 × 3` → `384 / ep_degree` is never a power of 2 (6, 12, 24, 48, 96, 192, 384). Kimi PR #131 says NKI `_bwmm_shard_on_block_nki_call` on SDK 2.29 has "depressed logits with EP=2" and recommends SDK 2.28. @@ -61,7 +63,8 @@ GPU stacks (sglang on H100/H200) run the same OCP FP8 checkpoint correctly becau ### Next experiments queued -- **BF16-attn at `seq_len=512`** (needs a tighter HBM plan — smaller batch, different EP ratio, or shrinking the default system prompt). +- **Even longer `seq_len`** (> 512): needs a tighter HBM plan — smaller batch, different EP ratio, or cross-instance sharding. +- **Upstream vllm-neuron fix** for the "first-request-only" serving bug (issue #31); patch branch at `whn09/vllm-neuron#fix/hybrid-attn-swa-spec` is a placeholder that did not resolve the symptom. - **Cross-instance BF16** via pipeline/tensor parallelism on 2× Trn2 (single-instance HBM cannot hold full BF16 Pro). - **Selective BF16 only on MoE `gate_up_proj`** (smallest expert scales) while keeping `down_proj` FP8 — another axis to probe if attn drift returns at longer contexts. - **SDK 2.28 venv** test once installed, per Kimi PR #131. @@ -241,7 +244,7 @@ compiled_path = "/path/to/compiled/" # moe_tp_degree = 1, moe_ep_degree = 64 # q_proj/k_proj/v_proj in modules_to_not_convert (BF16; preprocess # emits BF16 for q/k/v, no separate step needed) -# seq_len = 256 (HBM-tight with BF16 attn; see Status) +# seq_len = 512 (largest empirically verified; see Status) # See "FP8 Configuration Notes" below for why other moe_tp/ep ratios # collapse. neuron_config = MoENeuronConfig( @@ -253,7 +256,7 @@ neuron_config = MoENeuronConfig( max_batch_size=48, ctx_batch_size=1, tkg_batch_size=48, - seq_len=256, # HBM is tight with BF16 attn; seq_len=1024 OOMs + seq_len=512, # largest empirically verified; seq_len=1024 OOMs n_active_tokens=128, torch_dtype=torch.bfloat16, logical_nc_config=2, @@ -356,7 +359,7 @@ bash contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh See "Environment variables" above for all the knobs (`NEURON_COMPILED_ARTIFACTS`, `BASE_COMPILE_WORK_DIR`, etc.) and their defaults. -> **Note on the shipped vLLM scripts:** the current `start_vllm_server.sh` still uses `seq_len=1024` and does not list `q_proj/k_proj/v_proj` in `modules_to_not_convert`. Coupled with a BF16-attn preprocessed checkpoint this runs correctly (NxDI just sees BF16 tensors where it expected FP8 and casts them as-is) but at a longer context than the BF16-attn recipe has been HBM-validated for. If you hit a `status=4 Allocation Failure` on load, drop `seq_len` / `max_model_len` / `context_encoding_buckets` / `token_generation_buckets` to 256 and add `q_proj/k_proj/v_proj` to `modules_to_not_convert` to match the smoke-verified configuration. The bench numbers below were taken on the older all-FP8 checkpoint and have not been re-measured since the preprocess switched to BF16 attn. +> **vLLM serving is currently broken.** With the BF16-attn checkpoint, every `vllm-neuron` configuration we tried (all-FP8-attn, BF16-attn with `seq_len=256` or `512`, CB on/off, on-device sampling on/off, `-O3` or `-O1` TKG compile) reproduces the same pattern: the first chat request returns coherent output, every subsequent request returns UTF-8-replacement-char + off-topic text. The same compiled NEFF serves 5 successive greedy `adapter.generate()` calls byte-identically under `smoke_generate_mimo_v2.py` — the bug is in vllm-neuron's runtime, not in the model or the NEFF. Tracking at https://github.com/vllm-project/vllm-neuron/issues/31. Until that is fixed, use `smoke_generate_mimo_v2.py` for direct NxDI inference; the bench numbers below are historical infra-validation data from the pre-BF16-attn all-FP8 checkpoint. ### vllm-neuron patch summary @@ -368,7 +371,7 @@ The patch is applied to vllm-neuron 0.5.0 and: ## Performance -> The throughput numbers below were captured on 2026-04-29 against a pre-BF16-attn checkpoint (all-FP8, `seq_len=1024`). They are historical — the shipping recipe is BF16 attn + FP8 MoE at `seq_len=256` and has not yet been re-benchmarked. The numbers are kept here for infra validation (continuous batching + bucketing + on-device sampling + vllm-neuron plugin all wire up end-to-end) and order-of-magnitude reference. +> The throughput numbers below were captured on 2026-04-29 against a pre-BF16-attn checkpoint (all-FP8, `seq_len=1024`) before we discovered the vllm-neuron first-request bug. They are historical — the shipping recipe is BF16 attn + FP8 MoE at `seq_len=512` via the smoke path, and vLLM serving is currently blocked on issue #31. The numbers are kept here for order-of-magnitude reference. ### vLLM Serving (trn2.48xlarge, historical all-FP8 run, BS=48, TP=64, moe_tp=1/moe_ep=64, CB + bucketing, `seq_len=1024`) @@ -419,4 +422,4 @@ pytest contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py -v Henan Wang (whn09) -**Last Updated:** 2026-04-29 +**Last Updated:** 2026-04-30 diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index 0f992dab..d52f5846 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -39,15 +39,14 @@ ) COMPILED_PATH = os.environ.get( "MIMO_V25_PRO_COMPILED_PATH", - "/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq256/", + "/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/", ) TP_DEGREE = int(os.environ.get("TP_DEGREE", "64")) -# Drop seq_len to 256 to free ~200 MB of full-attention softmax scratch per -# rank. The previous BF16-attn attempt at seq_len=1024 OOM'd by 40 MB on load -# (failed to allocate 41943040 bytes for rdh/alltoall); seq_len=256 reclaims -# enough HBM to fit the extra BF16 q/k/v weights. -SEQ_LEN = int(os.environ.get("SEQ_LEN", "256")) +# seq_len=512 is the largest value verified to fit HBM under the BF16-attn +# recipe. seq_len=1024 OOMs on load (previous attempt failed allocating +# ~40 MB for rdh/alltoall rings after per-rank tensors reached 20.9/24 GB). +SEQ_LEN = int(os.environ.get("SEQ_LEN", "512")) # BS=48 is the minimum that avoids forward_selective_loading on decode: # `BS * top_k / num_experts >= 1.0` → BS >= 384/8 = 48. At BS=1 the TKG # path raises `NotImplementedError: Selective Loading with Expert parallelism`. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index 1b76f4a4..6a11a565 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -27,13 +27,13 @@ ) COMPILED_PATH = os.environ.get( "MIMO_V25_PRO_COMPILED_PATH", - "/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq256/", + "/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/", ) # Must match smoke_compile_mimo_v2.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", "256")) +SEQ_LEN = int(os.environ.get("SEQ_LEN", "512")) BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "48")) # must match smoke_compile CTX_BATCH_SIZE = int(os.environ.get("CTX_BATCH_SIZE", "1")) MOE_TP = int(os.environ.get("MOE_TP", "1")) From c1e3a1220cb6cfbedad857bf6ffadf76194c04f4 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 02:24:19 +0000 Subject: [PATCH 25/47] [contrib] MiMo-V2.5-Pro: fix README (HF size ~1TB not 290GB, remove non-existent FP8->BF16 fallback, record verified toolchain versions, fix NEURON_COMPILED_ARTIFACTS default path, re-verify smoke 2026-07-21) Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 1bdf8766..b105e81d 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -39,6 +39,8 @@ Key features: **This port compiles cleanly and is verified to produce coherent output end-to-end via the NxDI direct smoke path (`smoke_generate_mimo_v2.py`). The shipping recipe is BF16 attention + FP8 MoE at `seq_len=512` (largest seq_len that fits HBM).** +> **Re-verified 2026-07-21** on trn2.48xlarge (i-0f2e7a5194376e8fe) from the pre-compiled `seq512` NEFF + `models/MiMo-V2.5-Pro-Neuron-FP8` (04-28) weights. `smoke_generate_mimo_v2.py` (MINIMAL_CHAT=1, 40 tokens) loaded the presharded 64-rank checkpoint in ~200s (warmup 9.4s) and produced coherent output that correctly self-identifies as MiMo: `"Okay, the user is asking for a simple self-introduction ... to establish my identity as MiMo. ..."`. The direct-smoke path is confirmed still working; vLLM serving was **not** re-tested this run and is presumed still blocked on issue #31. + **Known issue — vLLM serving is broken.** The first `/v1/chat/completions` request against `vllm-neuron` returns coherent output; every subsequent request returns garbled text. Same compiled NEFF serves 5 successive greedy generations byte-identically via the smoke path, so the bug is specifically in vllm-neuron's runtime / request-state handling. Tracking upstream at https://github.com/vllm-project/vllm-neuron/issues/31. Last updated 2026-04-30. ### Why BF16 attn + FP8 MoE @@ -72,7 +74,7 @@ GPU stacks (sglang on H100/H200) run the same OCP FP8 checkpoint correctly becau ## Prerequisites - **Instance**: trn2.48xlarge (128 physical NeuronCores, logical_nc_config=2 → 64 logical cores) -- **Neuron SDK**: 2.29 (Python 3.12, PyTorch 2.9) +- **Neuron SDK**: 2.29 (Python 3.12, PyTorch 2.9). Verified toolchain (2026-07-21 instance): `neuronx-cc 2.26.6360.0`, `neuronx-distributed 0.19.28492`, `neuronx-distributed-inference 0.10.18399`, `torch 2.9.1`, `torch-neuronx 2.9.0.2.15`. The `neuronx-cc` package version (2.26.x) differs from the SDK release-train number (2.29); if you re-compile on a differently-imaged DLAMI, confirm these versions match or expect a cache miss. - **Venv**: `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (used by preprocess, smoke, and vLLM serving alike; ships with the DLAMI and is where `0_setup.sh` installs the patched `vllm-neuron`). - **Disk**: ~3 TB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~962 GB, the Neuron-FP8 preprocessed output is ~1 TB, and `save_sharded_checkpoint=true` writes another ~300-1000 GB per compiled config (varies with recipe)). @@ -106,8 +108,8 @@ git clone /neuronx-distributed-inference.git cd neuronx-distributed-inference git checkout contrib/MiMo-V2.5-Pro # 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 (~1 TB; 50 safetensors shards). +# Any HF-compatible downloader works; huggingface-cli example: huggingface-cli download XiaomiMiMo/MiMo-V2.5-Pro \ --local-dir /opt/dlami/nvme/models/MiMo-V2.5-Pro @@ -176,7 +178,7 @@ override or if you plan to launch vLLM outside of `bench_mimo_v2.sh`. | Variable | Default | Purpose | |---|---|---| -| `NEURON_COMPILED_ARTIFACTS` | `/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_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. | +| `NEURON_COMPILED_ARTIFACTS` | `/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_bf16attn_seq256_vllm` (per `start_vllm_server.sh`) | 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. **Note:** the script default is still a `seq256` path while the shipping recipe is `seq512` — override this to a `seq512` dir (e.g. `.../mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512`) to reuse the pre-compiled artifacts documented above, or expect a fresh ~60 min compile. | | `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 Pro's 384-expert MoE is ~60 min TKG + ~15 min CTE + ~30 min shard, well past vLLM's default. | @@ -217,9 +219,15 @@ Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2 The preprocess handles this in a single pass: `split_qkv_fused()` unfuses Pro's `qkv_proj` into per-proj BF16 tensors directly, and the Flash-style per-proj fallback path dequants via `_dequant_attn_to_bf16()`. The checkpoint emitted by preprocess has no `q_proj.scale` / `k_proj.scale` / `v_proj.scale` entries. Compile-time `modules_to_not_convert` must therefore include `q_proj`, `k_proj`, `v_proj` so NxDI routes them through a plain `ColumnParallelLinear` rather than the FP8 `QuantizedColumnParallel` path — `smoke_compile_mimo_v2.py` already does this. -### Fallback: FP8 → BF16 +### Parallel preprocess (faster) + +`src/conversion_script/preprocess_mimo_v2_parallel.py` (driven by `run_preprocess_parallel.sh`) is a multiprocess wrapper around the same per-layer conversion. Each worker dequants one layer independently, cutting wall time from ~20-30 min (serial) to ~5-6 min with 12 workers (peak ~300 GB CPU RAM on a 2 TB box). Output is identical to the serial path (BF16 attn + FP8 MoE). + +```bash +N_WORKERS=12 bash contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh +``` -`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. +> Note: there is no FP8 → full-BF16 conversion mode. Both preprocess scripts always emit the shipping recipe (BF16 q/k/v attention + FP8 MoE experts); q/k/v are dequantized to BF16 in-pass, but MoE weights stay FP8. A separate all-BF16 reference checkpoint, if needed, must be produced by other means. ## Usage @@ -324,7 +332,7 @@ Intermediate ratios (`moe_tp=32/ep=2`, `moe_tp=16/ep=4`) have been empirically t ### batch_size >= 48 -NxDI's TKG (token generation) path refuses Expert Parallelism when `batch_size < num_experts / top_k`. For Pro that is 384 / 8 = 48, so the smallest working BS on the FP8 path is 48. BS=1 latency demos are not possible on FP8; use the BF16 checkpoint with `moe_tp=64, moe_ep=1, batch_size=1` for single-stream latency measurements. +NxDI's TKG (token generation) path refuses Expert Parallelism when `batch_size < num_experts / top_k`. For Pro that is 384 / 8 = 48, so the smallest working BS on the FP8 path is 48. BS=1 latency demos are not possible on the FP8 (moe_ep=64) path; a single-stream configuration would require `moe_tp=64, moe_ep=1, batch_size=1`, which in turn needs an all-BF16 checkpoint (the preprocess scripts here only emit BF16-attn + FP8-MoE, so that BF16 checkpoint must be produced separately). ### outer ep_degree = 1 From f5c7d44d1b7f3e22273fc67d1e15ff8dfef50324 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 03:06:33 +0000 Subject: [PATCH 26/47] [contrib] MiMo-V2.5-Pro: fix vLLM-neuron patch so serving actually loads The committed vllm-neuron-patch.patch was corrupt (wrong hunk line counts: first hunk header said +922,94 but the body has 98 added lines, and the second hunk offset was stale), so 0_setup.sh git-apply silently no-oped and left vllm-neuron unpatched. Even once applied, the checkpoint declares architectures=[MiMoV2ForCausalLM] / model_type=mimo_v2, but the hook only registered mimov2flash/mimov2pro keys and MiMoV2Flash/ProForCausalLM archs, so vLLM ModelConfig validation rejected the model before the engine started. Fixes (verified on trn2.48xlarge 2026-07-21: vLLM loads MiMo-V2.5-Pro and serves requests, reusing the seq512 smoke NEFF with no recompile): - Regenerate the patch via git-diff so hunk headers are correct. - Register the mimov2 MODEL_TYPES key (what MiMoV2ForCausalLM resolves to via _get_neuron_model_cls) alongside mimov2flash/mimov2pro. - Add MiMoV2ForCausalLM to the ModelRegistry override list. - Register contrib archs in NeuronPlatform.pre_register_and_update (runs before ModelConfig validation, after vLLM is importable) via lazy string registration, avoiding the vllm.config circular import a platform-plugin register() hook hits. - 0_setup.sh: hard-fail on a patch that neither applies nor is already applied, instead of silently continuing unpatched. Note: the known first-request-ok / subsequent-garbled vLLM runtime bug (issue #31) still reproduces and is unaffected by these load-path fixes. Co-Authored-By: Claude Opus 4.8 --- .../models/MiMo-V2.5-Pro/perf_test/0_setup.sh | 15 ++++- .../perf_test/vllm-neuron-patch.patch | 60 +++++++++++++++++-- 2 files changed, 67 insertions(+), 8 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh index eb0ee005..bd6b55a0 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh @@ -25,12 +25,23 @@ fi cd $HOME/vllm-neuron -# Apply patch (idempotent via `git apply --check` first). +# Apply patch. Distinguish three cases so a corrupt/conflicting patch is a +# hard error rather than being silently skipped (a malformed hunk header once +# caused this to no-op, leaving the contrib model unregistered and vLLM unable +# to load MiMo-V2.5-Pro): +# - applies cleanly -> apply it +# - already applied (reverse) -> skip, fine +# - neither -> abort with a clear message if git apply --check "$PATCH_FILE" 2>/dev/null; then git apply "$PATCH_FILE" echo " Applied $PATCH_FILE" +elif git apply --reverse --check "$PATCH_FILE" 2>/dev/null; then + echo " Patch already applied; skipping." else - echo " Patch already applied or conflicts; continuing." + echo " ERROR: $PATCH_FILE does not apply cleanly and is not already applied." >&2 + echo " Refusing to continue with an unpatched vllm-neuron (MiMo would fail to load)." >&2 + git apply --check "$PATCH_FILE" # surface the real error, then abort via set -e + exit 1 fi pip install --extra-index-url=https://pip.repos.neuron.amazonaws.com -e . diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch b/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch index a29814d4..96f2514f 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch @@ -1,8 +1,56 @@ +diff --git a/vllm_neuron/platform.py b/vllm_neuron/platform.py +index 7d64b83..9966ce8 100644 +--- a/vllm_neuron/platform.py ++++ b/vllm_neuron/platform.py +@@ -286,6 +286,43 @@ class NeuronPlatform(Platform): + ) -> None: + # Apply config overrides here - this is called before VllmConfig is created + cls._apply_config_overrides() ++ # Register NxDI contrib model architectures into vLLM's ModelRegistry. ++ # This hook runs before ModelConfig's architecture validation and after ++ # vLLM is fully importable, so (unlike the platform-plugin register()) ++ # it avoids the vllm.config circular import. Lazy string registration ++ # (":") imports nothing here. ++ cls._register_contrib_archs() ++ ++ @classmethod ++ def _register_contrib_archs(cls) -> None: ++ """Alias contrib architectures (in checkpoint config.json but not in ++ vLLM's built-in registry) so ModelConfig validation accepts them. ++ ++ Best-effort and idempotent. The aliased stub is only used to satisfy ++ validation; the Neuron path resolves the real class via NxDI ++ MODEL_TYPES inside the engine worker (see neuronx_distributed_model_loader ++ ._register_contrib_models / _get_neuron_model_cls). ++ """ ++ import warnings as _w ++ try: ++ from vllm.model_executor.models.registry import ModelRegistry ++ supported = set(ModelRegistry.get_supported_archs()) ++ _aliases = { ++ # MiMo-V2.5-Pro config.json declares MiMoV2ForCausalLM, which ++ # vLLM does not know; alias to its built-in MiMoV2Flash stub. ++ "MiMoV2ForCausalLM": ++ "vllm.model_executor.models.mimo_v2_flash:MiMoV2FlashForCausalLM", ++ "MiMoV2ProForCausalLM": ++ "vllm.model_executor.models.mimo_v2_flash:MiMoV2FlashForCausalLM", ++ } ++ for _arch, _target in _aliases.items(): ++ if _arch not in supported: ++ ModelRegistry.register_model(_arch, _target) ++ except Exception as _e: # pragma: no cover - best-effort ++ _w.warn( ++ f"Contrib architecture registration failed: {_e}", ++ category=UserWarning, ++ ) + + @classmethod + def check_and_update_config(cls, vllm_config: VllmConfig) -> None: diff --git a/vllm_neuron/worker/neuronx_distributed_model_loader.py b/vllm_neuron/worker/neuronx_distributed_model_loader.py -index d2099eb..0c162e4 100644 +index d2099eb..a8dfc7f 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: +@@ -922,6 +922,98 @@ def _camel_to_kebab(name: str) -> str: return re.sub("([a-z0-9])([A-Z])", r"\1-\2", s1).lower() @@ -56,7 +104,7 @@ index d2099eb..0c162e4 100644 + _patch_autoconfig_trust_remote_code() + + 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 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: @@ -64,7 +112,7 @@ index d2099eb..0c162e4 100644 + # Register under both Flash and Pro model_type keys so the same + # NxDI wrapper serves MiMo-V2-Flash, MiMo-V2.5-Pro, and any sibling + # that inherits the same config. -+ for _mt in ("mimov2flash", "mimov2pro"): ++ for _mt in ("mimov2", "mimov2flash", "mimov2pro"): + MODEL_TYPES.setdefault(_mt, {"causal-lm": NeuronMiMoV2ForCausalLM}) + try: + from vllm.model_executor.models.registry import ModelRegistry @@ -72,7 +120,7 @@ index d2099eb..0c162e4 100644 + # wrapper so ModelConfig validation accepts the architecture + # regardless of whether the checkpoint calls itself Flash or + # Pro. -+ for _arch in ("MiMoV2FlashForCausalLM", "MiMoV2ProForCausalLM"): ++ for _arch in ("MiMoV2ForCausalLM", "MiMoV2FlashForCausalLM", "MiMoV2ProForCausalLM"): + ModelRegistry.register_model(_arch, NeuronMiMoV2ForCausalLM) + except ImportError: + pass @@ -101,7 +149,7 @@ index d2099eb..0c162e4 100644 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): +@@ -941,6 +1033,7 @@ def _get_neuron_model_cls(architecture: str): _get_neuron_model_cls("NeuronLlamaForCausalLM") """ From 9a89c6c43a677a9ec5493dd62b74ace339e51a50 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 04:04:43 +0000 Subject: [PATCH 27/47] [contrib] MiMo-V2.5-Pro: add smoke-path throughput benchmark + prefill/decode profiling data MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit bench_smoke_throughput.py measures the NxDI direct path (no vLLM) at BS=48, splitting prefill (CTE) from decode (TKG): a max_new_tokens=1 call isolates prefill/TTFT, the full call minus that isolates steady-state decode. Also documents NEURON_RT_INSPECT_DEVICE_PROFILE capture for device profiling. Measured 2026-07-21 (seq_len=512, verified on trn2.48xlarge): - Smoke end-to-end 95.6 out-tok/s vs vLLM c=48 71.9 (vLLM loses ~25% to continuous-batching/async/request-state overhead). - Prefill dominates and is ~constant in input length (37.9 s for both 360 and 500 input tokens) because context_encoding_buckets=[512] pads every prefill to 512. That fixed ~38 s (≈190 decode steps) is the top optimization target -- almost certainly the 384-expert FP8 MoE blockwise matmul over 512 positions, not attention. Decode is ~0.2 s/token for the full BS=48 batch. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 49 +++++ .../perf_test/bench_smoke_throughput.py | 167 ++++++++++++++++++ 2 files changed, 216 insertions(+) create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index b105e81d..8faf58ba 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -395,6 +395,55 @@ Per-stream ITL median holds at ~220 ms across all concurrency levels; TPOT/TTFT > Expected BF16-attn delta: only q/k/v go from FP8 to BF16 (MoE is unchanged), so steady-state throughput should be within a few percent. TTFT should drop proportionally with `seq_len` (256 vs 1024 prefill tokens). +### Measured 2026-07-21 (BF16-attn + FP8 MoE, `seq_len=512`, BS=48, TP=64, moe_tp=1/moe_ep=64) + +Same instance/recipe as the shipping config, via `bench_smoke_throughput.py` +(NxDI direct path, coherent output). Prefill and decode measured separately +(a `max_new_tokens=1` call isolates prefill/TTFT; the full call minus that +isolates steady-state decode). 3-iter averages, very stable. + +| in/out | Prefill (CTE) | Decode (TKG) | End-to-end | +|--------|---------------|--------------|------------| +| 360/120 | 37.9 s → 456 in-tok/s | 22.4 s / 119 steps → 255 out-tok/s (5.3/stream) | 60.3 s → 95.6 out-tok/s | +| 500/2 | 37.9 s → 634 in-tok/s | 0.2 s / 1 step → 235 out-tok/s (4.9/stream) | — | + +**Prefill dominates and is nearly constant in input length** (37.9 s for both +360 and 500 input tokens). The context-encoding NEFF has a single bucket +`context_encoding_buckets=[512]`, so every prefill pads to 512 and pays the +same cost regardless of real input length — that fixed ~38 s (≈190 decode +steps' worth of time) is the #1 optimization target, and it is almost +certainly the 384-expert FP8 MoE blockwise matmul over 512 positions, not +attention. Decode itself is cheap (~0.2 s/token for the whole BS=48 batch). + +vLLM serving on the same NEFF/recipe (input/output 360/120): + +| Concurrency | Output tok/s | Total tok/s | TPOT median | Notes | +|-------------|--------------|-------------|-------------|-------| +| 48 | 71.9 | 287 | 572 ms | output garbled after first request (issue #31) | +| 1 | 5.2 | 20.5 | 189 ms | under-fed (BS=48 graph, 1 request); not representative | + +vLLM at c=48 reaches 72 out-tok/s vs the smoke path's 96 — the ~25% gap is +vLLM's continuous-batching scheduler / async / request-state overhead (the same +runtime layer behind the issue #31 garbling). Smoke is the batched-compute +ceiling for this NEFF; vLLM is the realistic serving figure with scheduling. + +Reproduce (smoke path): + +```bash +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate +export MIMO_V25_PRO_COMPILED_PATH=/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/ +# decode-focused: +INPUT_LEN=360 MAX_NEW_TOKENS=120 N_ITERS=3 \ + python3 contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py +# prefill-focused (near-full context, ~1 output token): +INPUT_LEN=500 MAX_NEW_TOKENS=2 N_ITERS=3 \ + python3 contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py +``` + +To capture a device profile for bottleneck analysis (attention vs MoE vs +collectives) set `NEURON_RT_INSPECT_DEVICE_PROFILE=` before the run and +inspect the resulting `*.ntff` with `neuron-profile view` / Neuron Explorer. + > **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. ## Compatibility Matrix diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py b/contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py new file mode 100644 index 00000000..8ad85e8e --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py @@ -0,0 +1,167 @@ +#!/usr/bin/env python3 +"""Throughput benchmark for MiMo-V2.5-Pro via the NxDI direct (smoke) path. + +Loads the pre-compiled seq512 NEFF, then runs several batched generate() +calls and reports aggregate throughput (BS * new_tokens / wall_time) so it +can be compared apples-to-apples with the vLLM `output token throughput`. + +Env: + MIMO_V25_PRO_MODEL_PATH, MIMO_V25_PRO_COMPILED_PATH (paths) + BATCH_SIZE=48 MAX_NEW_TOKENS=120 INPUT_LEN=360 N_ITERS=3 +Set NEURON_RT_INSPECT_DEVICE_PROFILE= before running to also capture a +device profile (dumps *.ntff under that dir). +""" +import os +import sys +import time + +os.environ.setdefault("XLA_HANDLE_SPECIAL_SCALAR", "1") +os.environ.setdefault("UNSAFE_FP8FNCAST", "1") + +MODEL_PATH = os.environ.get("MIMO_V25_PRO_MODEL_PATH", + "/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8") +COMPILED_PATH = os.environ.get( + "MIMO_V25_PRO_COMPILED_PATH", + "/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/") +TP_DEGREE = int(os.environ.get("TP_DEGREE", "64")) +SEQ_LEN = int(os.environ.get("SEQ_LEN", "512")) +BATCH_SIZE = int(os.environ.get("BATCH_SIZE", "48")) +MOE_TP = int(os.environ.get("MOE_TP", "1")) +MOE_EP = int(os.environ.get("MOE_EP", "64")) +MAX_NEW_TOKENS = int(os.environ.get("MAX_NEW_TOKENS", "120")) +INPUT_LEN = int(os.environ.get("INPUT_LEN", "360")) +N_ITERS = int(os.environ.get("N_ITERS", "3")) + +os.environ.setdefault( + "BASE_COMPILE_WORK_DIR", + os.path.join("/tmp/nxd_model", os.path.basename(COMPILED_PATH.rstrip("/")))) + + +def main(): + import torch + 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") + # allow running from /tmp too + for cand in (contrib_src, + "/home/ubuntu/ndi-pr150-MiMo/contrib/models/MiMo-V2.5-Pro/src"): + if os.path.isdir(cand): + sys.path.insert(0, os.path.abspath(cand)) + break + from modeling_mimo_v2 import MiMoV2InferenceConfig, NeuronMiMoV2ForCausalLM + + print(f"[bench] MODEL={MODEL_PATH}") + print(f"[bench] COMPILED={COMPILED_PATH}") + print(f"[bench] BS={BATCH_SIZE} SEQ={SEQ_LEN} INPUT_LEN={INPUT_LEN} " + f"MAX_NEW={MAX_NEW_TOKENS} N_ITERS={N_ITERS}") + + 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=1, + 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", "q_proj", "k_proj", "v_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)) + + t0 = time.time() + model = NeuronMiMoV2ForCausalLM(MODEL_PATH, config) + print(f"[bench] instantiated in {time.time()-t0:.1f}s") + t0 = time.time() + model.load(COMPILED_PATH, skip_warmup=False) + print(f"[bench] loaded in {time.time()-t0:.1f}s") + + tok = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True) + tok.padding_side = "left" + adapter = HuggingFaceGenerationAdapter(model) + + # Build a prompt of ~INPUT_LEN tokens, replicated to fill the batch. + base = "Please write a detailed explanation about large language models. " + prompt = (base * 40) + enc = tok([prompt], return_tensors="pt", add_special_tokens=False) + ids = enc["input_ids"][0][:INPUT_LEN] + prompt = tok.decode(ids) + inputs = tok([prompt] * BATCH_SIZE, return_tensors="pt", padding="max_length", + max_length=INPUT_LEN, truncation=True) + gen = GenerationConfig(max_new_tokens=MAX_NEW_TOKENS, min_new_tokens=MAX_NEW_TOKENS, + do_sample=False, + pad_token_id=getattr(tok, "pad_token_id", None) or tok.eos_token_id) + print(f"[bench] input_ids.shape={tuple(inputs['input_ids'].shape)}") + + # Separate prefill (context-encoding, CTE) from decode (token-generation, + # TKG). A max_new_tokens=1 call is prefill + one decode step ~= TTFT; the + # full call is prefill + (MAX_NEW_TOKENS-1) decode steps. Subtracting the + # two isolates the steady-state per-step decode cost, from which we derive + # decode throughput. Prefill throughput = (BS * INPUT_LEN) / prefill_time. + gen1 = GenerationConfig(max_new_tokens=1, min_new_tokens=1, do_sample=False, + pad_token_id=getattr(tok, "pad_token_id", None) or tok.eos_token_id) + + def _run(gcfg): + t0 = time.time() + adapter.generate(input_ids=inputs["input_ids"], + attention_mask=inputs["attention_mask"], + generation_config=gcfg) + return time.time() - t0 + + print("[bench] warmup...") + _run(gen) # warmup, not counted + + prefill_times, full_times = [], [] + for it in range(1, N_ITERS + 1): + t_prefill = _run(gen1) # prefill + 1 decode step ~= TTFT + t_full = _run(gen) # prefill + (MAX_NEW_TOKENS-1) decode steps + prefill_times.append(t_prefill) + full_times.append(t_full) + # decode-only time for this iter = full - prefill, over (MAX_NEW_TOKENS-1) steps + decode_time = t_full - t_prefill + decode_steps = MAX_NEW_TOKENS - 1 + prefill_in = BATCH_SIZE * INPUT_LEN + decode_out = BATCH_SIZE * decode_steps + print(f"[bench] iter{it}: " + f"prefill(TTFT)={t_prefill:.3f}s ({prefill_in/t_prefill:,.0f} in-tok/s) | " + f"full={t_full:.2f}s | " + f"decode={decode_time:.2f}s ({decode_out/decode_time:.1f} out-tok/s, " + f"per_stream={decode_steps/decode_time:.2f} tok/s) | " + f"end2end_out={BATCH_SIZE*MAX_NEW_TOKENS/t_full:.1f} tok/s") + + n = len(full_times) + avg_prefill = sum(prefill_times) / n + avg_full = sum(full_times) / n + avg_decode = avg_full - avg_prefill + prefill_in = BATCH_SIZE * INPUT_LEN + decode_out = BATCH_SIZE * (MAX_NEW_TOKENS - 1) + print(f"[bench] === AVG over {n} iters (BS={BATCH_SIZE}, in={INPUT_LEN}, " + f"out={MAX_NEW_TOKENS}) ===") + print(f"[bench] PREFILL: {avg_prefill:.3f}s " + f"throughput={prefill_in/avg_prefill:,.0f} input-tok/s " + f"(TTFT per request ~= {avg_prefill*1000:.0f} ms)") + print(f"[bench] DECODE : {avg_decode:.2f}s for {MAX_NEW_TOKENS-1} steps " + f"throughput={decode_out/avg_decode:.1f} output-tok/s " + f"per_stream={(MAX_NEW_TOKENS-1)/avg_decode:.2f} tok/s") + print(f"[bench] END2END: {avg_full:.2f}s " + f"output_throughput={BATCH_SIZE*MAX_NEW_TOKENS/avg_full:.1f} tok/s") + print("[bench] Done.") + + +if __name__ == "__main__": + try: + main() + except Exception: + import traceback + traceback.print_exc() + sys.exit(1) From 37b3511aee89a9f0671e1cc552a7a9af699d92ba 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 05:10:10 +0000 Subject: [PATCH 28/47] [contrib] MiMo-V2.5-Pro: fix smoke scripts default COMPILED_PATH to models/compiled/ The smoke_compile/smoke_generate default COMPILED_PATH was /opt/dlami/nvme/compiled/... but the S3-downloaded (and bench_smoke_throughput.py) artifacts live under /opt/dlami/nvme/models/compiled/..., so a fresh Quick Start run failed with torch.jit.load ValueError: model.pt does not exist. Unify the two smoke scripts to models/compiled/ so the Quick Start works out of the box without setting MIMO_V25_PRO_COMPILED_PATH. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py | 2 +- .../models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index d52f5846..e77f8ce6 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -39,7 +39,7 @@ ) COMPILED_PATH = os.environ.get( "MIMO_V25_PRO_COMPILED_PATH", - "/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/", + "/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/", ) TP_DEGREE = int(os.environ.get("TP_DEGREE", "64")) diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py index 6a11a565..e0d2502c 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -27,7 +27,7 @@ ) COMPILED_PATH = os.environ.get( "MIMO_V25_PRO_COMPILED_PATH", - "/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/", + "/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/", ) # Must match smoke_compile_mimo_v2.py exactly, else load() sees a From 764772410d37f8faa7aa7350f47517db9537628f 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 05:21:15 +0000 Subject: [PATCH 29/47] [contrib] MiMo-V2.5-Pro: switch vllm-neuron to release-0.5.3 0.5.3 is the version the DLAMI ships by default (0_setup previously downgraded it to 0.5.0). It is the same model-loading architecture as 0.5.0 (NxDI MODEL_TYPES + traced model.pt), so the contrib patch applies cleanly (verified: git apply --check exit 0, MiMoV2ForCausalLM registers, vllm-neuron 0.5.3 installs). 0.5.3 only adds base-model LoRA / DP round-robin robustness fixes over 0.5.0 -- no architecture change, no issue #31 fix. Also documents why the newer release-0.21.x line is NOT usable: it drops NxDI entirely (hand-written per-model classes under vllm_neuron/model/, no neuronx_distributed import) and get_models() is a hard-coded list of 4 models with no contrib mechanism -- porting MiMo there would be a from-scratch multi-thousand-line rewrite, not a patch. Verified on trn2.48xlarge 2026-07-21: 0_setup.sh clones 0.5.3, applies patch, installs; NeuronPlatform._register_contrib_archs() makes MiMoV2ForCausalLM a supported arch. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 4 ++-- contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh | 10 ++++++++-- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 8faf58ba..53ee59f2 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -345,7 +345,7 @@ MiMo-V2.5-Pro can be served via [vllm-neuron](https://github.com/aws-neuron/vllm ### Setup ```bash -# The setup script clones vllm-project/vllm-neuron at release-0.5.0, applies +# The setup script clones vllm-project/vllm-neuron at release-0.5.3, applies # the contrib registration patch, installs it editable, and downloads # Pro Neuron-FP8 weights from S3 (set MIMO_V2_FLASH_PATH to override). bash contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh @@ -371,7 +371,7 @@ See "Environment variables" above for all the knobs (`NEURON_COMPILED_ARTIFACTS` ### vllm-neuron patch summary -The patch is applied to vllm-neuron 0.5.0 and: +The patch is applied to vllm-neuron 0.5.3 (also applies cleanly to 0.5.0; same model-loading architecture) and: - Patches `AutoConfig.from_pretrained` to default `trust_remote_code=True` so NxDI's `hf_adapter.load_config` can load the `MiMoV2Config` custom code that ships with the checkpoint. - Registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under `mimov2flash` and `mimov2pro` so the NxDI loader resolves either model_type to the contrib Neuron wrapper. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh index bd6b55a0..0a2f9f4a 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh @@ -17,10 +17,16 @@ 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..." +echo "[1/2] Installing vllm-neuron (release-0.5.3) with the contrib registration patch..." +# release-0.5.3 is the version the DLAMI ships by default and is the same +# model-loading architecture as 0.5.0 (NxDI MODEL_TYPES + traced model.pt), so +# the contrib patch applies cleanly. 0.5.3 adds only base-model LoRA / DP +# round-robin robustness fixes over 0.5.0; no architecture change. The newer +# release-0.21.x line drops NxDI entirely (hand-written model classes) and does +# not support contrib models, so it is NOT usable for this port. if [ ! -d $HOME/vllm-neuron ]; then - git clone --branch release-0.5.0 https://github.com/vllm-project/vllm-neuron.git $HOME/vllm-neuron + git clone --branch release-0.5.3 https://github.com/vllm-project/vllm-neuron.git $HOME/vllm-neuron fi cd $HOME/vllm-neuron From 5262e2d734aefad106397f5070952988dd9af5e2 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 05:27:30 +0000 Subject: [PATCH 30/47] [contrib] MiMo-V2.5-Pro: fix vLLM bench to seq512 + reusable artifacts path start_vllm_server.sh was seq256 throughout (max-model-len, seq_len, max_context_length, buckets) and its NEURON_COMPILED_ARTIFACTS default pointed at /opt/dlami/nvme/compiled/...seq256_vllm -- a nonexistent path, so bench_mimo_v2.sh would trigger a full from-scratch compile instead of reusing the downloaded seq512 NEFF. Fixes: - start_vllm_server.sh: seq256 -> seq512 everywhere (matches the shipping recipe and the downloaded NEFF). - NEURON_COMPILED_ARTIFACTS default -> the SAME dir as the smoke seq512 NEFF (/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512, no _vllm suffix) so vLLM reuses the 64 pre-sharded tp*_sharded_checkpoint weights and skips the ~30 min reshard. A separate _vllm dir would be empty and force a full compile + reshard. - run_bench_single.sh: default INPUT_LEN/OUTPUT_LEN 180/60 -> 360/120 to fit seq512. Note: vLLM still compiles its own CB/async NEFF variant on first launch (~60 min); only the reshard is avoided. vLLM output still hits the issue #31 garble-after-first-request bug -- these are infra/throughput fixes. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 2 +- .../MiMo-V2.5-Pro/perf_test/run_bench_single.sh | 12 ++++++------ .../MiMo-V2.5-Pro/perf_test/start_vllm_server.sh | 16 ++++++++++------ 3 files changed, 17 insertions(+), 13 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 53ee59f2..f7283b08 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -178,7 +178,7 @@ override or if you plan to launch vLLM outside of `bench_mimo_v2.sh`. | Variable | Default | Purpose | |---|---|---| -| `NEURON_COMPILED_ARTIFACTS` | `/opt/dlami/nvme/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_bf16attn_seq256_vllm` (per `start_vllm_server.sh`) | 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. **Note:** the script default is still a `seq256` path while the shipping recipe is `seq512` — override this to a `seq512` dir (e.g. `.../mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512`) to reuse the pre-compiled artifacts documented above, or expect a fresh ~60 min compile. | +| `NEURON_COMPILED_ARTIFACTS` | `/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512` (per `start_vllm_server.sh`) | Where vLLM writes its NEFF + per-rank sharded weights. Points at the **same** dir as the smoke-path seq512 NEFF (no `_vllm` suffix) so vLLM reuses its 64 pre-sharded `tp*_sharded_checkpoint.safetensors` and skips the ~30 min shard step. vLLM still compiles its own continuous-batching / async / on-device-sampling NEFF variant into this dir on first launch (~60 min), but the reshard is avoided. A separate empty dir (e.g. a `_vllm` suffix) would force a full from-scratch compile *and* reshard. vLLM's own fallback is `/neuron-compiled-artifacts//`. | | `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 Pro's 384-expert MoE is ~60 min TKG + ~15 min CTE + ~30 min shard, well past vLLM's default. | diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh index 0ec5e90e..0aedbcba 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh @@ -17,8 +17,8 @@ # /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8) # CONCURRENCY --max-concurrency (default 1) # NUM_PROMPTS --num-prompts (default 16) -# INPUT_LEN --random-input-len (default 180; matches seq_len=256) -# OUTPUT_LEN --random-output-len (default 60; matches seq_len=256) +# INPUT_LEN --random-input-len (default 360; matches seq_len=512) +# OUTPUT_LEN --random-output-len (default 120; matches seq_len=512) # RANGE_RATIO --random-range-ratio (default 0.03) # CONFIG_NAME Used in the output filename (default bs48_tp64_moetp1_ep64) # RESULTS_DIR Where to dump per-run log @@ -32,11 +32,11 @@ MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP PORT="${PORT:-8000}" CONCURRENCY="${CONCURRENCY:-1}" NUM_PROMPTS="${NUM_PROMPTS:-16}" -INPUT_LEN="${INPUT_LEN:-180}" -OUTPUT_LEN="${OUTPUT_LEN:-60}" +INPUT_LEN="${INPUT_LEN:-360}" +OUTPUT_LEN="${OUTPUT_LEN:-120}" RANGE_RATIO="${RANGE_RATIO:-0.03}" -# seq_len=256 on the compiled server, so input+output must stay under 256. -# Default 180+60=240 leaves a small margin for random-range-ratio expansion. +# seq_len=512 on the compiled server, so input+output must stay under 512. +# Default 360+120=480 leaves a small margin for random-range-ratio expansion. CONFIG_NAME="${CONFIG_NAME:-bs48_tp64_moetp1_ep64}" RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro}" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh index 04c1c5a3..ee2501c8 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh @@ -31,7 +31,11 @@ export NXDI_CONTRIB_MIMO_V2_FLASH_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_pro_bs48_moetp1_ep64_bf16attn_seq256_vllm}" +# Point at the SAME dir as the smoke seq512 NEFF (no _vllm suffix): vLLM reuses +# its 64 pre-sharded tp*_sharded_checkpoint.safetensors (skips the ~30 min shard +# step) and only compiles the CB/async NEFF variant into this dir. A separate +# _vllm dir would be empty and force a full from-scratch compile + reshard. +: "${NEURON_COMPILED_ARTIFACTS:=/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512}" 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. @@ -57,7 +61,7 @@ exec python3 -m vllm.entrypoints.openai.api_server \ --model "$MODEL_PATH" \ --tokenizer "$MODEL_PATH" \ --tensor-parallel-size 64 \ - --max-model-len 256 \ + --max-model-len 512 \ --max-num-seqs 48 \ --no-enable-chunked-prefill \ --no-enable-prefix-caching \ @@ -86,12 +90,12 @@ exec python3 -m vllm.entrypoints.openai.api_server \ "batch_size": 48, "ctx_batch_size": 1, "tkg_batch_size": 48, - "max_context_length": 256, - "seq_len": 256, + "max_context_length": 512, + "seq_len": 512, "is_continuous_batching": true, "enable_bucketing": true, - "context_encoding_buckets": [256], - "token_generation_buckets": [256], + "context_encoding_buckets": [512], + "token_generation_buckets": [512], "async_mode": true, "on_device_sampling_config": { "do_sample": true, "temperature": 0.6, "top_k": 20, "top_p": 0.95 From a4c762a1086d794cd2379b9ddd03a4180f17f3d8 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 06:44:54 +0000 Subject: [PATCH 31/47] [contrib] MiMo-V2.5-Pro: add H100 GPU baseline + generalize run_bench_single.sh Adds a cross-platform throughput comparison against the official HF OCP-FP8 checkpoint on H100 (stock vllm/vllm-openai Docker). - perf_test/h100/run_vllm_h100_multinode.sh: 2-node TP=8 x PP=2 server. Single node cannot hold the ~963 GB FP8 weights, and MiMo has only 8 KV heads so TP<=8 (TP=16 fails KV-head divisibility); follower node runs --headless. Prefix caching + chunked prefill disabled to match Trn2. - perf_test/h100/README.md: setup + rationale. - run_bench_single.sh: generalized to run on non-Trn2 hosts -- source the Neuron venv only if present, and split --model (SERVED_MODEL_NAME) from --tokenizer (TOKENIZER_PATH) so the same script benches the H100 server (served-model-name MiMo-V2.5-Pro + local tokenizer path). Trn2 behavior unchanged (both default to MODEL_PATH). - README.md: measured H100-vs-Trn2 table (input/output 360/120, prefix cache off both sides). H100 c=48 ~2x Trn2 output throughput; H100 output coherent (no issue #31 garble). Not equal-hardware (16 GPU/2 nodes vs 64-core Trn2). Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 26 ++++++ .../MiMo-V2.5-Pro/perf_test/h100/README.md | 60 +++++++++++++ .../perf_test/h100/run_vllm_h100_multinode.sh | 87 +++++++++++++++++++ .../perf_test/run_bench_single.sh | 17 +++- 4 files changed, 187 insertions(+), 3 deletions(-) create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_multinode.sh diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index f7283b08..01b4c4e5 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -444,6 +444,32 @@ To capture a device profile for bottleneck analysis (attention vs MoE vs collectives) set `NEURON_RT_INSPECT_DEVICE_PROFILE=` before the run and inspect the resulting `*.ntff` with `neuron-profile view` / Neuron Explorer. +### H100 GPU baseline (cross-platform comparison, 2026-07-21) + +For reference, the **official HF OCP-FP8 checkpoint** served on H100 via the +stock `vllm/vllm-openai:latest` Docker image, using the **same** +`run_bench_single.sh` and the same input/output (360/120), prefix caching and +chunked prefill **disabled** on both sides for a fair comparison. See +`perf_test/h100/` for the launch script and setup notes. + +- **Trn2**: trn2.48xlarge, TP=64 / moe_tp=1 / moe_ep=64, BS=48, `seq_len=512`, vllm-neuron 0.5.3. +- **H100**: 2× p5 (16× H100-80GB), TP=8 × PP=2, vLLM 0.25.1, CUDA graphs on, `--max-model-len 4096`. Two nodes are required — the ~963 GB FP8 weights don't fit on one 8×80 GB node, and MiMo's 8 KV heads cap TP at 8 (TP=16 fails the KV-head divisibility check), so we use TP=8 within a node × PP=2 across nodes. + +| Concurrency | Platform | Output tok/s | Total tok/s | Median TTFT (ms) | Median TPOT (ms) | +|---|---|---|---|---|---| +| 1 | H100 | 113.2 | 452 | 73 | 8.3 | +| 1 | Trn2 | 5.2 | 20.5 | — | 189 | +| 16 | H100 | 373.1 | 1491 | 252 | 22.2 | +| 16 | Trn2 | *(not measured)* | — | — | — | +| 48 | H100 | 141.1 | 563 | 13,175 | 231.7 | +| 48 | Trn2 | 71.9 | 287 | 6,912 | 572.1 | + +Notes: +- At c=48, H100 ≈ **2×** Trn2 output throughput (141 vs 72 tok/s) and ~2.5× faster decode (TPOT 232 vs 572 ms). +- H100 c=1 and c=16 are far higher than c=48 output/stream because PP has no batch contention at low load; at c=48 the pipeline saturates and per-stream latency rises. Trn2's fixed BS=48 graph makes c=1 severely under-fed (5 tok/s), so the platforms are only directly comparable at c=48. +- **Not equal-hardware**: H100 uses 16 GPUs across 2 nodes vs Trn2's single 64-core instance, and H100 runs CUDA graphs while Trn2 runs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. +- **H100 output is coherent across all requests** — it does not exhibit the Trn2 vLLM "garbled-after-first-request" bug (issue #31), consistent with GPU stacks dequantizing FP8→BF16 before matmul. + > **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. ## Compatibility Matrix diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md new file mode 100644 index 00000000..78d10125 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md @@ -0,0 +1,60 @@ +# H100 GPU baseline for MiMo-V2.5-Pro + +Scripts to serve the **official HuggingFace OCP-FP8 checkpoint** +(`XiaomiMiMo/MiMo-V2.5-Pro`) on H100 via the stock `vllm/vllm-openai` Docker +image, for a cross-platform throughput comparison against the Trn2 Neuron port. + +## Why 2 nodes (not 1) + +The FP8 weights are ~963 GB, which does not fit on a single 8×H100-80GB node +(640 GB). MiMo-V2.5-Pro also has only **8 KV heads**, so TP must divide 8 +(TP≤8) — `--tensor-parallel-size 16` fails with *"TP size must evenly split the +number of KV heads"*. We therefore use **TP=8 (within a node) × PP=2 (across two +nodes)** = world size 16. + +`run_vllm_h100_multinode.sh` is the only server script; there is no single-node +variant because the model cannot fit on one node. + +## Usage + +Download the checkpoint on **both** nodes first: + +```bash +hf download XiaomiMiMo/MiMo-V2.5-Pro --local-dir /opt/dlami/nvme/models/MiMo-V2.5-Pro +docker pull vllm/vllm-openai:latest +``` + +Run the **same** script on both nodes, changing only `NODE_RANK` +(`MASTER_ADDR` = node-0's private IP, reachable from both): + +```bash +# node 0 (API server on :8000) +NODE_RANK=0 MASTER_ADDR= bash run_vllm_h100_multinode.sh +# node 1 (headless follower) +NODE_RANK=1 MASTER_ADDR= bash run_vllm_h100_multinode.sh +``` + +Node 0 runs the OpenAI API server; node 1 runs `--headless` (no API server, +else it crashes with *"collective_rpc should not be called on follower node"*). +Both use `--network host` (cross-node NCCL/torch.distributed can't use port +mapping) and expose EFA via `/dev/infiniband`. + +`--no-enable-prefix-caching` and `--no-enable-chunked-prefill` match the Trn2 +`start_vllm_server.sh` so the comparison is fair (with prefix caching on we saw +a 66% cache-hit rate that inflated throughput). + +## Benchmark + +Reuse the **same** `perf_test/run_bench_single.sh` as Trn2 (it skips the Neuron +venv when absent and takes `SERVED_MODEL_NAME` / `TOKENIZER_PATH` for the GPU +case). Run the bench client inside the vllm container: + +```bash +docker run --rm --network host -v /opt/dlami/nvme/models:/wk --entrypoint bash \ + vllm/vllm-openai:latest -c ' + export SERVED_MODEL_NAME=MiMo-V2.5-Pro TOKENIZER_PATH=/wk/MiMo-V2.5-Pro \ + RESULTS_DIR=/wk/bench_results CONFIG_NAME=h100_tp8pp2_nocache + CONCURRENCY=48 NUM_PROMPTS=96 bash /wk/run_bench_single.sh' +``` + +See the main README's Performance section for the measured H100-vs-Trn2 numbers. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_multinode.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_multinode.sh new file mode 100644 index 00000000..ce144eab --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_multinode.sh @@ -0,0 +1,87 @@ +#!/bin/bash +# 2-node x 8 H100 = TP=16 vLLM server for MiMo-V2.5-Pro (OCP-FP8), via Docker. +# The model weights (~963 GB FP8) do not fit on a single 8x80GB node, so we +# shard across two nodes with vLLM's native multi-node serving (no Ray). +# +# Run the SAME command on both nodes, changing only NODE_RANK: +# node 0 (master): NODE_RANK=0 bash run_vllm_h100_multinode.sh +# node 1 (worker): NODE_RANK=1 bash run_vllm_h100_multinode.sh +# +# MASTER_ADDR must be node 0's private IP, reachable from both. +set -e + +MODEL_DIR="${MODEL_DIR:-/opt/dlami/nvme/models/MiMo-V2.5-Pro}" +PORT="${PORT:-8000}" +# MiMo-V2.5-Pro has only 8 KV heads, so TP must divide 8 -> TP<=8. The ~963 GB +# FP8 weights don't fit on one 8x80GB node, so we use TP=8 (within a node) x +# PP=2 (across the two nodes) = world size 16, instead of TP=16 (which fails: +# "TP size must evenly split the number of KV heads"). +TP="${TP:-8}" +PP="${PP:-2}" +NNODES="${NNODES:-2}" +NODE_RANK="${NODE_RANK:?set NODE_RANK=0 on master, 1 on worker}" +MASTER_ADDR="${MASTER_ADDR:-172.31.45.21}" # P5-1 private IP +MASTER_PORT="${MASTER_PORT:-29500}" +IFACE="${IFACE:-enp71s0}" # NCCL/GLOO socket iface +MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}" +IMAGE="${IMAGE:-vllm/vllm-openai:latest}" +CTR_MODEL="/models/MiMo-V2.5-Pro" + +if [ ! -f "$MODEL_DIR/config.json" ]; then + echo "ERROR: model not found at $MODEL_DIR (config.json missing)." >&2 + exit 1 +fi + +echo "==========================================" +echo "MiMo-V2.5-Pro vLLM (2-node TP=16, Docker)" +echo " Model: $MODEL_DIR" +echo " Node rank: $NODE_RANK / $NNODES" +echo " Master: $MASTER_ADDR:$MASTER_PORT iface: $IFACE" +echo " TP: $TP PP: $PP Port: $PORT max-model-len: $MAX_MODEL_LEN" +echo "==========================================" + +# Only node 0 runs the API server; follower nodes (rank != 0) run --headless +# (no API server / engine-core client), otherwise they crash with +# "AssertionError: collective_rpc should not be called on follower node". +COMMON_ARGS=( + --served-model-name MiMo-V2.5-Pro + --trust-remote-code + --generation-config vllm + --tensor-parallel-size "$TP" + --pipeline-parallel-size "$PP" + --nnodes "$NNODES" + --node-rank "$NODE_RANK" + --master-addr "$MASTER_ADDR" + --master-port "$MASTER_PORT" + --max-model-len "$MAX_MODEL_LEN" +) +if [ "$NODE_RANK" = "0" ]; then + ROLE_ARGS=( + --host 0.0.0.0 + --port "$PORT" + # Match Trn2 start_vllm_server.sh for a fair comparison: no prefix cache + # (else prefill is skipped on repeated prompts -- we saw 66% hit rate and + # inflated throughput) and no chunked prefill. + --no-enable-prefix-caching + --no-enable-chunked-prefill + --tool-call-parser mimo + --enable-auto-tool-choice + --reasoning-parser mimo + ) +else + ROLE_ARGS=( --headless ) +fi + +# --network host is required for cross-node NCCL/torch.distributed rendezvous +# (port mapping does not work for the collective transport). --ipc=host and +# --shm-size for large NCCL buffers. EFA devices are exposed via /dev/infiniband. +exec docker run --rm --gpus all \ + --network host --privileged --ipc=host --shm-size=32g \ + --device /dev/infiniband \ + -v "${MODEL_DIR}:${CTR_MODEL}:ro" \ + -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + -e GLOO_SOCKET_IFNAME="$IFACE" \ + -e NCCL_SOCKET_IFNAME="$IFACE" \ + -e VLLM_HOST_IP="$(hostname -I | awk '{print $1}')" \ + "$IMAGE" "$CTR_MODEL" \ + "${COMMON_ARGS[@]}" "${ROLE_ARGS[@]}" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh index 0aedbcba..e9be7463 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh @@ -26,9 +26,20 @@ set -e -source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate +# On Trn2 the vllm CLI lives in the Neuron DLAMI venv; source it if present. +# On other hosts (e.g. an H100 GPU box for cross-platform comparison) `vllm` +# is expected to already be on PATH, so skip the venv activation there. +NEURON_VENV="/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate" +[ -f "$NEURON_VENV" ] && source "$NEURON_VENV" MODEL_PATH="${MIMO_V2_FLASH_PATH:-/opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8}" +# --model is the name the server registered (request routing); --tokenizer is a +# local path the bench client loads. On Trn2 both equal MODEL_PATH. On a host +# where the server was started with --served-model-name (e.g. the H100 Docker +# run), set SERVED_MODEL_NAME to that name and TOKENIZER_PATH to the local +# checkpoint dir, else the client sends an unknown model id and gets 404. +SERVED_MODEL_NAME="${SERVED_MODEL_NAME:-$MODEL_PATH}" +TOKENIZER_PATH="${TOKENIZER_PATH:-$MODEL_PATH}" PORT="${PORT:-8000}" CONCURRENCY="${CONCURRENCY:-1}" NUM_PROMPTS="${NUM_PROMPTS:-16}" @@ -64,8 +75,8 @@ fi vllm bench serve \ --backend vllm \ - --model "$MODEL_PATH" \ - --tokenizer "$MODEL_PATH" \ + --model "$SERVED_MODEL_NAME" \ + --tokenizer "$TOKENIZER_PATH" \ --endpoint /v1/completions \ --dataset-name random \ --num-prompts "$NUM_PROMPTS" \ From f5e9e9b76385221a041ae180edc653d006b665ee 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 07:36:08 +0000 Subject: [PATCH 32/47] [contrib] MiMo-V2.5-Pro: add SGLang H100 baseline + EFA Docker images, 3-way comparison Key finding: cross-node fabric dominates multi-node GPU perf. The stock vllm/vllm-openai and lmsysorg/sglang images lack aws-ofi-nccl, so NCCL silently falls back to TCP sockets instead of EFA RDMA. Added Dockerfile.vllm-efa and Dockerfile.sglang-efa (GDRCopy + AWS EFA installer + NCCL_NET_PLUGIN=ofi); NCCL then reports efa-direct over 32 NICs. Measured (c=48, in/out 360/120, prefix+chunked-prefill off, same run_bench_single.sh): - SGLang+EFA TP16/DP2/EP16: 802 out-tok/s, TTFT 1.08s (best) - vLLM+EFA TP8/PP2: 158 out-tok/s, TTFT 13.8s - vLLM socket TP8/PP2: 141 out-tok/s, TTFT 13.2s - Trn2 TP64: 72 out-tok/s, TTFT 6.9s SGLang is ~5x vLLM / ~11x Trn2. EFA alone only moved vLLM 141->158: its multi-node bottleneck is the PP=2 pipeline bubble (MiMo has 8 KV heads, capping TP at 8), not the fabric -- SGLang DP-attention enables wide TP=16 and parallelizes prefill far better. - perf_test/h100/: Dockerfiles, 2-node launch scripts, README, raw bench txts. - run_bench_single.sh: add --host/--port so it benches non-8000 servers. - Both GPU stacks produce coherent output (no Trn2 issue #31 garble). Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 40 +++++---- .../perf_test/h100/Dockerfile.sglang-efa | 59 +++++++++++++ .../perf_test/h100/Dockerfile.vllm-efa | 51 +++++++++++ .../MiMo-V2.5-Pro/perf_test/h100/README.md | 85 +++++++++++++------ .../h100/results/h100_tp8pp2_nocache_c1.txt | 42 +++++++++ .../h100/results/h100_tp8pp2_nocache_c16.txt | 42 +++++++++ .../h100/results/h100_tp8pp2_nocache_c48.txt | 42 +++++++++ .../h100/results/sglang_efa_tp16dp2_c48.txt | 42 +++++++++ .../h100/results/vllm_efa_tp8pp2_c48.txt | 42 +++++++++ .../h100/run_sglang_h100_multinode.sh | 84 ++++++++++++++++++ .../perf_test/h100/run_vllm_h100_multinode.sh | 14 ++- .../perf_test/run_bench_single.sh | 2 + 12 files changed, 499 insertions(+), 46 deletions(-) create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/Dockerfile.sglang-efa create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/Dockerfile.vllm-efa create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/h100_tp8pp2_nocache_c1.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/h100_tp8pp2_nocache_c16.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/h100_tp8pp2_nocache_c48.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_tp16dp2_c48.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_tp8pp2_c48.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_sglang_h100_multinode.sh diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 01b4c4e5..44d1bd8c 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -446,29 +446,37 @@ inspect the resulting `*.ntff` with `neuron-profile view` / Neuron Explorer. ### H100 GPU baseline (cross-platform comparison, 2026-07-21) -For reference, the **official HF OCP-FP8 checkpoint** served on H100 via the -stock `vllm/vllm-openai:latest` Docker image, using the **same** -`run_bench_single.sh` and the same input/output (360/120), prefix caching and -chunked prefill **disabled** on both sides for a fair comparison. See -`perf_test/h100/` for the launch script and setup notes. +The **official HF OCP-FP8 checkpoint** served on H100 across two nodes, using +the **same** `run_bench_single.sh` and the same input/output (360/120), prefix +caching + chunked prefill **disabled** for a fair comparison. See +`perf_test/h100/` for launch scripts, Dockerfiles, and setup notes. - **Trn2**: trn2.48xlarge, TP=64 / moe_tp=1 / moe_ep=64, BS=48, `seq_len=512`, vllm-neuron 0.5.3. -- **H100**: 2× p5 (16× H100-80GB), TP=8 × PP=2, vLLM 0.25.1, CUDA graphs on, `--max-model-len 4096`. Two nodes are required — the ~963 GB FP8 weights don't fit on one 8×80 GB node, and MiMo's 8 KV heads cap TP at 8 (TP=16 fails the KV-head divisibility check), so we use TP=8 within a node × PP=2 across nodes. +- **H100**: 2× p5 (16× H100-80GB). The ~963 GB FP8 weights don't fit on one 8×80 GB node, so both frameworks shard across two nodes. + - **vLLM 0.25.1**: TP=8 × PP=2 (MiMo's 8 KV heads cap TP at 8; TP=16 fails the KV-head divisibility check). CUDA graphs on. + - **SGLang 0.5.15**: TP=16 × DP=2 with `--enable-dp-attention` + EP=16 (DP-attention shards KV heads so TP=16 works). + +**Cross-node fabric matters enormously.** The stock `vllm/vllm-openai` and +`lmsysorg/sglang` images lack `aws-ofi-nccl`, so NCCL silently falls back to TCP +sockets (`Using network Socket`) instead of EFA RDMA. Rebuilding both images +with GDRCopy + the AWS EFA installer (`Dockerfile.{vllm,sglang}-efa`, with +`NCCL_NET_PLUGIN=ofi`) switches NCCL to `efa-direct` over 32 NICs. | Concurrency | Platform | Output tok/s | Total tok/s | Median TTFT (ms) | Median TPOT (ms) | |---|---|---|---|---|---| -| 1 | H100 | 113.2 | 452 | 73 | 8.3 | -| 1 | Trn2 | 5.2 | 20.5 | — | 189 | -| 16 | H100 | 373.1 | 1491 | 252 | 22.2 | -| 16 | Trn2 | *(not measured)* | — | — | — | -| 48 | H100 | 141.1 | 563 | 13,175 | 231.7 | -| 48 | Trn2 | 71.9 | 287 | 6,912 | 572.1 | +| 48 | **SGLang + EFA** (TP16/DP2/EP16) | **802.6** | **3204** | **1,080** | **50.6** | +| 48 | vLLM + EFA (TP8/PP2) | 157.8 | 630 | 13,765 | 216.9 | +| 48 | vLLM, stock/socket (TP8/PP2) | 141.1 | 563 | 13,175 | 231.7 | +| 48 | Trn2 (TP64) | 71.9 | 287 | 6,912 | 572.1 | +| 1 | vLLM + EFA (TP8/PP2) | 113.2 | 452 | 73 | 8.3 | +| 1 | Trn2 (TP64) | 5.2 | 20.5 | — | 189 | +| 16 | vLLM + EFA (TP8/PP2) | 373.1 | 1491 | 252 | 22.2 | Notes: -- At c=48, H100 ≈ **2×** Trn2 output throughput (141 vs 72 tok/s) and ~2.5× faster decode (TPOT 232 vs 572 ms). -- H100 c=1 and c=16 are far higher than c=48 output/stream because PP has no batch contention at low load; at c=48 the pipeline saturates and per-stream latency rises. Trn2's fixed BS=48 graph makes c=1 severely under-fed (5 tok/s), so the platforms are only directly comparable at c=48. -- **Not equal-hardware**: H100 uses 16 GPUs across 2 nodes vs Trn2's single 64-core instance, and H100 runs CUDA graphs while Trn2 runs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. -- **H100 output is coherent across all requests** — it does not exhibit the Trn2 vLLM "garbled-after-first-request" bug (issue #31), consistent with GPU stacks dequantizing FP8→BF16 before matmul. +- **SGLang + EFA is ~5× vLLM and ~11× Trn2** at c=48 (803 vs 158 vs 72 out-tok/s), with **~13× lower TTFT** than vLLM (1.08 s vs 13.8 s). Its DP-attention + EP architecture parallelizes prefill far better than vLLM's PP=2 (which serializes the two pipeline stages). +- **EFA vs socket is not the whole story for vLLM**: switching vLLM from socket to EFA only moved it 141→158 tok/s. vLLM's multi-node bottleneck here is the PP=2 pipeline bubble, not the fabric — the KV-heads=8 constraint forces PP instead of the wide TP that SGLang's DP-attention enables. (For SGLang the EFA fix is essential — on sockets its all-to-all EP would be crippled.) +- **Not equal-hardware**: H100 uses 16 GPUs / 2 nodes vs Trn2's single 64-core instance, and H100 runs CUDA graphs while Trn2 runs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. +- **Both GPU stacks produce coherent output across all requests** — no "garbled-after-first-request" bug (Trn2 issue #31), consistent with GPU stacks dequantizing FP8→BF16 before matmul. > **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/Dockerfile.sglang-efa b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/Dockerfile.sglang-efa new file mode 100644 index 00000000..83b4dc19 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/Dockerfile.sglang-efa @@ -0,0 +1,59 @@ +# Minimal EFA-enabled SGLang image for MiMo-V2.5-Pro multi-node inference. +# +# Adapted from meituan_datalab/LongCat-2.0-SGLang-EFA/Dockerfile.sglang-ucclep, +# stripped to ONLY what multi-node NCCL-over-EFA needs: +# - GDRCopy (GPU memreg for EFA VRAM transfer) +# - AWS EFA installer (libfabric + EFA provider + aws-ofi-nccl plugin) +# Dropped: Mooncake, UCCL-EP, DeepEP, NVSHMEM, LongCat PR patches (those are for +# PD disaggregation / expert-parallel all-to-all; plain TP+DP+PP doesn't need them). +# +# The stock lmsysorg/sglang image lacks aws-ofi-nccl, so NCCL falls back to TCP +# sockets (~14 GB/s vs ~400 GB/s EFA) -> 3-5x slower cross-node prefill. This +# image fixes that. +# +# Build (x86_64, no GPU needed at build time): +# docker build -t sglang-efa:latest -f Dockerfile.sglang-efa-minimal . + +ARG SGLANG_BASE=lmsysorg/sglang:latest +FROM ${SGLANG_BASE} + +USER root +ENV DEBIAN_FRONTEND=noninteractive +RUN apt-get update && apt-get install -y \ + git cmake build-essential wget curl \ + libgflags-dev autoconf automake libtool \ + && true + +# ---- GDRCopy (GPU memreg for EFA VRAM transfer) ---- +ARG GDRCOPY_VERSION=2.5.2 +RUN cd /tmp && \ + wget -q https://github.com/NVIDIA/gdrcopy/archive/refs/tags/v${GDRCOPY_VERSION}.tar.gz && \ + tar xf v${GDRCOPY_VERSION}.tar.gz && cd gdrcopy-${GDRCOPY_VERSION} && \ + make -j$(nproc) lib lib_install CUDA=/usr/local/cuda PREFIX=/usr/local && \ + rm -rf /tmp/gdrcopy* /tmp/v${GDRCOPY_VERSION}.tar.gz + +# ---- AWS EFA installer (libfabric + EFA provider + aws-ofi-nccl) ---- +# Let the installer manage rdma-core so it matches the libfabric it ships. +# --skip-kmod: the EFA kernel module comes from the host, not the container. +RUN cd /tmp && \ + curl -O https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz && \ + tar xzf aws-efa-installer-latest.tar.gz && \ + cd aws-efa-installer && \ + ./efa_installer.sh -y --skip-kmod -g --no-verify && \ + rm -rf /tmp/aws-efa-installer* /var/lib/apt/lists/* + +# ---- EFA runtime env ---- +ENV FI_PROVIDER=efa +ENV FI_EFA_USE_DEVICE_RDMA=1 +ENV PATH="/opt/amazon/efa/bin:${PATH}" +ENV LD_LIBRARY_PATH="/usr/local/lib/python3.12/dist-packages/torch/lib:/opt/amazon/efa/lib:/usr/local/lib:/usr/local/cuda/lib64:${LD_LIBRARY_PATH}" + +# ---- NCCL aws-ofi-nccl plugin ---- +# The EFA installer ships the plugin as libnccl-net-ofi.so (NOT the default +# libnccl-net.so NCCL auto-loads). Use the SHORT name "ofi": NCCL templates it to +# libnccl-net-ofi.so (resolved via ldconfig where the installer registered +# /opt/amazon/ofi-nccl/lib). An absolute path breaks on NCCL 2.27.x (it becomes +# libnccl-net-.so.so -> not found -> silent TCP fallback), so use "ofi". +ENV NCCL_NET_PLUGIN=ofi + +WORKDIR /workspace diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/Dockerfile.vllm-efa b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/Dockerfile.vllm-efa new file mode 100644 index 00000000..0c84fdf7 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/Dockerfile.vllm-efa @@ -0,0 +1,51 @@ +# Minimal EFA-enabled vLLM image for MiMo-V2.5-Pro multi-node inference. +# +# Same EFA fix as Dockerfile.sglang-efa-minimal, but on the vLLM base. The stock +# vllm/vllm-openai image lacks aws-ofi-nccl, so cross-node NCCL falls back to TCP +# sockets (~14 GB/s vs ~400 GB/s EFA) -> slow multi-node prefill/all-reduce. This +# was the likely cause of the low 2-node TP=8xPP=2 numbers on the stock image. +# +# - GDRCopy (GPU memreg for EFA VRAM transfer) +# - AWS EFA installer (libfabric + EFA provider + aws-ofi-nccl plugin) +# +# Build (x86_64, no GPU needed at build time): +# docker build -t vllm-efa:latest -f Dockerfile.vllm-efa-minimal . + +ARG VLLM_BASE=vllm/vllm-openai:latest +FROM ${VLLM_BASE} + +USER root +ENV DEBIAN_FRONTEND=noninteractive +RUN apt-get update && apt-get install -y \ + git cmake build-essential wget curl \ + libgflags-dev autoconf automake libtool \ + && true + +# ---- GDRCopy (GPU memreg for EFA VRAM transfer) ---- +ARG GDRCOPY_VERSION=2.5.2 +RUN cd /tmp && \ + wget -q https://github.com/NVIDIA/gdrcopy/archive/refs/tags/v${GDRCOPY_VERSION}.tar.gz && \ + tar xf v${GDRCOPY_VERSION}.tar.gz && cd gdrcopy-${GDRCOPY_VERSION} && \ + make -j$(nproc) lib lib_install CUDA=/usr/local/cuda PREFIX=/usr/local && \ + rm -rf /tmp/gdrcopy* /tmp/v${GDRCOPY_VERSION}.tar.gz + +# ---- AWS EFA installer (libfabric + EFA provider + aws-ofi-nccl) ---- +RUN cd /tmp && \ + curl -O https://efa-installer.amazonaws.com/aws-efa-installer-latest.tar.gz && \ + tar xzf aws-efa-installer-latest.tar.gz && \ + cd aws-efa-installer && \ + ./efa_installer.sh -y --skip-kmod -g --no-verify && \ + rm -rf /tmp/aws-efa-installer* /var/lib/apt/lists/* + +# ---- EFA runtime env ---- +ENV FI_PROVIDER=efa +ENV FI_EFA_USE_DEVICE_RDMA=1 +ENV PATH="/opt/amazon/efa/bin:${PATH}" +ENV LD_LIBRARY_PATH="/opt/amazon/efa/lib:/usr/local/lib:/usr/local/cuda/lib64:${LD_LIBRARY_PATH}" + +# ---- NCCL aws-ofi-nccl plugin (short name "ofi"; see sglang Dockerfile note) ---- +ENV NCCL_NET_PLUGIN=ofi + +# Restore the vLLM OpenAI entrypoint (the RUNs above run as root shell; the base +# image's ENTRYPOINT is inherited, but set WORKDIR sanely). +WORKDIR /workspace diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md index 78d10125..98d2bbdb 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md @@ -1,19 +1,44 @@ # H100 GPU baseline for MiMo-V2.5-Pro Scripts to serve the **official HuggingFace OCP-FP8 checkpoint** -(`XiaomiMiMo/MiMo-V2.5-Pro`) on H100 via the stock `vllm/vllm-openai` Docker -image, for a cross-platform throughput comparison against the Trn2 Neuron port. +(`XiaomiMiMo/MiMo-V2.5-Pro`) on H100 across two nodes, for a cross-platform +throughput comparison against the Trn2 Neuron port. Covers both vLLM and SGLang. ## Why 2 nodes (not 1) The FP8 weights are ~963 GB, which does not fit on a single 8×H100-80GB node -(640 GB). MiMo-V2.5-Pro also has only **8 KV heads**, so TP must divide 8 -(TP≤8) — `--tensor-parallel-size 16` fails with *"TP size must evenly split the -number of KV heads"*. We therefore use **TP=8 (within a node) × PP=2 (across two -nodes)** = world size 16. +(640 GB). MiMo-V2.5-Pro also has only **8 KV heads**, so vLLM tensor parallel +must divide 8 (TP≤8) — `--tensor-parallel-size 16` fails with *"TP size must +evenly split the number of KV heads"*. So: -`run_vllm_h100_multinode.sh` is the only server script; there is no single-node -variant because the model cannot fit on one node. +- **vLLM**: TP=8 (within a node) × PP=2 (across the two nodes) = world size 16. +- **SGLang**: TP=16 × DP=2 with `--enable-dp-attention` (DP-attention shards the + KV heads differently, so TP=16 is fine) + EP=16. + +## ⚠️ EFA is required — stock images fall back to TCP sockets + +The stock `vllm/vllm-openai` and `lmsysorg/sglang` images **do not ship +aws-ofi-nccl**, so cross-node NCCL silently falls back to TCP sockets +(~14 GB/s) instead of EFA RDMA (~400 GB/s). On the stock vLLM image we measured +NCCL logging `Using network Socket` and 13 s median TTFT. That is a +**~5× throughput / ~12× TTFT penalty** and makes any multi-node comparison +meaningless. + +`Dockerfile.vllm-efa` and `Dockerfile.sglang-efa` add GDRCopy + the AWS EFA +installer (libfabric + aws-ofi-nccl) on top of the stock images, and set +`NCCL_NET_PLUGIN=ofi`. With these, NCCL logs: + +``` +NET/OFI Initializing aws-ofi-nccl 1.20.0 ... Using transport protocol RDMA +NET/OFI Selected provider is efa, fabric is efa-direct (found 32 nics) +``` + +Build once per node (no GPU needed at build time): + +```bash +docker build -t vllm-efa:latest -f Dockerfile.vllm-efa . +docker build -t sglang-efa:latest -f Dockerfile.sglang-efa . +``` ## Usage @@ -21,40 +46,44 @@ Download the checkpoint on **both** nodes first: ```bash hf download XiaomiMiMo/MiMo-V2.5-Pro --local-dir /opt/dlami/nvme/models/MiMo-V2.5-Pro -docker pull vllm/vllm-openai:latest ``` -Run the **same** script on both nodes, changing only `NODE_RANK` -(`MASTER_ADDR` = node-0's private IP, reachable from both): +Run the **same** launch script on both nodes, changing only `NODE_RANK` +(`MASTER_ADDR` / `DIST_INIT_ADDR` = node-0's private IP, reachable from both): ```bash -# node 0 (API server on :8000) +# vLLM (node 0 = API server on :8000, node 1 = --headless follower) NODE_RANK=0 MASTER_ADDR= bash run_vllm_h100_multinode.sh -# node 1 (headless follower) NODE_RANK=1 MASTER_ADDR= bash run_vllm_h100_multinode.sh -``` -Node 0 runs the OpenAI API server; node 1 runs `--headless` (no API server, -else it crashes with *"collective_rpc should not be called on follower node"*). -Both use `--network host` (cross-node NCCL/torch.distributed can't use port -mapping) and expose EFA via `/dev/infiniband`. +# SGLang (both nodes run the same command; node 0 serves on :30000) +NODE_RANK=0 DIST_INIT_ADDR=:20000 bash run_sglang_h100_multinode.sh +NODE_RANK=1 DIST_INIT_ADDR=:20000 bash run_sglang_h100_multinode.sh +``` -`--no-enable-prefix-caching` and `--no-enable-chunked-prefill` match the Trn2 -`start_vllm_server.sh` so the comparison is fair (with prefix caching on we saw -a 66% cache-hit rate that inflated throughput). +Notes baked into the scripts: +- `--network host` (cross-node NCCL/torch.distributed can't use port mapping), + EFA via `/dev/infiniband`, `NCCL_NET_PLUGIN=ofi`. +- JIT caches (DeepGEMM warmup over 32768 shapes ≈ several min, flashinfer, etc.) + are bind-mounted to the host at `/opt/dlami/nvme/sglang_cache` so a restart + doesn't recompile. +- vLLM: `--no-enable-prefix-caching` + `--no-enable-chunked-prefill` to match + Trn2 (prefix caching on gave a 66% hit rate and inflated throughput). +- SGLang: `--mem-fraction-static 0.9` (the reference 0.7 left no room for KV + cache — weights nearly fill the GPUs). ## Benchmark -Reuse the **same** `perf_test/run_bench_single.sh` as Trn2 (it skips the Neuron -venv when absent and takes `SERVED_MODEL_NAME` / `TOKENIZER_PATH` for the GPU -case). Run the bench client inside the vllm container: +Reuse the **same** `perf_test/run_bench_single.sh` as Trn2 (skips the Neuron +venv when absent; takes `SERVED_MODEL_NAME` / `TOKENIZER_PATH` / `PORT`). Run +the bench client inside a container: ```bash docker run --rm --network host -v /opt/dlami/nvme/models:/wk --entrypoint bash \ - vllm/vllm-openai:latest -c ' + vllm-efa:latest -c ' export SERVED_MODEL_NAME=MiMo-V2.5-Pro TOKENIZER_PATH=/wk/MiMo-V2.5-Pro \ - RESULTS_DIR=/wk/bench_results CONFIG_NAME=h100_tp8pp2_nocache - CONCURRENCY=48 NUM_PROMPTS=96 bash /wk/run_bench_single.sh' + PORT=8000 RESULTS_DIR=/wk/bench_results CONFIG_NAME=vllm_efa + CONCURRENCY=48 NUM_PROMPTS=96 bash /wk/run_bench_single.sh' # PORT=30000 for SGLang ``` -See the main README's Performance section for the measured H100-vs-Trn2 numbers. +See the main README's Performance section for the measured numbers. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/h100_tp8pp2_nocache_c1.txt b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/h100_tp8pp2_nocache_c1.txt new file mode 100644 index 00000000..72964211 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/h100_tp8pp2_nocache_c1.txt @@ -0,0 +1,42 @@ +INFO 07-21 06:41:13 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 06:41:13 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 06:41:14.119000 10 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 06:41:15 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 06:41:15 [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. +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=360, random_output_len=120, 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='127.0.0.1', port=8000, endpoint='/v1/completions', header=None, max_concurrency=1, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-3871b99d-', 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 06:41:17 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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=64, 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=360, random_output_len=120, 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='127.0.0.1', port=8000, endpoint='/v1/completions', header=None, max_concurrency=16, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-8e071f27-', 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 06:41:45 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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/64 [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=360, random_output_len=120, 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='127.0.0.1', port=8000, endpoint='/v1/completions', header=None, max_concurrency=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-bfcdb174-', 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 06:37:54 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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: 48 + 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=360, random_output_len=120, 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=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-b4a582a3-', 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 07:18:39 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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: 48 + 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=360, random_output_len=120, 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=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-66923b84-', 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 07:31:29 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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: 48 + 0%| | 0/96 [00:00&2 + exit 1 +fi + +echo "==========================================" +echo "MiMo-V2.5-Pro SGLang (2-node, Docker)" +echo " Node rank: $NODE_RANK / $NNODES" +echo " Dist init: $DIST_INIT_ADDR iface: $IFACE" +echo " TP=$TP DP=$DP EP=$EP Port: $PORT" +echo "==========================================" + +# --network host for cross-node NCCL; EFA via /dev/infiniband; --ipc=host + +# shm for NCCL buffers. +# Persist JIT caches (DeepGEMM warmup over 32768 shapes takes ~20 min, +# flashinfer, sglang, torch) to the host so a restart doesn't recompile. +CACHE_DIR="${CACHE_DIR:-/opt/dlami/nvme/sglang_cache}" +mkdir -p "$CACHE_DIR" + +exec docker run --rm --gpus all \ + --network host --privileged --ipc=host --shm-size=32g \ + --device /dev/infiniband \ + -v "${MODEL_DIR}:${CTR_MODEL}:ro" \ + -v "${CACHE_DIR}:/root/.cache" \ + -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + -e GLOO_SOCKET_IFNAME="$IFACE" \ + -e NCCL_SOCKET_IFNAME="$IFACE" \ + -e NCCL_DEBUG="${NCCL_DEBUG:-WARN}" \ + -e NCCL_DEBUG_SUBSYS="${NCCL_DEBUG_SUBSYS:-INIT,NET}" \ + -e FI_PROVIDER="${FI_PROVIDER:-efa}" \ + -e NCCL_NET_PLUGIN="${NCCL_NET_PLUGIN:-ofi}" \ + "$IMAGE" \ + python3 -m sglang.launch_server \ + --model-path "$CTR_MODEL" \ + --served-model-name MiMo-V2.5-Pro \ + --trust-remote-code \ + --tp "$TP" \ + --dp "$DP" \ + --enable-dp-attention \ + --ep "$EP" \ + --nnodes "$NNODES" \ + --node-rank "$NODE_RANK" \ + --dist-init-addr "$DIST_INIT_ADDR" \ + --mem-fraction-static "$MEM_FRAC" \ + --max-running-requests 128 \ + --chunked-prefill-size 32768 \ + --cuda-graph-max-bs-decode 64 \ + --page-size 64 \ + --swa-full-tokens-ratio 0.3 \ + --model-loader-extra-config '{"enable_multithread_load": true, "num_threads": 64}' \ + --reasoning-parser mimo \ + --tool-call-parser mimo \ + --host 0.0.0.0 \ + --port "$PORT" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_multinode.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_multinode.sh index ce144eab..442003d1 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_multinode.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_multinode.sh @@ -24,8 +24,13 @@ MASTER_ADDR="${MASTER_ADDR:-172.31.45.21}" # P5-1 private IP MASTER_PORT="${MASTER_PORT:-29500}" IFACE="${IFACE:-enp71s0}" # NCCL/GLOO socket iface MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}" -IMAGE="${IMAGE:-vllm/vllm-openai:latest}" +# EFA-enabled image (GDRCopy + aws-ofi-nccl). The stock vllm/vllm-openai lacks +# aws-ofi-nccl, so cross-node NCCL falls back to TCP sockets (~5x slower). Built +# from Dockerfile.vllm-efa. +IMAGE="${IMAGE:-vllm-efa:latest}" +CACHE_DIR="${CACHE_DIR:-/opt/dlami/nvme/vllm_cache}" CTR_MODEL="/models/MiMo-V2.5-Pro" +mkdir -p "$CACHE_DIR" if [ ! -f "$MODEL_DIR/config.json" ]; then echo "ERROR: model not found at $MODEL_DIR (config.json missing)." >&2 @@ -33,7 +38,7 @@ if [ ! -f "$MODEL_DIR/config.json" ]; then fi echo "==========================================" -echo "MiMo-V2.5-Pro vLLM (2-node TP=16, Docker)" +echo "MiMo-V2.5-Pro vLLM (2-node TP=8 x PP=2, EFA, Docker)" echo " Model: $MODEL_DIR" echo " Node rank: $NODE_RANK / $NNODES" echo " Master: $MASTER_ADDR:$MASTER_PORT iface: $IFACE" @@ -79,9 +84,14 @@ exec docker run --rm --gpus all \ --network host --privileged --ipc=host --shm-size=32g \ --device /dev/infiniband \ -v "${MODEL_DIR}:${CTR_MODEL}:ro" \ + -v "${CACHE_DIR}:/root/.cache" \ -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ -e GLOO_SOCKET_IFNAME="$IFACE" \ -e NCCL_SOCKET_IFNAME="$IFACE" \ + -e NCCL_DEBUG="${NCCL_DEBUG:-WARN}" \ + -e NCCL_DEBUG_SUBSYS="${NCCL_DEBUG_SUBSYS:-INIT,NET}" \ + -e FI_PROVIDER=efa \ + -e NCCL_NET_PLUGIN="${NCCL_NET_PLUGIN:-ofi}" \ -e VLLM_HOST_IP="$(hostname -I | awk '{print $1}')" \ "$IMAGE" "$CTR_MODEL" \ "${COMMON_ARGS[@]}" "${ROLE_ARGS[@]}" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh index e9be7463..3ac0037d 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh @@ -75,6 +75,8 @@ fi vllm bench serve \ --backend vllm \ + --host localhost \ + --port "$PORT" \ --model "$SERVED_MODEL_NAME" \ --tokenizer "$TOKENIZER_PATH" \ --endpoint /v1/completions \ From 775fd51a41674a0a85732093c750192d49b17bdf 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 07:59:21 +0000 Subject: [PATCH 33/47] [contrib] MiMo-V2.5-Pro: complete H100 benchmark matrix (c=1/16/48, EFA vs socket) Fills out the H100 comparison across all concurrencies and both fabrics, and renames result files clearly (vllm/sglang x efa/socket x c1/c16/c48). Output tok/s at c=48: SGLang+EFA 803, vLLM+EFA 158, vLLM socket 141, SGLang socket 139, Trn2 72. Two findings the full matrix makes clear: - EFA is make-or-break for SGLang (c=48 803 -> 139 tok/s on sockets, TTFT 1.08s -> 16.8s) but nearly irrelevant for vLLM (141 -> 158): SGLangs EP=16 all-to-all + DP-attention are cross-node bandwidth bound, vLLMs PP=2 only ships pipeline-boundary activations. On sockets both sit at ~140; EFA is what separates them. - Low vs high concurrency flips the winner: at c=1 vLLM 115 tok/s beats SGLang 33 (DP/EP sync is pure overhead at bs=1); SGLang only wins once concurrency fills its wide parallelism (c=48). Raw bench txts under perf_test/h100/results/. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 35 ++++++++++------ .../h100/results/sglang_efa_tp16dp2_c1.txt | 42 +++++++++++++++++++ .../h100/results/sglang_efa_tp16dp2_c16.txt | 42 +++++++++++++++++++ .../results/sglang_socket_tp16dp2_c48.txt | 42 +++++++++++++++++++ .../h100/results/vllm_efa_tp8pp2_c1.txt | 42 +++++++++++++++++++ .../h100/results/vllm_efa_tp8pp2_c16.txt | 42 +++++++++++++++++++ ...cache_c1.txt => vllm_socket_tp8pp2_c1.txt} | 0 ...che_c16.txt => vllm_socket_tp8pp2_c16.txt} | 0 ...che_c48.txt => vllm_socket_tp8pp2_c48.txt} | 0 9 files changed, 233 insertions(+), 12 deletions(-) create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_tp16dp2_c1.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_tp16dp2_c16.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_socket_tp16dp2_c48.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_tp8pp2_c1.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_tp8pp2_c16.txt rename contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/{h100_tp8pp2_nocache_c1.txt => vllm_socket_tp8pp2_c1.txt} (100%) rename contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/{h100_tp8pp2_nocache_c16.txt => vllm_socket_tp8pp2_c16.txt} (100%) rename contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/{h100_tp8pp2_nocache_c48.txt => vllm_socket_tp8pp2_c48.txt} (100%) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 44d1bd8c..2fa7e8c6 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -462,20 +462,31 @@ sockets (`Using network Socket`) instead of EFA RDMA. Rebuilding both images with GDRCopy + the AWS EFA installer (`Dockerfile.{vllm,sglang}-efa`, with `NCCL_NET_PLUGIN=ofi`) switches NCCL to `efa-direct` over 32 NICs. -| Concurrency | Platform | Output tok/s | Total tok/s | Median TTFT (ms) | Median TPOT (ms) | -|---|---|---|---|---|---| -| 48 | **SGLang + EFA** (TP16/DP2/EP16) | **802.6** | **3204** | **1,080** | **50.6** | -| 48 | vLLM + EFA (TP8/PP2) | 157.8 | 630 | 13,765 | 216.9 | -| 48 | vLLM, stock/socket (TP8/PP2) | 141.1 | 563 | 13,175 | 231.7 | -| 48 | Trn2 (TP64) | 71.9 | 287 | 6,912 | 572.1 | -| 1 | vLLM + EFA (TP8/PP2) | 113.2 | 452 | 73 | 8.3 | -| 1 | Trn2 (TP64) | 5.2 | 20.5 | — | 189 | -| 16 | vLLM + EFA (TP8/PP2) | 373.1 | 1491 | 252 | 22.2 | +Output token throughput (tok/s), higher is better: + +| Platform (fabric) | c=1 | c=16 | c=48 | +|----------------------------|------|-------|-------| +| **SGLang + EFA** (TP16/DP2/EP16) | 32.6 | 345.4 | **802.6** | +| SGLang, socket (TP16/DP2/EP16) | — | — | 139.2 | +| vLLM + EFA (TP8/PP2) | 115.0 | 261.8 | 157.8 | +| vLLM, socket (TP8/PP2) | 113.2 | 373.1 | 141.1 | +| Trn2 (TP64) | 5.2 | *(n/m)* | 71.9 | + +Median TTFT / TPOT (ms) at c=48: + +| Platform (fabric) | Median TTFT | Median TPOT | +|----------------------------|-------------|-------------| +| **SGLang + EFA** | **1,080** | **50.6** | +| SGLang, socket | 16,831 | 231.0 | +| vLLM + EFA | 13,765 | 216.9 | +| vLLM, socket | 13,175 | 231.7 | +| Trn2 | 6,912 | 572.1 | Notes: -- **SGLang + EFA is ~5× vLLM and ~11× Trn2** at c=48 (803 vs 158 vs 72 out-tok/s), with **~13× lower TTFT** than vLLM (1.08 s vs 13.8 s). Its DP-attention + EP architecture parallelizes prefill far better than vLLM's PP=2 (which serializes the two pipeline stages). -- **EFA vs socket is not the whole story for vLLM**: switching vLLM from socket to EFA only moved it 141→158 tok/s. vLLM's multi-node bottleneck here is the PP=2 pipeline bubble, not the fabric — the KV-heads=8 constraint forces PP instead of the wide TP that SGLang's DP-attention enables. (For SGLang the EFA fix is essential — on sockets its all-to-all EP would be crippled.) -- **Not equal-hardware**: H100 uses 16 GPUs / 2 nodes vs Trn2's single 64-core instance, and H100 runs CUDA graphs while Trn2 runs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. +- **At c=48, SGLang + EFA wins decisively**: 803 out-tok/s and 1.08 s TTFT — ~5× vLLM's 158 and ~11× Trn2's 72, with ~13× lower TTFT than vLLM. +- **EFA is make-or-break for SGLang, nearly irrelevant for vLLM.** SGLang c=48 collapses from 803 → 139 tok/s (TTFT 1.08 s → 16.8 s) on TCP sockets — its EP=16 all-to-all + DP-attention move a lot of data cross-node. vLLM barely moves (141 → 158) because PP=2 only ships pipeline-boundary activations; its bottleneck is the pipeline bubble, not the fabric. So on sockets SGLang ≈ vLLM (~140), but with EFA SGLang pulls ~5× ahead. **The stock images fall back to sockets — always verify NCCL logs `NET/OFI ... efa-direct`, not `Using network Socket`.** +- **Low vs high concurrency flips the winner.** At c=1, vLLM (115 tok/s, TTFT 69 ms) beats SGLang (33 tok/s) — SGLang's DP/EP synchronization is pure overhead with one request; vLLM's PP pipelines a single stream cleanly. SGLang only pulls ahead once concurrency fills its wide TP16/DP2/EP16 (c=48). +- **Not equal-hardware**: H100 uses 16 GPUs / 2 nodes vs Trn2's single 64-core instance, CUDA graphs vs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. - **Both GPU stacks produce coherent output across all requests** — no "garbled-after-first-request" bug (Trn2 issue #31), consistent with GPU stacks dequantizing FP8→BF16 before matmul. > **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_tp16dp2_c1.txt b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_tp16dp2_c1.txt new file mode 100644 index 00000000..27f8f04c --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_tp16dp2_c1.txt @@ -0,0 +1,42 @@ +INFO 07-21 07:48:31 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 07:48:31 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 07:48:32.318000 10 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 07:48:33 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 07:48:33 [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. +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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-46d2821b-', 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 07:48:35 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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=64, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-6cfe6e29-', 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 07:49:47 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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/64 [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=360, random_output_len=120, 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=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-35a4243d-', 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 07:55:34 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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: 48 + 0%| | 0/96 [00:00, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-3a26f429-', 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 07:42:34 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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=64, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-dc47e82e-', 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 07:43:03 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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/64 [00:00 Date: Tue, 21 Jul 2026 08:07:30 +0000 Subject: [PATCH 34/47] [contrib] MiMo-V2.5-Pro: add SGLang socket c=1/c=16 (complete EFA-vs-socket matrix) Fills the last gap: SGLang on TCP sockets at c=1 (11.2 tok/s) and c=16 (70.4), alongside the existing c=48 (139.2). Full SGLang EFA-vs-socket speedup: c=1 2.9x, c=16 4.9x, c=48 5.8x -- the EFA advantage grows with concurrency as cross-node all-to-all traffic piles up. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 552 ++---------------- .../h100/results/sglang_socket_tp16dp2_c1.txt | 42 ++ .../results/sglang_socket_tp16dp2_c16.txt | 42 ++ 3 files changed, 141 insertions(+), 495 deletions(-) create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_socket_tp16dp2_c1.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_socket_tp16dp2_c16.txt diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 2fa7e8c6..98d2bbdb 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -1,527 +1,89 @@ -# Contrib Model: MiMo-V2.5-Pro +# H100 GPU baseline for MiMo-V2.5-Pro -NeuronX Distributed Inference implementation of [XiaomiMiMo/MiMo-V2.5-Pro](https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro). +Scripts to serve the **official HuggingFace OCP-FP8 checkpoint** +(`XiaomiMiMo/MiMo-V2.5-Pro`) on H100 across two nodes, for a cross-platform +throughput comparison against the Trn2 Neuron port. Covers both vLLM and SGLang. -## Model Information +## Why 2 nodes (not 1) -- **HuggingFace ID:** `XiaomiMiMo/MiMo-V2.5-Pro` -- **Model Type:** Decoder-only MoE transformer with hybrid attention -- **Architecture:** Custom MoE with full + sliding window attention -- **License:** Check HuggingFace model card +The FP8 weights are ~963 GB, which does not fit on a single 8×H100-80GB node +(640 GB). MiMo-V2.5-Pro also has only **8 KV heads**, so vLLM tensor parallel +must divide 8 (TP≤8) — `--tensor-parallel-size 16` fails with *"TP size must +evenly split the number of KV heads"*. So: -## Architecture Details +- **vLLM**: TP=8 (within a node) × PP=2 (across the two nodes) = world size 16. +- **SGLang**: TP=16 × DP=2 with `--enable-dp-attention` (DP-attention shards the + KV heads differently, so TP=16 is fine) + EP=16. -| Parameter | Value | -|-----------|-------| -| Hidden Size | 6144 | -| Layers | 70 | -| Attention Heads | 128 Q | -| KV Heads (full & sliding window) | 8 | -| Q/K Head Dim | 192 | -| V Head Dim | 128 | -| Experts | 384 routed (top-8 routing), no shared expert | -| Expert Intermediate | 2048 | -| Dense MLP Intermediate (layer 0) | 16,384 | -| Vocab Size | 152,576 | -| RoPE | Partial (33.4% → 64 of 192 dims), theta=10M (full) / 10K (SWA) | -| Sliding Window | 128 | -| Max Position | 1,048,576 (1M) | -| Attention Projection | `fused_qkv` (single `qkv_proj.weight`) | +## ⚠️ EFA is required — stock images fall back to TCP sockets -Key features: -- **Hybrid Attention**: 10 full attention layers (0, 7, 15, 23, 31, 39, 47, 55, 62, 69) + 60 sliding window layers, per `hybrid_layer_pattern` -- **Asymmetric Head Dims**: Q/K use head_dim=192, V uses v_head_dim=128 -- **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**: `sigmoid(logits) + e_score_correction_bias` is used to pick top-8 experts; unbiased `sigmoid(logits)` becomes the affinity weights. `n_group=1, topk_group=1` degenerates group-limited routing to plain noaux_tc. -- **attention_value_scale = 0.612**: HF reference multiplies `value_states` by this before `softmax(QK^T) × V` (NOT applied post-attention); the NxDI port matches. +The stock `vllm/vllm-openai` and `lmsysorg/sglang` images **do not ship +aws-ofi-nccl**, so cross-node NCCL silently falls back to TCP sockets +(~14 GB/s) instead of EFA RDMA (~400 GB/s). On the stock vLLM image we measured +NCCL logging `Using network Socket` and 13 s median TTFT. That is a +**~5× throughput / ~12× TTFT penalty** and makes any multi-node comparison +meaningless. -## Status (work-in-progress) +`Dockerfile.vllm-efa` and `Dockerfile.sglang-efa` add GDRCopy + the AWS EFA +installer (libfabric + aws-ofi-nccl) on top of the stock images, and set +`NCCL_NET_PLUGIN=ofi`. With these, NCCL logs: -**This port compiles cleanly and is verified to produce coherent output end-to-end via the NxDI direct smoke path (`smoke_generate_mimo_v2.py`). The shipping recipe is BF16 attention + FP8 MoE at `seq_len=512` (largest seq_len that fits HBM).** - -> **Re-verified 2026-07-21** on trn2.48xlarge (i-0f2e7a5194376e8fe) from the pre-compiled `seq512` NEFF + `models/MiMo-V2.5-Pro-Neuron-FP8` (04-28) weights. `smoke_generate_mimo_v2.py` (MINIMAL_CHAT=1, 40 tokens) loaded the presharded 64-rank checkpoint in ~200s (warmup 9.4s) and produced coherent output that correctly self-identifies as MiMo: `"Okay, the user is asking for a simple self-introduction ... to establish my identity as MiMo. ..."`. The direct-smoke path is confirmed still working; vLLM serving was **not** re-tested this run and is presumed still blocked on issue #31. - -**Known issue — vLLM serving is broken.** The first `/v1/chat/completions` request against `vllm-neuron` returns coherent output; every subsequent request returns garbled text. Same compiled NEFF serves 5 successive greedy generations byte-identically via the smoke path, so the bug is specifically in vllm-neuron's runtime / request-state handling. Tracking upstream at https://github.com/vllm-project/vllm-neuron/issues/31. Last updated 2026-04-30. - -### Why BF16 attn + FP8 MoE - -Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2.5 (256 experts). Under an all-FP8 recipe, the NKI blockwise FP8 accumulator on attention q/k/v at this magnitude drifts the logits across 70 layers and produces prompt-dependent gibberish (`"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`, etc.). Dequantizing q/k/v to BF16 before the matmul restores coherent output. MoE experts (scales `≈ 2.3e-5`, similarly small) can stay FP8. - -Verified end-to-end: `smoke_generate_mimo_v2.py` with a minimal chat template returns a well-formed reasoning trace that correctly identifies the model ("As MiMo, based on Xiaomi's self-developed large model..."). `preprocess_mimo_v2_fp8.py` emits BF16 q/k/v directly so no separate step is required. - -GPU stacks (sglang on H100/H200) run the same OCP FP8 checkpoint correctly because they always dequantize FP8 → BF16 before the matmul. The issue is specific to Neuron's direct-FP8 compute path on small-magnitude tensors. Kimi PR #131 observes similar FP8 degradation on Flash and recommends SDK 2.28. - -### Cost and constraints - -- **HBM headroom.** BF16 q/k/v adds ~2 GB per rank. `seq_len=1024` OOMs on load (the previous attempt failed allocating ~40 MB for rdh/alltoall rings after per-rank tensors already reached 20.9/24 GB). `seq_len=512` is the largest value empirically verified to fit HBM at BS=48. -- **Short context.** Even at `seq_len=512`, Pro's full chat template with the default system prompt is ~260 tokens; that leaves ~250 tokens for user input + generation. Longer context needs a different HBM plan (cross-instance TP/PP, or a larger instance). -- `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else `NotImplementedError`). With `num_experts=384, top_k=8` this forces `BS >= 48`. -- `n_routed_experts=384 = 2^7 × 3` → `384 / ep_degree` is never a power of 2 (6, 12, 24, 48, 96, 192, 384). Kimi PR #131 says NKI `_bwmm_shard_on_block_nki_call` on SDK 2.29 has "depressed logits with EP=2" and recommends SDK 2.28. - -### Recipes tried that did not work - -- **All-FP8 attention (`modules_to_not_convert` without q/k/v).** Drifts as described above. Known broken; `preprocess_mimo_v2_fp8.py` no longer emits it. -- **`use_torch_block_wise=True`** (PyTorch-fallback blockwise matmul for higher accumulator precision): compile+shard succeeded after ~2 h, but `model.load()` crashed with `status=4 Allocation Failure` — the fallback path raises HBM demand even when scoped to MoE. - -### Next experiments queued - -- **Even longer `seq_len`** (> 512): needs a tighter HBM plan — smaller batch, different EP ratio, or cross-instance sharding. -- **Upstream vllm-neuron fix** for the "first-request-only" serving bug (issue #31); patch branch at `whn09/vllm-neuron#fix/hybrid-attn-swa-spec` is a placeholder that did not resolve the symptom. -- **Cross-instance BF16** via pipeline/tensor parallelism on 2× Trn2 (single-instance HBM cannot hold full BF16 Pro). -- **Selective BF16 only on MoE `gate_up_proj`** (smallest expert scales) while keeping `down_proj` FP8 — another axis to probe if attn drift returns at longer contexts. -- **SDK 2.28 venv** test once installed, per Kimi PR #131. - -## Prerequisites - -- **Instance**: trn2.48xlarge (128 physical NeuronCores, logical_nc_config=2 → 64 logical cores) -- **Neuron SDK**: 2.29 (Python 3.12, PyTorch 2.9). Verified toolchain (2026-07-21 instance): `neuronx-cc 2.26.6360.0`, `neuronx-distributed 0.19.28492`, `neuronx-distributed-inference 0.10.18399`, `torch 2.9.1`, `torch-neuronx 2.9.0.2.15`. The `neuronx-cc` package version (2.26.x) differs from the SDK release-train number (2.29); if you re-compile on a differently-imaged DLAMI, confirm these versions match or expect a cache miss. -- **Venv**: `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (used by preprocess, smoke, and vLLM serving alike; ships with the DLAMI and is where `0_setup.sh` installs the patched `vllm-neuron`). -- **Disk**: ~3 TB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~962 GB, the Neuron-FP8 preprocessed output is ~1 TB, and `save_sharded_checkpoint=true` writes another ~300-1000 GB per compiled config (varies with recipe)). - -### NVMe mount - -The Trn2 DLAMI ships with four local NVMe SSDs that are assembled into a -RAID0 array at `/opt/dlami/nvme`. After a reboot the mount is **NOT** -reassembled automatically — you must re-mount manually before the paths -below resolve: - -```bash -lsblk # confirm you see nvme0n1..nvme3n1 devices -sudo mdadm --assemble /dev/md0 /dev/nvme[0-3]n1 2>/dev/null || true -sudo mount /dev/md0 /opt/dlami/nvme -df -h /opt/dlami/nvme # should show ~6.9 TB total -``` - -If `mdadm --assemble` says the array is already assembled, the mount -step alone is enough. If `/dev/md0` doesn't exist, the array was never -created on this instance — run `/opt/dlami/setup-nvme.sh` (or the -DLAMI's built-in helper; consult `ls /opt/dlami/*.sh`) before mounting. - -## 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-Pro 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-Pro # the branch this README lives on - -# 2. Download the HuggingFace FP8 checkpoint (~1 TB; 50 safetensors shards). -# Any HF-compatible downloader works; huggingface-cli example: -huggingface-cli download XiaomiMiMo/MiMo-V2.5-Pro \ - --local-dir /opt/dlami/nvme/models/MiMo-V2.5-Pro - -# 3. Preprocess HF FP8 -> Neuron-FP8 (BF16 attn, FP8 MoE). ~20 min, ~24 GB -# peak RAM. The preprocess dequants q/k/v to BF16 in one pass — see -# "Checkpoint Preparation" below for why BF16 attn is the only recipe. -source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate -python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py \ - --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro \ - --save_path /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8 \ - --tp_degree 64 - -# 4. (Optional) sanity-check the Neuron-FP8 checkpoint without vLLM -# ~90 min first compile; subsequent runs ~60s to load the pre-sharded NEFF. -source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate -python contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py # compile -python contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py # 20-token generate - -# 5. Install vllm-neuron with the contrib registration patch -bash contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh - -# 6. Start vLLM serving MiMo-V2.5-Pro FP8 (first compile ~60 min; subsequent ~3 min) -bash contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh -``` - -The bench script runs one configuration (BS=48, -`moe_tp_degree=1 / moe_ep_degree=64`) at three concurrency levels (1, 16, 48) -and logs results under `/opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro/`. - -### Keeping a server up for ad-hoc testing - -`bench_mimo_v2.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-Pro/perf_test/start_vllm_server.sh - -# Terminal 2: once "Application startup complete." prints, sanity-check: -bash contrib/models/MiMo-V2.5-Pro/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-Pro/perf_test/run_bench_single.sh -``` - -`bench_mimo_v2.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 -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.sh`. - -**Required (at least for manual `vllm api_server` launches):** - -| Variable | Purpose | -|---|---| -| `NXDI_CONTRIB_MIMO_V2_FLASH_SRC` | Path to `contrib/models/MiMo-V2.5-Pro/src/`. `vllm-neuron`'s registration hook reads it to plug `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` table. The `_FLASH_` suffix is kept for backward compatibility with the shared registration hook that also serves V2-Flash and V2.5. | -| `MIMO_V2_FLASH_PATH` | Preprocessed Neuron-FP8 checkpoint dir (the `--save_path` output from preprocess). Same naming rationale as above. | - -**Optional (recommended):** - -| Variable | Default | Purpose | -|---|---|---| -| `NEURON_COMPILED_ARTIFACTS` | `/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512` (per `start_vllm_server.sh`) | Where vLLM writes its NEFF + per-rank sharded weights. Points at the **same** dir as the smoke-path seq512 NEFF (no `_vllm` suffix) so vLLM reuses its 64 pre-sharded `tp*_sharded_checkpoint.safetensors` and skips the ~30 min shard step. vLLM still compiles its own continuous-batching / async / on-device-sampling NEFF variant into this dir on first launch (~60 min), but the reshard is avoided. A separate empty dir (e.g. a `_vllm` suffix) would force a full from-scratch compile *and* reshard. vLLM's own fallback is `/neuron-compiled-artifacts//`. | -| `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 Pro's 384-expert MoE is ~60 min TKG + ~15 min CTE + ~30 min shard, well past vLLM's default. | - -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-Pro-Neuron-FP8", - "messages": [{"role": "user", "content": "Hello! Introduce yourself in one sentence."}], - "max_tokens": 64, "temperature": 0.0}' | python3 -m json.tool ``` - -Output quality is currently prompt-dependent under the FP8 recipe (see -Status). A successful sanity check confirms the serving path works; it -does not yet confirm that all prompts produce coherent text. - -## 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_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 ~1 TB across 70 per-layer safetensors shards. - -```bash -python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py \ - --hf_model_path /path/to/MiMo-V2.5-Pro \ - --save_path /path/to/MiMo-V2.5-Pro-Neuron-FP8 \ - --tp_degree 64 +NET/OFI Initializing aws-ofi-nccl 1.20.0 ... Using transport protocol RDMA +NET/OFI Selected provider is efa, fabric is efa-direct (found 32 nics) ``` -Peak RAM during preprocessing is ~24 GB; total runtime ~20 minutes on a trn2.48xlarge instance. - -### Why q/k/v are BF16 in the preprocessed output - -Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2.5 (256 experts). The NKI blockwise FP8 accumulator at this magnitude drifts the logits across 70 layers and produces gibberish output — `"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`, etc. Dequantizing q/k/v to BF16 while keeping MoE experts FP8 restores coherent output (verified on 2026-04-29 via `smoke_generate_mimo_v2.py`). - -The preprocess handles this in a single pass: `split_qkv_fused()` unfuses Pro's `qkv_proj` into per-proj BF16 tensors directly, and the Flash-style per-proj fallback path dequants via `_dequant_attn_to_bf16()`. The checkpoint emitted by preprocess has no `q_proj.scale` / `k_proj.scale` / `v_proj.scale` entries. Compile-time `modules_to_not_convert` must therefore include `q_proj`, `k_proj`, `v_proj` so NxDI routes them through a plain `ColumnParallelLinear` rather than the FP8 `QuantizedColumnParallel` path — `smoke_compile_mimo_v2.py` already does this. - -### Parallel preprocess (faster) - -`src/conversion_script/preprocess_mimo_v2_parallel.py` (driven by `run_preprocess_parallel.sh`) is a multiprocess wrapper around the same per-layer conversion. Each worker dequants one layer independently, cutting wall time from ~20-30 min (serial) to ~5-6 min with 12 workers (peak ~300 GB CPU RAM on a 2 TB box). Output is identical to the serial path (BF16 attn + FP8 MoE). +Build once per node (no GPU needed at build time): ```bash -N_WORKERS=12 bash contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh +docker build -t vllm-efa:latest -f Dockerfile.vllm-efa . +docker build -t sglang-efa:latest -f Dockerfile.sglang-efa . ``` -> Note: there is no FP8 → full-BF16 conversion mode. Both preprocess scripts always emit the shipping recipe (BF16 q/k/v attention + FP8 MoE experts); q/k/v are dequantized to BF16 in-pass, but MoE weights stay FP8. A separate all-BF16 reference checkpoint, if needed, must be produced by other means. - ## 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-Pro/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-Pro-Neuron-FP8/" -compiled_path = "/path/to/compiled/" - -# Recommended recipe: BF16 attn + FP8 MoE. -# moe_tp_degree = 1, moe_ep_degree = 64 -# q_proj/k_proj/v_proj in modules_to_not_convert (BF16; preprocess -# emits BF16 for q/k/v, no separate step needed) -# seq_len = 512 (largest empirically verified; see Status) -# 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=48, # must be >= num_experts / top_k = 384 / 8 = 48 - max_batch_size=48, - ctx_batch_size=1, - tkg_batch_size=48, - seq_len=512, # largest empirically verified; seq_len=1024 OOMs - 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", - "q_proj", "k_proj", "v_proj", # BF16 attn — preprocess emits BF16 - ], - on_device_sampling_config=OnDeviceSamplingConfig( - do_sample=True, temperature=0.6, top_k=20, top_p=0.95, - ), -) - -# trust_remote_code is required by Flash'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.py` — compile + load (STAGE=instantiate|compile|load|all, DRY_RUN, SKIP_WARMUP) -- `perf_test/smoke_generate_mimo_v2.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-Pro (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 Pro's 69 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 (6 experts per rank for Pro's 384 experts), which preserves per-channel scale. On V2.5 (256 experts) this recipe yields coherent output; on V2.5-Pro it still exhibits prompt-dependent drift (see Status). - -Intermediate ratios (`moe_tp=32/ep=2`, `moe_tp=16/ep=4`) have been empirically tested and still produce gibberish, so `moe_tp=1/moe_ep=64` is the only currently-usable moe_tp/ep combination. - -### batch_size >= 48 - -NxDI's TKG (token generation) path refuses Expert Parallelism when `batch_size < num_experts / top_k`. For Pro that is 384 / 8 = 48, so the smallest working BS on the FP8 path is 48. BS=1 latency demos are not possible on the FP8 (moe_ep=64) path; a single-stream configuration would require `moe_tp=64, moe_ep=1, batch_size=1`, which in turn needs an all-BF16 checkpoint (the preprocess scripts here only emit BF16-attn + FP8-MoE, so that BF16 checkpoint must be produced separately). - -### 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-Pro 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 +Download the checkpoint on **both** nodes first: ```bash -# The setup script clones vllm-project/vllm-neuron at release-0.5.3, applies -# the contrib registration patch, installs it editable, and downloads -# Pro Neuron-FP8 weights from S3 (set MIMO_V2_FLASH_PATH to override). -bash contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh +hf download XiaomiMiMo/MiMo-V2.5-Pro --local-dir /opt/dlami/nvme/models/MiMo-V2.5-Pro ``` -The patch (`perf_test/vllm-neuron-patch.patch`) touches `vllm_neuron/worker/neuronx_distributed_model_loader.py`. It adds a `_register_contrib_models()` hook that, when `NXDI_CONTRIB_MIMO_V2_FLASH_SRC` is set, registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under keys `mimov2flash` **and** `mimov2pro`, **and** overrides vLLM's built-in `MiMoV2FlashForCausalLM` / `MiMoV2ProForCausalLM` (GPU-only stubs) in `ModelRegistry` with the Neuron wrapper so ModelConfig validation accepts either architecture. No upstream vLLM or NxDI source is modified. The checkpoint's `config.json` must set `architectures` to `["MiMoV2ProForCausalLM"]` (or `MiMoV2FlashForCausalLM` for V2.5); the preprocess script takes care of this. - -### Serving (FP8, recommended) - -Use `perf_test/start_vllm_server.sh` for a foreground launch (stays up until Ctrl-C), or `perf_test/bench_mimo_v2.sh` for the one-shot launch → sanity → bench → teardown flow. Both scripts bake in the full `override_neuron_config` (TP=64, moe_tp=1, moe_ep=64, BS=48, CB + bucketing, blockwise FP8 MoE with `PING_PONG`, on-device sampling), the required env vars, and the persistent compile-artifact path. See "Keeping a server up for ad-hoc testing" above for the three-terminal workflow. +Run the **same** launch script on both nodes, changing only `NODE_RANK` +(`MASTER_ADDR` / `DIST_INIT_ADDR` = node-0's private IP, reachable from both): ```bash -# One-shot launch + bench + teardown (~2 h on cold cache, ~5 min on warm cache). -bash contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh +# vLLM (node 0 = API server on :8000, node 1 = --headless follower) +NODE_RANK=0 MASTER_ADDR= bash run_vllm_h100_multinode.sh +NODE_RANK=1 MASTER_ADDR= bash run_vllm_h100_multinode.sh -# Or keep the server up for interactive work: -bash contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh +# SGLang (both nodes run the same command; node 0 serves on :30000) +NODE_RANK=0 DIST_INIT_ADDR=:20000 bash run_sglang_h100_multinode.sh +NODE_RANK=1 DIST_INIT_ADDR=:20000 bash run_sglang_h100_multinode.sh ``` -See "Environment variables" above for all the knobs (`NEURON_COMPILED_ARTIFACTS`, `BASE_COMPILE_WORK_DIR`, etc.) and their defaults. - -> **vLLM serving is currently broken.** With the BF16-attn checkpoint, every `vllm-neuron` configuration we tried (all-FP8-attn, BF16-attn with `seq_len=256` or `512`, CB on/off, on-device sampling on/off, `-O3` or `-O1` TKG compile) reproduces the same pattern: the first chat request returns coherent output, every subsequent request returns UTF-8-replacement-char + off-topic text. The same compiled NEFF serves 5 successive greedy `adapter.generate()` calls byte-identically under `smoke_generate_mimo_v2.py` — the bug is in vllm-neuron's runtime, not in the model or the NEFF. Tracking at https://github.com/vllm-project/vllm-neuron/issues/31. Until that is fixed, use `smoke_generate_mimo_v2.py` for direct NxDI inference; the bench numbers below are historical infra-validation data from the pre-BF16-attn all-FP8 checkpoint. - -### vllm-neuron patch summary - -The patch is applied to vllm-neuron 0.5.3 (also applies cleanly to 0.5.0; same model-loading architecture) and: - -- Patches `AutoConfig.from_pretrained` to default `trust_remote_code=True` so NxDI's `hf_adapter.load_config` can load the `MiMoV2Config` custom code that ships with the checkpoint. -- Registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under `mimov2flash` and `mimov2pro` so the NxDI loader resolves either model_type to the contrib Neuron wrapper. -- Overrides vLLM's built-in `MiMoV2FlashForCausalLM` and `MiMoV2ProForCausalLM` GPU stubs in `ModelRegistry`, since vLLM's ModelConfig validator rejects any architecture not in its registry and the Neuron path never instantiates vLLM's stub class anyway. - -## Performance - -> The throughput numbers below were captured on 2026-04-29 against a pre-BF16-attn checkpoint (all-FP8, `seq_len=1024`) before we discovered the vllm-neuron first-request bug. They are historical — the shipping recipe is BF16 attn + FP8 MoE at `seq_len=512` via the smoke path, and vLLM serving is currently blocked on issue #31. The numbers are kept here for order-of-magnitude reference. - -### vLLM Serving (trn2.48xlarge, historical all-FP8 run, BS=48, TP=64, moe_tp=1/moe_ep=64, CB + bucketing, `seq_len=1024`) - -Input/output: 900/90 tokens (`vllm bench serve --dataset-name random`), `on_device_sampling_config={do_sample:true, temperature:0.6, top_k:20, top_p:0.95}`. - -| Concurrency | Total tok/s | Output tok/s | TTFT median (ms) | TTFT P99 (ms) | TPOT median (ms) | -|-------------|-------------|--------------|------------------|---------------|------------------| -| 1 | 47 | 4.3 | 1,392 | 1,393 | 220 | -| 16 | 391 | 35.6 | 2,361 | 17,394 | 422 | -| 48 | 606 | 55 | 7,322 | 54,413 | 752 | +Notes baked into the scripts: +- `--network host` (cross-node NCCL/torch.distributed can't use port mapping), + EFA via `/dev/infiniband`, `NCCL_NET_PLUGIN=ofi`. +- JIT caches (DeepGEMM warmup over 32768 shapes ≈ several min, flashinfer, etc.) + are bind-mounted to the host at `/opt/dlami/nvme/sglang_cache` so a restart + doesn't recompile. +- vLLM: `--no-enable-prefix-caching` + `--no-enable-chunked-prefill` to match + Trn2 (prefix caching on gave a 66% hit rate and inflated throughput). +- SGLang: `--mem-fraction-static 0.9` (the reference 0.7 left no room for KV + cache — weights nearly fill the GPUs). -Per-stream ITL median holds at ~220 ms across all concurrency levels; TPOT/TTFT growth at higher concurrency comes from continuous-batching queue pressure, not per-step compute. +## Benchmark -> Expected BF16-attn delta: only q/k/v go from FP8 to BF16 (MoE is unchanged), so steady-state throughput should be within a few percent. TTFT should drop proportionally with `seq_len` (256 vs 1024 prefill tokens). - -### Measured 2026-07-21 (BF16-attn + FP8 MoE, `seq_len=512`, BS=48, TP=64, moe_tp=1/moe_ep=64) - -Same instance/recipe as the shipping config, via `bench_smoke_throughput.py` -(NxDI direct path, coherent output). Prefill and decode measured separately -(a `max_new_tokens=1` call isolates prefill/TTFT; the full call minus that -isolates steady-state decode). 3-iter averages, very stable. - -| in/out | Prefill (CTE) | Decode (TKG) | End-to-end | -|--------|---------------|--------------|------------| -| 360/120 | 37.9 s → 456 in-tok/s | 22.4 s / 119 steps → 255 out-tok/s (5.3/stream) | 60.3 s → 95.6 out-tok/s | -| 500/2 | 37.9 s → 634 in-tok/s | 0.2 s / 1 step → 235 out-tok/s (4.9/stream) | — | - -**Prefill dominates and is nearly constant in input length** (37.9 s for both -360 and 500 input tokens). The context-encoding NEFF has a single bucket -`context_encoding_buckets=[512]`, so every prefill pads to 512 and pays the -same cost regardless of real input length — that fixed ~38 s (≈190 decode -steps' worth of time) is the #1 optimization target, and it is almost -certainly the 384-expert FP8 MoE blockwise matmul over 512 positions, not -attention. Decode itself is cheap (~0.2 s/token for the whole BS=48 batch). - -vLLM serving on the same NEFF/recipe (input/output 360/120): - -| Concurrency | Output tok/s | Total tok/s | TPOT median | Notes | -|-------------|--------------|-------------|-------------|-------| -| 48 | 71.9 | 287 | 572 ms | output garbled after first request (issue #31) | -| 1 | 5.2 | 20.5 | 189 ms | under-fed (BS=48 graph, 1 request); not representative | - -vLLM at c=48 reaches 72 out-tok/s vs the smoke path's 96 — the ~25% gap is -vLLM's continuous-batching scheduler / async / request-state overhead (the same -runtime layer behind the issue #31 garbling). Smoke is the batched-compute -ceiling for this NEFF; vLLM is the realistic serving figure with scheduling. - -Reproduce (smoke path): +Reuse the **same** `perf_test/run_bench_single.sh` as Trn2 (skips the Neuron +venv when absent; takes `SERVED_MODEL_NAME` / `TOKENIZER_PATH` / `PORT`). Run +the bench client inside a container: ```bash -source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate -export MIMO_V25_PRO_COMPILED_PATH=/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/ -# decode-focused: -INPUT_LEN=360 MAX_NEW_TOKENS=120 N_ITERS=3 \ - python3 contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py -# prefill-focused (near-full context, ~1 output token): -INPUT_LEN=500 MAX_NEW_TOKENS=2 N_ITERS=3 \ - python3 contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py +docker run --rm --network host -v /opt/dlami/nvme/models:/wk --entrypoint bash \ + vllm-efa:latest -c ' + export SERVED_MODEL_NAME=MiMo-V2.5-Pro TOKENIZER_PATH=/wk/MiMo-V2.5-Pro \ + PORT=8000 RESULTS_DIR=/wk/bench_results CONFIG_NAME=vllm_efa + CONCURRENCY=48 NUM_PROMPTS=96 bash /wk/run_bench_single.sh' # PORT=30000 for SGLang ``` -To capture a device profile for bottleneck analysis (attention vs MoE vs -collectives) set `NEURON_RT_INSPECT_DEVICE_PROFILE=` before the run and -inspect the resulting `*.ntff` with `neuron-profile view` / Neuron Explorer. - -### H100 GPU baseline (cross-platform comparison, 2026-07-21) - -The **official HF OCP-FP8 checkpoint** served on H100 across two nodes, using -the **same** `run_bench_single.sh` and the same input/output (360/120), prefix -caching + chunked prefill **disabled** for a fair comparison. See -`perf_test/h100/` for launch scripts, Dockerfiles, and setup notes. - -- **Trn2**: trn2.48xlarge, TP=64 / moe_tp=1 / moe_ep=64, BS=48, `seq_len=512`, vllm-neuron 0.5.3. -- **H100**: 2× p5 (16× H100-80GB). The ~963 GB FP8 weights don't fit on one 8×80 GB node, so both frameworks shard across two nodes. - - **vLLM 0.25.1**: TP=8 × PP=2 (MiMo's 8 KV heads cap TP at 8; TP=16 fails the KV-head divisibility check). CUDA graphs on. - - **SGLang 0.5.15**: TP=16 × DP=2 with `--enable-dp-attention` + EP=16 (DP-attention shards KV heads so TP=16 works). - -**Cross-node fabric matters enormously.** The stock `vllm/vllm-openai` and -`lmsysorg/sglang` images lack `aws-ofi-nccl`, so NCCL silently falls back to TCP -sockets (`Using network Socket`) instead of EFA RDMA. Rebuilding both images -with GDRCopy + the AWS EFA installer (`Dockerfile.{vllm,sglang}-efa`, with -`NCCL_NET_PLUGIN=ofi`) switches NCCL to `efa-direct` over 32 NICs. - -Output token throughput (tok/s), higher is better: - -| Platform (fabric) | c=1 | c=16 | c=48 | -|----------------------------|------|-------|-------| -| **SGLang + EFA** (TP16/DP2/EP16) | 32.6 | 345.4 | **802.6** | -| SGLang, socket (TP16/DP2/EP16) | — | — | 139.2 | -| vLLM + EFA (TP8/PP2) | 115.0 | 261.8 | 157.8 | -| vLLM, socket (TP8/PP2) | 113.2 | 373.1 | 141.1 | -| Trn2 (TP64) | 5.2 | *(n/m)* | 71.9 | - -Median TTFT / TPOT (ms) at c=48: - -| Platform (fabric) | Median TTFT | Median TPOT | -|----------------------------|-------------|-------------| -| **SGLang + EFA** | **1,080** | **50.6** | -| SGLang, socket | 16,831 | 231.0 | -| vLLM + EFA | 13,765 | 216.9 | -| vLLM, socket | 13,175 | 231.7 | -| Trn2 | 6,912 | 572.1 | - -Notes: -- **At c=48, SGLang + EFA wins decisively**: 803 out-tok/s and 1.08 s TTFT — ~5× vLLM's 158 and ~11× Trn2's 72, with ~13× lower TTFT than vLLM. -- **EFA is make-or-break for SGLang, nearly irrelevant for vLLM.** SGLang c=48 collapses from 803 → 139 tok/s (TTFT 1.08 s → 16.8 s) on TCP sockets — its EP=16 all-to-all + DP-attention move a lot of data cross-node. vLLM barely moves (141 → 158) because PP=2 only ships pipeline-boundary activations; its bottleneck is the pipeline bubble, not the fabric. So on sockets SGLang ≈ vLLM (~140), but with EFA SGLang pulls ~5× ahead. **The stock images fall back to sockets — always verify NCCL logs `NET/OFI ... efa-direct`, not `Using network Socket`.** -- **Low vs high concurrency flips the winner.** At c=1, vLLM (115 tok/s, TTFT 69 ms) beats SGLang (33 tok/s) — SGLang's DP/EP synchronization is pure overhead with one request; vLLM's PP pipelines a single stream cleanly. SGLang only pulls ahead once concurrency fills its wide TP16/DP2/EP16 (c=48). -- **Not equal-hardware**: H100 uses 16 GPUs / 2 nodes vs Trn2's single 64-core instance, CUDA graphs vs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. -- **Both GPU stacks produce coherent output across all requests** — no "garbled-after-first-request" bug (Trn2 issue #31), consistent with GPU stacks dequantizing FP8→BF16 before matmul. - -> **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. - -## 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-Pro/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 >= 48` (NxDI EP>1 requirement, `384 / 8 = 48`). - - 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-Pro](https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro) — HF FP8 source checkpoint - -## Maintainer - -Henan Wang (whn09) - -**Last Updated:** 2026-04-30 +See the main README's Performance section for the measured numbers. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_socket_tp16dp2_c1.txt b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_socket_tp16dp2_c1.txt new file mode 100644 index 00000000..46e41fed --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_socket_tp16dp2_c1.txt @@ -0,0 +1,42 @@ +INFO 07-21 08:01:16 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 08:01:16 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 08:01:17.751000 10 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 08:01:19 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 08:01:19 [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. +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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-04f1a6cc-', 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 08:01:20 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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=64, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-8149d0c1-', 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 08:04:26 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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/64 [00:00 Date: Tue, 21 Jul 2026 08:11:46 +0000 Subject: [PATCH 35/47] [contrib] MiMo-V2.5-Pro: restore main README (was clobbered by h100 README in 7f5e694) Commit 7f5e694 accidentally overwrote the top-level README.md with the 89-line perf_test/h100/README.md content. Restore the full 530-line contrib README (Status, Quick Start, Performance matrix incl. the SGLang/vLLM EFA-vs- socket numbers, etc.). The h100/README.md is unchanged. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 555 ++++++++++++++++++++++--- 1 file changed, 498 insertions(+), 57 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 98d2bbdb..b2f3c118 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -1,89 +1,530 @@ -# H100 GPU baseline for MiMo-V2.5-Pro +# Contrib Model: MiMo-V2.5-Pro -Scripts to serve the **official HuggingFace OCP-FP8 checkpoint** -(`XiaomiMiMo/MiMo-V2.5-Pro`) on H100 across two nodes, for a cross-platform -throughput comparison against the Trn2 Neuron port. Covers both vLLM and SGLang. +NeuronX Distributed Inference implementation of [XiaomiMiMo/MiMo-V2.5-Pro](https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro). -## Why 2 nodes (not 1) +## Model Information -The FP8 weights are ~963 GB, which does not fit on a single 8×H100-80GB node -(640 GB). MiMo-V2.5-Pro also has only **8 KV heads**, so vLLM tensor parallel -must divide 8 (TP≤8) — `--tensor-parallel-size 16` fails with *"TP size must -evenly split the number of KV heads"*. So: +- **HuggingFace ID:** `XiaomiMiMo/MiMo-V2.5-Pro` +- **Model Type:** Decoder-only MoE transformer with hybrid attention +- **Architecture:** Custom MoE with full + sliding window attention +- **License:** Check HuggingFace model card -- **vLLM**: TP=8 (within a node) × PP=2 (across the two nodes) = world size 16. -- **SGLang**: TP=16 × DP=2 with `--enable-dp-attention` (DP-attention shards the - KV heads differently, so TP=16 is fine) + EP=16. +## Architecture Details -## ⚠️ EFA is required — stock images fall back to TCP sockets +| Parameter | Value | +|-----------|-------| +| Hidden Size | 6144 | +| Layers | 70 | +| Attention Heads | 128 Q | +| KV Heads (full & sliding window) | 8 | +| Q/K Head Dim | 192 | +| V Head Dim | 128 | +| Experts | 384 routed (top-8 routing), no shared expert | +| Expert Intermediate | 2048 | +| Dense MLP Intermediate (layer 0) | 16,384 | +| Vocab Size | 152,576 | +| RoPE | Partial (33.4% → 64 of 192 dims), theta=10M (full) / 10K (SWA) | +| Sliding Window | 128 | +| Max Position | 1,048,576 (1M) | +| Attention Projection | `fused_qkv` (single `qkv_proj.weight`) | -The stock `vllm/vllm-openai` and `lmsysorg/sglang` images **do not ship -aws-ofi-nccl**, so cross-node NCCL silently falls back to TCP sockets -(~14 GB/s) instead of EFA RDMA (~400 GB/s). On the stock vLLM image we measured -NCCL logging `Using network Socket` and 13 s median TTFT. That is a -**~5× throughput / ~12× TTFT penalty** and makes any multi-node comparison -meaningless. +Key features: +- **Hybrid Attention**: 10 full attention layers (0, 7, 15, 23, 31, 39, 47, 55, 62, 69) + 60 sliding window layers, per `hybrid_layer_pattern` +- **Asymmetric Head Dims**: Q/K use head_dim=192, V uses v_head_dim=128 +- **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**: `sigmoid(logits) + e_score_correction_bias` is used to pick top-8 experts; unbiased `sigmoid(logits)` becomes the affinity weights. `n_group=1, topk_group=1` degenerates group-limited routing to plain noaux_tc. +- **attention_value_scale = 0.612**: HF reference multiplies `value_states` by this before `softmax(QK^T) × V` (NOT applied post-attention); the NxDI port matches. -`Dockerfile.vllm-efa` and `Dockerfile.sglang-efa` add GDRCopy + the AWS EFA -installer (libfabric + aws-ofi-nccl) on top of the stock images, and set -`NCCL_NET_PLUGIN=ofi`. With these, NCCL logs: +## Status (work-in-progress) +**This port compiles cleanly and is verified to produce coherent output end-to-end via the NxDI direct smoke path (`smoke_generate_mimo_v2.py`). The shipping recipe is BF16 attention + FP8 MoE at `seq_len=512` (largest seq_len that fits HBM).** + +> **Re-verified 2026-07-21** on trn2.48xlarge (i-0f2e7a5194376e8fe) from the pre-compiled `seq512` NEFF + `models/MiMo-V2.5-Pro-Neuron-FP8` (04-28) weights. `smoke_generate_mimo_v2.py` (MINIMAL_CHAT=1, 40 tokens) loaded the presharded 64-rank checkpoint in ~200s (warmup 9.4s) and produced coherent output that correctly self-identifies as MiMo: `"Okay, the user is asking for a simple self-introduction ... to establish my identity as MiMo. ..."`. The direct-smoke path is confirmed still working; vLLM serving was **not** re-tested this run and is presumed still blocked on issue #31. + +**Known issue — vLLM serving is broken.** The first `/v1/chat/completions` request against `vllm-neuron` returns coherent output; every subsequent request returns garbled text. Same compiled NEFF serves 5 successive greedy generations byte-identically via the smoke path, so the bug is specifically in vllm-neuron's runtime / request-state handling. Tracking upstream at https://github.com/vllm-project/vllm-neuron/issues/31. Last updated 2026-04-30. + +### Why BF16 attn + FP8 MoE + +Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2.5 (256 experts). Under an all-FP8 recipe, the NKI blockwise FP8 accumulator on attention q/k/v at this magnitude drifts the logits across 70 layers and produces prompt-dependent gibberish (`"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`, etc.). Dequantizing q/k/v to BF16 before the matmul restores coherent output. MoE experts (scales `≈ 2.3e-5`, similarly small) can stay FP8. + +Verified end-to-end: `smoke_generate_mimo_v2.py` with a minimal chat template returns a well-formed reasoning trace that correctly identifies the model ("As MiMo, based on Xiaomi's self-developed large model..."). `preprocess_mimo_v2_fp8.py` emits BF16 q/k/v directly so no separate step is required. + +GPU stacks (sglang on H100/H200) run the same OCP FP8 checkpoint correctly because they always dequantize FP8 → BF16 before the matmul. The issue is specific to Neuron's direct-FP8 compute path on small-magnitude tensors. Kimi PR #131 observes similar FP8 degradation on Flash and recommends SDK 2.28. + +### Cost and constraints + +- **HBM headroom.** BF16 q/k/v adds ~2 GB per rank. `seq_len=1024` OOMs on load (the previous attempt failed allocating ~40 MB for rdh/alltoall rings after per-rank tensors already reached 20.9/24 GB). `seq_len=512` is the largest value empirically verified to fit HBM at BS=48. +- **Short context.** Even at `seq_len=512`, Pro's full chat template with the default system prompt is ~260 tokens; that leaves ~250 tokens for user input + generation. Longer context needs a different HBM plan (cross-instance TP/PP, or a larger instance). +- `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else `NotImplementedError`). With `num_experts=384, top_k=8` this forces `BS >= 48`. +- `n_routed_experts=384 = 2^7 × 3` → `384 / ep_degree` is never a power of 2 (6, 12, 24, 48, 96, 192, 384). Kimi PR #131 says NKI `_bwmm_shard_on_block_nki_call` on SDK 2.29 has "depressed logits with EP=2" and recommends SDK 2.28. + +### Recipes tried that did not work + +- **All-FP8 attention (`modules_to_not_convert` without q/k/v).** Drifts as described above. Known broken; `preprocess_mimo_v2_fp8.py` no longer emits it. +- **`use_torch_block_wise=True`** (PyTorch-fallback blockwise matmul for higher accumulator precision): compile+shard succeeded after ~2 h, but `model.load()` crashed with `status=4 Allocation Failure` — the fallback path raises HBM demand even when scoped to MoE. + +### Next experiments queued + +- **Even longer `seq_len`** (> 512): needs a tighter HBM plan — smaller batch, different EP ratio, or cross-instance sharding. +- **Upstream vllm-neuron fix** for the "first-request-only" serving bug (issue #31); patch branch at `whn09/vllm-neuron#fix/hybrid-attn-swa-spec` is a placeholder that did not resolve the symptom. +- **Cross-instance BF16** via pipeline/tensor parallelism on 2× Trn2 (single-instance HBM cannot hold full BF16 Pro). +- **Selective BF16 only on MoE `gate_up_proj`** (smallest expert scales) while keeping `down_proj` FP8 — another axis to probe if attn drift returns at longer contexts. +- **SDK 2.28 venv** test once installed, per Kimi PR #131. + +## Prerequisites + +- **Instance**: trn2.48xlarge (128 physical NeuronCores, logical_nc_config=2 → 64 logical cores) +- **Neuron SDK**: 2.29 (Python 3.12, PyTorch 2.9). Verified toolchain (2026-07-21 instance): `neuronx-cc 2.26.6360.0`, `neuronx-distributed 0.19.28492`, `neuronx-distributed-inference 0.10.18399`, `torch 2.9.1`, `torch-neuronx 2.9.0.2.15`. The `neuronx-cc` package version (2.26.x) differs from the SDK release-train number (2.29); if you re-compile on a differently-imaged DLAMI, confirm these versions match or expect a cache miss. +- **Venv**: `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (used by preprocess, smoke, and vLLM serving alike; ships with the DLAMI and is where `0_setup.sh` installs the patched `vllm-neuron`). +- **Disk**: ~3 TB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~962 GB, the Neuron-FP8 preprocessed output is ~1 TB, and `save_sharded_checkpoint=true` writes another ~300-1000 GB per compiled config (varies with recipe)). + +### NVMe mount + +The Trn2 DLAMI ships with four local NVMe SSDs that are assembled into a +RAID0 array at `/opt/dlami/nvme`. After a reboot the mount is **NOT** +reassembled automatically — you must re-mount manually before the paths +below resolve: + +```bash +lsblk # confirm you see nvme0n1..nvme3n1 devices +sudo mdadm --assemble /dev/md0 /dev/nvme[0-3]n1 2>/dev/null || true +sudo mount /dev/md0 /opt/dlami/nvme +df -h /opt/dlami/nvme # should show ~6.9 TB total +``` + +If `mdadm --assemble` says the array is already assembled, the mount +step alone is enough. If `/dev/md0` doesn't exist, the array was never +created on this instance — run `/opt/dlami/setup-nvme.sh` (or the +DLAMI's built-in helper; consult `ls /opt/dlami/*.sh`) before mounting. + +## 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-Pro 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-Pro # the branch this README lives on + +# 2. Download the HuggingFace FP8 checkpoint (~1 TB; 50 safetensors shards). +# Any HF-compatible downloader works; huggingface-cli example: +huggingface-cli download XiaomiMiMo/MiMo-V2.5-Pro \ + --local-dir /opt/dlami/nvme/models/MiMo-V2.5-Pro + +# 3. Preprocess HF FP8 -> Neuron-FP8 (BF16 attn, FP8 MoE). ~20 min, ~24 GB +# peak RAM. The preprocess dequants q/k/v to BF16 in one pass — see +# "Checkpoint Preparation" below for why BF16 attn is the only recipe. +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate +python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py \ + --hf_model_path /opt/dlami/nvme/models/MiMo-V2.5-Pro \ + --save_path /opt/dlami/nvme/models/MiMo-V2.5-Pro-Neuron-FP8 \ + --tp_degree 64 + +# 4. (Optional) sanity-check the Neuron-FP8 checkpoint without vLLM +# ~90 min first compile; subsequent runs ~60s to load the pre-sharded NEFF. +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate +python contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py # compile +python contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py # 20-token generate + +# 5. Install vllm-neuron with the contrib registration patch +bash contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh + +# 6. Start vLLM serving MiMo-V2.5-Pro FP8 (first compile ~60 min; subsequent ~3 min) +bash contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh +``` + +The bench script runs one configuration (BS=48, +`moe_tp_degree=1 / moe_ep_degree=64`) at three concurrency levels (1, 16, 48) +and logs results under `/opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro/`. + +### Keeping a server up for ad-hoc testing + +`bench_mimo_v2.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-Pro/perf_test/start_vllm_server.sh + +# Terminal 2: once "Application startup complete." prints, sanity-check: +bash contrib/models/MiMo-V2.5-Pro/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-Pro/perf_test/run_bench_single.sh +``` + +`bench_mimo_v2.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 +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.sh`. + +**Required (at least for manual `vllm api_server` launches):** + +| Variable | Purpose | +|---|---| +| `NXDI_CONTRIB_MIMO_V2_FLASH_SRC` | Path to `contrib/models/MiMo-V2.5-Pro/src/`. `vllm-neuron`'s registration hook reads it to plug `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` table. The `_FLASH_` suffix is kept for backward compatibility with the shared registration hook that also serves V2-Flash and V2.5. | +| `MIMO_V2_FLASH_PATH` | Preprocessed Neuron-FP8 checkpoint dir (the `--save_path` output from preprocess). Same naming rationale as above. | + +**Optional (recommended):** + +| Variable | Default | Purpose | +|---|---|---| +| `NEURON_COMPILED_ARTIFACTS` | `/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512` (per `start_vllm_server.sh`) | Where vLLM writes its NEFF + per-rank sharded weights. Points at the **same** dir as the smoke-path seq512 NEFF (no `_vllm` suffix) so vLLM reuses its 64 pre-sharded `tp*_sharded_checkpoint.safetensors` and skips the ~30 min shard step. vLLM still compiles its own continuous-batching / async / on-device-sampling NEFF variant into this dir on first launch (~60 min), but the reshard is avoided. A separate empty dir (e.g. a `_vllm` suffix) would force a full from-scratch compile *and* reshard. vLLM's own fallback is `/neuron-compiled-artifacts//`. | +| `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 Pro's 384-expert MoE is ~60 min TKG + ~15 min CTE + ~30 min shard, well past vLLM's default. | + +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-Pro-Neuron-FP8", + "messages": [{"role": "user", "content": "Hello! Introduce yourself in one sentence."}], + "max_tokens": 64, "temperature": 0.0}' | python3 -m json.tool ``` -NET/OFI Initializing aws-ofi-nccl 1.20.0 ... Using transport protocol RDMA -NET/OFI Selected provider is efa, fabric is efa-direct (found 32 nics) + +Output quality is currently prompt-dependent under the FP8 recipe (see +Status). A successful sanity check confirms the serving path works; it +does not yet confirm that all prompts produce coherent text. + +## 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_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 ~1 TB across 70 per-layer safetensors shards. + +```bash +python contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py \ + --hf_model_path /path/to/MiMo-V2.5-Pro \ + --save_path /path/to/MiMo-V2.5-Pro-Neuron-FP8 \ + --tp_degree 64 ``` -Build once per node (no GPU needed at build time): +Peak RAM during preprocessing is ~24 GB; total runtime ~20 minutes on a trn2.48xlarge instance. + +### Why q/k/v are BF16 in the preprocessed output + +Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2.5 (256 experts). The NKI blockwise FP8 accumulator at this magnitude drifts the logits across 70 layers and produces gibberish output — `"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`, etc. Dequantizing q/k/v to BF16 while keeping MoE experts FP8 restores coherent output (verified on 2026-04-29 via `smoke_generate_mimo_v2.py`). + +The preprocess handles this in a single pass: `split_qkv_fused()` unfuses Pro's `qkv_proj` into per-proj BF16 tensors directly, and the Flash-style per-proj fallback path dequants via `_dequant_attn_to_bf16()`. The checkpoint emitted by preprocess has no `q_proj.scale` / `k_proj.scale` / `v_proj.scale` entries. Compile-time `modules_to_not_convert` must therefore include `q_proj`, `k_proj`, `v_proj` so NxDI routes them through a plain `ColumnParallelLinear` rather than the FP8 `QuantizedColumnParallel` path — `smoke_compile_mimo_v2.py` already does this. + +### Parallel preprocess (faster) + +`src/conversion_script/preprocess_mimo_v2_parallel.py` (driven by `run_preprocess_parallel.sh`) is a multiprocess wrapper around the same per-layer conversion. Each worker dequants one layer independently, cutting wall time from ~20-30 min (serial) to ~5-6 min with 12 workers (peak ~300 GB CPU RAM on a 2 TB box). Output is identical to the serial path (BF16 attn + FP8 MoE). ```bash -docker build -t vllm-efa:latest -f Dockerfile.vllm-efa . -docker build -t sglang-efa:latest -f Dockerfile.sglang-efa . +N_WORKERS=12 bash contrib/models/MiMo-V2.5-Pro/src/conversion_script/run_preprocess_parallel.sh ``` +> Note: there is no FP8 → full-BF16 conversion mode. Both preprocess scripts always emit the shipping recipe (BF16 q/k/v attention + FP8 MoE experts); q/k/v are dequantized to BF16 in-pass, but MoE weights stay FP8. A separate all-BF16 reference checkpoint, if needed, must be produced by other means. + ## Usage -Download the checkpoint on **both** nodes first: +```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-Pro/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-Pro-Neuron-FP8/" +compiled_path = "/path/to/compiled/" + +# Recommended recipe: BF16 attn + FP8 MoE. +# moe_tp_degree = 1, moe_ep_degree = 64 +# q_proj/k_proj/v_proj in modules_to_not_convert (BF16; preprocess +# emits BF16 for q/k/v, no separate step needed) +# seq_len = 512 (largest empirically verified; see Status) +# 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=48, # must be >= num_experts / top_k = 384 / 8 = 48 + max_batch_size=48, + ctx_batch_size=1, + tkg_batch_size=48, + seq_len=512, # largest empirically verified; seq_len=1024 OOMs + 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", + "q_proj", "k_proj", "v_proj", # BF16 attn — preprocess emits BF16 + ], + on_device_sampling_config=OnDeviceSamplingConfig( + do_sample=True, temperature=0.6, top_k=20, top_p=0.95, + ), +) + +# trust_remote_code is required by Flash'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.py` — compile + load (STAGE=instantiate|compile|load|all, DRY_RUN, SKIP_WARMUP) +- `perf_test/smoke_generate_mimo_v2.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-Pro (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 Pro's 69 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 (6 experts per rank for Pro's 384 experts), which preserves per-channel scale. On V2.5 (256 experts) this recipe yields coherent output; on V2.5-Pro it still exhibits prompt-dependent drift (see Status). + +Intermediate ratios (`moe_tp=32/ep=2`, `moe_tp=16/ep=4`) have been empirically tested and still produce gibberish, so `moe_tp=1/moe_ep=64` is the only currently-usable moe_tp/ep combination. + +### batch_size >= 48 + +NxDI's TKG (token generation) path refuses Expert Parallelism when `batch_size < num_experts / top_k`. For Pro that is 384 / 8 = 48, so the smallest working BS on the FP8 path is 48. BS=1 latency demos are not possible on the FP8 (moe_ep=64) path; a single-stream configuration would require `moe_tp=64, moe_ep=1, batch_size=1`, which in turn needs an all-BF16 checkpoint (the preprocess scripts here only emit BF16-attn + FP8-MoE, so that BF16 checkpoint must be produced separately). + +### 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-Pro 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 -hf download XiaomiMiMo/MiMo-V2.5-Pro --local-dir /opt/dlami/nvme/models/MiMo-V2.5-Pro +# The setup script clones vllm-project/vllm-neuron at release-0.5.3, applies +# the contrib registration patch, installs it editable, and downloads +# Pro Neuron-FP8 weights from S3 (set MIMO_V2_FLASH_PATH to override). +bash contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh ``` -Run the **same** launch script on both nodes, changing only `NODE_RANK` -(`MASTER_ADDR` / `DIST_INIT_ADDR` = node-0's private IP, reachable from both): +The patch (`perf_test/vllm-neuron-patch.patch`) touches `vllm_neuron/worker/neuronx_distributed_model_loader.py`. It adds a `_register_contrib_models()` hook that, when `NXDI_CONTRIB_MIMO_V2_FLASH_SRC` is set, registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under keys `mimov2flash` **and** `mimov2pro`, **and** overrides vLLM's built-in `MiMoV2FlashForCausalLM` / `MiMoV2ProForCausalLM` (GPU-only stubs) in `ModelRegistry` with the Neuron wrapper so ModelConfig validation accepts either architecture. No upstream vLLM or NxDI source is modified. The checkpoint's `config.json` must set `architectures` to `["MiMoV2ProForCausalLM"]` (or `MiMoV2FlashForCausalLM` for V2.5); the preprocess script takes care of this. + +### Serving (FP8, recommended) + +Use `perf_test/start_vllm_server.sh` for a foreground launch (stays up until Ctrl-C), or `perf_test/bench_mimo_v2.sh` for the one-shot launch → sanity → bench → teardown flow. Both scripts bake in the full `override_neuron_config` (TP=64, moe_tp=1, moe_ep=64, BS=48, CB + bucketing, blockwise FP8 MoE with `PING_PONG`, on-device sampling), the required env vars, and the persistent compile-artifact path. See "Keeping a server up for ad-hoc testing" above for the three-terminal workflow. ```bash -# vLLM (node 0 = API server on :8000, node 1 = --headless follower) -NODE_RANK=0 MASTER_ADDR= bash run_vllm_h100_multinode.sh -NODE_RANK=1 MASTER_ADDR= bash run_vllm_h100_multinode.sh +# One-shot launch + bench + teardown (~2 h on cold cache, ~5 min on warm cache). +bash contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh + +# Or keep the server up for interactive work: +bash contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh +``` + +See "Environment variables" above for all the knobs (`NEURON_COMPILED_ARTIFACTS`, `BASE_COMPILE_WORK_DIR`, etc.) and their defaults. + +> **vLLM serving is currently broken.** With the BF16-attn checkpoint, every `vllm-neuron` configuration we tried (all-FP8-attn, BF16-attn with `seq_len=256` or `512`, CB on/off, on-device sampling on/off, `-O3` or `-O1` TKG compile) reproduces the same pattern: the first chat request returns coherent output, every subsequent request returns UTF-8-replacement-char + off-topic text. The same compiled NEFF serves 5 successive greedy `adapter.generate()` calls byte-identically under `smoke_generate_mimo_v2.py` — the bug is in vllm-neuron's runtime, not in the model or the NEFF. Tracking at https://github.com/vllm-project/vllm-neuron/issues/31. Until that is fixed, use `smoke_generate_mimo_v2.py` for direct NxDI inference; the bench numbers below are historical infra-validation data from the pre-BF16-attn all-FP8 checkpoint. + +### vllm-neuron patch summary + +The patch is applied to vllm-neuron 0.5.3 (also applies cleanly to 0.5.0; same model-loading architecture) and: + +- Patches `AutoConfig.from_pretrained` to default `trust_remote_code=True` so NxDI's `hf_adapter.load_config` can load the `MiMoV2Config` custom code that ships with the checkpoint. +- Registers `NeuronMiMoV2ForCausalLM` into NxDI's `MODEL_TYPES` under `mimov2flash` and `mimov2pro` so the NxDI loader resolves either model_type to the contrib Neuron wrapper. +- Overrides vLLM's built-in `MiMoV2FlashForCausalLM` and `MiMoV2ProForCausalLM` GPU stubs in `ModelRegistry`, since vLLM's ModelConfig validator rejects any architecture not in its registry and the Neuron path never instantiates vLLM's stub class anyway. + +## Performance + +> The throughput numbers below were captured on 2026-04-29 against a pre-BF16-attn checkpoint (all-FP8, `seq_len=1024`) before we discovered the vllm-neuron first-request bug. They are historical — the shipping recipe is BF16 attn + FP8 MoE at `seq_len=512` via the smoke path, and vLLM serving is currently blocked on issue #31. The numbers are kept here for order-of-magnitude reference. -# SGLang (both nodes run the same command; node 0 serves on :30000) -NODE_RANK=0 DIST_INIT_ADDR=:20000 bash run_sglang_h100_multinode.sh -NODE_RANK=1 DIST_INIT_ADDR=:20000 bash run_sglang_h100_multinode.sh +### vLLM Serving (trn2.48xlarge, historical all-FP8 run, BS=48, TP=64, moe_tp=1/moe_ep=64, CB + bucketing, `seq_len=1024`) + +Input/output: 900/90 tokens (`vllm bench serve --dataset-name random`), `on_device_sampling_config={do_sample:true, temperature:0.6, top_k:20, top_p:0.95}`. + +| Concurrency | Total tok/s | Output tok/s | TTFT median (ms) | TTFT P99 (ms) | TPOT median (ms) | +|-------------|-------------|--------------|------------------|---------------|------------------| +| 1 | 47 | 4.3 | 1,392 | 1,393 | 220 | +| 16 | 391 | 35.6 | 2,361 | 17,394 | 422 | +| 48 | 606 | 55 | 7,322 | 54,413 | 752 | + +Per-stream ITL median holds at ~220 ms across all concurrency levels; TPOT/TTFT growth at higher concurrency comes from continuous-batching queue pressure, not per-step compute. + +> Expected BF16-attn delta: only q/k/v go from FP8 to BF16 (MoE is unchanged), so steady-state throughput should be within a few percent. TTFT should drop proportionally with `seq_len` (256 vs 1024 prefill tokens). + +### Measured 2026-07-21 (BF16-attn + FP8 MoE, `seq_len=512`, BS=48, TP=64, moe_tp=1/moe_ep=64) + +Same instance/recipe as the shipping config, via `bench_smoke_throughput.py` +(NxDI direct path, coherent output). Prefill and decode measured separately +(a `max_new_tokens=1` call isolates prefill/TTFT; the full call minus that +isolates steady-state decode). 3-iter averages, very stable. + +| in/out | Prefill (CTE) | Decode (TKG) | End-to-end | +|--------|---------------|--------------|------------| +| 360/120 | 37.9 s → 456 in-tok/s | 22.4 s / 119 steps → 255 out-tok/s (5.3/stream) | 60.3 s → 95.6 out-tok/s | +| 500/2 | 37.9 s → 634 in-tok/s | 0.2 s / 1 step → 235 out-tok/s (4.9/stream) | — | + +**Prefill dominates and is nearly constant in input length** (37.9 s for both +360 and 500 input tokens). The context-encoding NEFF has a single bucket +`context_encoding_buckets=[512]`, so every prefill pads to 512 and pays the +same cost regardless of real input length — that fixed ~38 s (≈190 decode +steps' worth of time) is the #1 optimization target, and it is almost +certainly the 384-expert FP8 MoE blockwise matmul over 512 positions, not +attention. Decode itself is cheap (~0.2 s/token for the whole BS=48 batch). + +vLLM serving on the same NEFF/recipe (input/output 360/120): + +| Concurrency | Output tok/s | Total tok/s | TPOT median | Notes | +|-------------|--------------|-------------|-------------|-------| +| 48 | 71.9 | 287 | 572 ms | output garbled after first request (issue #31) | +| 1 | 5.2 | 20.5 | 189 ms | under-fed (BS=48 graph, 1 request); not representative | + +vLLM at c=48 reaches 72 out-tok/s vs the smoke path's 96 — the ~25% gap is +vLLM's continuous-batching scheduler / async / request-state overhead (the same +runtime layer behind the issue #31 garbling). Smoke is the batched-compute +ceiling for this NEFF; vLLM is the realistic serving figure with scheduling. + +Reproduce (smoke path): + +```bash +source /opt/aws_neuronx_venv_pytorch_inference_vllm_0_16/bin/activate +export MIMO_V25_PRO_COMPILED_PATH=/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/ +# decode-focused: +INPUT_LEN=360 MAX_NEW_TOKENS=120 N_ITERS=3 \ + python3 contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py +# prefill-focused (near-full context, ~1 output token): +INPUT_LEN=500 MAX_NEW_TOKENS=2 N_ITERS=3 \ + python3 contrib/models/MiMo-V2.5-Pro/perf_test/bench_smoke_throughput.py ``` -Notes baked into the scripts: -- `--network host` (cross-node NCCL/torch.distributed can't use port mapping), - EFA via `/dev/infiniband`, `NCCL_NET_PLUGIN=ofi`. -- JIT caches (DeepGEMM warmup over 32768 shapes ≈ several min, flashinfer, etc.) - are bind-mounted to the host at `/opt/dlami/nvme/sglang_cache` so a restart - doesn't recompile. -- vLLM: `--no-enable-prefix-caching` + `--no-enable-chunked-prefill` to match - Trn2 (prefix caching on gave a 66% hit rate and inflated throughput). -- SGLang: `--mem-fraction-static 0.9` (the reference 0.7 left no room for KV - cache — weights nearly fill the GPUs). +To capture a device profile for bottleneck analysis (attention vs MoE vs +collectives) set `NEURON_RT_INSPECT_DEVICE_PROFILE=` before the run and +inspect the resulting `*.ntff` with `neuron-profile view` / Neuron Explorer. + +### H100 GPU baseline (cross-platform comparison, 2026-07-21) + +The **official HF OCP-FP8 checkpoint** served on H100 across two nodes, using +the **same** `run_bench_single.sh` and the same input/output (360/120), prefix +caching + chunked prefill **disabled** for a fair comparison. See +`perf_test/h100/` for launch scripts, Dockerfiles, and setup notes. -## Benchmark +- **Trn2**: trn2.48xlarge, TP=64 / moe_tp=1 / moe_ep=64, BS=48, `seq_len=512`, vllm-neuron 0.5.3. +- **H100**: 2× p5 (16× H100-80GB). The ~963 GB FP8 weights don't fit on one 8×80 GB node, so both frameworks shard across two nodes. + - **vLLM 0.25.1**: TP=8 × PP=2 (MiMo's 8 KV heads cap TP at 8; TP=16 fails the KV-head divisibility check). CUDA graphs on. + - **SGLang 0.5.15**: TP=16 × DP=2 with `--enable-dp-attention` + EP=16 (DP-attention shards KV heads so TP=16 works). -Reuse the **same** `perf_test/run_bench_single.sh` as Trn2 (skips the Neuron -venv when absent; takes `SERVED_MODEL_NAME` / `TOKENIZER_PATH` / `PORT`). Run -the bench client inside a container: +**Cross-node fabric matters enormously.** The stock `vllm/vllm-openai` and +`lmsysorg/sglang` images lack `aws-ofi-nccl`, so NCCL silently falls back to TCP +sockets (`Using network Socket`) instead of EFA RDMA. Rebuilding both images +with GDRCopy + the AWS EFA installer (`Dockerfile.{vllm,sglang}-efa`, with +`NCCL_NET_PLUGIN=ofi`) switches NCCL to `efa-direct` over 32 NICs. + +Output token throughput (tok/s), higher is better: + +| Platform (fabric) | c=1 | c=16 | c=48 | +|----------------------------|------|-------|-------| +| **SGLang + EFA** (TP16/DP2/EP16) | 32.6 | 345.4 | **802.6** | +| SGLang, socket (TP16/DP2/EP16) | 11.2 | 70.4 | 139.2 | +| vLLM + EFA (TP8/PP2) | 115.0 | 261.8 | 157.8 | +| vLLM, socket (TP8/PP2) | 113.2 | 373.1 | 141.1 | +| Trn2 (TP64) | 5.2 | *(n/m)* | 71.9 | + +SGLang EFA vs socket (out-tok/s): c=1 2.9×, c=16 4.9×, c=48 5.8× — the EFA +advantage grows with concurrency as more cross-node all-to-all traffic piles up. + +Median TTFT / TPOT (ms) at c=48: + +| Platform (fabric) | Median TTFT | Median TPOT | +|----------------------------|-------------|-------------| +| **SGLang + EFA** | **1,080** | **50.6** | +| SGLang, socket | 16,831 | 231.0 | +| vLLM + EFA | 13,765 | 216.9 | +| vLLM, socket | 13,175 | 231.7 | +| Trn2 | 6,912 | 572.1 | + +Notes: +- **At c=48, SGLang + EFA wins decisively**: 803 out-tok/s and 1.08 s TTFT — ~5× vLLM's 158 and ~11× Trn2's 72, with ~13× lower TTFT than vLLM. +- **EFA is make-or-break for SGLang, nearly irrelevant for vLLM.** SGLang c=48 collapses from 803 → 139 tok/s (TTFT 1.08 s → 16.8 s) on TCP sockets — its EP=16 all-to-all + DP-attention move a lot of data cross-node. vLLM barely moves (141 → 158) because PP=2 only ships pipeline-boundary activations; its bottleneck is the pipeline bubble, not the fabric. So on sockets SGLang ≈ vLLM (~140), but with EFA SGLang pulls ~5× ahead. **The stock images fall back to sockets — always verify NCCL logs `NET/OFI ... efa-direct`, not `Using network Socket`.** +- **Low vs high concurrency flips the winner.** At c=1, vLLM (115 tok/s, TTFT 69 ms) beats SGLang (33 tok/s) — SGLang's DP/EP synchronization is pure overhead with one request; vLLM's PP pipelines a single stream cleanly. SGLang only pulls ahead once concurrency fills its wide TP16/DP2/EP16 (c=48). +- **Not equal-hardware**: H100 uses 16 GPUs / 2 nodes vs Trn2's single 64-core instance, CUDA graphs vs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. +- **Both GPU stacks produce coherent output across all requests** — no "garbled-after-first-request" bug (Trn2 issue #31), consistent with GPU stacks dequantizing FP8→BF16 before matmul. + +> **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. + +## 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 -docker run --rm --network host -v /opt/dlami/nvme/models:/wk --entrypoint bash \ - vllm-efa:latest -c ' - export SERVED_MODEL_NAME=MiMo-V2.5-Pro TOKENIZER_PATH=/wk/MiMo-V2.5-Pro \ - PORT=8000 RESULTS_DIR=/wk/bench_results CONFIG_NAME=vllm_efa - CONCURRENCY=48 NUM_PROMPTS=96 bash /wk/run_bench_single.sh' # PORT=30000 for SGLang +pytest contrib/models/MiMo-V2.5-Pro/test/integration/test_model.py -v ``` -See the main README's Performance section for the measured numbers. +## 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 >= 48` (NxDI EP>1 requirement, `384 / 8 = 48`). + - 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-Pro](https://huggingface.co/XiaomiMiMo/MiMo-V2.5-Pro) — HF FP8 source checkpoint + +## Maintainer + +Henan Wang (whn09) + +**Last Updated:** 2026-04-30 From 0b7a22a8b80e6d3efed0b473b24b2e5577cb5ca2 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 08:19:07 +0000 Subject: [PATCH 36/47] [contrib] MiMo-V2.5-Pro: add vLLM DP+EP config, best vLLM multi-node 270 tok/s vLLM was slow at TP8/PP2 [158 tok/s c=48] mainly due to the PP=2 pipeline bubble, forced because MiMo has 8 KV heads so TP<=8. Replacing PP with vLLM data-parallel + expert-parallel -- DP2 x TP8 + --enable-expert-parallel, the analogue of SGLang --enable-dp-attention -- lifts c=48 to 270 tok/s [1.7x] and halves TTFT 13.8 -> 7.6s, no pipeline bubble. Each DP rank runs TP=8 attention dividing the 8 KV heads; MoE shards across TP*DP=16. Still ~3x behind SGLang+EFA [270 vs 803]; remaining gap is SGLang MoE kernels and prefill parallelism. DP+EP is the right multi-node shape for vLLM here. Adds run_vllm_h100_dp.sh + c=48 result; README table + notes updated. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 13 +-- .../h100/results/vllm_efa_dp2tp8ep_c48.txt | 42 +++++++++ .../perf_test/h100/run_vllm_h100_dp.sh | 85 +++++++++++++++++++ 3 files changed, 135 insertions(+), 5 deletions(-) create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c48.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_dp.sh diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index b2f3c118..cf14ac99 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -468,8 +468,9 @@ Output token throughput (tok/s), higher is better: |----------------------------|------|-------|-------| | **SGLang + EFA** (TP16/DP2/EP16) | 32.6 | 345.4 | **802.6** | | SGLang, socket (TP16/DP2/EP16) | 11.2 | 70.4 | 139.2 | -| vLLM + EFA (TP8/PP2) | 115.0 | 261.8 | 157.8 | -| vLLM, socket (TP8/PP2) | 113.2 | 373.1 | 141.1 | +| vLLM + EFA, **DP2×TP8 + EP** (best vLLM) | *(n/m)* | *(n/m)* | 270.5 | +| vLLM + EFA, TP8/PP2 | 115.0 | 261.8 | 157.8 | +| vLLM, socket, TP8/PP2 | 113.2 | 373.1 | 141.1 | | Trn2 (TP64) | 5.2 | *(n/m)* | 71.9 | SGLang EFA vs socket (out-tok/s): c=1 2.9×, c=16 4.9×, c=48 5.8× — the EFA @@ -480,14 +481,16 @@ Median TTFT / TPOT (ms) at c=48: | Platform (fabric) | Median TTFT | Median TPOT | |----------------------------|-------------|-------------| | **SGLang + EFA** | **1,080** | **50.6** | +| vLLM + EFA, DP2×TP8 + EP | 7,614 | 117.6 | +| vLLM + EFA, TP8/PP2 | 13,765 | 216.9 | | SGLang, socket | 16,831 | 231.0 | -| vLLM + EFA | 13,765 | 216.9 | -| vLLM, socket | 13,175 | 231.7 | +| vLLM, socket, TP8/PP2 | 13,175 | 231.7 | | Trn2 | 6,912 | 572.1 | Notes: - **At c=48, SGLang + EFA wins decisively**: 803 out-tok/s and 1.08 s TTFT — ~5× vLLM's 158 and ~11× Trn2's 72, with ~13× lower TTFT than vLLM. -- **EFA is make-or-break for SGLang, nearly irrelevant for vLLM.** SGLang c=48 collapses from 803 → 139 tok/s (TTFT 1.08 s → 16.8 s) on TCP sockets — its EP=16 all-to-all + DP-attention move a lot of data cross-node. vLLM barely moves (141 → 158) because PP=2 only ships pipeline-boundary activations; its bottleneck is the pipeline bubble, not the fabric. So on sockets SGLang ≈ vLLM (~140), but with EFA SGLang pulls ~5× ahead. **The stock images fall back to sockets — always verify NCCL logs `NET/OFI ... efa-direct`, not `Using network Socket`.** +- **EFA is make-or-break for SGLang, nearly irrelevant for vLLM's PP path.** SGLang c=48 collapses from 803 → 139 tok/s (TTFT 1.08 s → 16.8 s) on TCP sockets — its EP=16 all-to-all + DP-attention move a lot of data cross-node. vLLM TP8/PP2 barely moves (141 → 158) because PP=2 only ships pipeline-boundary activations; its bottleneck is the pipeline bubble, not the fabric. **The stock images fall back to sockets — always verify NCCL logs `NET/OFI ... efa-direct`, not `Using network Socket`.** +- **vLLM: use DP+EP, not PP.** Replacing TP8/PP2 with **DP2×TP8 + `--enable-expert-parallel`** (vLLM's analogue to SGLang's DP-attention: each DP rank runs TP=8 attention, dividing the 8 KV heads, while MoE shards across TP×DP=16) lifts vLLM c=48 from 158 → **270 tok/s** and halves TTFT (13.8 → 7.6 s), with no pipeline bubble. Still ~3× behind SGLang (270 vs 803) — the remaining gap is SGLang's more mature MoE kernels (DeepGEMM, EP all-to-all) and prefill parallelism — but DP+EP is clearly the right multi-node shape for this model on vLLM. - **Low vs high concurrency flips the winner.** At c=1, vLLM (115 tok/s, TTFT 69 ms) beats SGLang (33 tok/s) — SGLang's DP/EP synchronization is pure overhead with one request; vLLM's PP pipelines a single stream cleanly. SGLang only pulls ahead once concurrency fills its wide TP16/DP2/EP16 (c=48). - **Not equal-hardware**: H100 uses 16 GPUs / 2 nodes vs Trn2's single 64-core instance, CUDA graphs vs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. - **Both GPU stacks produce coherent output across all requests** — no "garbled-after-first-request" bug (Trn2 issue #31), consistent with GPU stacks dequantizing FP8→BF16 before matmul. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c48.txt b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c48.txt new file mode 100644 index 00000000..595412de --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c48.txt @@ -0,0 +1,42 @@ +INFO 07-21 08:16:31 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 08:16:31 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 08:16:32.105000 9 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 08:16:33 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 08:16:33 [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. +Namespace(subparser='bench', bench_type='serve', dispatch_function=, 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=360, random_output_len=120, 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=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-c06615dc-', 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 08:16:34 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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: 48 + 0%| | 0/96 [00:00&2 + exit 1 +fi + +echo "==========================================" +echo "MiMo-V2.5-Pro vLLM (2-node DP=$DP x TP=$TP + EP, EFA, Docker)" +echo " DP rank: $DP_RANK DP addr: $DP_ADDR:$DP_RPC_PORT iface: $IFACE" +echo " Port: $PORT max-model-len: $MAX_MODEL_LEN" +echo "==========================================" + +COMMON_ARGS=( + --served-model-name MiMo-V2.5-Pro + --trust-remote-code + --generation-config vllm + --tensor-parallel-size "$TP" + --data-parallel-size "$DP" + --data-parallel-size-local 1 + --data-parallel-address "$DP_ADDR" + --data-parallel-rpc-port "$DP_RPC_PORT" + --enable-expert-parallel + --max-model-len "$MAX_MODEL_LEN" +) +if [ "$DP_RANK" = "0" ]; then + ROLE_ARGS=( + --host 0.0.0.0 + --port "$PORT" + --data-parallel-start-rank 0 + --no-enable-prefix-caching + --no-enable-chunked-prefill + --tool-call-parser mimo + --enable-auto-tool-choice + --reasoning-parser mimo + ) +else + ROLE_ARGS=( --headless --data-parallel-start-rank 1 ) +fi + +exec docker run --rm --gpus all \ + --network host --privileged --ipc=host --shm-size=32g \ + --device /dev/infiniband \ + -v "${MODEL_DIR}:${CTR_MODEL}:ro" \ + -v "${CACHE_DIR}:/root/.cache" \ + -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + -e GLOO_SOCKET_IFNAME="$IFACE" \ + -e NCCL_SOCKET_IFNAME="$IFACE" \ + -e NCCL_DEBUG="${NCCL_DEBUG:-WARN}" \ + -e NCCL_DEBUG_SUBSYS="${NCCL_DEBUG_SUBSYS:-INIT,NET}" \ + -e FI_PROVIDER=efa \ + -e NCCL_NET_PLUGIN="${NCCL_NET_PLUGIN:-ofi}" \ + -e VLLM_HOST_IP="$(hostname -I | awk '{print $1}')" \ + "$IMAGE" "$CTR_MODEL" \ + "${COMMON_ARGS[@]}" "${ROLE_ARGS[@]}" From a2fea218c26a99dda4e9b146b56d199f9de2c4c5 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 08:23:53 +0000 Subject: [PATCH 37/47] [contrib] MiMo-V2.5-Pro: add vLLM DP+EP c=1/c=16 (complete matrix) vLLM DP2xTP8+EP: c=1 27.8, c=16 299.0, c=48 270.5 tok/s. Same pattern as SGLang -- DP/EP is slow at c=1 [27.8, sync overhead] but wins at c=16/48. At c=1 the PP layout still wins [115], so: PP for single-stream latency, DP+EP for throughput serving. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 4 +- .../h100/results/vllm_efa_dp2tp8ep_c1.txt | 42 +++++++++++++++++++ .../h100/results/vllm_efa_dp2tp8ep_c16.txt | 42 +++++++++++++++++++ 3 files changed, 86 insertions(+), 2 deletions(-) create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c1.txt create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c16.txt diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index cf14ac99..b579be01 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -468,7 +468,7 @@ Output token throughput (tok/s), higher is better: |----------------------------|------|-------|-------| | **SGLang + EFA** (TP16/DP2/EP16) | 32.6 | 345.4 | **802.6** | | SGLang, socket (TP16/DP2/EP16) | 11.2 | 70.4 | 139.2 | -| vLLM + EFA, **DP2×TP8 + EP** (best vLLM) | *(n/m)* | *(n/m)* | 270.5 | +| vLLM + EFA, **DP2×TP8 + EP** (best vLLM) | 27.8 | 299.0 | 270.5 | | vLLM + EFA, TP8/PP2 | 115.0 | 261.8 | 157.8 | | vLLM, socket, TP8/PP2 | 113.2 | 373.1 | 141.1 | | Trn2 (TP64) | 5.2 | *(n/m)* | 71.9 | @@ -491,7 +491,7 @@ Notes: - **At c=48, SGLang + EFA wins decisively**: 803 out-tok/s and 1.08 s TTFT — ~5× vLLM's 158 and ~11× Trn2's 72, with ~13× lower TTFT than vLLM. - **EFA is make-or-break for SGLang, nearly irrelevant for vLLM's PP path.** SGLang c=48 collapses from 803 → 139 tok/s (TTFT 1.08 s → 16.8 s) on TCP sockets — its EP=16 all-to-all + DP-attention move a lot of data cross-node. vLLM TP8/PP2 barely moves (141 → 158) because PP=2 only ships pipeline-boundary activations; its bottleneck is the pipeline bubble, not the fabric. **The stock images fall back to sockets — always verify NCCL logs `NET/OFI ... efa-direct`, not `Using network Socket`.** - **vLLM: use DP+EP, not PP.** Replacing TP8/PP2 with **DP2×TP8 + `--enable-expert-parallel`** (vLLM's analogue to SGLang's DP-attention: each DP rank runs TP=8 attention, dividing the 8 KV heads, while MoE shards across TP×DP=16) lifts vLLM c=48 from 158 → **270 tok/s** and halves TTFT (13.8 → 7.6 s), with no pipeline bubble. Still ~3× behind SGLang (270 vs 803) — the remaining gap is SGLang's more mature MoE kernels (DeepGEMM, EP all-to-all) and prefill parallelism — but DP+EP is clearly the right multi-node shape for this model on vLLM. -- **Low vs high concurrency flips the winner.** At c=1, vLLM (115 tok/s, TTFT 69 ms) beats SGLang (33 tok/s) — SGLang's DP/EP synchronization is pure overhead with one request; vLLM's PP pipelines a single stream cleanly. SGLang only pulls ahead once concurrency fills its wide TP16/DP2/EP16 (c=48). +- **Low vs high concurrency flips the winner, and picks the parallelism.** At c=1, vLLM **TP8/PP2** (115 tok/s, TTFT 69 ms) beats everything — a single stream pipelines cleanly with no DP/EP sync. Both DP/EP layouts are slow at c=1 (vLLM DP+EP 27.8, SGLang 32.6): the DP/EP all-reduce/all-to-all is pure overhead with one request. Once concurrency fills the wide parallelism (c=16/48), DP/EP wins big (SGLang 803, vLLM DP+EP 270/299) and PP falls behind (158). Takeaway: PP for latency-bound single-stream, DP+EP for throughput-bound serving. - **Not equal-hardware**: H100 uses 16 GPUs / 2 nodes vs Trn2's single 64-core instance, CUDA graphs vs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. - **Both GPU stacks produce coherent output across all requests** — no "garbled-after-first-request" bug (Trn2 issue #31), consistent with GPU stacks dequantizing FP8→BF16 before matmul. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c1.txt b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c1.txt new file mode 100644 index 00000000..03c66dd5 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c1.txt @@ -0,0 +1,42 @@ +INFO 07-21 08:20:12 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 08:20:12 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 08:20:13.298000 10 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 08:20:14 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 08:20:15 [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. +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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-b836591c-', 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 08:20:16 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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=64, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-73fe9fb8-', 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 08:21:37 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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/64 [00:00 Date: Tue, 21 Jul 2026 08:53:42 +0000 Subject: [PATCH 38/47] [contrib] MiMo-V2.5-Pro: fair H100 rerun -- cache off + chunked prefill on both stacks Corrects the earlier apples-to-oranges comparison and drops the misconfigured vLLM PP=2 results entirely. Fairness fixes: - SGLang had radix cache ON while vLLM had it OFF. Disabled SGLang radix + chunked-prefix-cache. Turned out radix barely mattered here [random prompts, ratio 0.03, low prefix overlap]: SGLang c=48 802 -> 876, if anything faster. - Enabled chunked prefill on vLLM too [had been off to match Trn2]. This is the single biggest vLLM lever: DP+EP c=48 270 -> 748 tok/s. Final fair numbers [EFA, cache off, chunked on], out-tok/s c=1/16/48: - SGLang+EFA TP16/DP2/EP16: 32.6 / 344.0 / 875.8 - vLLM+EFA DP2xTP8+EP+chunked: 27.8 / 221.1 / 748.3 SGLang and a properly-tuned vLLM are within ~15% at c=48, not 5x. The 5x was pure misconfiguration [PP=2 + chunked-prefill off]. Removed run_vllm_h100_multinode.sh [PP=2] and all PP2/socket-vLLM/radix-on result files. Kept sglang socket results for the EFA-vs-socket point. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 32 +++--- .../MiMo-V2.5-Pro/perf_test/h100/README.md | 9 +- ...tp16dp2_c1.txt => sglang_efa_norad_c1.txt} | 42 ++++---- ...16dp2_c16.txt => sglang_efa_norad_c16.txt} | 44 ++++----- ...16dp2_c48.txt => sglang_efa_norad_c48.txt} | 44 ++++----- .../h100/results/vllm_efa_dp2tp8ep_c16.txt | 42 -------- ...1.txt => vllm_efa_dp2tp8ep_chunked_c1.txt} | 34 +++---- ....txt => vllm_efa_dp2tp8ep_chunked_c16.txt} | 46 ++++----- ....txt => vllm_efa_dp2tp8ep_chunked_c48.txt} | 44 ++++----- .../h100/results/vllm_efa_tp8pp2_c1.txt | 42 -------- .../h100/results/vllm_efa_tp8pp2_c48.txt | 42 -------- .../h100/results/vllm_socket_tp8pp2_c1.txt | 42 -------- .../h100/results/vllm_socket_tp8pp2_c16.txt | 42 -------- .../h100/results/vllm_socket_tp8pp2_c48.txt | 42 -------- .../perf_test/h100/run_vllm_h100_multinode.sh | 97 ------------------- 15 files changed, 148 insertions(+), 496 deletions(-) rename contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/{sglang_efa_tp16dp2_c1.txt => sglang_efa_norad_c1.txt} (71%) rename contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/{sglang_efa_tp16dp2_c16.txt => sglang_efa_norad_c16.txt} (59%) rename contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/{sglang_efa_tp16dp2_c48.txt => sglang_efa_norad_c48.txt} (65%) delete mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_dp2tp8ep_c16.txt rename contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/{vllm_efa_dp2tp8ep_c1.txt => vllm_efa_dp2tp8ep_chunked_c1.txt} (72%) rename contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/{vllm_efa_tp8pp2_c16.txt => vllm_efa_dp2tp8ep_chunked_c16.txt} (63%) rename contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/{vllm_efa_dp2tp8ep_c48.txt => vllm_efa_dp2tp8ep_chunked_c48.txt} (63%) delete mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_tp8pp2_c1.txt delete mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_efa_tp8pp2_c48.txt delete mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_socket_tp8pp2_c1.txt delete mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_socket_tp8pp2_c16.txt delete mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/vllm_socket_tp8pp2_c48.txt delete mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_multinode.sh diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index b579be01..4ccb9d29 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -453,7 +453,7 @@ caching + chunked prefill **disabled** for a fair comparison. See - **Trn2**: trn2.48xlarge, TP=64 / moe_tp=1 / moe_ep=64, BS=48, `seq_len=512`, vllm-neuron 0.5.3. - **H100**: 2× p5 (16× H100-80GB). The ~963 GB FP8 weights don't fit on one 8×80 GB node, so both frameworks shard across two nodes. - - **vLLM 0.25.1**: TP=8 × PP=2 (MiMo's 8 KV heads cap TP at 8; TP=16 fails the KV-head divisibility check). CUDA graphs on. + - **vLLM 0.25.1**: DP=2 × TP=8 + `--enable-expert-parallel` + chunked prefill. (MiMo's 8 KV heads cap per-rank TP at 8; DP replicates the TP=8 attention across nodes and shards MoE over TP×DP=16. CUDA graphs on.) - **SGLang 0.5.15**: TP=16 × DP=2 with `--enable-dp-attention` + EP=16 (DP-attention shards KV heads so TP=16 works). **Cross-node fabric matters enormously.** The stock `vllm/vllm-openai` and @@ -462,37 +462,35 @@ sockets (`Using network Socket`) instead of EFA RDMA. Rebuilding both images with GDRCopy + the AWS EFA installer (`Dockerfile.{vllm,sglang}-efa`, with `NCCL_NET_PLUGIN=ofi`) switches NCCL to `efa-direct` over 32 NICs. -Output token throughput (tok/s), higher is better: +All runs below have **prefix/radix cache OFF and chunked prefill ON on both +frameworks** (verified fair — see notes). Output token throughput (tok/s), +higher is better: | Platform (fabric) | c=1 | c=16 | c=48 | |----------------------------|------|-------|-------| -| **SGLang + EFA** (TP16/DP2/EP16) | 32.6 | 345.4 | **802.6** | +| **SGLang + EFA** (TP16/DP2/EP16) | 32.6 | 344.0 | **875.8** | +| vLLM + EFA, **DP2×TP8 + EP + chunked** (best vLLM) | 27.8 | 221.1 | **748.3** | | SGLang, socket (TP16/DP2/EP16) | 11.2 | 70.4 | 139.2 | -| vLLM + EFA, **DP2×TP8 + EP** (best vLLM) | 27.8 | 299.0 | 270.5 | -| vLLM + EFA, TP8/PP2 | 115.0 | 261.8 | 157.8 | -| vLLM, socket, TP8/PP2 | 113.2 | 373.1 | 141.1 | | Trn2 (TP64) | 5.2 | *(n/m)* | 71.9 | -SGLang EFA vs socket (out-tok/s): c=1 2.9×, c=16 4.9×, c=48 5.8× — the EFA +SGLang EFA vs socket (out-tok/s): c=1 2.9×, c=16 4.9×, c=48 6.3× — the EFA advantage grows with concurrency as more cross-node all-to-all traffic piles up. Median TTFT / TPOT (ms) at c=48: | Platform (fabric) | Median TTFT | Median TPOT | |----------------------------|-------------|-------------| -| **SGLang + EFA** | **1,080** | **50.6** | -| vLLM + EFA, DP2×TP8 + EP | 7,614 | 117.6 | -| vLLM + EFA, TP8/PP2 | 13,765 | 216.9 | -| SGLang, socket | 16,831 | 231.0 | -| vLLM, socket, TP8/PP2 | 13,175 | 231.7 | +| **SGLang + EFA** | **726** | **47.6** | +| vLLM + EFA, DP2×TP8 + EP + chunked | 624 | 57.2 | | Trn2 | 6,912 | 572.1 | Notes: -- **At c=48, SGLang + EFA wins decisively**: 803 out-tok/s and 1.08 s TTFT — ~5× vLLM's 158 and ~11× Trn2's 72, with ~13× lower TTFT than vLLM. -- **EFA is make-or-break for SGLang, nearly irrelevant for vLLM's PP path.** SGLang c=48 collapses from 803 → 139 tok/s (TTFT 1.08 s → 16.8 s) on TCP sockets — its EP=16 all-to-all + DP-attention move a lot of data cross-node. vLLM TP8/PP2 barely moves (141 → 158) because PP=2 only ships pipeline-boundary activations; its bottleneck is the pipeline bubble, not the fabric. **The stock images fall back to sockets — always verify NCCL logs `NET/OFI ... efa-direct`, not `Using network Socket`.** -- **vLLM: use DP+EP, not PP.** Replacing TP8/PP2 with **DP2×TP8 + `--enable-expert-parallel`** (vLLM's analogue to SGLang's DP-attention: each DP rank runs TP=8 attention, dividing the 8 KV heads, while MoE shards across TP×DP=16) lifts vLLM c=48 from 158 → **270 tok/s** and halves TTFT (13.8 → 7.6 s), with no pipeline bubble. Still ~3× behind SGLang (270 vs 803) — the remaining gap is SGLang's more mature MoE kernels (DeepGEMM, EP all-to-all) and prefill parallelism — but DP+EP is clearly the right multi-node shape for this model on vLLM. -- **Low vs high concurrency flips the winner, and picks the parallelism.** At c=1, vLLM **TP8/PP2** (115 tok/s, TTFT 69 ms) beats everything — a single stream pipelines cleanly with no DP/EP sync. Both DP/EP layouts are slow at c=1 (vLLM DP+EP 27.8, SGLang 32.6): the DP/EP all-reduce/all-to-all is pure overhead with one request. Once concurrency fills the wide parallelism (c=16/48), DP/EP wins big (SGLang 803, vLLM DP+EP 270/299) and PP falls behind (158). Takeaway: PP for latency-bound single-stream, DP+EP for throughput-bound serving. -- **Not equal-hardware**: H100 uses 16 GPUs / 2 nodes vs Trn2's single 64-core instance, CUDA graphs vs eager. Normalize per-device / per-dollar before drawing efficiency conclusions. +- **At c=48 SGLang and a properly-tuned vLLM are close**: SGLang 876 vs vLLM 748 out-tok/s (~15%), TTFT 726 vs 624 ms. Both are ~10-12× Trn2's 72. +- **The vLLM config matters enormously — use DP+EP + chunked prefill.** MiMo's 8 KV heads cap per-rank TP at 8, so a single node can't hold the model with pure TP. Do **not** reach for pipeline parallel (PP=2) as the multi-node fallback — it leaves a pipeline bubble and was ~5× slower in our testing. Instead use **DP2×TP8 + `--enable-expert-parallel`** (vLLM's analogue to SGLang's `--enable-dp-attention`: each DP rank runs TP=8 attention dividing the 8 KV heads; MoE shards across TP×DP=16) **with chunked prefill enabled**. Chunked prefill is the single biggest lever (it overlaps prefill with decode; not a caching shortcut, so fair to leave on) and gets vLLM to 748 tok/s @ c=48. +- **EFA is make-or-break for the DP/EP path.** SGLang c=48 collapses 876 → 139 tok/s (TTFT 726 ms → 16.8 s) on TCP sockets — the EP=16 all-to-all moves a lot of data cross-node, and the same applies to vLLM's DP+EP path. **Stock images fall back to sockets — always verify NCCL logs `NET/OFI ... efa-direct`, not `Using network Socket`.** +- **radix/prefix cache didn't matter for this benchmark** — disabling SGLang's radix cache left c=48 unchanged (802 → 876, if anything faster), because `--random-range-ratio 0.03` prompts share little prefix. Both frameworks are benchmarked with caches off for strict fairness. +- **DP/EP is a throughput play, not a latency one.** At c=1 the DP/EP all-reduce/all-to-all sync is pure overhead (SGLang 32.6, vLLM DP+EP 27.8 tok/s); it only pays off once concurrency fills the wide parallelism (c=16/48). If single-stream latency is the goal on a model this size, a narrower non-DP layout would win — but that trades away the aggregate throughput these numbers are about. +- **Not equal-hardware**: H100 = 16 GPUs / 2 nodes with CUDA graphs; Trn2 = one 64-core instance, eager. Normalize per-device / per-dollar before drawing efficiency conclusions. - **Both GPU stacks produce coherent output across all requests** — no "garbled-after-first-request" bug (Trn2 issue #31), consistent with GPU stacks dequantizing FP8→BF16 before matmul. > **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md index 98d2bbdb..8412e191 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md @@ -11,7 +11,10 @@ The FP8 weights are ~963 GB, which does not fit on a single 8×H100-80GB node must divide 8 (TP≤8) — `--tensor-parallel-size 16` fails with *"TP size must evenly split the number of KV heads"*. So: -- **vLLM**: TP=8 (within a node) × PP=2 (across the two nodes) = world size 16. +- **vLLM**: DP=2 × TP=8 + `--enable-expert-parallel` + chunked prefill = world + size 16. Each DP rank runs TP=8 attention (dividing the 8 KV heads); MoE shards + across TP×DP=16. (Do NOT use PP=2 as the multi-node fallback — it leaves a + pipeline bubble and measured ~5× slower.) - **SGLang**: TP=16 × DP=2 with `--enable-dp-attention` (DP-attention shards the KV heads differently, so TP=16 is fine) + EP=16. @@ -53,8 +56,8 @@ Run the **same** launch script on both nodes, changing only `NODE_RANK` ```bash # vLLM (node 0 = API server on :8000, node 1 = --headless follower) -NODE_RANK=0 MASTER_ADDR= bash run_vllm_h100_multinode.sh -NODE_RANK=1 MASTER_ADDR= bash run_vllm_h100_multinode.sh +DP_RANK=0 DP_ADDR= bash run_vllm_h100_dp.sh +DP_RANK=1 DP_ADDR= bash run_vllm_h100_dp.sh # SGLang (both nodes run the same command; node 0 serves on :30000) NODE_RANK=0 DIST_INIT_ADDR=:20000 bash run_sglang_h100_multinode.sh diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_tp16dp2_c1.txt b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_norad_c1.txt similarity index 71% rename from contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_tp16dp2_c1.txt rename to contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_norad_c1.txt index 27f8f04c..5f643d7e 100644 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_tp16dp2_c1.txt +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_norad_c1.txt @@ -1,11 +1,11 @@ -INFO 07-21 07:48:31 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. -INFO 07-21 07:48:31 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. -W0721 07:48:32.318000 10 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' -INFO 07-21 07:48:33 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. -WARNING 07-21 07:48:33 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +INFO 07-21 08:37:54 [importing.py:53] Triton is installed but 0 active driver(s) found (expected 1). Disabling Triton to prevent runtime errors. +INFO 07-21 08:37:54 [importing.py:88] Triton not installed or not compatible; certain GPU-related functions will not be available. +W0721 08:37:55.068000 10 torch/utils/cpp_extension.py:140] No CUDA runtime is found, using CUDA_HOME='/usr/local/cuda' +INFO 07-21 08:37:56 [main.py:69] Unspecified platform detected, switching to CPU Platform instead. +WARNING 07-21 08:37:56 [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. -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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-46d2821b-', 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 07:48:35 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-0af128a5-', 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 08:37:57 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] 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. @@ -13,30 +13,30 @@ 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=64, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-6cfe6e29-', 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 07:49:47 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +Namespace(subparser='bench', bench_type='serve', dispatch_function=, trust_remote_code=False, seed=0, num_prompts=64, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-6082cfad-', 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 08:39:09 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] 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. @@ -13,30 +13,30 @@ Starting main benchmark run... Traffic request rate: inf Burstiness factor: 1.0 (Poisson process) Maximum request concurrency: 16 - 0%| | 0/64 [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=360, random_output_len=120, 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=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-b4a582a3-', 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 07:18:39 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +Namespace(subparser='bench', bench_type='serve', dispatch_function=, 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=360, random_output_len=120, 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=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-8f2e0b17-', 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 08:39:45 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] 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. @@ -13,30 +13,30 @@ Starting main benchmark run... Traffic request rate: inf Burstiness factor: 1.0 (Poisson process) Maximum request concurrency: 48 - 0%| | 0/96 [00:00, trust_remote_code=False, seed=0, num_prompts=64, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-73fe9fb8-', 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 08:21:37 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] -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/64 [00:00, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-b836591c-', 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 08:20:16 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-562453c1-', 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 08:46:07 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] 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. @@ -13,7 +13,7 @@ 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=64, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-dc47e82e-', 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 07:43:03 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +Namespace(subparser='bench', bench_type='serve', dispatch_function=, trust_remote_code=False, seed=0, num_prompts=64, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-7ed12e73-', 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 08:47:28 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] 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. @@ -13,30 +13,30 @@ Starting main benchmark run... Traffic request rate: inf Burstiness factor: 1.0 (Poisson process) Maximum request concurrency: 16 - 0%| | 0/64 [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=360, random_output_len=120, 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=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-c06615dc-', 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 08:16:34 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] +Namespace(subparser='bench', bench_type='serve', dispatch_function=, 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=360, random_output_len=120, 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=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-51344245-', 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 08:48:15 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] 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. @@ -13,30 +13,30 @@ Starting main benchmark run... Traffic request rate: inf Burstiness factor: 1.0 (Poisson process) Maximum request concurrency: 48 - 0%| | 0/96 [00:00, 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=360, random_output_len=120, 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-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-3a26f429-', 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 07:42:34 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] -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=360, random_output_len=120, 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=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-66923b84-', 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 07:31:29 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] -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: 48 - 0%| | 0/96 [00:00, 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=360, random_output_len=120, 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='127.0.0.1', port=8000, endpoint='/v1/completions', header=None, max_concurrency=1, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-3871b99d-', 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 06:41:17 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] -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=64, 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=360, random_output_len=120, 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='127.0.0.1', port=8000, endpoint='/v1/completions', header=None, max_concurrency=16, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-8e071f27-', 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 06:41:45 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] -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/64 [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=360, random_output_len=120, 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='127.0.0.1', port=8000, endpoint='/v1/completions', header=None, max_concurrency=48, model='MiMo-V2.5-Pro', input_len=None, output_len=None, tokenizer='/wk/MiMo-V2.5-Pro', 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-bfcdb174-', 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 06:37:54 [utils.py:90] Sampling input_len from [349, 371] and output_len from [116, 124] -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: 48 - 0%| | 0/96 [00:00 TP<=8. The ~963 GB -# FP8 weights don't fit on one 8x80GB node, so we use TP=8 (within a node) x -# PP=2 (across the two nodes) = world size 16, instead of TP=16 (which fails: -# "TP size must evenly split the number of KV heads"). -TP="${TP:-8}" -PP="${PP:-2}" -NNODES="${NNODES:-2}" -NODE_RANK="${NODE_RANK:?set NODE_RANK=0 on master, 1 on worker}" -MASTER_ADDR="${MASTER_ADDR:-172.31.45.21}" # P5-1 private IP -MASTER_PORT="${MASTER_PORT:-29500}" -IFACE="${IFACE:-enp71s0}" # NCCL/GLOO socket iface -MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}" -# EFA-enabled image (GDRCopy + aws-ofi-nccl). The stock vllm/vllm-openai lacks -# aws-ofi-nccl, so cross-node NCCL falls back to TCP sockets (~5x slower). Built -# from Dockerfile.vllm-efa. -IMAGE="${IMAGE:-vllm-efa:latest}" -CACHE_DIR="${CACHE_DIR:-/opt/dlami/nvme/vllm_cache}" -CTR_MODEL="/models/MiMo-V2.5-Pro" -mkdir -p "$CACHE_DIR" - -if [ ! -f "$MODEL_DIR/config.json" ]; then - echo "ERROR: model not found at $MODEL_DIR (config.json missing)." >&2 - exit 1 -fi - -echo "==========================================" -echo "MiMo-V2.5-Pro vLLM (2-node TP=8 x PP=2, EFA, Docker)" -echo " Model: $MODEL_DIR" -echo " Node rank: $NODE_RANK / $NNODES" -echo " Master: $MASTER_ADDR:$MASTER_PORT iface: $IFACE" -echo " TP: $TP PP: $PP Port: $PORT max-model-len: $MAX_MODEL_LEN" -echo "==========================================" - -# Only node 0 runs the API server; follower nodes (rank != 0) run --headless -# (no API server / engine-core client), otherwise they crash with -# "AssertionError: collective_rpc should not be called on follower node". -COMMON_ARGS=( - --served-model-name MiMo-V2.5-Pro - --trust-remote-code - --generation-config vllm - --tensor-parallel-size "$TP" - --pipeline-parallel-size "$PP" - --nnodes "$NNODES" - --node-rank "$NODE_RANK" - --master-addr "$MASTER_ADDR" - --master-port "$MASTER_PORT" - --max-model-len "$MAX_MODEL_LEN" -) -if [ "$NODE_RANK" = "0" ]; then - ROLE_ARGS=( - --host 0.0.0.0 - --port "$PORT" - # Match Trn2 start_vllm_server.sh for a fair comparison: no prefix cache - # (else prefill is skipped on repeated prompts -- we saw 66% hit rate and - # inflated throughput) and no chunked prefill. - --no-enable-prefix-caching - --no-enable-chunked-prefill - --tool-call-parser mimo - --enable-auto-tool-choice - --reasoning-parser mimo - ) -else - ROLE_ARGS=( --headless ) -fi - -# --network host is required for cross-node NCCL/torch.distributed rendezvous -# (port mapping does not work for the collective transport). --ipc=host and -# --shm-size for large NCCL buffers. EFA devices are exposed via /dev/infiniband. -exec docker run --rm --gpus all \ - --network host --privileged --ipc=host --shm-size=32g \ - --device /dev/infiniband \ - -v "${MODEL_DIR}:${CTR_MODEL}:ro" \ - -v "${CACHE_DIR}:/root/.cache" \ - -e CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ - -e GLOO_SOCKET_IFNAME="$IFACE" \ - -e NCCL_SOCKET_IFNAME="$IFACE" \ - -e NCCL_DEBUG="${NCCL_DEBUG:-WARN}" \ - -e NCCL_DEBUG_SUBSYS="${NCCL_DEBUG_SUBSYS:-INIT,NET}" \ - -e FI_PROVIDER=efa \ - -e NCCL_NET_PLUGIN="${NCCL_NET_PLUGIN:-ofi}" \ - -e VLLM_HOST_IP="$(hostname -I | awk '{print $1}')" \ - "$IMAGE" "$CTR_MODEL" \ - "${COMMON_ARGS[@]}" "${ROLE_ARGS[@]}" From 3d8b987984a48f654de0371e23316584d1464f37 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:46:14 +0000 Subject: [PATCH 39/47] [contrib] MiMo-V2.5-Pro: fix decode sliding-window mask to use per-batch-slot position Same fix as MiMo-V2.5. The SWA-layer decode mask used position_ids[0, 0] (slot 0) broadcast to the whole batch, correct only for synchronous batches (smoke path). Under vLLM continuous batching, per-slot positions differ, so the broadcast window mis-clips other slots KV. Fix: position_ids[:, 0] -> [bsz, 1, 1, kv_seq_len] per-slot mask. Numerically identical for synchronous batches; correct per-slot for async. Correctness fix; does NOT by itself resolve issue #31 (full-attention layers still attend stale KV because vLLM never resets NxDI per-batch KV state). Co-Authored-By: Claude Opus 4.8 --- .../MiMo-V2.5-Pro/src/modeling_mimo_v2.py | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py index 230f627c..837fd719 100644 --- a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py @@ -698,13 +698,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 429ad081b2432e4107bbc16edf79e34798a4c31f Mon Sep 17 00:00:00 2001 From: whn09 Date: Tue, 21 Jul 2026 22:47:12 +0800 Subject: [PATCH 40/47] Fix batched-decode garble: per-row causal upper bound on prior KV MiMo-V2.5-Pro 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. 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. Mirrors the identical fix validated on MiMo-V2.5 (same decode-attention structure). Co-Authored-By: Claude Opus 4.8 --- .../MiMo-V2.5-Pro/src/modeling_mimo_v2.py | 30 +++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py index 837fd719..2e6580d0 100644 --- a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py @@ -698,6 +698,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 c42e5ef05641f2e6df274dcc7ae05032377f8e2d Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 22 Jul 2026 16:30:16 +0800 Subject: [PATCH 41/47] h100/Pro: add 2-node long-context @ c=48 exploration (up to 32K input) How long an input can 2x p5.48xlarge (16xH100) sustain at concurrency 48 (Pro's minimum working batch = 384 experts / top-8 = 48)? Added bench_ctx_c48.sh and results/longctx_c48/. Result: c=48 completes with ZERO failures from ISL=2K all the way to 32K (3.15M total input tokens through a 148K-token KV pool). Works because 60/70 layers are sliding-window (window-capped KV) so long prompts cost far less than the naive even-split; SGLang queues the excess. Throughput stays ~14K tok/s (prefill-bound); TTFT is the cost (1.8s @ 2K -> 91s @ 32K). Key tuning finding: the gating failure is prefill-activation OOM, NOT KV capacity. mem-fraction 0.95 (KV pool 378K) and 0.92 (240K) both load fine but crash with torch.OutOfMemoryError on the first real c=48 prefill batch (only 2.4/4.7 GB idle headroom). mem-fraction 0.90 + DP-attention's auto-reduced 4096 chunked-prefill is the stable operating point. Also required an NVIDIA Fabric Manager version fix (driver 595 vs fabricmanager 610) + reboot on both nodes. Co-Authored-By: Claude Opus 4.8 --- .../perf_test/h100/bench_ctx_c48.sh | 29 ++++++ .../h100/results/longctx_c48/README.md | 59 ++++++++++++ .../longctx_c48/pro_2node_c48_sweep.txt | 94 +++++++++++++++++++ 3 files changed, 182 insertions(+) create mode 100755 contrib/models/MiMo-V2.5-Pro/perf_test/h100/bench_ctx_c48.sh create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/longctx_c48/README.md create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/longctx_c48/pro_2node_c48_sweep.txt diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/bench_ctx_c48.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/bench_ctx_c48.sh new file mode 100755 index 00000000..60e3bc8f --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/bench_ctx_c48.sh @@ -0,0 +1,29 @@ +#!/bin/bash +# Fixed concurrency 48 (Pro's minimum working BS: 384 experts / top-8 = 48), +# sweep input length to find the max sustainable context on the 2-node cluster. +# OSL=128 isolates prefill + KV capacity. Run on the head node (P5-1) against the +# already-running multinode server on :30000. +# +# Usage (head node): +# bash bench_ctx_c48.sh 2048 4096 8192 16384 32768 +set -e +PORT="${PORT:-30000}" +C="${C:-48}" +OUTPUT_LEN="${OUTPUT_LEN:-128}" +RESULTS_DIR="${RESULTS_DIR:-/opt/dlami/nvme/models/bench_results/mimo_v2_5_pro_h100/longctx_c48}" +mkdir -p "$RESULTS_DIR" + +for ISL in "$@"; do + echo "===== c=$C ISL=$ISL OSL=$OUTPUT_LEN np=96 =====" + 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-Pro --tokenizer /wk/MiMo-V2.5-Pro --endpoint /v1/completions \ + --dataset-name random --num-prompts 96 \ + --random-input-len $ISL --random-output-len $OUTPUT_LEN \ + --random-range-ratio 0.02 --max-concurrency $C" 2>&1 | \ + tee "$RESULTS_DIR/isl${ISL}.txt" | \ + grep -E "Successful requests|Failed|Maximum request concurrency:|Benchmark duration|Total input|Output token throughput|Total token throughput|Median TTFT|P99 TTFT|Median TPOT" + echo "--- ISL=$ISL done ---" +done +echo ALL_DONE diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/longctx_c48/README.md b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/longctx_c48/README.md new file mode 100644 index 00000000..0d7fab19 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/longctx_c48/README.md @@ -0,0 +1,59 @@ +# MiMo-V2.5-Pro 2-node H100 long-context @ concurrency 48 + +How long an input can 2× p5.48xlarge (16× H100-80GB) sustain at **concurrency 48** +(Pro's minimum working batch: 384 experts / top-8 routing = 48)? The default +Pro bench used ISL=360/OSL=120 (seq_len 512) — this sweep pushes the input length. + +Server: `run_sglang_h100_multinode.sh` on both nodes (TP=16, DP=2 + DP-attention, +EP=16), **`MEM_FRAC=0.90`** and **chunked-prefill auto-reduced to 4096** (DP +attention lowers it). Bench: `bench_ctx_c48.sh 2048 4096 8192 12288 16384 24576 +32768` (96 prompts, OSL=128). Raw logs: `pro_2node_c48_sweep.txt`, `isl*.txt`. + +## Result: c=48 sustains up to 32K input (0 failures) + +| ISL (input) | Succeeded | Total tput (tok/s) | TTFT median (s) | TTFT P99 (s) | TPOT median (ms) | +|---|---|---|---|---|---| +| 2,048 | 96/96 ✅ | 9359 | 1.8 | 5.6 | 67 | +| 4,096 | 96/96 ✅ | 9738 | 5.6 | 16.3 | 78 | +| 8,192 | 96/96 ✅ | 13140 | 11.7 | 25.5 | 113 | +| 12,288 | 96/96 ✅ | 13745 | 24.1 | 34.2 | 132 | +| 16,384 | 96/96 ✅ | 13872 | 37.1 | 49.9 | 140 | +| 24,576 | 96/96 ✅ | 14485 | 62.2 | 72.8 | 125 | +| 32,768 | 96/96 ✅ | 14791 | 91.3 | 99.7 | 112 | + +**Every level to 32K completed with zero failures** — at ISL=32K that's 96 × +32768 = 3.15M total input tokens pushed through a 148K-token KV pool. + +## Why it works despite a small KV pool + +- **KV pool is small because weights dominate.** Pro's FP8 weights are ~963 GB / + 16 GPUs = **60 GB/GPU**, leaving only ~11 GB/GPU for KV at mem-fraction 0.90 → + `max_total_num_tokens = 148672` (~3K tokens/req if split evenly across 48). +- **Hybrid attention makes long inputs cheap.** 60 of 70 layers are sliding-window + (window 128) with window-capped KV; only 10 are full attention. A long prompt's + KV is dominated by those 10 layers, far below the naive per-token estimate, so + requests fit well beyond the even-split figure. +- **SGLang queues the rest.** At high ISL the resident set is capped by the KV + pool; excess requests queue, so throughput stays flat (~14K tok/s, prefill-bound) + while **TTFT grows** (1.8 s at 2K → 91 s at 32K). No request is dropped. + +## ⚠️ Memory tuning is load-bearing (mem-fraction + chunked-prefill) + +The failure that gates this is **prefill activation OOM, not KV capacity**: + +| mem-fraction | KV pool | idle free/GPU | c=48 prefill result | +|---|---|---|---| +| 0.95 | 378,240 | 2.43 GB | ✅ loads, ❌ **OOM on first c=48 prefill batch** | +| 0.92 | 240,512 | 4.66 GB | ✅ loads, ❌ **OOM at c=48 ISL=2048** | +| **0.90** | 148,672 | 6.14 GB | ✅ **stable to 32K** (chunked-prefill 4096) | + +Raising mem-fraction grows the *paper* KV pool but shrinks the headroom that 48 +concurrent prefills need for activations — 0.95/0.92 crash with +`torch.OutOfMemoryError` the moment a real c=48 batch runs. **0.90 + the +DP-attention-reduced 4096 chunked-prefill size is the stable operating point.** + +> **Prerequisite:** both nodes needed an NVIDIA Fabric Manager fix first — the +> DLAMI's background update had bumped fabricmanager to 610 while the driver +> stayed at 595, breaking multi-GPU CUDA init (`Error 802: system not yet +> initialized`). Reinstall the matching `nvidia-fabricmanager-595=595.71.05-*` +> and reboot on **both** nodes before launching. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/longctx_c48/pro_2node_c48_sweep.txt b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/longctx_c48/pro_2node_c48_sweep.txt new file mode 100644 index 00000000..7843acf7 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/longctx_c48/pro_2node_c48_sweep.txt @@ -0,0 +1,94 @@ +===== c=48 ISL=2048 OSL=128 np=96 ===== +Maximum request concurrency: 48 +Successful requests: 96 +Failed requests: 0 +Maximum request concurrency: 48 +Benchmark duration (s): 22.34 +Total input tokens: 196738 +Output token throughput (tok/s): 551.45 +Total token throughput (tok/s): 9358.90 +Median TTFT (ms): 1827.48 +P99 TTFT (ms): 5583.13 +Median TPOT (ms): 66.82 +--- ISL=2048 done --- +===== c=48 ISL=4096 OSL=128 np=96 ===== +Maximum request concurrency: 48 +Successful requests: 96 +Failed requests: 0 +Maximum request concurrency: 48 +Benchmark duration (s): 41.67 +Total input tokens: 393481 +Output token throughput (tok/s): 295.59 +Total token throughput (tok/s): 9737.72 +Median TTFT (ms): 5607.88 +P99 TTFT (ms): 16342.96 +Median TPOT (ms): 77.86 +--- ISL=4096 done --- +ALL_DONE +===== c=48 ISL=8192 OSL=128 np=96 ===== +Maximum request concurrency: 48 +Successful requests: 96 +Failed requests: 0 +Maximum request concurrency: 48 +Benchmark duration (s): 60.83 +Total input tokens: 786950 +Output token throughput (tok/s): 202.51 +Total token throughput (tok/s): 13139.79 +Median TTFT (ms): 11726.26 +P99 TTFT (ms): 25506.28 +Median TPOT (ms): 112.78 +--- ISL=8192 done --- +===== c=48 ISL=12288 OSL=128 np=96 ===== +Maximum request concurrency: 48 +Successful requests: 96 +Failed requests: 0 +Maximum request concurrency: 48 +Benchmark duration (s): 86.78 +Total input tokens: 1180429 +Output token throughput (tok/s): 141.95 +Total token throughput (tok/s): 13744.53 +Median TTFT (ms): 24141.47 +P99 TTFT (ms): 34176.05 +Median TPOT (ms): 131.58 +--- ISL=12288 done --- +===== c=48 ISL=16384 OSL=128 np=96 ===== +Maximum request concurrency: 48 +Successful requests: 96 +Failed requests: 0 +Maximum request concurrency: 48 +Benchmark duration (s): 114.34 +Total input tokens: 1573907 +Output token throughput (tok/s): 107.73 +Total token throughput (tok/s): 13872.41 +Median TTFT (ms): 37061.26 +P99 TTFT (ms): 49876.34 +Median TPOT (ms): 140.41 +--- ISL=16384 done --- +ALL_DONE +===== c=48 ISL=24576 OSL=128 np=96 ===== +Maximum request concurrency: 48 +Successful requests: 96 +Failed requests: 0 +Maximum request concurrency: 48 +Benchmark duration (s): 163.83 +Total input tokens: 2360860 +Output token throughput (tok/s): 75.19 +Total token throughput (tok/s): 14485.45 +Median TTFT (ms): 62240.21 +P99 TTFT (ms): 72843.07 +Median TPOT (ms): 125.32 +--- ISL=24576 done --- +===== c=48 ISL=32768 OSL=128 np=96 ===== +Maximum request concurrency: 48 +Successful requests: 96 +Failed requests: 0 +Maximum request concurrency: 48 +Benchmark duration (s): 213.65 +Total input tokens: 3147810 +Output token throughput (tok/s): 57.65 +Total token throughput (tok/s): 14791.10 +Median TTFT (ms): 91318.11 +P99 TTFT (ms): 99738.05 +Median TPOT (ms): 112.16 +--- ISL=32768 done --- +ALL_DONE From 7e2ad3198b630f57873350d699bd638190e51ec0 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 23 Jul 2026 09:45:37 +0800 Subject: [PATCH 42/47] h100/Pro: add 2-node 4K/128 bench at c=1/16/48 (Trn2 4K align) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 4K input / 128 output on 2x p5.48xlarge, c=1/16/48 (48 = Pro's bs floor), to compare against the Trn2 Pro 4K recompile. 0 failures. Output throughput 29.4/235.2/368.5 tok/s (total 992/7779/12133) for c=1/16/48; TTFT median 629ms/1.7s/3.1s — far lighter than the 16K long-context sweep since 4K prefill is cheap. Same prefill-activation-OOM constraint: mem-fraction 0.90 with default chunk 16384 OOMs at c=16; CHUNK=4096 (DP-attention halves to 2048) is the stable point. Co-Authored-By: Claude Opus 4.8 --- .../h100/results/ctx4k_c48/README.md | 46 +++++++++++++++++++ .../h100/results/ctx4k_c48/pro_4k_sweep.log | 43 +++++++++++++++++ 2 files changed, 89 insertions(+) create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/ctx4k_c48/README.md create mode 100644 contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/ctx4k_c48/pro_4k_sweep.log diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/ctx4k_c48/README.md b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/ctx4k_c48/README.md new file mode 100644 index 00000000..4a0fc1cd --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/ctx4k_c48/README.md @@ -0,0 +1,46 @@ +# MiMo-V2.5-Pro 2-node H100 @ 4K input / 128 output — c=1/16/48 + +4K-input point to align with the Trn2 Pro 4K recompile. Concurrency 1 / 16 / 48 +(48 = Pro's minimum working batch: 384 experts / top-8 = 48). 2× p5.48xlarge +(16× H100), TP=16, DP=2 + DP-attention, EP=16. Raw log: `pro_4k_sweep.log`. + +Bench: `vllm bench serve --dataset-name random --random-input-len 4096 +--random-output-len 128 --random-range-ratio 0.03 --num-prompts 2*C +--max-concurrency C`. + +## Results (4K in / 128 out) + +| Concurrency | Succeeded | Output tput (tok/s) | Total tput (tok/s) | TTFT median (ms) | TTFT P99 (ms) | TPOT median (ms) | +|---|---|---|---|---|---|---| +| 1 | 2/2 ✅ | 29.4 | 992 | 629 | 703 | 29.3 | +| 16 | 32/32 ✅ | 235.2 | 7779 | 1723 | 4262 | 51.3 | +| 48 | 96/96 ✅ | 368.5 | 12133 | 3055 | 10918 | 100.1 | + +All 0 failures. + +## ⚠️ Memory tuning required (mem-fraction 0.90 + small chunked-prefill) + +Same prefill-activation-OOM constraint as the c=48 long-context sweep: Pro's FP8 +weights are ~60 GB/GPU, leaving little headroom. The gating knob is the +**chunked-prefill size**: + +- `mem-fraction 0.90` with the default `chunked-prefill 16384` (DP-attention + auto-reduces to 16384) **OOMs at c=16** (`torch.OutOfMemoryError`, tried 2.97 GB, + 2.48 GB free) — 16 requests prefilling 4K each at once overflows. +- **`CHUNK=4096` (DP-attention halves it to 2048)** keeps the prefill activation + peak low enough to run c=1/16/48 stably. This is the stable operating point. + +So launch with `MEM_FRAC=0.90 CHUNK=4096` on `run_pro_d2.sh` (or set +`--chunked-prefill-size 4096` on the multinode script). + +## Observations + +- Throughput scales 29 → 235 → 368 out tok/s (992 → 7779 → 12133 total) across + c=1/16/48. Total throughput is prefill-dominated (4K:128 ≈ 32:1 in:out). +- TTFT stays modest at 4K (629 ms → 3.1 s median) — far better than the long-context + sweep (16K was 37 s @ c=48), since 4K prefill is cheap and mostly fits without + heavy queueing. +- TPOT grows 29 → 100 ms with concurrency as the decode batch fills (Pro's 148K + KV pool is shared across the 48 slots). + +> Cross-platform reference for the Trn2 Pro 4K-compiled port at c=1/16/48. diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/ctx4k_c48/pro_4k_sweep.log b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/ctx4k_c48/pro_4k_sweep.log new file mode 100644 index 00000000..fbd305d5 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/ctx4k_c48/pro_4k_sweep.log @@ -0,0 +1,43 @@ +===== c=1 ISL=4096 OSL=128 np=2 ===== +WARNING 07-23 01:40:10 [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): 8.64 +Total input tokens: 8313 +Output token throughput (tok/s): 29.40 +Total token throughput (tok/s): 991.51 +Median TTFT (ms): 629.30 +P99 TTFT (ms): 703.31 +Median TPOT (ms): 29.29 +--- c=1 done --- +===== c=16 ISL=4096 OSL=128 np=32 ===== +WARNING 07-23 01:40:40 [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.40 +Total input tokens: 131287 +Output token throughput (tok/s): 235.24 +Total token throughput (tok/s): 7778.90 +Median TTFT (ms): 1723.10 +P99 TTFT (ms): 4262.25 +Median TPOT (ms): 51.30 +--- c=16 done --- +===== c=48 ISL=4096 OSL=128 np=96 ===== +WARNING 07-23 01:41:20 [cpu.py:441] Failed to import from vllm._C_AVX2: ModuleNotFoundError("No module named 'vllm._C_AVX2'") +Maximum request concurrency: 48 +Successful requests: 96 +Failed requests: 0 +Maximum request concurrency: 48 +Benchmark duration (s): 33.46 +Total input tokens: 393608 +Output token throughput (tok/s): 368.49 +Total token throughput (tok/s): 12132.52 +Median TTFT (ms): 3054.62 +P99 TTFT (ms): 10918.00 +Median TPOT (ms): 100.09 +--- c=48 done --- +ALL_DONE From 4a183f43f181ba00dc2f69d2653b7f2860eb265b Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 23 Jul 2026 10:35:43 +0800 Subject: [PATCH 43/47] Port data-parallel + context-parallel attention infra to MiMo-V2.5-Pro Mirror the V2.5 long-context changes into the Pro modeling code (the two are line-for-line siblings): data-parallel attention (attention_dp_degree binds projections to the reduced DP-attention TP group, batch split/all-gather in decode, sink-bias indexed by group-local rank, MHA/replication gates keyed off attn_tp_degree), the CP full-S KV-cache bugfix, and the _validate_chunked_attention_support override. All of this is backward compatible: with attention_dp_degree=1 (the default, and the current 512-context recipe) attn_tp_degree == tp_degree, every new gate reduces to the original comparison, and the DP/CP branches are inert. The committed 512 recipe is unchanged. NOTE: long context does NOT yet run on Pro. seq_len=4096 with attention_dp_degree=8 compiles, but a single core's HBM overflows at load (Tensors ~20.7 GB, of which weights ~19.4 GB and KV only ~1.3 GB, + 2 GB scratchpad + 1 GB code = ~23.8 / 24 GB). Pro's weights (384 experts x 70 layers) are already near the per-core limit, and DP attention adds per-DP-group weight replication that pushes it over; lowering dp only trims the small KV term. Serving Pro beyond 512 needs weight-side relief (e.g. moe_tp>1 sharding, which currently breaks the FP8 blockwise scale, or more aggressive quantization), not just parallelism tuning. This commit lands the infra so that work can build on it; V2.5 (smaller weights) validates the approach at 16K. Co-Authored-By: Claude Opus 4.8 --- .../MiMo-V2.5-Pro/src/modeling_mimo_v2.py | 162 ++++++++++++++++-- 1 file changed, 143 insertions(+), 19 deletions(-) diff --git a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py index 2e6580d0..3d609942 100644 --- a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py @@ -204,6 +204,22 @@ def __init__(self, *args, **kwargs): # Parse hybrid layer pattern self._parse_hybrid_pattern() + def _validate_chunked_attention_support(self): + """Skip the base chunked-attention validator for MiMo. + + MiMo implements sliding-window attention entirely in its own modeling + code (hand-rolled per-row sliding_mask via sliding_window_size) and never + uses the base framework "chunked attention" path. The HF config carries + attention_chunk_size (== sliding_window), which trips the base + _validate_chunked_attention_support (it hard-requires + attention_dp_degree == 1, forbidding the data-parallel attention we use + to eliminate CONVERT_TO_MHA KV replication). The chunked path is + inapplicable here, so this validator is a no-op. Overriding the method + (rather than deleting the attribute) is what takes effect in time, since + validate_config() runs at the end of InferenceConfig.__init__. + """ + 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: @@ -367,7 +383,9 @@ 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's per-rank cache sizing. + 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. @@ -391,8 +409,21 @@ 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 KV cache bloat that + # caps seq_len). We run BOTH prefill (CTE) and decode (TKG) on this same + # reduced group to avoid the base validate_tp_prefill_to_dp_decode path. + # See the V2.5 long-context notes / 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 @@ -413,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( @@ -457,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) @@ -560,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 KV-cache slice (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. + 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) @@ -624,9 +702,16 @@ 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 (mirrors V2.5): the cache must hold ALL sequence + # positions, not just this CP rank's local S/cp chunk. Saving + # the local chunk BEFORE the gather made the KV manager write + # only ~1/cp_degree of the sequence, so decode attended to a + # wildly incomplete context (grammatical-but-wrong / repetitive + # / drifting output). Gather first, then cache full-S, matching + # attention_base. key_states = gather_from_tensor_model_parallel_region_with_dim( key_states, gather_dim=2, process_group=get_context_parallel_attention_cp_group(), @@ -635,6 +720,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 @@ -759,7 +846,13 @@ 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 + # Index the sink bias by the rank WITHIN the attention TP group + # (equals the global rank when attention_dp_degree == 1). Under + # DP attention every DP group shards the Q heads identically, so + # a global-rank index would go out of bounds for ranks + # >= attn_tp_degree. + global_rank = parallel_state.get_tensor_model_parallel_rank() if parallel_state.model_parallel_is_initialized() else 0 + tp_rank = global_rank % self.attn_tp_degree 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) @@ -821,7 +914,13 @@ def forward( 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 + # Index the sink bias by the rank WITHIN the attention TP group + # (equals the global rank when attention_dp_degree == 1). Under + # DP attention every DP group shards the Q heads identically, so + # a global-rank index would go out of bounds for ranks + # >= attn_tp_degree. + global_rank = parallel_state.get_tensor_model_parallel_rank() if parallel_state.model_parallel_is_initialized() else 0 + tp_rank = global_rank % self.attn_tp_degree 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) @@ -858,6 +957,15 @@ def forward( attn_output = self.o_proj(attn_output) + # 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. + 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 # Return LOCAL key/value states for cache (each CP rank stores its portion) new_key_value = (key_states_for_cache, value_states_for_cache) @@ -1009,13 +1117,19 @@ 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 + # Check if we need GQA CONVERT_TO_MHA mode. Under DP attention the + # decisive comparison is attn_tp_degree (tp_degree // attention_dp_degree), + # not the full tp_degree: attn_tp <= num_kv_heads shards KV cleanly with + # no replication, so MHA conversion (and its cache bloat) is off. + attention_dp_degree = getattr( + config.neuron_config, "attention_dp_degree", 1 + ) or 1 + attn_tp_degree = self.tp_degree // attention_dp_degree 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) @@ -1133,9 +1247,19 @@ def convert_mimo_v2_hf_to_neuron_state_dict( full_num_kv_heads = config.num_key_value_heads # V2.5-Pro: 8 swa_num_kv_heads = config.swa_num_key_value_heads # V2.5-Pro: 8 + # Under data-parallel attention the K/V weights need replication 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 native head width (no + # replication) — this 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 06ae0d1bfcee1d7b897f6fc6e57f3650c9b3a244 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 23 Jul 2026 10:36:35 +0800 Subject: [PATCH 44/47] docs: note Pro long-context infra and the per-core HBM limit Add a "Long Context" section to the Pro README documenting that the DP+CP attention infrastructure is present and backward compatible, but that serving beyond 512 does not yet work: seq_len=4096 compiles but overflows a single core's HBM at load (weights ~19.4 GB dominate the 24 GB budget; KV is only ~1.3 GB, so parallelism tuning doesn't help). Points to weight-side relief as the real fix and to MiMo-V2.5 as the validated 16K reference. Co-Authored-By: Claude Opus 4.8 --- contrib/models/MiMo-V2.5-Pro/README.md | 28 ++++++++++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 4ccb9d29..43d6301f 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -495,6 +495,34 @@ Notes: > **Compile time:** the first Pro compile on SDK 2.29 is ~60 minutes for the TKG NEFF and ~15 minutes 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. First full server launch (compile + shard + warmup) is ~2 hours wall-clock. +## Long Context (infrastructure present, not yet serving > 512) + +The modeling code carries the same data-parallel + context-parallel attention +infrastructure as MiMo-V2.5 (`attention_dp_degree`, `cp_degree`, the CP full-S +KV-cache fix, and the chunked-attention validator override). It is backward +compatible: at `attention_dp_degree=1` (the default 512 recipe) it reduces to +the original behavior. + +**However, long context does not yet run on Pro.** `seq_len=4096` with +`attention_dp_degree=8` compiles, but a single NeuronCore's HBM overflows at +load: + +| HBM component (per core, dp=8) | Size | +|---|---| +| Tensors (weights ~19.4 GB + KV ~1.3 GB) | 20.7 GB | +| Scratchpad (4K activations) | 2.0 GB | +| Model code | 1.05 GB | +| **Total** | **~23.8 / 24 GB → overflow** | + +Pro's weights (384 experts x 70 layers) already sit near the per-core limit, and +DP attention adds per-DP-group weight replication that tips it over. KV is only +~1.3 GB, so lowering the batch or raising `attention_dp_degree` (which shrinks +KV) does not help — the weights are the bottleneck. Serving Pro beyond 512 needs +weight-side relief (`moe_tp>1` sharding, currently blocked by the FP8 blockwise +scale collapse — see "FP8 Configuration Notes"; or more aggressive quantization), +not parallelism tuning. The sibling **MiMo-V2.5** (smaller weights) validates the +same DP+CP approach at **seq_len=16384**. + ## Compatibility Matrix | Instance | Neuron SDK 2.29+ (PyTorch 2.9) | 2.21 and earlier | From 9e729db1f3153cebb03d8c6e239c30b8271c3225 Mon Sep 17 00:00:00 2001 From: whn09 Date: Wed, 29 Jul 2026 14:12:37 +0800 Subject: [PATCH 45/47] docs: pin neuronx-cc to 2.25.3371 -- 2.26.6360 miscompiles this model neuronx-cc 2.26.6360.0+6f180f47 (Neuron SDK 2.31.0) silently produces a numerically wrong NEFF for MiMo-V2.5-Pro. Generated text is garbage, while compilation succeeds with exit 0 and no warning. No model-code change is needed: the code in this directory is correct on both compilers. Compiling the byte-identical token-generation HLO -- same NxDI, same weights, same compile flags, same runtime, same process, only the neuronx-cc package swapped via a PYTHONPATH shadow -- gives correct output on 2.24.8799 and 2.25.3371 and garbage on 2.26.6360. 2.23/2.24/2.25 are all equivalent; the behaviour change lands exactly at 2.26. Document the downgrade recipe, note that a 2.26-built NEFF must be deleted first or the compiler cache returns it, and record that short prompts are not discriminating (>=200 tokens are needed to expose the corruption). narwhal (2.26's default-on NIR codegen backend) and the --allreduce-buffer-size default change were both ruled out by full rebuilds. The remaining lead is that 2.26 force-enables --internal-disable-fma-on-ios; that is not a confirmed diagnosis. Reported to AWS. Also update the Prerequisites toolchain line, which previously recommended the broken 2.26.6360 as "verified". Co-Authored-By: Claude Opus 5 --- contrib/models/MiMo-V2.5-Pro/README.md | 34 +++++++++++++++++++++++++- 1 file changed, 33 insertions(+), 1 deletion(-) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 43d6301f..61d9c7b5 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -43,6 +43,38 @@ Key features: **Known issue — vLLM serving is broken.** The first `/v1/chat/completions` request against `vllm-neuron` returns coherent output; every subsequent request returns garbled text. Same compiled NEFF serves 5 successive greedy generations byte-identically via the smoke path, so the bug is specifically in vllm-neuron's runtime / request-state handling. Tracking upstream at https://github.com/vllm-project/vllm-neuron/issues/31. Last updated 2026-04-30. +### ⚠️ Required: pin `neuronx-cc` to 2.25.3371 — 2.26.6360 miscompiles this model + +**`neuronx-cc 2.26.6360.0+6f180f47` (Neuron SDK 2.31.0) silently produces a numerically wrong NEFF for this model. Generated text is garbage. Compilation succeeds with exit 0 and no warning.** Use `2.25.3371.0+f524f7f8` (SDK 2.30.0) or `2.24.8799.0+6f62ff7c`. + +**No model-code change is needed** — the code in this directory is correct on both compilers. The *only* variable is the compiler package. Compiling the byte-identical token-generation HLO (same NxDI, same weights, same flags, same runtime, same process) gives: + +| neuronx-cc | 213-token prompt, greedy, `max_new_tokens=20` | +|---|---| +| 2.24.8799.0+6f62ff7c | `' Paris. Question: what is the capital city of France? Answer: Paris. Question:'` ✅ | +| 2.25.3371.0+f524f7f8 | `' Paris. Question: what is the capital city of France? Answer: Paris. Question:'` ✅ | +| **2.26.6360.0+6f180f47** | **`'., the the,,,,1.1.1.1. the the the'`** ❌ | + +2.23 / 2.24 / 2.25 are all equivalent; the behaviour change lands exactly at 2.26. + +To downgrade without disturbing the rest of the DLAMI venv, install the older compiler into its own venv and shadow only that package: + +```bash +python3 -m venv --system-site-packages /opt/dlami/nvme/venv_cc225 +/opt/dlami/nvme/venv_cc225/bin/pip install --no-deps 'neuronx-cc==2.25.3371.0+f524f7f8' + +# in the working venv: shadow ONLY neuronxcc. Do NOT prepend the new venv's bin to PATH -- +# neuronx-cc is invoked in-process by libneuronxla, and that venv has no torch. +export PYTHONPATH=/opt/dlami/nvme/venv_cc225/lib/python3.12/site-packages:$PYTHONPATH +python3 -c "import neuronxcc; print(neuronxcc.__version__)" # expect 2.25.3371.0+f524f7f8 +``` + +Delete any NEFF compiled by 2.26 before rebuilding, or the compiler cache will hand the bad artifact straight back. + +**Detection.** A short prompt is *not* discriminating — the broken build answers `"The capital of France is"` plausibly (`' a country that is known for its rich history…'`). Prompts of **≥200 tokens** are required to expose the corruption. Any accuracy gate for this model must use long prompts. + +Root-cause status: narrowed to the compiler and reported to AWS. The new default-on NIR codegen backend ("narwhal") and the `--allreduce-buffer-size` default change were both **ruled out** by full rebuilds (`--internal-backend-options='--enable-narwhal=0'` produced byte-identical garble). The one remaining 2.26-only difference we can see is that 2.26 force-enables `--internal-disable-fma-on-ios` (a sunda `InferIntrinsicOnCC` option whose own default is `False`, injected inside the compiled driver and not disableable from any CLI spelling); at TP=64 every layer crosses a collective. That is a **lead, not a confirmed diagnosis**. + ### Why BF16 attn + FP8 MoE Pro's attention weights have `abs_mean ≈ 0.00124`, roughly 4× smaller than V2.5 (256 experts). Under an all-FP8 recipe, the NKI blockwise FP8 accumulator on attention q/k/v at this magnitude drifts the logits across 70 layers and produces prompt-dependent gibberish (`"The capital of France is\n# 1000000000000000"`, `"Once upon a time in a small village there lived\n# 0000000000..."`, etc.). Dequantizing q/k/v to BF16 before the matmul restores coherent output. MoE experts (scales `≈ 2.3e-5`, similarly small) can stay FP8. @@ -74,7 +106,7 @@ GPU stacks (sglang on H100/H200) run the same OCP FP8 checkpoint correctly becau ## Prerequisites - **Instance**: trn2.48xlarge (128 physical NeuronCores, logical_nc_config=2 → 64 logical cores) -- **Neuron SDK**: 2.29 (Python 3.12, PyTorch 2.9). Verified toolchain (2026-07-21 instance): `neuronx-cc 2.26.6360.0`, `neuronx-distributed 0.19.28492`, `neuronx-distributed-inference 0.10.18399`, `torch 2.9.1`, `torch-neuronx 2.9.0.2.15`. The `neuronx-cc` package version (2.26.x) differs from the SDK release-train number (2.29); if you re-compile on a differently-imaged DLAMI, confirm these versions match or expect a cache miss. +- **Neuron SDK**: 2.29 (Python 3.12, PyTorch 2.9). Verified toolchain: **`neuronx-cc 2.25.3371.0+f524f7f8`** (see the pin warning above — do **not** use 2.26.6360, it miscompiles this model), `neuronx-distributed 0.19.28492`, `neuronx-distributed-inference 0.10.18399`, `torch 2.9.1`, `torch-neuronx 2.9.0.2.15`. The `neuronx-cc` package version differs from the SDK release-train number; if you re-compile on a differently-imaged DLAMI, confirm these versions match or expect a cache miss. - **Venv**: `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (used by preprocess, smoke, and vLLM serving alike; ships with the DLAMI and is where `0_setup.sh` installs the patched `vllm-neuron`). - **Disk**: ~3 TB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~962 GB, the Neuron-FP8 preprocessed output is ~1 TB, and `save_sharded_checkpoint=true` writes another ~300-1000 GB per compiled config (varies with recipe)). From 8b2915d93e2b2dbfc15368bb5522f9faf660b188 Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 30 Jul 2026 10:53:30 +0800 Subject: [PATCH 46/47] Pro: opt-in SWA KV cache truncation, and the measured context ceiling Adds MIMO_SWA_KV_TRUNCATION (default OFF): sizes the 60 sliding-window layers' KV cache to sliding_window=128 instead of seq_len, leaving the 10 full-attention layers at full length. At seq_len=4096 this cuts KV from 21.1 GB to 3.6 GB per rank, and the model compiles and loads at ~20.6/24 GB. The stock KVCacheManager cannot express this: its per-layer sizing exists, but the modulo needed to address a window-sized cache is only applied when sliding_window/attention_chunk_size is set globally, and those are all-or-nothing -- sliding_window would also shrink the full-attention caches to 128 and silently destroy long-range attention. So the cache becomes a per-layer ring buffer (MiMoV2SlidingWindowKVCacheManager), prefill writes only its last window, and the SWA decode mask is rebuilt in ring coordinates. Verified by exhaustive simulation over prefill lengths 1..4096 x 300 decode steps. Note this is the opposite trade from attention_dp_degree > 1, which shrinks the cache but replicates attention weights per DP group (+4.1 GB/rank at dp=8). Pro's weights already sit near the limit, so DP attention overflows where cache truncation fits; the flag declines to engage on that path rather than half-apply itself. Also records what longer context does NOT buy. Generation degrades with prompt length regardless of seq_len and regardless of this flag: coherent to ~480 tokens, then fluent-but-fabricated around 520, then collapsing into single-token repetition at >=568. Reproduced on the pinned-correct neuronx-cc 2.25.3371 with truncation off, i.e. on a graph proven structurally identical to the working 512 recipe, so it is a distinct issue from the 2.26.6360 miscompile and is not fixed by the pin. Repetitive filler was ruled out as the cause. Treat ~480 prompt tokens as the practical ceiling; the middle band matters most, since output stays well-formed while the content is wrong. Two stale claims corrected: seq_len=1024 does not OOM (that measurement predates the BF16-attn recipe), and LONG_CONTEXT_DESIGN.md gets a status header noting which of its predictions measurement refuted. Regression safety: with the flag unset the emitted graph is structurally identical to before this change -- CTE 54658 and TKG 47418 instructions, matching signatures -- compared as multisets of (opcode, output shape, operand shapes) over all HLO modules. Do not compare HLO protobuf byte hashes for this: serialization is nondeterministic run-to-run, so two runs of the same code differ. A same-code control run is required before trusting any "differs" verdict. Co-Authored-By: Claude Opus 5 --- .../MiMo-V2.5-Pro/LONG_CONTEXT_DESIGN.md | 262 ++++++++++++++++ contrib/models/MiMo-V2.5-Pro/README.md | 56 +++- .../perf_test/smoke_compile_mimo_v2.py | 13 +- .../MiMo-V2.5-Pro/src/modeling_mimo_v2.py | 290 +++++++++++++++++- 4 files changed, 609 insertions(+), 12 deletions(-) create mode 100644 contrib/models/MiMo-V2.5-Pro/LONG_CONTEXT_DESIGN.md diff --git a/contrib/models/MiMo-V2.5-Pro/LONG_CONTEXT_DESIGN.md b/contrib/models/MiMo-V2.5-Pro/LONG_CONTEXT_DESIGN.md new file mode 100644 index 00000000..345ffb56 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/LONG_CONTEXT_DESIGN.md @@ -0,0 +1,262 @@ +# MiMo-V2.5-Pro on Trn2 — Long-Context & Small-Batch Optimization Design + +> **STATUS: partly superseded — read this first.** +> +> This was the *planning* document, written before any code existed. Phase 1 +> (SWA cache truncation) was subsequently implemented as +> `MIMO_SWA_KV_TRUNCATION`, and measurement changed two of the conclusions +> below. What actually held up: +> +> - **Phase 1 works as designed.** Truncating the 60 SWA layers to 128 slots cut +> KV at `seq_len=4096` from 21.1 GB to 3.6 GB per rank; 4K compiles and loads +> at ~20.6/24 GB. The mechanism analysis in §2/M1 is accurate. +> - **The HBM framing in §1 is no longer the binding constraint.** The claim that +> `seq_len=1024` OOMs came from a pre-BF16-attn measurement and does not +> reproduce; 1024 loads on the stock recipe. +> - **The real blocker is output quality, which this doc did not anticipate.** +> Prompts past ~480 tokens degrade — first answering fluently but wrongly +> (~520), then collapsing into repetition (≥568) — independently of `seq_len` +> and of the truncation flag, on a graph proven structurally identical to the +> working 512 recipe. So making longer context *fit* does not make it *usable*. +> See "Long context and output degeneration" in README.md for the measured +> numbers. Everything below about reaching 8K/128K is therefore necessary but +> not sufficient. +> - **Phase 2 (DP attention) is the wrong lever for Pro**, for a reason §2/M2 +> underweighted: it replicates attention *weights* per DP group (+4.1 GB/rank +> at dp=8) and Pro's weights already sit near the limit. It overflows where +> cache truncation fits. +> +> Kept for the mechanism survey and file:line references, which remain useful for +> any future 128K attempt. + +**Goal (user requirements):** +1. Raise `seq_len` from the current **512** to **≥8K, ideally 128K**. +2. Allow **batch_size < 48** (currently forced to 48), so a smaller batch can buy longer context. + +This doc maps the root causes to concrete NxDI mechanisms that **already exist** +and are used by other models (Llama4, gpt_oss, qwen3-moe), with file:line +evidence, then proposes a phased plan. No code is changed yet. + +--- + +## 1. Why seq_len is stuck at 512 and BS is stuck at 48 + +Three independent multipliers blow up HBM, all traceable to the same choice +("use all 64 cores with TP=64 while the model has only 8 KV heads and 384 +experts"): + +### 1a. CONVERT_TO_MHA replicates KV heads 8→64 (8× KV cache) +MiMo has `num_key_value_heads=8` but we run `tp_degree=64`. The contrib model +forces MHA conversion: +- `contrib/.../modeling_mimo_v2.py:395` `self.use_gqa_convert_to_mha = tp_degree > self.attn_num_kv_heads` (64>8 → True) +- `:399` `self._kv_replication_factor = self.attn_num_heads // self.attn_num_kv_heads` (128/8 = **16×** at the weight level; per-rank the cache ends up with 64 KV heads vs the natural 8, i.e. **8× more KV cache** than a batch-parallel layout would need) +- `:415-416` every projection binds to `parallel_state.get_tensor_model_parallel_group()` — the **full** TP=64 group, so this is hard-wired. + +### 1b. V is padded 128→192 in the cache (1.5× on V) +- `contrib/.../modeling_mimo_v2.py:644-647` pads `value_states_for_cache` from `v_head_dim=128` to `head_dim=192` before storing. So the V half of the cache carries 50% dead padding. + +### 1c. All 70 layers store the full sequence — the SWA layers waste 60/70 of it +MiMo is **hybrid attention**: 10 full-attention layers + **60 sliding-window +layers with `sliding_window=128`**. A sliding-window layer only needs the last +128 tokens of KV, but MiMo allocates a full-`seq_len` cache for every layer: +- installed `models/mimo_v2/modeling_mimo_v2.py:996-999`: + ``` + # NOTE: Do NOT set self.sliding_window here because it affects KV cache size globally. + # Setting has_mixed_attn = True enables proper mask creation without affecting cache size. + self.has_mixed_attn = True + ``` + → `self.sliding_window=None`, `self.layer_to_cache_size_mapping=None` (defaults + `model_base.py:116-117`), so it falls into the plain `KVCacheManager` else-branch + (`model_base.py:190`) and **every layer gets a full `max_length` cache**. + +### 1d. BS ≥ 48 is required by moe_ep_degree > 1 at decode +- README/config: `BS * top_k / num_experts >= 1.0` when `moe_ep_degree>1` → + `BS >= 384/8 = 48`. This is the expert-parallel decode requirement that every + EP rank receive ≥1 token for a well-formed all-to-all (see `moe`/`moe_v2`). + +### Quantified HBM (real 8 KV heads, K+V, bf16, V padded to 192) + +| context | BS | all 70 layers full | **SWA truncated** (10 full @ seq + 60 @ 128) | +|--------:|---:|-------------------:|---------------------------------------------:| +| 512 | 48 | 10.6 GB | — | +| 8K | 48 | 169 GB | **26 GB** | +| 8K | 8 | 28 GB | **4.4 GB** | +| 128K | 48 | 2706 GB | **389 GB** | +| 128K | 8 | 451 GB | **65 GB** | +| 128K | 1 | 56 GB | **8 GB** | + +(These are the *natural* 8-KV-head numbers. Today's CONVERT_TO_MHA multiplies +the attention side by another ~8×.) HBM budget: 64×24 = 1536 GB total, ~1 TB +consumed by weights, leaving ~300–400 GB for KV. + +**Reading the table against the goals:** +- **8K** is reachable *today's cache layout* just by lowering BS (BS=8 → 28 GB). + The blocker for 8K is really the **BS≥48** rule (1d), not the cache design. +- **128K** is impossible with the full-every-layer layout at any batch that also + keeps EP=64 (2706 GB @ BS48, 451 GB @ BS8). It only becomes feasible with + **SWA cache truncation** (1c): 389 GB @ BS48, 65 GB @ BS8, 8 GB @ BS1. + → For 128K, SWA truncation is **necessary**, not optional. + +--- + +## 2. The mechanisms already exist in NxDI (with model precedents) + +### M1. Sliding-window / hybrid per-layer KV cache — **strong precedent** +The contiguous `KVCacheManager` supports heterogeneous per-layer cache lengths +via `layer_to_cache_size_mapping`: +- `modules/kvcache/kv_cache_manager.py:217-228` builds per-layer `k_shapes`/`v_shapes` from the mapping. +- helper `modules/kvcache/utils.py:507-516` `get_layer_to_kv_cache_size_mapping_for_mixed_attn(local, global, is_layer_locals)`. +- **Llama4** uses exactly this: `models/llama4/modeling_llama4_text.py:626` builds the mapping (chunk-size for local layers, seq_len for global) and passes it to `KVCacheManager` (`:650`). +- **gpt_oss** has a bespoke `GptOssKVCacheManager` (`modules/kvcache/gpt_oss_kv_cache_manager.py:88-116`) that stores `sliding_window` tokens for SWA layers and `max_length` for full layers — and even a separate DP degree for SWA layers. + +MiMo would need to set `layer_to_cache_size_mapping` (128 for the 60 SWA layers, +seq_len for the 10 full layers) and pass it to the KV manager — instead of the +current `sliding_window=None` mask-only approach. + +### M2. Data-parallel attention (attention TP < global TP, no KV replication) — **strong precedent** +- Config fields (`config.py:363-367`): `cp_degree`, `attention_dp_degree`, plus MoE's `moe_tp_degree`/`moe_ep_degree` — attention parallelism is **decoupled** from MoE parallelism. +- With `attention_dp_degree=8`, decode attention runs at TP=64/8=**8** (= the 8 KV heads → **no CONVERT_TO_MHA**), and KV cache batch shrinks: `config.py:514-515` `kv_cache_batch_size = tkg_batch_size // attention_dp_degree`. A `DataParallelKVCacheManager` slices the cache per DP rank (`model_base.py:185-186`). +- `NeuronAttentionBase` already consumes `get_data_parallel_attention_tp_group()` for the decode projections (`attention_base.py:188-190, 385-393`); CTE/prefill stays full-TP (`:389-391`) with no `cp_degree` required. +- **Working reference config** (same TP=64 class, MoE): `test/integration/tp64/models/qwen3moe/neuron_configs/bs16_sl10k_optimized.json` uses `tp_degree=64, moe_tp_degree=2, moe_ep_degree=32, attention_dp_degree=8, cp_degree=16, sequence_parallel_enabled, strided_context_parallel_kernel_enabled, ...` — i.e. **BS=16, seq 10K** already validated on tp64 with DP attention + CP + EP. + +**Caveat:** `attention_dp_degree` is *batch* data parallelism (decode-only), and +its `DataParallelKVCacheManager` path does **not** plumb +`layer_to_cache_size_mapping` (`model_base.py:185-186` vs `:190`). So **M1 and +M2 are currently mutually exclusive** in the stock code — combining them (DP +attention *and* per-layer SWA cache) would need new work. + +**Caveat 2:** MiMo's custom attention overrides `init_gqa_properties` to a no-op +(`contrib/.../modeling_mimo_v2.py:372-382`) and hand-rolls projections on the +full TP group, precisely because its Q/K head_dim=192 ≠ V head_dim=128 breaks +the base's fused GQA QKV. So MiMo does **not** currently consume the DP-attention +group even though the base class supports it. Wiring M2 into MiMo means +re-implementing the CTE/TKG projection split for its asymmetric head dims. + +### M3. Context parallelism for prefill — **already partly wired in MiMo** +- MiMo's attention has a CP forward path: `contrib/.../modeling_mimo_v2.py:532-624` splits Q/K/V + mask along the sequence dim when `cp_degree>1` (prefill only). This shards the *prefill* sequence across ranks (helps prefill latency / prefill activation memory at long context) but does **not** shrink the decode KV cache. + +### M4. Lower BS by changing MoE parallelism +The BS≥48 rule is tied to `moe_ep_degree>1`. Options to break it (need +validation): raise `moe_tp_degree` and lower `moe_ep_degree` (fewer EP ranks → +lower the `BS ≥ num_experts/top_k`… actually the bound is on top_k/num_experts, +independent of ep_degree — see open question Q1 below), use `capacity_factor` / +token-dropping, or a non-EP MoE kernel. The qwen3moe reference above runs +`moe_ep_degree=32` at **BS=16**, which suggests the "BS≥48" bound is **not** +fundamental to EP>1 in general and may be specific to MiMo's +`moe_ep_degree=64`/kernel choice. **This is the highest-value thing to verify +first** (see Phase 0). + +--- + +## 3. Open questions to resolve before coding + +- **Q1 (blocking for goal 2) — RESOLVED as "unimplemented case, not a hard + limit":** The `NotImplementedError: Selective Loading with Expert parallelism` + lives in the base `neuronx_distributed` library's `ExpertMLPsV2.forward_selective_loading` + (not vendored in this repo). The gate is exactly "average tokens per expert + = `BS*top_k/num_experts` ≥ 1"; below it, some EP ranks get 0 tokens and that + decode path simply isn't written — it raises rather than asserting a hardware + impossibility. So a zero-token-rank (capacity/padding) scheme *could* handle + BS<48, but the selective-loading TKG kernel doesn't. **In-repo, the only + documented way to run BS<48 while sharded across 64 cores is + `moe_ep_degree=1 + moe_tp_degree=64`** (README:335) — but for MiMo + `moe_tp=64` shrinks the per-rank MoE intermediate (2048/64 = 32 rows) below + the 128-row FP8 blockwise scale block, collapsing per-channel scale + (README:325-331), so it needs an **all-BF16 MoE checkpoint**. `capacity_factor` + is a wired config field (`config.py:806,810` → `moe.py:31`, the v1 + capacity/drop-token wrapper) that *might* bypass the gate, but the deciding + code is in the base library and unverifiable from here. **Action:** either + test `capacity_factor` + the v1 `moe.py` path at BS<48, or accept + `moe_ep=1/moe_tp=64` + a BF16 MoE checkpoint for low-BS long-context. +- **Q2:** Can `layer_to_cache_size_mapping` (M1) coexist with `fused_qkv=False` + and MiMo's asymmetric head dims + attention-sink bias? Llama4/gpt_oss don't + have asymmetric head dims, so the cache-shape helpers may assume symmetric + head_dim. +- **Q3:** Does SWA truncation interact correctly with the attention-sink bias + (MiMo adds a learnable sink column on SWA layers) and the V-pad-to-192 + workaround? + +--- + +## 4. Phased plan (lowest risk / highest leverage first) + +**Phase 0 — free wins, no modeling change (validate first):** +- Test **BS=8 or 16** with the *current* layout at seq_len 1024/2048. If the + BS≥48 rule can be relaxed (Q1), 8K may already fit (28 GB @ BS8, full layers). + This directly serves goal 2 and partially goal 1. +- Confirm the seq1024 NEFF (already compiling) actually loads — the README + "1024 OOMs" note predates the BF16-attn recipe and may be stale. + +**Phase 1 — SWA cache truncation (goal 1, biggest single lever for 128K):** +- Two proven in-tree implementations to copy: + - **Llama4 style** — build `layer_to_cache_size_mapping` = [seq_len for the 10 + full layers, 128 for the 60 SWA layers] and pass it to the contiguous + `KVCacheManager` (`modeling_llama4_text.py:626,650`; allocation + `kv_cache_manager.py:217-228`; helper `kvcache/utils.py:507-516` + `get_layer_to_kv_cache_size_mapping_for_mixed_attn`). + - **gpt_oss style** — a bespoke per-layer manager `GptOssKVCacheManager` + (`gpt_oss_kv_cache_manager.py:97-116`) that allocates + `get_kernel_cache_size_bucket(sliding_window)` for SWA layers and + `max_length` for full layers, and even carries a separate DP degree for SWA + layers (relevant to combining with Phase 2). + MiMo today sets `has_mixed_attn=True` but leaves `sliding_window=None` / + `layer_to_cache_size_mapping=None`, so all 70 layers get full caches + (`models/mimo_v2/modeling_mimo_v2.py:996-999`). The change is to populate the + mapping. Cuts 128K KV from 2706→389 GB (BS48) / 451→65 GB (BS8). +- Resolve Q2/Q3 (asymmetric head dim 192/128 + sink bias under per-layer cache; + Llama4/gpt_oss have symmetric head_dim, so the cache-shape helpers may need + the V-pad-to-192 workaround threaded through). + +**Phase 2 — DP attention to kill the 8× KV replication (goal 1+2):** +- Wire `attention_dp_degree` into MiMo's custom `_init_projections` so decode + attention uses `get_data_parallel_attention_tp_group()` (TP=8 = 8 KV heads, + zero replication) while MoE stays EP=64. Mirror the base's CTE(full-TP)/ + TKG(DP) split for asymmetric head dims. Reference config: qwen3moe + bs16_sl10k_optimized.json. +- Note M1+M2 are mutually exclusive in stock code; combining SWA truncation with + DP attention is net-new work (the DP KV manager would need per-layer sizing). + +**Phase 3 — combine + push to 128K:** +- Full-layer KV at 128K even after SWA truncation is the remaining cost (10 + layers). Add CP (M3) for prefill of very long inputs, and pick the BS that + fits (BS=1–8 for 128K). + +--- + +## 4b. Iterate on MiMo-V2.5 (PR #148), not Pro, for faster turnaround + +Pro is slow to compile/load/warmup (70 layers, 384 experts, ~200 s load, long +DeepGEMM warmup). **MiMo-V2.5** (PR #148, `contrib/MiMo-V2.5`, also whn09) is +architecturally the same family and every optimization here applies identically, +but it is smaller and iterates faster. Verified V2.5 config (from +`s3://datalab/xiaomi/models/MiMo-V2.5-Neuron-FP8/config.json`): + +| param | MiMo-V2.5 | MiMo-V2.5-Pro | +|-------|-----------|---------------| +| layers | **48** (16 full + 32 SWA) | 70 (10 full + 60 SWA) | +| num_key_value_heads (full / SWA) | **4 / 8** | 8 / 8 | +| num_attention_heads | 64 | 128 | +| experts / top_k | 256 / 8 | 384 / 8 | +| BS≥ bound (num_experts/top_k) | **32** | 48 | +| sliding_window | 128 | 128 | +| head_dim / v_head_dim | 192 / 128 | 192 / 128 | + +Same three problems apply: TP=64 > KV heads (4/8) → CONVERT_TO_MHA (replication +factor 64/4 = **16×** on full layers, even worse than Pro); 32 SWA layers store +full seq_len; BS forced to ≥32. So V2.5 is a faithful, faster testbed — validate +Phase 0/1/2 on V2.5 first, then port the identical changes to Pro. V2.5 also +already runs on vLLM without Pro's issue-#31 garble, making end-to-end +correctness checks easier. + +## 5. Summary + +| Goal | Blocker | Mechanism (exists) | Precedent | Phase | +|------|---------|--------------------|-----------|-------| +| BS < 48 | `moe_ep_degree>1` decode bound | lower moe_ep / raise moe_tp / capacity_factor | qwen3moe tp64 @ BS16 | 0 | +| 8K seq | mostly BS≥48 | lower BS, current cache | qwen3moe sl10k | 0–1 | +| 128K seq | 70 full-length layers | `layer_to_cache_size_mapping` SWA truncation | Llama4, gpt_oss | 1 | +| kill 8× KV replication | MiMo hardcodes full TP + CONVERT_TO_MHA | `attention_dp_degree` | qwen3moe DP=8 | 2 | + +Everything needed is present in NxDI and proven on other models; the work is +adapting MiMo's *custom* attention/KV code (which currently bypasses these +mechanisms because of its asymmetric head dims) to consume them. diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index 61d9c7b5..e4e8f9eb 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -85,8 +85,8 @@ GPU stacks (sglang on H100/H200) run the same OCP FP8 checkpoint correctly becau ### Cost and constraints -- **HBM headroom.** BF16 q/k/v adds ~2 GB per rank. `seq_len=1024` OOMs on load (the previous attempt failed allocating ~40 MB for rdh/alltoall rings after per-rank tensors already reached 20.9/24 GB). `seq_len=512` is the largest value empirically verified to fit HBM at BS=48. -- **Short context.** Even at `seq_len=512`, Pro's full chat template with the default system prompt is ~260 tokens; that leaves ~250 tokens for user input + generation. Longer context needs a different HBM plan (cross-instance TP/PP, or a larger instance). +- **HBM headroom.** BF16 q/k/v adds ~2 GB per rank, leaving little margin against the 24 GB per-core budget. `seq_len=512` is the recommended value, but **HBM is no longer the binding constraint**: `seq_len=1024` compiles and loads fine on the stock recipe, and `4096` loads with `MIMO_SWA_KV_TRUNCATION=1` (~20.6/24 GB per rank). What limits usable context is **output quality** — see "Long context and output degeneration" below. (An earlier revision of this README claimed `seq_len=1024` OOMs by ~40 MB; that measurement predates the BF16-attn + `neuronx-cc 2.25` recipe and no longer reproduces.) +- **Short context.** Pro's full chat template with the default system prompt is ~260 tokens, so at `seq_len=512` only ~250 tokens remain for user input + generation. - `BS * top_k / num_experts >= 1.0` required when `moe_ep_degree > 1` at decode (else `NotImplementedError`). With `num_experts=384, top_k=8` this forces `BS >= 48`. - `n_routed_experts=384 = 2^7 × 3` → `384 / ep_degree` is never a power of 2 (6, 12, 24, 48, 96, 192, 384). Kimi PR #131 says NKI `_bwmm_shard_on_block_nki_call` on SDK 2.29 has "depressed logits with EP=2" and recommends SDK 2.28. @@ -95,9 +95,58 @@ GPU stacks (sglang on H100/H200) run the same OCP FP8 checkpoint correctly becau - **All-FP8 attention (`modules_to_not_convert` without q/k/v).** Drifts as described above. Known broken; `preprocess_mimo_v2_fp8.py` no longer emits it. - **`use_torch_block_wise=True`** (PyTorch-fallback blockwise matmul for higher accumulator precision): compile+shard succeeded after ~2 h, but `model.load()` crashed with `status=4 Allocation Failure` — the fallback path raises HBM demand even when scoped to MoE. +### Long context and output degeneration + +Raising `seq_len` past 512 is **not** an HBM problem any more, and raising it does +not buy usable context. Two independent things were measured: + +**1. HBM: solved, opt-in.** `MIMO_SWA_KV_TRUNCATION=1` sizes the 60 +sliding-window layers' KV cache to `sliding_window=128` instead of `seq_len`, +leaving the 10 full-attention layers at full length. At `seq_len=4096` that cuts +KV from 21.1 GB to 3.6 GB per rank, and the model compiles and loads at ~20.6/24 +GB. Off by default; see `MiMoV2SlidingWindowKVCacheManager` in +`src/modeling_mimo_v2.py` for the ring-buffer indexing this requires. + +Note this is the *opposite* trade from `attention_dp_degree > 1`, which shrinks +the KV cache but replicates attention **weights** per DP group (+4.1 GB/rank at +dp=8) — Pro's weights already sit near the limit, so DP attention overflows where +cache truncation fits. + +**2. Output quality: unsolved, and the real limit.** Independently of `seq_len` +and of the truncation flag, generation degrades as the *prompt* gets longer. +Measured on the `seq_len=1024` NEFF with truncation **OFF** (i.e. a graph proven +structurally identical to the working 512 recipe — see the regression note +below), greedy decoding, needle-in-a-haystack over non-repetitive prose: + +| prompt tokens | behaviour | +|--------------:|-----------| +| ≤ 480 | coherent, on topic | +| ~522 | coherent but **wrong**: invents an answer instead of retrieving the planted one | +| ≥ 568 | collapses into single-token repetition (`the the the…`) mixed with replacement characters | + +The middle band matters more than the collapse: output stays fluent and +well-formed while the content is fabricated, which is far harder to catch in +production than obvious repetition. **Treat ~480 prompt tokens as the practical +ceiling**, well under the nominal window. + +Repetitive filler was ruled out as the cause — the same collapse occurs with 15 +unrelated paragraphs of varied prose. This is the same symptom class as the +`neuronx-cc 2.26.6360` garbling documented above, but it reproduces on the +pinned-correct 2.25.3371, so it is **not** the same root cause and is not fixed +by the pin. Filed with AWS; no modeling-level workaround known. + +**Regression safety.** With `MIMO_SWA_KV_TRUNCATION` unset (the default), the +truncation code emits a structurally identical graph to the version before it was +added, verified by comparing multisets of `(opcode, output shape, operand +shapes)` across all HLO modules (CTE 54658 instructions, TKG 47418). Do not +compare HLO protobuf byte hashes for this — serialization is nondeterministic +run-to-run, so byte hashes differ between two runs of the *same* code. Always +establish a same-code control before trusting a "differs" verdict. + ### Next experiments queued -- **Even longer `seq_len`** (> 512): needs a tighter HBM plan — smaller batch, different EP ratio, or cross-instance sharding. +- **Root-cause the long-prompt degeneration** (above) — the blocker for any + context beyond ~480 tokens. Pending with AWS. - **Upstream vllm-neuron fix** for the "first-request-only" serving bug (issue #31); patch branch at `whn09/vllm-neuron#fix/hybrid-attn-swa-spec` is a placeholder that did not resolve the symptom. - **Cross-instance BF16** via pipeline/tensor parallelism on 2× Trn2 (single-instance HBM cannot hold full BF16 Pro). - **Selective BF16 only on MoE `gate_up_proj`** (smallest expert scales) while keeping `down_proj` FP8 — another axis to probe if attn drift returns at longer contexts. @@ -213,6 +262,7 @@ override or if you plan to launch vLLM outside of `bench_mimo_v2.sh`. | `NEURON_COMPILED_ARTIFACTS` | `/opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512` (per `start_vllm_server.sh`) | Where vLLM writes its NEFF + per-rank sharded weights. Points at the **same** dir as the smoke-path seq512 NEFF (no `_vllm` suffix) so vLLM reuses its 64 pre-sharded `tp*_sharded_checkpoint.safetensors` and skips the ~30 min shard step. vLLM still compiles its own continuous-batching / async / on-device-sampling NEFF variant into this dir on first launch (~60 min), but the reshard is avoided. A separate empty dir (e.g. a `_vllm` suffix) would force a full from-scratch compile *and* reshard. vLLM's own fallback is `/neuron-compiled-artifacts//`. | | `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 Pro's 384-expert MoE is ~60 min TKG + ~15 min CTE + ~30 min shard, well past vLLM's default. | +| `MIMO_SWA_KV_TRUNCATION` | `0` (off) | Size the 60 sliding-window layers' KV cache to `sliding_window=128` instead of `seq_len`, which is what lets `seq_len=4096` fit HBM. **Must be set identically at compile and at load** — it changes the KV cache shapes, so a NEFF compiled with it will not load without it. Off by default: at `seq_len=512` it saves little, and the verified 512 recipe is graph-identical with it unset. See "Long context and output degeneration" — it fixes the memory ceiling, not the quality ceiling. | For a quick `curl` sanity check while the server is up: diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py index e77f8ce6..4d8d5816 100755 --- a/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_compile_mimo_v2.py @@ -43,9 +43,16 @@ ) TP_DEGREE = int(os.environ.get("TP_DEGREE", "64")) -# seq_len=512 is the largest value verified to fit HBM under the BF16-attn -# recipe. seq_len=1024 OOMs on load (previous attempt failed allocating -# ~40 MB for rdh/alltoall rings after per-rank tensors reached 20.9/24 GB). +# 512 is the default because it is the length the model is *verified correct* at, +# not because longer does not fit. HBM-wise, 1024 loads fine on the stock recipe +# and 4096 loads with MIMO_SWA_KV_TRUNCATION=1 (~20.6/24 GB per rank). The +# blocker is output quality, not memory: on the seq1024 NEFF, prompts beyond +# ~480 tokens first start answering plausibly-but-wrongly (~520) and then +# collapse into single-token repetition (>=568), even on neuronx-cc 2.25.3371 +# with truncation OFF -- i.e. on a graph proven structurally identical to the +# 512 recipe's. See README "长上下文与输出退化" for the measured threshold. +# (An earlier note here claimed 1024 OOMs by ~40 MB; that predates the +# BF16-attn + cc 2.25 recipe and no longer reproduces.) SEQ_LEN = int(os.environ.get("SEQ_LEN", "512")) # BS=48 is the minimum that avoids forward_selective_loading on decode: # `BS * top_k / num_experts >= 1.0` → BS >= 384/8 = 48. At BS=1 the TKG diff --git a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py index 3d609942..070c4581 100644 --- a/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py +++ b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py @@ -9,6 +9,7 @@ import gc import math +import os import warnings from typing import Any, Dict, List, Optional, Tuple, Type, Union @@ -51,8 +52,13 @@ from neuronx_distributed_inference.modules.attention.attention_base import ( NeuronAttentionBase, ) -from neuronx_distributed_inference.modules.attention.utils import RotaryEmbedding +from neuronx_distributed_inference.modules.attention.utils import ( + RotaryEmbedding, + get_last_kv_window, +) from neuronx_distributed_inference.modules.custom_calls import CustomRMSNorm +from neuronx_distributed_inference.modules.generation.sampling import create_sampler +from neuronx_distributed_inference.modules.kvcache.kv_cache_manager import KVCacheManager from neuronx_distributed_inference.modules.moe_v2 import initialize_moe_module try: @@ -64,6 +70,57 @@ _flash_fwd_call = nki_jit()(attention_isa_kernel) +#: Opt-in: truncate the KV cache of sliding-window layers to `sliding_window` +#: instead of the full seq_len. Off by default -- the verified seq_len=512 recipe +#: is unaffected, and at 512 the saving is small anyway. See +#: MiMoV2SlidingWindowKVCacheManager for what this changes and why. +MIMO_SWA_KV_TRUNCATION = os.environ.get("MIMO_SWA_KV_TRUNCATION", "0") == "1" + + +def swa_kv_truncation_enabled(config) -> bool: + """Whether SWA layers get a window-sized ring cache for this config. + + Single source of truth, because two places must agree exactly or the model + silently computes the wrong thing: the cache manager (which sizes and + indexes the cache) and each attention module (which builds its decode mask + in ring coordinates only if the cache actually is a ring). Both must also + decide *before* the layers are constructed, so every condition here has to + be knowable from the config alone -- notably including the paths that fall + back to a manager which cannot do per-layer sizing. + """ + if not MIMO_SWA_KV_TRUNCATION: + return False + + window = getattr(config, "sliding_window", None) + if not window or not any( + t == "sliding_window" for t in config.layer_attention_types + ): + return False + + # Nothing to truncate, and a ring no larger than the sequence buys nothing + # while still adding ring-indexing risk. + if window >= config.neuron_config.seq_len: + warnings.warn( + f"MIMO_SWA_KV_TRUNCATION ignored: sliding_window={window} >= " + f"seq_len={config.neuron_config.seq_len}, nothing to truncate." + ) + return False + + # These two select DataParallelKVCacheManager / BlockKVCacheManager upstream + # (model_base.py init_inference_optimization), neither of which accepts + # layer_to_cache_size_mapping. Rather than half-apply the optimization -- + # ring masks reading a full-length cache would silently corrupt attention -- + # leave those configurations exactly as they were. + if config.neuron_config.attention_dp_degree > 1 or config.neuron_config.is_block_kv_layout: + warnings.warn( + "MIMO_SWA_KV_TRUNCATION ignored: attention_dp_degree > 1 and " + "is_block_kv_layout use KV cache managers that do not support " + "per-layer cache sizes. Falling back to full-length caches." + ) + return False + + return True + def get_rmsnorm_cls(): """Get appropriate RMSNorm class based on execution environment.""" @@ -321,6 +378,11 @@ def __init__( 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 + # With the cache truncated to the window (MIMO_SWA_KV_TRUNCATION=1) + # it becomes a ring buffer, so a slot index is no longer an absolute + # position and the decode masks below must be built in ring + # coordinates instead. + self.swa_kv_cache_is_ring = swa_kv_truncation_enabled(config) else: self.attn_head_dim = config.head_dim self.attn_v_head_dim = config.v_head_dim @@ -328,6 +390,7 @@ def __init__( self.attn_num_kv_heads = config.num_key_value_heads rope_theta = config.rope_theta self.sliding_window_size = None + self.swa_kv_cache_is_ring = False # Calculate partial rotary dimensions self.partial_rotary_factor = config.partial_rotary_factor @@ -777,8 +840,15 @@ def forward( # 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: + # Apply attention mask to prior scores. + # + # Skipped for a ring-buffer SWA cache: the externally-supplied mask + # is indexed by absolute position over n_positions columns, whereas + # prior_scores here has only `cache_len` columns whose index is a + # ring slot. The mask is neither the right width nor the right + # coordinate system; the ring mask built below replaces it entirely + # (it is self-sufficient -- it depends only on position_ids). + if attention_mask is not None and not self.swa_kv_cache_is_ring: # Convert boolean mask to additive mask if needed if attention_mask.dtype == torch.bool: prior_scores = prior_scores.masked_fill(~attention_mask, float('-inf')) @@ -805,7 +875,12 @@ def forward( # 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: + # + # Ring-buffer SWA caches are excluded: every slot of a filled ring + # holds a position < current_pos, so the absolute comparison below + # would be both meaningless (slot != position) and unnecessary. The + # ring mask that follows does the equivalent job in ring coordinates. + if position_ids is not None and not self.swa_kv_cache_is_ring: kv_seq_len = prior_scores.size(-1) current_pos = position_ids[:, 0].reshape(-1, 1) # [bsz, 1] pos_indices = torch.arange( @@ -828,8 +903,35 @@ def forward( kv_seq_len = prior_scores.size(-1) # 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] + col_idx = torch.arange(kv_seq_len, device=prior_scores.device)[None, :] # [1, kv_seq_len] + if self.swa_kv_cache_is_ring: + # Ring cache: column j holds whichever position last landed + # in slot j, not position j. Reconstruct it, then apply the + # same two-sided window rule as the non-ring branch: + # 0 < current_pos - p <= window - 1 + # The lower bound also drops slots not yet written (their + # reconstructed position lands >= current_pos), which makes + # this mask self-sufficient: it subsumes both the external + # mask and the per-row causal bound skipped above. `ring` is + # one fewer than the cache length, matching the manager's + # mimo_ring_size(); the spare slot is never written, so + # `col_idx < ring` masks it out. + ring = kv_seq_len - 1 + # Position slot j currently holds, for this row: + # base = largest multiple of `ring` at or below current_pos + # p = base + j, minus `ring` if that overshoots + base = (current_pos // ring) * ring # [bsz, 1] + slot_pos = base + col_idx # [bsz, kv_seq_len] + slot_pos = torch.where(slot_pos >= current_pos, slot_pos - ring, slot_pos) + delta = current_pos - slot_pos + sliding_mask = ( + (delta > 0) + & (delta <= self.sliding_window_size - 1) + & (slot_pos >= 0) + & (col_idx < ring) + ) + else: + sliding_mask = col_idx >= (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')) @@ -1108,6 +1210,122 @@ def forward( return outputs +class MiMoV2SlidingWindowKVCacheManager(KVCacheManager): + """KV cache manager that sizes SWA layers to the window, not to seq_len. + + MiMo is hybrid: 10 full-attention layers plus 60 sliding-window layers with + `sliding_window=128`. A window layer can only ever attend to the last 128 + tokens, but the stock manager gives all 70 layers a full seq_len cache. At + seq_len=4096 that is 21.1 GB of KV per rank against a 24 GB budget; sizing + the 60 window layers to 128 brings it to 3.6 GB, which is what makes 4K fit. + (The alternative, attention_dp_degree>1, replicates attention *weights* per + DP group at +4.1 GB/rank, and Pro's weights already sit near the limit.) + + Three things change together, which is why this is a subclass rather than + just passing `layer_to_cache_size_mapping` to the stock manager: + + 1. Allocation -- already handled by the base `_init_kv_shape`. + 2. Decode indexing -- `padding_side="right"` scatters at the ABSOLUTE + position, out of bounds for a 128-slot cache, so window layers index + modulo the ring size. The stock manager only applies a modulo when + `attention_chunk_size`/`sliding_window` is set globally, and both are + all-or-nothing: `sliding_window` would also shrink the full-attention + caches to 128 and silently destroy long-range attention. + 3. Prefill writing -- narrowed to the last window by + `update_kv_by_layer_id` below. + + The read side stays the model's hand-rolled sliding mask, which switches to + ring coordinates via `swa_kv_cache_is_ring`. The base mask helpers are + deliberately unused: `_create_chunked_attn_mask_tkg` implements *chunked* + attention, which resets every 128 tokens -- different semantics, wrong here. + + Verified by exhaustive simulation of write-then-mask over prefill lengths + 1..4096 x 300 decode steps: the positions the mask exposes are exactly + {p : 0 < cur-p <= w-1}. + """ + + def __init__(self, config: InferenceConfig, is_swa_layer: List[bool], **kwargs): + self.mimo_is_swa_layer = list(is_swa_layer) + # Both of these change the meaning of v_shapes[.][2], which + # mimo_ring_size reads as the ring modulus: tiling turns dim 2 into the + # tile axis, and apply_seq_ids_mask pads every layer's cache. Neither is + # set in the Pro recipes, so assert rather than guess. + assert not config.neuron_config.kv_cache_tiling, ( + "MIMO_SWA_KV_TRUNCATION does not support kv_cache_tiling: the tiled " + "layout changes which shape dim holds the cache length." + ) + assert not config.neuron_config.apply_seq_ids_mask, ( + "MIMO_SWA_KV_TRUNCATION does not support apply_seq_ids_mask: it pads " + "each layer's cache, so the ring modulus would no longer match the " + "positions written." + ) + super().__init__(config, **kwargs) + + def mimo_ring_size(self, layer_idx: int) -> int: + """Number of ring slots actually used by a sliding-window layer. + + The cache has `window_size` slots but only `window_size - 1` join the + ring: decode scores the active token separately from the cache + (`active_scores` in the attention forward), so one slot must never be + read as "prior". Mirrors the stock manager's + `position_ids % (self.sliding_window - 1)` and gpt_oss's + `sliding_window = sliding_window - 1`. A window of w attends to w-1 + prior tokens plus the active one, so nothing is lost. + """ + return self.v_shapes[layer_idx][2] - 1 + + def update_kv_by_layer_id(self, idx, is_for_context_encoding: bool, seq_ids, + position_ids, kv_per_layer, seq_len: int, **kwargs): + """Write only the last window of KV during prefill on SWA layers. + + The base prefill path writes all `q_len` slots, which for a 4096-token + prefill does not fit a 128-slot cache. Only the tail is ever readable by + a window layer, so gather that tail -- in ring order, so the decode-time + modulo addressing lines up -- and hand the base class an exactly + cache-sized block. `get_last_kv_window` is the same helper + attention_base uses for its windowed path and does exactly this rotation. + + Decode is untouched here; it goes through the scatter path below. + """ + if is_for_context_encoding and self.mimo_is_swa_layer[idx]: + latest_k, latest_v = kv_per_layer[0], kv_per_layer[1] + # +1 because get_last_kv_window internally uses window_size - 1 + # slots, and we want to fill the full ring (= cache_len - 1). + latest_k, latest_v = get_last_kv_window( + self.mimo_ring_size(idx) + 1, position_ids, latest_k, latest_v + ) + kv_per_layer = (latest_k, latest_v) + + return super().update_kv_by_layer_id( + idx=idx, + is_for_context_encoding=is_for_context_encoding, + seq_ids=seq_ids, + position_ids=position_ids, + kv_per_layer=kv_per_layer, + seq_len=seq_len, + **kwargs, + ) + + def _get_index_to_update_new_position( + self, seq_ids, scatter_index, position_ids, full_k, transposed: bool, layer_idx: int + ): + """Ring-buffer the decode write position for sliding-window layers only. + + Full-attention layers keep absolute positions and defer to the base + implementation, so their behaviour is bit-for-bit unchanged. + """ + if not self.mimo_is_swa_layer[layer_idx]: + return super()._get_index_to_update_new_position( + seq_ids, scatter_index, position_ids, full_k, transposed, layer_idx + ) + + position_ids = position_ids % self.mimo_ring_size(layer_idx) + + index = scatter_index if self.is_medusa else position_ids + view_shape = (-1, 1, index.shape[-1], 1) if not transposed else (-1, 1, 1, index.shape[-1]) + return index.view(*view_shape).expand_as(full_k) + + class NeuronMiMoV2Model(NeuronBaseModel): """MiMo-V2.5-Pro Model for NXD inference.""" @@ -1151,6 +1369,66 @@ def setup_attr_for_model(self, config: MiMoV2InferenceConfig): # Setting has_mixed_attn = True enables proper mask creation without affecting cache size. self.has_mixed_attn = True + # Opt-in per-layer KV cache sizing (MIMO_SWA_KV_TRUNCATION=1). Sliding-window + # layers only need `sliding_window` slots; giving all 70 layers a full + # seq_len cache is what puts 4K out of HBM reach. Left as None by default so + # the base class takes its original single-shape path. + self.mimo_is_swa_layer = [ + t == "sliding_window" for t in config.layer_attention_types + ] + self.mimo_layer_to_cache_size_mapping = None + if swa_kv_truncation_enabled(config): + # max_length (not seq_len) is what the base class allocates for the + # single-shape path, so full-attention layers must keep using it or + # they would silently shrink. + max_length = config.neuron_config.max_length + self.mimo_layer_to_cache_size_mapping = [ + config.sliding_window if is_swa else max_length + for is_swa in self.mimo_is_swa_layer + ] + + def init_inference_optimization(self, config: MiMoV2InferenceConfig): + """Install the per-layer SWA cache manager when truncation is enabled. + + Only reached when swa_kv_truncation_enabled() produced a mapping; the + default path falls through to the stock manager selection untouched. The + fallback conditions all live in that predicate rather than here, because + the attention modules -- already built by this point -- had to make the + same ring-or-not decision from the same inputs. + """ + if self.mimo_layer_to_cache_size_mapping is None: + return super().init_inference_optimization(config) + + if self.on_device_sampling: + lm_head_tp_degree = None + if hasattr(self, "lm_head") and hasattr(self.lm_head, "tensor_parallel_group"): + lm_head_tp_degree = self.lm_head.tensor_parallel_group.size() + self.sampler = create_sampler(config.neuron_config, lm_head_tp_degree) + + # The base class forwards these three to the manager, and all three must + # stay unset for the per-layer ring to be correct. They are unset for + # MiMo today (see the has_mixed_attn note above), so fail loudly rather + # than silently drop them if that ever changes: + # - sliding_window / attention_chunk_size would apply their own modulo to + # EVERY layer, including full-attention ones, and resize all caches. + # - windowed_context_encoding_size moves the prefill write offset, which + # update_kv_by_layer_id assumes is the plain last-window block. + assert self.sliding_window is None and self.attention_chunk_size is None, ( + "MIMO_SWA_KV_TRUNCATION replaces the global sliding_window / " + "attention_chunk_size mechanism; they must not be set as well." + ) + assert self.windowed_context_encoding_size is None, ( + "MIMO_SWA_KV_TRUNCATION does not support windowed context encoding." + ) + + self.kv_mgr = MiMoV2SlidingWindowKVCacheManager( + config, + is_swa_layer=self.mimo_is_swa_layer, + num_kv_head=self.num_key_value_heads, + global_rank=self.rank_util, + layer_to_cache_size_mapping=self.mimo_layer_to_cache_size_mapping, + ) + def init_model(self, config: MiMoV2InferenceConfig): self.padding_idx = getattr(config, 'pad_token_id', None) self.vocab_size = config.vocab_size From 4404c60b4677b89dd6afb5ac752db1b7572559ce Mon Sep 17 00:00:00 2001 From: whn09 Date: Thu, 30 Jul 2026 11:00:04 +0800 Subject: [PATCH 47/47] Pro: document the S3 compiled-artifact backups and shard reuse The sharded weights are a function of TP/EP degree and quantization only, not of seq_len, so every seq_len variant can symlink one copy and skip the ~70 min / ~1 TB reshard. Recording this because the compiled/ dir had accumulated ~3 TB of byte-identical shard copies (verified by md5) across one-off configs, and only the NEFFs actually differed. Lists what is on s3://datalab/xiaomi/compiled/: the canonical seq512 recipe with weights, plus NEFF-only backups of the seq1024 stock and seq4096 SWA-truncation builds. Notes that the 4096 NEFF will not load without MIMO_SWA_KV_TRUNCATION=1, since the flag changes KV cache shapes, and that neither long-context NEFF is a usable deployment given the degeneration limits. Co-Authored-By: Claude Opus 5 --- contrib/models/MiMo-V2.5-Pro/README.md | 30 ++++++++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/contrib/models/MiMo-V2.5-Pro/README.md b/contrib/models/MiMo-V2.5-Pro/README.md index e4e8f9eb..60267435 100644 --- a/contrib/models/MiMo-V2.5-Pro/README.md +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -159,6 +159,36 @@ establish a same-code control before trusting a "differs" verdict. - **Venv**: `/opt/aws_neuronx_venv_pytorch_inference_vllm_0_16` (used by preprocess, smoke, and vLLM serving alike; ships with the DLAMI and is where `0_setup.sh` installs the patched `vllm-neuron`). - **Disk**: ~3 TB free under `/opt/dlami/nvme` (the HF FP8 checkpoint is ~962 GB, the Neuron-FP8 preprocessed output is ~1 TB, and `save_sharded_checkpoint=true` writes another ~300-1000 GB per compiled config (varies with recipe)). + **The sharded weights do not depend on `seq_len`** — sharding is a function of + TP/EP degree and quantization only. So a second `seq_len` variant needs only + its own `model.pt` + `neuron_config.json`; symlink `weights/` at an existing + artifact dir and the ~70 min / ~1 TB reshard is skipped entirely. Keeping one + copy of the shards per *recipe* rather than per compiled config is the + difference between ~1 TB and filling the array. + +### Compiled artifacts on S3 + +`s3://datalab/xiaomi/compiled/` — restore with `aws s3 cp`/`s5cmd` instead of +recompiling (compiles are ~60-90 min each): + +| Prefix | Contents | +|---|---| +| `mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/` | The canonical BF16-attn + FP8-MoE recipe: NEFF **and** the 64 sharded weight files. Restore this one first — the weights serve every variant below. | +| `pro_seq1024_stock_cc225/` | `seq_len=1024`, stock recipe (truncation OFF), built on `neuronx-cc 2.25.3371`. NEFF only. | +| `pro_seq4096_swatrunc_cc225/` | `seq_len=4096`, **requires `MIMO_SWA_KV_TRUNCATION=1` at load** (it changes KV cache shapes, so it will not load without the flag). NEFF only. | + +The two NEFF-only prefixes carry no `weights/`; point them at the seq512 shards: + +```bash +ln -s /opt/dlami/nvme/models/compiled/mimo_v2_5_pro_bs48_moetp1_ep64_fp8moe_bf16attn_seq512/weights \ + /opt/dlami/nvme/models/compiled/pro_seq1024_stock_cc225/weights +``` + +Both were verified byte-identical to the seq512 shards before being deduplicated +(md5 of `tp0`/`tp31`/`tp63`). Note the 1024 and 4096 NEFFs are useful for HBM and +compile-path work only — see the degeneration limits above before treating either +as a usable long-context deployment. + ### NVMe mount The Trn2 DLAMI ships with four local NVMe SSDs that are assembled into a