[contrib] Add MiMo-V2.5-Pro (Xiaomi, 384 experts MoE, FP8 on Trn2)

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"

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"

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=<dir> 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)

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-<path>.so.so -> not found -> silent TCP fallback), so use "ofi".
+ENV NCCL_NET_PLUGIN=ofi
+
+WORKDIR /workspace