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 new file mode 100644 index 00000000..60267435 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/README.md @@ -0,0 +1,671 @@ +# 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. + +## 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. + +### ⚠️ 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. + +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, 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. + +### 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. + +### 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 + +- **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. +- **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: **`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)). + + **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 +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. | +| `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: + +```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 +``` + +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 +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 + +```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 +# 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 +``` + +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 +# 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. + +### 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 +``` + +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**: 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 +`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. + +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 | 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 | +| 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 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** | **726** | **47.6** | +| vLLM + EFA, DP2×TP8 + EP + chunked | 624 | 57.2 | +| Trn2 | 6,912 | 572.1 | + +Notes: +- **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. + +## 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 | +|----------|--------------------------------|------------------| +| 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 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..0a2f9f4a --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/0_setup.sh @@ -0,0 +1,78 @@ +#!/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_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.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.3 https://github.com/vllm-project/vllm-neuron.git $HOME/vllm-neuron +fi + +cd $HOME/vllm-neuron + +# 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 " 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 . +pip install s5cmd + +python3 -c "import vllm_neuron; print('vllm-neuron installed:', vllm_neuron.__file__)" + +echo "" +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-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 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-Neuron-FP8/**" "$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..23c3f59d --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/bench_mimo_v2.sh @@ -0,0 +1,87 @@ +#!/bin/bash +set -e + +# MiMo-V2.5-Pro FP8 vLLM benchmark on Trn2. One-shot wrapper: +# launch server -> sanity check -> bench at c=1,16,48 -> stop server. +# +# 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` +# +# 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" + +mkdir -p "$RESULTS_DIR" + +# 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 on port $PORT (up to 2 h for first compile)..." + local interval=10 + 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 after $((SECONDS - start))s." + return 0 + fi + 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 +} + +stop_server() { + echo " Stopping vLLM server..." + pkill -f "vllm.entrypoints.openai.api_server" 2>/dev/null || true + sleep 5 +} + +echo "==========================================" +echo "MiMo-V2.5-Pro FP8 Performance Benchmark" +echo "==========================================" +echo "Port: $PORT" +echo "Results: $RESULTS_DIR" +echo "" + +# 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 + +# One-shot sanity check (curl the chat endpoint). +PORT="$PORT" bash "$SCRIPT_DIR/sanity_check.sh" || true + +# Three concurrency levels. run_bench_single.sh reads knobs from the +# environment; see its header for all the options. +PORT="$PORT" RESULTS_DIR="$RESULTS_DIR" CONFIG_NAME="$CONFIG_NAME" \ + CONCURRENCY=1 NUM_PROMPTS=16 bash "$SCRIPT_DIR/run_bench_single.sh" +PORT="$PORT" RESULTS_DIR="$RESULTS_DIR" CONFIG_NAME="$CONFIG_NAME" \ + CONCURRENCY=16 NUM_PROMPTS=128 bash "$SCRIPT_DIR/run_bench_single.sh" +PORT="$PORT" RESULTS_DIR="$RESULTS_DIR" CONFIG_NAME="$CONFIG_NAME" \ + CONCURRENCY=48 NUM_PROMPTS=192 bash "$SCRIPT_DIR/run_bench_single.sh" + +echo "==========================================" +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/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) 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 new file mode 100644 index 00000000..8412e191 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/README.md @@ -0,0 +1,92 @@ +# H100 GPU baseline for 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. + +## 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 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: + +- **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. + +## ⚠️ 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 + +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 +``` + +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 +# vLLM (node 0 = API server on :8000, node 1 = --headless follower) +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 +NODE_RANK=1 DIST_INIT_ADDR=:20000 bash run_sglang_h100_multinode.sh +``` + +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 (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-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 +``` + +See the main README's Performance section for the measured numbers. 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/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 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 diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_norad_c1.txt b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_norad_c1.txt new file mode 100644 index 00000000..5f643d7e --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/results/sglang_efa_norad_c1.txt @@ -0,0 +1,42 @@ +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-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. +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-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. +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-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. +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=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, 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-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. +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-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. +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-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. +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_dp.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_dp.sh new file mode 100644 index 00000000..e4672638 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/h100/run_vllm_h100_dp.sh @@ -0,0 +1,85 @@ +#!/bin/bash +# 2-node x 8 H100 vLLM server for MiMo-V2.5-Pro (OCP-FP8) via DATA PARALLEL + +# EXPERT PARALLEL, instead of pipeline parallel. +# +# Why DP+EP over TP8xPP2: MiMo has only 8 KV heads (TP<=8), and the ~963 GB FP8 +# weights don't fit on one node. The PP=2 approach works but leaves a pipeline +# bubble (measured 158 out-tok/s @ c=48). DP=2 x TP=8 instead replicates the +# attention path per DP rank (each is TP=8, dividing 8 KV heads) and shards MoE +# experts across TP*DP=16 (equivalent to EP=16) -- the same shape SGLang uses +# with --enable-dp-attention, and with NO pipeline bubble. This is vLLM's +# closest analogue to SGLang's DP-attention. +# +# Run the SAME command on both nodes, changing only DP_RANK: +# node 0 (head, API server on :8000): DP_RANK=0 bash run_vllm_h100_dp.sh +# node 1 (headless): DP_RANK=1 bash run_vllm_h100_dp.sh +# DP_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}" +TP="${TP:-8}" # per-node tensor parallel (divides 8 KV heads) +DP="${DP:-2}" # data parallel across the 2 nodes +DP_RANK="${DP_RANK:?set DP_RANK=0 on head, 1 on worker}" +DP_ADDR="${DP_ADDR:-172.31.45.21}" # P5-1 private IP +DP_RPC_PORT="${DP_RPC_PORT:-29550}" +IFACE="${IFACE:-enp71s0}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-4096}" +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" >&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[@]}" 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..3ac0037d --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/run_bench_single.sh @@ -0,0 +1,92 @@ +#!/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 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 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 +# (default /opt/dlami/nvme/logs/bench_results/mimo_v2_5_pro) + +set -e + +# 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}" +INPUT_LEN="${INPUT_LEN:-360}" +OUTPUT_LEN="${OUTPUT_LEN:-120}" +RANGE_RATIO="${RANGE_RATIO:-0.03}" +# 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}" + +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 \ + --host localhost \ + --port "$PORT" \ + --model "$SERVED_MODEL_NAME" \ + --tokenizer "$TOKENIZER_PATH" \ + --endpoint /v1/completions \ + --dataset-name random \ + --num-prompts "$NUM_PROMPTS" \ + --random-input-len "$INPUT_LEN" \ + --random-output-len "$OUTPUT_LEN" \ + --random-range-ratio "$RANGE_RATIO" \ + --max-concurrency "$CONCURRENCY" \ + 2>&1 | tee "$RESULTS_DIR/${CONFIG_NAME}_c${CONCURRENCY}.txt" + +echo "" +echo "Saved to: $RESULTS_DIR/${CONFIG_NAME}_c${CONCURRENCY}.txt" diff --git a/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh new file mode 100755 index 00000000..a80e85c1 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/sanity_check.sh @@ -0,0 +1,88 @@ +#!/bin/bash +# Quick sanity check against an already-running vLLM server. +# +# 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/chat/completions on port $PORT" +echo " Model: $MODEL_PATH" +echo " System: $SYSTEM" +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 start_vllm_server.sh' (or bench_mimo_v2.sh)" + echo "first and wait for 'Application startup complete.'" + exit 1 +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, or restart the server with `do_sample=false` in +# start_vllm_server.sh to force deterministic greedy decoding. +python3 < 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) + +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 +as vllm serving; both NxDI direct and vllm-neuron are installed there). +""" + +import os +import sys +import time +import traceback + +# 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", + "/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")) +# 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 +# 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=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() + +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. 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], + # 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. 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", + "norm", + "router", + "o_proj", + "q_proj", + "k_proj", + "v_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..e0d2502c --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/smoke_generate_mimo_v2.py @@ -0,0 +1,244 @@ +#!/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 + +# 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") + +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/", +) + +# 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", "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")) +MOE_EP = int(os.environ.get("MOE_EP", "64")) + +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", + "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) + ) + + 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) + # 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 + # 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/start_vllm_server.sh b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh new file mode 100644 index 00000000..ee2501c8 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/start_vllm_server.sh @@ -0,0 +1,104 @@ +#!/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//. +# 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. +: "${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 512 \ + --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", "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": 512, + "seq_len": 512, + "is_continuous_batching": true, + "enable_bucketing": true, + "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 + } + } + }' 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..96f2514f --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/perf_test/vllm-neuron-patch.patch @@ -0,0 +1,159 @@ +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..a8dfc7f 100644 +--- a/vllm_neuron/worker/neuronx_distributed_model_loader.py ++++ b/vllm_neuron/worker/neuronx_distributed_model_loader.py +@@ -922,6 +922,98 @@ 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 "mimov2" not in MODEL_TYPES: ++ if mimo_src not in _sys.path: ++ _sys.path.insert(0, mimo_src) ++ try: ++ from modeling_mimo_v2 import NeuronMiMoV2ForCausalLM ++ # 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 ("mimov2", "mimov2flash", "mimov2pro"): ++ MODEL_TYPES.setdefault(_mt, {"causal-lm": NeuronMiMoV2ForCausalLM}) ++ try: ++ from vllm.model_executor.models.registry import ModelRegistry ++ # 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 ("MiMoV2ForCausalLM", "MiMoV2FlashForCausalLM", "MiMoV2ProForCausalLM"): ++ ModelRegistry.register_model(_arch, NeuronMiMoV2ForCausalLM) ++ except ImportError: ++ pass ++ except Exception as e: ++ _w.warn( ++ f"Failed to register MiMo-V2 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 +1033,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..e63b5913 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/src/conversion_script/preprocess_mimo_v2_fp8.py @@ -0,0 +1,680 @@ +""" +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, 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 + 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 _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], + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + v_head_dim: int, +) -> Dict[str, torch.Tensor]: + """Split Pro's pre-fused qkv_proj into q/k/v (BF16 output). + + 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. + + 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 + 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 source path (rare — most Pro checkpoints ship as FP8+scale). + w = qkv_weight.view(num_kv_heads, real_rows_per_group, in_features) + 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})" + ) + + 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 + + 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] + + w = (w_padded * s_exp)[:, :real_rows_per_group, :] + + 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_bf16, "k_proj": k_bf16, "v_proj": v_bf16} + + +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_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, w_bf16 in split.items(): + out[f"{out_prefix}self_attn.{proj}.weight"] = w_bf16 + else: + # 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") + 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 + # 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: + # 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 + # 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/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..623946c0 --- /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_inference_vllm_0_16} +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" 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..070c4581 --- /dev/null +++ b/contrib/models/MiMo-V2.5-Pro/src/modeling_mimo_v2.py @@ -0,0 +1,2127 @@ +# 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 os +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, + 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: + 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) + +#: 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.""" + 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 + + # 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 + + # 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 _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: + 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 + # 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 + 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 + self.swa_kv_cache_is_ring = False + + # 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.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)) + + # 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) + ) + # 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. + + 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 + + # 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 + 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(): + 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( + 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 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 // 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) + 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 + ) + + # 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) + value_states = self.v_proj(hidden_states) + + # 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. + 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. 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(), + ) + value_states = gather_from_tensor_model_parallel_region_with_dim( + 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 + 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. + # + # 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')) + 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. + # + # 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( + 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) + # 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) + # position_ids: [bsz, q_len] -> per-slot current position [bsz, 1] + current_pos = position_ids[:, 0].reshape(-1, 1) # [bsz, 1] + 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')) + + 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: + # 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) + + # 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 + # 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) + 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) + + # 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) + + 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. + # `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 + + 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 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.""" + + 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. 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 = 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) + 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 + + # 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 + + 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 # 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 = 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 + 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