From 83e63d07385c43b4052b167d536ff7eb8a592b25 Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Mon, 29 Jun 2026 18:29:15 -0500 Subject: [PATCH 01/12] Add DeepSeek-V4-Pro disaggregated P/D inference recipe (scripts/DeepSeekV4) Multi-node prefill/decode disaggregated serving/benchmark harness for DeepSeek-V4-Pro (ATOM+mooncake primary; SGLang+MoRI experimental), ported from InferenceX. Adds a Blueprints row to the top-level README. --- README.md | 1 + scripts/DeepSeekV4/README.md | 59 + .../DeepSeekV4/atom_disagg/accuracy_ctx.py | 73 + scripts/DeepSeekV4/atom_disagg/bench.sh | 47 + scripts/DeepSeekV4/atom_disagg/env_atom.sh | 71 + scripts/DeepSeekV4/atom_disagg/launch.sh | 130 ++ scripts/DeepSeekV4/atom_disagg/server_atom.sh | 585 ++++++++ scripts/DeepSeekV4/atom_disagg/setup_deps.sh | 11 + scripts/DeepSeekV4/cluster.yaml | 28 + scripts/DeepSeekV4/docs/proxy_and_disagg.md | 58 + scripts/DeepSeekV4/lib/cfg.py | 34 + scripts/DeepSeekV4/lib/check_accuracy.py | 59 + scripts/DeepSeekV4/lib/clean_node.sh | 38 + scripts/DeepSeekV4/lib/lib_inferencex.sh | 1281 +++++++++++++++++ scripts/DeepSeekV4/lib/prompts.json | 32 + scripts/DeepSeekV4/lib/run_disagg.sh | 200 +++ scripts/DeepSeekV4/lib/topo.py | 52 + scripts/DeepSeekV4/model.yaml | 72 + scripts/DeepSeekV4/run_atom_disagg.sh | 5 + scripts/DeepSeekV4/run_sglang_disagg.sh | 5 + .../utils/bench_serving/KNOWN_LIMITATION.md | 7 + .../bench_serving/backend_request_func.py | 574 ++++++++ .../utils/bench_serving/benchmark_serving.py | 1276 ++++++++++++++++ .../utils/bench_serving/benchmark_utils.py | 39 + .../utils/bench_serving/encoding_dsv4.py | 744 ++++++++++ 25 files changed, 5481 insertions(+) create mode 100644 scripts/DeepSeekV4/README.md create mode 100644 scripts/DeepSeekV4/atom_disagg/accuracy_ctx.py create mode 100755 scripts/DeepSeekV4/atom_disagg/bench.sh create mode 100755 scripts/DeepSeekV4/atom_disagg/env_atom.sh create mode 100755 scripts/DeepSeekV4/atom_disagg/launch.sh create mode 100755 scripts/DeepSeekV4/atom_disagg/server_atom.sh create mode 100755 scripts/DeepSeekV4/atom_disagg/setup_deps.sh create mode 100644 scripts/DeepSeekV4/cluster.yaml create mode 100644 scripts/DeepSeekV4/docs/proxy_and_disagg.md create mode 100755 scripts/DeepSeekV4/lib/cfg.py create mode 100755 scripts/DeepSeekV4/lib/check_accuracy.py create mode 100755 scripts/DeepSeekV4/lib/clean_node.sh create mode 100755 scripts/DeepSeekV4/lib/lib_inferencex.sh create mode 100644 scripts/DeepSeekV4/lib/prompts.json create mode 100755 scripts/DeepSeekV4/lib/run_disagg.sh create mode 100644 scripts/DeepSeekV4/lib/topo.py create mode 100644 scripts/DeepSeekV4/model.yaml create mode 100755 scripts/DeepSeekV4/run_atom_disagg.sh create mode 100755 scripts/DeepSeekV4/run_sglang_disagg.sh create mode 100644 scripts/DeepSeekV4/utils/bench_serving/KNOWN_LIMITATION.md create mode 100644 scripts/DeepSeekV4/utils/bench_serving/backend_request_func.py create mode 100644 scripts/DeepSeekV4/utils/bench_serving/benchmark_serving.py create mode 100644 scripts/DeepSeekV4/utils/bench_serving/benchmark_utils.py create mode 100644 scripts/DeepSeekV4/utils/bench_serving/encoding_dsv4.py diff --git a/README.md b/README.md index 2a972a38..e99868b3 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ Below are blueprints of supported models along with their documentation. | [**vLLM disaggregated P/D inference**](scripts/vllm_dissag/README.MD) | Distributed Inference P/D disaggregation with vLLM (Default, MoRI EP, DeepEP) | DeepSeek-R1, DeepSeek-V3, DeepSeek-V3-5layer, amd-Llama-3.3-70B-Instruct-FP8-KV, Llama-3.1-405B-Instruct-FP8-KV, gpt-oss-120b | | [**SGLang disaggregated P/D inference**](scripts/sglang_disagg/README.MD) | Distributed Inference P/D disaggregation with SGLang (MoRI IO, Mooncake) | Llama-3.1-8B, Qwen3-32B, Llama-3.3-70B-FP8, Llama-3.1-405B-FP8, Mixtral-8x7B, DeepSeek-V3, DeepSeek-R1 | | [**SGLang disaggregated P/D inference with WideEP/LargeEP**](scripts/sglang_disagg/README.MD) | Distributed Inference P/D disaggregation with SGLang with WideEP/LargeEP | DeepSeek-V3, DeepSeek-R1 | +| [**DeepSeek-V4-Pro disaggregated P/D inference**](scripts/DeepSeekV4/README.md) | Distributed Inference P/D disaggregation for DeepSeek-V4-Pro (ATOM+mooncake / SGLang+MoRI) | DeepSeek-V4-Pro | | [**KVCache Transfer Bench**](scripts/kvcache_transfer_bench/README.md) | Inter-node Transfer Benchmark | no specific models | ## Table of Contents diff --git a/scripts/DeepSeekV4/README.md b/scripts/DeepSeekV4/README.md new file mode 100644 index 00000000..7569f3eb --- /dev/null +++ b/scripts/DeepSeekV4/README.md @@ -0,0 +1,59 @@ +# DeepSeek-V4-Pro — disaggregated (P/D) serving/benchmark harness + +Multi-node **prefill/decode disaggregated** serving for DeepSeek-V4-Pro on AMD Instinct +(gfx950), ported from InferenceX. + +## Why disagg +DSv4-Pro is a large MoE (61 layers, 384 experts, hidden 7168, moe_intermediate 3072; FP8 block-quant ++ FP4 experts). It's served **disaggregated**: prefill workers (kv_producer) and decode workers +(kv_consumer) on separate nodes, KV streamed over RDMA, fronted by ONE router. Bench drives the +router → a TRUE aggregate number (no ×N extrapolation). TP8 is valid here (3072/8=384 ≥ 128). + +## Engines & proxy/router (see docs/proxy_and_disagg.md for detail) +| Engine | Router | KV transfer | Status | +|--------|--------|-------------|--------| +| **atom-disagg** | **atomesh** | **mooncake** (RDMA) | PRIMARY — InferenceX-validated DSv4 config | +| **sglang-disagg** | sglang_router | MoRI-IO | EXPERIMENTAL — adapted from DeepSeek-R1 recipe | + +(vLLM disagg uses MoRI-IO/Nixl; not scaffolded — its `server_vllm.sh` wasn't in our snapshot.) + +## Files +| File | Role | +|------|------| +| `cluster.yaml` | nodes, gpus/node, reservation, router/prefill/decode/handshake ports, paths | +| `model.yaml` | DSv4 def + shared `defaults{}` + per-engine image/router/kv + `topologies{}` | +| `run_atom_disagg.sh` / `run_sglang_disagg.sh` | thin wrappers → `lib/run_disagg.sh` | +| `lib/run_disagg.sh` | orchestrator: clean nodes → launch prefill/decode servers → start router → bench → teardown | +| `lib/topo.py` | resolve a topology (1p1d / 2p1d_dpa) → node-role placement (prefill/decode, ports) | +| `lib/cfg.py` | yaml reader (shared with qwen_v2) | +| `lib/clean_node.sh` | pre-flight GPU clean (shared with qwen_v2) | +| `utils/bench_serving/` | vendored InferenceX bench client | +| `docs/proxy_and_disagg.md` | full per-engine proxy/router/KV mechanics + topologies | + +## Topologies (from dsv4-fp4-mi355x-atom-disagg) +| name | nodes | layout | shapes | conc | +|------|-------|--------|--------|------| +| `1p1d` | 2 | 1 prefill TP8 + 1 decode TP8 | 1024/1024, 8192/1024 | 4–128 | +| `2p1d_dpa` | 3 | 2 prefill + 1 decode, TP8 + DP-attn | 8192/1024 | 256–2048 | + +## Usage +```bash +# ATOM disagg, 1-prefill/1-decode (2 nodes), full bench sweep through the router: +MODEL=dsv4-pro TOPO=1p1d ./run_atom_disagg.sh +MODEL=dsv4-pro TOPO=2p1d_dpa ./run_atom_disagg.sh # 3 nodes, high-conc + +# Just serve (hold the endpoint up; drive it yourself): +MODEL=dsv4-pro TOPO=1p1d ACTION=serve ./run_atom_disagg.sh + +# SGLang disagg (experimental, R1-template): +MODEL=dsv4-pro TOPO=1p1d ./run_sglang_disagg.sh +``` +Bench drives `http://:8000` (OpenAI-compatible). Results land in `results_dsv4/`. + +## Status / TODO +- ✅ Weights: DSv4-Pro 64/64 shards integrity-verified on NFS; staging to local scratch (12 nodes). +- ⏳ ATOM disagg: scaffolded faithfully from server_atom.sh; **needs a live dry-run** (RDMA/mooncake + handshake, atomesh router) once weights are staged. +- ⏳ SGLang disagg: scaffolded from R1 template; DSv4 model support + sglang_router flags TBC. +- ATOM DSv4 env applied by launcher: `ATOM_MOE_GU_ITLV=1`, `AITER_BF16_FP8_MOE_BOUND=0` (+ TBO path + `GPU_MAX_HW_QUEUES=5`, `ATOM_CPU_AFFINITY=1`). diff --git a/scripts/DeepSeekV4/atom_disagg/accuracy_ctx.py b/scripts/DeepSeekV4/atom_disagg/accuracy_ctx.py new file mode 100644 index 00000000..9db66477 --- /dev/null +++ b/scripts/DeepSeekV4/atom_disagg/accuracy_ctx.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# accuracy_ctx.py — context-varied accuracy probe for the DSv4 ATOM disagg pipeline. +# +# Why: a PD-disaggregated pipeline can be correct at short context but corrupt the KV that +# mooncake transfers prefill->decode at long context. A 5-short-prompt gate misses that. +# This embeds a known "needle" fact inside filler of increasing size and checks the router +# (which spans prefill+decode+mooncake) returns it verbatim — exercising the FULL pipeline +# at each context length. +# +# Usage: ROUTER_URL, MODEL, OUT(json), [CTX_SIZES="512 2048 8192 16384"] via env. +# Exits 0 if all context sizes pass, 1 otherwise. +import json, os, sys, time, urllib.request + +ROUTER = os.environ.get("ROUTER_URL", "http://0.0.0.0:8000") +MODEL = os.environ["MODEL"] +OUT = os.environ.get("OUT", "/out/accuracy_ctx.json") +# approx tokens of filler context to test at (the needle adds a few more) +CTX_SIZES = [int(x) for x in os.environ.get("CTX_SIZES", "512 2048 8192 16384").split()] + +# A deterministic needle: a unique code the model must echo back. Greedy (temp 0). +SECRET = "PLATYPUS-7731-AZURE" +FILLER_SENTENCE = ("The quarterly logistics report notes that warehouse throughput " + "remained within nominal parameters across all regional hubs. ") + +def make_prompt(approx_tokens): + # ~1.3 tokens/word for English; build filler then plant the needle in the MIDDLE + words_needed = int(approx_tokens / 1.3) + filler_words = FILLER_SENTENCE.split() + reps = max(1, words_needed // len(filler_words)) + block = (" ".join(filler_words) + " ") * reps + half = len(block) // 2 + needle = f" IMPORTANT: the access code is {SECRET}. Remember it exactly. " + ctx = block[:half] + needle + block[half:] + q = ("\n\nQuestion: What is the exact access code mentioned above? " + "Reply with ONLY the code, nothing else.") + return ctx + q + +def ask(prompt): + body = json.dumps({ + "model": MODEL, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": 32, "temperature": 0.0, + }).encode() + req = urllib.request.Request(f"{ROUTER}/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=300) as r: + d = json.load(r) + return d["choices"][0]["message"]["content"], time.time() - t0 + +results = []; allpass = True +for ctx in CTX_SIZES: + try: + prompt = make_prompt(ctx) + text, dt = ask(prompt) + found = SECRET in text.upper() + ok = found + results.append({"ctx_tokens": ctx, "ok": ok, "found_needle": found, + "reply": text.strip()[:80], "latency_s": round(dt, 2), + "prompt_chars": len(prompt)}) + print(f" ctx~{ctx:>6}tok: {'PASS' if ok else 'FAIL'} " + f"(needle={'yes' if found else 'NO'}, {dt:.1f}s) reply='{text.strip()[:40]}'") + allpass = allpass and ok + except Exception as e: + results.append({"ctx_tokens": ctx, "ok": False, "error": str(e)[:120]}) + print(f" ctx~{ctx:>6}tok: ERROR {str(e)[:80]}") + allpass = False + +os.makedirs(os.path.dirname(OUT), exist_ok=True) +json.dump({"model": MODEL, "router": ROUTER, "all_pass": allpass, "results": results}, + open(OUT, "w"), indent=2) +print(f"[accuracy_ctx] {'ALL PASS' if allpass else 'FAILURES'} -> {OUT}") +sys.exit(0 if allpass else 1) diff --git a/scripts/DeepSeekV4/atom_disagg/bench.sh b/scripts/DeepSeekV4/atom_disagg/bench.sh new file mode 100755 index 00000000..53442f10 --- /dev/null +++ b/scripts/DeepSeekV4/atom_disagg/bench.sh @@ -0,0 +1,47 @@ +#!/bin/bash +# bench.sh — driven by server_atom.sh on node 0 (router node), AFTER atomesh is /v1/models-ready. +# Reproduces the InferenceX call signature exactly: +# bash bench.sh \ +# "" +# +# Drives the OpenAI-compatible ATOM router at localhost:$ROUTER_PORT with the vendored +# InferenceX benchmark_serving client (same one used for the Qwen sweep) — one result JSON +# per concurrency. This is a TRUE aggregate (one endpoint, real PD + mooncake KV), no x N. +set -uo pipefail +xP="$1"; yD="$2"; PF_GPUS="$3"; DEC_GPUS="$4" +MODEL_DIR="$5"; MODEL_NAME="$6"; LOG_DIR="$7" +ISL="$8"; OSL="$9"; CONC_LIST="${10}"; REQ_RATE="${11}"; RANGE_RATIO="${12}"; PROMPTS_MULT="${13}" +# server_atom.sh passes the conc list 'x'-separated (e.g. 4x8x16). Normalise to spaces. +CONC_LIST=$(echo "$CONC_LIST" | tr 'x' ' ') + +ROUTER_PORT="${ROUTER_PORT:-8000}" +BENCH_CLIENT="${BENCH_CLIENT:-/scripts/utils/bench_serving/benchmark_serving.py}" +OUT="${LOG_DIR}"; mkdir -p "$OUT" 2>/dev/null || true + +echo "[bench] router=localhost:${ROUTER_PORT} model=${MODEL_NAME} isl=${ISL} osl=${OSL} conc='${CONC_LIST}' (xP=${xP} yD=${yD})" + +# --- ACCURACY GATE (context-varied) before perf ------------------------------------------- +# Validate the FULL PD pipeline (prefill+mooncake KV+decode) returns correct output at varied +# context sizes. Runs ONCE per launch (guarded by a marker), not per-shape. Non-fatal: logs +# result but does not abort the perf sweep (so we still get numbers even if a long-ctx probe fails). +ACC_OUT="${OUT}/accuracy_ctx.json" +if [[ ! -f "$ACC_OUT" ]]; then + echo "[accuracy] context-varied probe via router (sizes: ${CTX_SIZES:-512 2048 8192 16384})" + ROUTER_URL="http://localhost:${ROUTER_PORT}" MODEL="${MODEL_DIR}/${MODEL_NAME}" \ + OUT="$ACC_OUT" CTX_SIZES="${CTX_SIZES:-512 2048 8192 16384}" \ + python3 /ws/accuracy_ctx.py || echo "[accuracy] WARN: some context sizes failed (see $ACC_OUT)" +fi + +for c in $CONC_LIST; do + echo "[bench] --- concurrency $c ---" + python3 "$BENCH_CLIENT" \ + --model "${MODEL_DIR}/${MODEL_NAME}" --backend vllm \ + --base-url "http://localhost:${ROUTER_PORT}" --dataset-name random \ + --random-input-len "$ISL" --random-output-len "$OSL" --random-range-ratio "$RANGE_RATIO" \ + --num-prompts "$(( c * PROMPTS_MULT ))" --num-warmups "$(( c * 2 ))" \ + --max-concurrency "$c" --request-rate "$REQ_RATE" --ignore-eos --save-result \ + --percentile-metrics ttft,tpot,itl,e2el --result-dir "$OUT" \ + --result-filename "disagg_${xP}p${yD}d_isl${ISL}_osl${OSL}_c${c}.json" \ + || echo "[bench] FAILED isl=$ISL osl=$OSL conc=$c" +done +echo "[bench] done -> $OUT" diff --git a/scripts/DeepSeekV4/atom_disagg/env_atom.sh b/scripts/DeepSeekV4/atom_disagg/env_atom.sh new file mode 100755 index 00000000..78c7d495 --- /dev/null +++ b/scripts/DeepSeekV4/atom_disagg/env_atom.sh @@ -0,0 +1,71 @@ +#!/bin/bash +# ATOM/mooncake-specific environment setup for multi-node disaggregated serving. +# +# Sourced by server_atom.sh in place of env.sh (which is SGLang/MoRI-specific). +# +# REQUIRED ENVIRONMENT VARIABLES: +# IBDEVICES - RDMA/InfiniBand device names (e.g., ionic_0,ionic_1,...) +# Set by runner or auto-detected from hostname. + +set -x + +export PYTHONUNBUFFERED=1 +export PYTHONDONTWRITEBYTECODE=1 + +# ============================================================================= +# IBDEVICES detection (same as env.sh) +# ============================================================================= + +if [[ -z "$IBDEVICES" ]]; then + DETECTED=$(ibv_devinfo 2>/dev/null | grep "hca_id:" | awk '{print $2}' | paste -sd',') + if [[ -n "$DETECTED" ]]; then + export IBDEVICES="$DETECTED" + echo "[INFO] Auto-detected IBDEVICES=$IBDEVICES via ibv_devinfo on $(hostname -s)" + else + # ATOM uses mooncake proxy_ip/handshake_port for KV transfer — IBDEVICES is + # not passed as a server argument (unlike SGLang --disaggregation-ib-device). + # Log a warning but do not fail; mooncake will use its own RDMA device selection. + echo "[WARN] Unable to detect RDMA devices via ibv_devinfo; IBDEVICES unset (non-fatal for ATOM/mooncake)" >&2 + fi +else + echo "[INFO] Using IBDEVICES=$IBDEVICES (set by runner or environment)" +fi +export IBDEVICES + +export SAFETENSORS_FAST_GPU=1 +export VLLM_LOG_LEVEL=WARNING +export ATOM_LOG_LEVEL=WARNING +export AITER_LOG_LEVEL=WARNING +export LOG_LEVEL=WARNING +export LOGLEVEL=WARNING + +# ============================================================================= +# ATOM/mooncake-specific environment +# ============================================================================= + +# mooncake RDMA KV transfer library path +# FIX (verified 2026-06-28): image rocm/atom-dev:nightly_202606101403 ships Python 3.12, NOT 3.10. +# Original InferenceX path (.../python3.10/...) does not exist here -> mooncake native libs fail to +# load. Resolve the actual site-packages dir dynamically so this is image-version agnostic. +_MOONCAKE_DIR=$(python3 -c "import os,mooncake;print(os.path.dirname(mooncake.__file__))" 2>/dev/null) +export LD_LIBRARY_PATH=${_MOONCAKE_DIR:-/opt/venv/lib/python3.12/site-packages/mooncake}:/opt/rocm/lib:${LD_LIBRARY_PATH:-} + + +# ATOM_HOST_IP is set per-node in server_atom.sh (= host_ip, used as handshake IP) + +# aiter logging (WARNING to reduce noise; use DEBUG for troubleshooting) +export AITER_LOG_LEVEL=WARNING + +if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then + # ATOM MoE gather/scatter interleave optimization + export ATOM_MOE_GU_ITLV=1 + # Disable bf16->fp8 MoE bound (only for DeepSeek-V4-Pro) + export AITER_BF16_FP8_MOE_BOUND=0 +fi + +# Clear stale ATOM cache on startup (server_atom.sh handles this via rm -rf) +# No env var needed; documented here for reference. + +set +x + +echo "[INFO] ATOM env: IBDEVICES=$IBDEVICES LD_LIBRARY_PATH includes mooncake" diff --git a/scripts/DeepSeekV4/atom_disagg/launch.sh b/scripts/DeepSeekV4/atom_disagg/launch.sh new file mode 100755 index 00000000..d47d242e --- /dev/null +++ b/scripts/DeepSeekV4/atom_disagg/launch.sh @@ -0,0 +1,130 @@ +#!/bin/bash +# launch.sh — multi-node ATOM PD-disagg launcher for DSv4-Pro (the `job.slurm`-equivalent). +# Drives the VENDORED, verbatim InferenceX server_atom.sh across N nodes via docker+srun. +# +# Faithful to the InferenceX chain: +# dsv4_..._atom-disagg.sh -> submit.sh -> job.slurm -> server_atom.sh +# We replace submit.sh+job.slurm (the slurm plumbing) with this; server_atom.sh runs unmodified. +# +# Env (with defaults): +# TOPO=1p1d|2p1d topology (sets xP/yD/dp-attn) [default 1p1d] +# ISL, OSL bench shape [default 1024 1024] +# CONC_LIST space-separated concurrencies [topology default] +# ACTION=bench|serve|dry bench (default) | hold endpoint | print cmds only +# +# Topologies (from dsv4-fp4-mi355x-atom-disagg): +# 1p1d : xP=1 yD=1, 2 nodes, TP8, dp-attn OFF (conc 4-128; shapes 1024/1024 & 8192/1024) +# 2p1d : xP=2 yD=1, 3 nodes, TP8, dp-attn ON (conc 256-2048; shape 8192/1024) +set -uo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # scripts_dsv4/atom_disagg +ROOT="$(cd "$HERE/.." && pwd)" # scripts_dsv4 +cfg(){ python3 "$ROOT/lib/cfg.py" "$@"; } +CLUSTER="$ROOT/cluster.yaml" + +RESV=$(cfg "$CLUSTER" cluster reservation) +MROOT=$(cfg "$CLUSTER" cluster models_root) +RROOT=$(cfg "$CLUSTER" cluster results_root) +RPORT=$(cfg "$CLUSTER" cluster router_port) +PPORT=$(cfg "$CLUSTER" cluster prefill_port) +DPORT=$(cfg "$CLUSTER" cluster decode_port) +HPORT=$(cfg "$CLUSTER" cluster handshake_port) +IMG=$(cfg "$ROOT/model.yaml" engine atom-disagg image) +MODEL_NAME="DeepSeek-V4-Pro" + +TOPO="${TOPO:-1p1d}"; ACTION="${ACTION:-bench}" +ISL="${ISL:-1024}"; OSL="${OSL:-1024}" +case "$TOPO" in + 1p1d) xP=1; yD=1; DPATTN=0; DEF_CONC="4 8 16 32 64 128" ;; + 2p1d) xP=2; yD=1; DPATTN=1; DEF_CONC="256 512 768 1024 2048" ;; + *) echo "ERROR: TOPO must be 1p1d|2p1d (got $TOPO)"; exit 1 ;; +esac +CONC_LIST="${CONC_LIST:-$DEF_CONC}" +# server_atom.sh expects BENCH_MAX_CONCURRENCY 'x'-separated (it does `tr 'x' '\n'|sort -n|tail -1` +# to derive the decode server's --max-num-seqs). Our bench.sh wants space-separated. Keep both. +CONC_X=$(echo "$CONC_LIST" | tr ' ' 'x') +TP=8 # TP8 throughout (DSv4 valid: 3072/8=384>=128) +NNODES=$(( xP + yD )) # 1 node/worker at TP8 on 8-GPU nodes + +# --- node selection: explicit NODES= override (REQUIRED for parallel launches to avoid +# two launches racing to grab the same 'first free' nodes), else auto-pick free ones. ---- +if [[ -n "${NODES:-}" ]]; then + IFS=',' read -ra NODES <<< "$NODES" + [[ ${#NODES[@]} -lt $NNODES ]] && { echo "ERROR: NODES= has ${#NODES[@]}, need $NNODES"; exit 1; } + NODES=("${NODES[@]:0:$NNODES}") +else + mapfile -t ALL < <(scontrol show hostnames "$(scontrol show res "$RESV" 2>/dev/null | grep -oE 'Nodes=[^ ]+' | head -1 | cut -d= -f2)") + BUSY=$(squeue -u "$USER" -h -o "%N" 2>/dev/null | scontrol show hostnames 2>/dev/null | sort -u) + NODES=(); for n in "${ALL[@]}"; do grep -qx "$n" <<<"$BUSY" || NODES+=("$n"); [[ ${#NODES[@]} -ge $NNODES ]] && break; done + [[ ${#NODES[@]} -lt $NNODES ]] && { echo "ERROR: need $NNODES free nodes, have ${#NODES[@]}"; exit 1; } +fi + +STAMP=$(python3 -c "import time;print(time.strftime('%Y%m%d_%H%M%S'))") +RUN="$RROOT/dsv4_atom_${TOPO}_isl${ISL}_${STAMP}"; mkdir -p "$RUN"; chmod 777 "$RUN" +echo "[launch] TOPO=$TOPO xP=$xP yD=$yD nodes=${NODES[*]} shape=${ISL}/${OSL} conc='$CONC_LIST'" +echo "[launch] image=$IMG run=$RUN" + +# --- resolve node IPs (IPADDRS = prefill nodes first, then decode; server_atom.sh slices it) - +declare -a IPS=() +for n in "${NODES[@]}"; do + ip=$(srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=1G \ + bash -c "ip route get 1.1.1.1 2>/dev/null | awk '/src/{print \$7}'" 2>/dev/null \ + | grep -oE '([0-9]+\.){3}[0-9]+' | head -1) + IPS+=("$ip"); echo " $n -> $ip" +done +IPADDRS=$(IFS=,; echo "${IPS[*]}") +NODE0="${NODES[0]}"; NODE0_IP="${IPS[0]}" + +# --- pre-flight clean every node ------------------------------------------------------------ +for n in "${NODES[@]}"; do + srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=0 \ + bash "$ROOT/lib/clean_node.sh" > "$RUN/clean_${n}.log" 2>&1 && echo " [clean] $n OK" || echo " [clean] $n WARN" +done + +DRY=0; [[ "$ACTION" == "dry" ]] && DRY=1 + +# --- RDMA device flags (THE fix for mooncake 30s timeouts) ---------------------------------- +# VERIFIED 2026-06-28: the disagg image (rocm/atom-dev:nightly_202606101403) ALREADY has a +# working bnxt_re RDMA stack (its native libbnxt_re-rdmav34.so enumerates all 8 NICs PORT_ACTIVE). +# The earlier failure was missing --device=/dev/infiniband + --ulimit memlock=-1, NOT missing +# host libs. DO NOT bind-mount host libibverbs over the container's: the host (CentOS, +# libibverbs.so.1.15) and container (Ubuntu, .1.14) have incompatible IBVERBS_PRIVATE ABIs, so +# mounting host libs BREAKS the container's RDMA tools. Expose devices + memlock only. +RDMA_DEVICES="--device=/dev/infiniband --ulimit memlock=-1 --ulimit stack=67108864 --init" +RDMA_MOUNT_SNIPPET='RM=""' # no host-lib mounts (container stack is self-sufficient) + +# --- launch server_atom.sh on each node, NODE_RANK = its index ------------------------------ +# server_atom.sh logic: rank 0 = prefill0 + atomesh router (+ bench); rank=NODE_OFFSET = decode. We pass the SAME env to every node; the script self-selects by rank. +rank=0 +for n in "${NODES[@]}"; do + rdir="$RUN/node${rank}_${n}"; mkdir -p "$rdir"; chmod 777 "$rdir" + srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --exclusive bash -c " + docker rm -f dsv4_atom_n${rank} >/dev/null 2>&1 || true + $RDMA_MOUNT_SNIPPET # builds \$RM (host RDMA lib bind-mounts) on THIS node + echo \"[rdma] mounts: \$RM\" + docker run --rm --name dsv4_atom_n${rank} --entrypoint bash \ + --device=/dev/kfd --device=/dev/dri $RDMA_DEVICES \$RM -v /sys:/sys \ + --group-add video --ipc host --shm-size 128G \ + --network host --privileged --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e ATOM_WS_PATH=/ws -e NODE_RANK=${rank} -e NODE0_ADDR=${NODE0_IP} -e IPADDRS='${IPADDRS}' \ + -e MODEL_DIR=${MROOT} -e MODEL_NAME=${MODEL_NAME} \ + -e xP=${xP} -e yD=${yD} -e GPUS_PER_NODE=8 \ + -e PREFILL_TP_SIZE=${TP} -e DECODE_TP_SIZE=${TP} \ + -e PREFILL_ENABLE_DP=$([[ $DPATTN == 1 ]] && echo true || echo false) -e PREFILL_ENABLE_EP=1 \ + -e DECODE_ENABLE_DP=$([[ $DPATTN == 1 ]] && echo true || echo false) -e DECODE_ENABLE_EP=1 \ + -e PREFILL_PORT=${PPORT} -e DECODE_PORT=${DPORT} -e ROUTER_PORT=${RPORT} -e HANDSHAKE_PORT=${HPORT} \ + -e MEM_FRAC_STATIC=0.85 -e KV_CACHE_DTYPE=fp8 -e BLOCK_SIZE=16 -e MAX_NUM_SEQS=256 \ + -e BENCH_INPUT_LEN=${ISL} -e BENCH_OUTPUT_LEN=${OSL} -e BENCH_MAX_CONCURRENCY='${CONC_X}' \ + -e BENCH_REQUEST_RATE=inf -e BENCH_RANDOM_RANGE_RATIO=1 -e BENCH_NUM_PROMPTS_MULTIPLIER=10 \ + -e DRY_RUN=${DRY} -e SLURM_JOB_ID=${STAMP} \ + -v ${MROOT}:${MROOT} -v ${HERE}:/ws -v ${ROOT}:/scripts -v ${rdir}:/run_logs/slurm_job-${STAMP} \ + ${IMG} /ws/server_atom.sh + " > "$rdir/srun.log" 2>&1 & + echo " [launch] node${rank} ($n) server_atom.sh dispatched" + rank=$(( rank + 1 )) +done + +echo "[launch] all $NNODES nodes dispatched. node0=$NODE0 hosts atomesh router :$RPORT + drives bench." +echo "[launch] tail node0: $RUN/node0_${NODE0}/srun.log" +wait +echo "[launch] done -> $RUN (results: disagg_*.json under node0_*/)" diff --git a/scripts/DeepSeekV4/atom_disagg/server_atom.sh b/scripts/DeepSeekV4/atom_disagg/server_atom.sh new file mode 100755 index 00000000..ccc86403 --- /dev/null +++ b/scripts/DeepSeekV4/atom_disagg/server_atom.sh @@ -0,0 +1,585 @@ +#!/bin/bash +# ATOM Disaggregated Server Launcher +# ============================================================================= +# Uses atom.entrypoints.openai_server with mooncake RDMA KV transfer. +# Mirrors server_sglang.sh topology (dynamic xP/yD) but adapts to ATOM's +# explicit kv-transfer-config and atomesh router. +# +# Key differences from server_sglang.sh: +# - Engine: atom.entrypoints.openai_server (not sglang.launch_server) +# - KV transfer: mooncake (--kv-transfer-config JSON) +# - Router: atomesh (not sglang_router) +# - Prefill port: $PREFILL_PORT (default 8010) / Decode port: $DECODE_PORT (default 8020) +# - Router port: $ROUTER_PORT (default 8000) +# ============================================================================= + +# ============================================================================= +# Environment Configuration +# ============================================================================= + +NODE0_ADDR="${NODE0_ADDR:-localhost}" +NODE_RANK="${NODE_RANK:-0}" +MODEL_DIR="${MODEL_DIR:-}" +MODEL_NAME="${MODEL_NAME:-}" + +xP="${xP:-1}" +yD="${yD:-1}" + +IPADDRS="${IPADDRS:-localhost}" + +# Parallelism +PREFILL_TP_SIZE="${PREFILL_TP_SIZE:-8}" +PREFILL_ENABLE_EP="${PREFILL_ENABLE_EP}" +PREFILL_ENABLE_DP="${PREFILL_ENABLE_DP}" +DECODE_TP_SIZE="${DECODE_TP_SIZE:-8}" +DECODE_ENABLE_EP="${DECODE_ENABLE_EP}" +DECODE_ENABLE_DP="${DECODE_ENABLE_DP}" + +# MTP +SPEC_DECODING="${SPEC_DECODING:-}" +DECODE_MTP_SIZE="${DECODE_MTP_SIZE:-1}" + +# ATOM server ports (different from SGLang which uses 8000 for all) +PREFILL_PORT="${PREFILL_PORT:-8010}" +DECODE_PORT="${DECODE_PORT:-8020}" +ROUTER_PORT="${ROUTER_PORT:-8000}" +HANDSHAKE_PORT="${HANDSHAKE_PORT:-6301}" + +# ATOM server tuning (from reference script defaults) +MEM_FRAC_STATIC="${MEM_FRAC_STATIC:-0.85}" +KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" +BLOCK_SIZE="${BLOCK_SIZE:-16}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-256}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-}" +MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-}" +EXTRA_SERVER_ARGS="${EXTRA_SERVER_ARGS:-}" + +# Benchmark Configuration +BENCH_INPUT_LEN="${BENCH_INPUT_LEN:-1024}" +BENCH_OUTPUT_LEN="${BENCH_OUTPUT_LEN:-1024}" +BENCH_RANDOM_RANGE_RATIO="${BENCH_RANDOM_RANGE_RATIO:-1}" +BENCH_REQUEST_RATE="${BENCH_REQUEST_RATE:-inf}" +BENCH_NUM_PROMPTS_MULTIPLIER="${BENCH_NUM_PROMPTS_MULTIPLIER:-10}" +BENCH_MAX_CONCURRENCY="${BENCH_MAX_CONCURRENCY:-512}" + +DRY_RUN="${DRY_RUN:-0}" +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" + +# ============================================================================= +# Dependencies and Environment Setup +# ============================================================================= + +source $ATOM_WS_PATH/setup_deps.sh +source $ATOM_WS_PATH/env_atom.sh + +host_ip=$(ip route get 1.1.1.1 2>/dev/null | awk '/src/ {print $7}') +if [[ -z "$host_ip" ]]; then + host_ip=$(hostname -I 2>/dev/null | awk '{print $1}') +fi +host_name=$(hostname) + +# ============================================================================= +# Cluster Topology Configuration +# ============================================================================= + +IFS=',' read -ra IP_ARRAY <<< "$IPADDRS" + +PREFILL_NODES_PER_WORKER=$(((PREFILL_TP_SIZE + GPUS_PER_NODE - 1) / GPUS_PER_NODE)) +DECODE_NODES_PER_WORKER=$(((DECODE_TP_SIZE + GPUS_PER_NODE - 1) / GPUS_PER_NODE)) +NODE_OFFSET=$((PREFILL_NODES_PER_WORKER * xP)) + +# Build prefill IP list and atomesh --prefill args +PREFILL_ARGS="" +PREFILL_IPS=() +for i in $(seq 0 $((xP - 1))); do + idx=$((i * PREFILL_NODES_PER_WORKER)) + PREFILL_IPS[$i]="${IP_ARRAY[$idx]}" + PREFILL_ARGS="$PREFILL_ARGS --prefill http://${IP_ARRAY[$idx]}:${PREFILL_PORT}" +done + +# Build decode IP list and atomesh --decode args +DECODE_ARGS="" +DECODE_IPS=() +for i in $(seq 0 $((yD - 1))); do + idx=$((i * DECODE_NODES_PER_WORKER + NODE_OFFSET)) + DECODE_IPS[$i]="${IP_ARRAY[$idx]}" + DECODE_ARGS="$DECODE_ARGS --decode http://${IP_ARRAY[$idx]}:${DECODE_PORT}" +done + +PREFILL_ENABLE_EP="${PREFILL_ENABLE_EP}" +PREFILL_ENABLE_DP="${PREFILL_ENABLE_DP}" +DECODE_ENABLE_EP="${DECODE_ENABLE_EP}" +DECODE_ENABLE_DP="${DECODE_ENABLE_DP}" + +# Parallel args +PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE") #TP +if [ "$PREFILL_ENABLE_DP" = "true" ]; then + if [ "$PREFILL_ENABLE_EP" -gt 1 ]; then #DPA+EP + PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" --enable-expert-parallel --enable-dp-attention ) + else #TP+DPA+TBO + if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then + PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" --enable-dp-attention --enable-tbo ) + export GPU_MAX_HW_QUEUES=5 + export ATOM_CPU_AFFINITY=1 + else #TP+DPA + PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" --enable-dp-attention ) + fi + fi +fi + +# (srok), split DPA & TBO cases +DECODE_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE") #TP +if [ "$DECODE_ENABLE_DP" = "true" ]; then + if [ "$DECODE_ENABLE_EP" -gt 1 ]; then #DPA+EP + DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" --enable-expert-parallel --enable-dp-attention ) + else #TP+DPA+TBO + if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then + DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" --enable-dp-attention --enable-tbo ) + export GPU_MAX_HW_QUEUES=5 + export ATOM_CPU_AFFINITY=1 + else #TP+DPA + DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" --enable-dp-attention ) + fi + fi +fi + +# MTP args +SPEC_ARGS=() #TP +if [ "$SPEC_DECODING" = "mtp" ]; then + SPEC_ARGS=(--method mtp --num-speculative-tokens "$DECODE_MTP_SIZE") +fi + +# HF overrides (single-quoted JSON preserved through eval) +HF_OVERRIDES_ARG="" +if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then + HF_OVERRIDES_ARG="--hf-overrides '{\"use_index_cache\":true,\"index_topk_freq\":4}'" +fi + +# KV cache dtype (skip if unset or 'auto') +KV_CACHE_ARG="" +if [[ -n "$KV_CACHE_DTYPE" && "$KV_CACHE_DTYPE" != "auto" ]]; then + KV_CACHE_ARG="--kv_cache_dtype ${KV_CACHE_DTYPE}" +fi + +# Optional model length / batched-token cap +MODEL_LEN_ARGS="" +if [[ -n "$MAX_MODEL_LEN" ]]; then + MODEL_LEN_ARGS="${MODEL_LEN_ARGS} --max-model-len ${MAX_MODEL_LEN}" +fi +if [[ -n "$MAX_NUM_BATCHED_TOKENS" ]]; then + MODEL_LEN_ARGS="${MODEL_LEN_ARGS} --max-num-batched-tokens ${MAX_NUM_BATCHED_TOKENS}" +fi + +if [[ "$MODEL_NAME" != "DeepSeek-V4-Pro" ]]; then + export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +fi + +cat < prefill node 0 + router +# rank 1 .. (NODE_OFFSET-1) -> remaining prefill nodes +# rank NODE_OFFSET .. -> decode nodes +# ============================================================================= +if [ "$NODE_RANK" -eq 0 ]; then + # ────────────────────────────────────────────────────────────────────────── + # Node 0: prefill server (producer) + atomesh router + # ────────────────────────────────────────────────────────────────────────── + echo "NODE INFO =======================================" + echo "${host_name}:${host_ip} is Prefill Node 0 + Router" + echo "Prefill TP=${PREFILL_TP_SIZE}, Decode TP=${DECODE_TP_SIZE}" + echo "Prefill servers: ${PREFILL_ARGS}" + echo "Decode servers: ${DECODE_ARGS}" + echo "================================================" + + PREFILL_CMD="python3 -m atom.entrypoints.openai_server \ + --model ${MODEL_DIR}/${MODEL_NAME} \ + --host 0.0.0.0 --server-port ${PREFILL_PORT} \ + --trust-remote-code \ + ${PREFILL_PARALLEL_ARGS[*]} \ + ${SPEC_ARGS[*]} \ + ${KV_CACHE_ARG} \ + --block-size ${BLOCK_SIZE} \ + --gpu-memory-utilization ${MEM_FRAC_STATIC} \ + --max-num-seqs ${MAX_NUM_SEQS} \ + ${MODEL_LEN_ARGS} \ + --no-enable_prefix_caching \ + ${HF_OVERRIDES_ARG} \ + --kv-transfer-config '{\"kv_role\":\"kv_producer\",\"kv_connector\":\"mooncake\",\"proxy_ip\":\"${host_ip}\",\"handshake_port\":${HANDSHAKE_PORT}}' \ + ${EXTRA_SERVER_ARGS}" + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $PREFILL_CMD" + else + set -x + eval "$PREFILL_CMD" \ + 2>&1 | tee /run_logs/slurm_job-${SLURM_JOB_ID}/prefill0_${host_name}.log & + set +x + prefill0_pid=$! + fi + + # Wait for all prefill and decode servers to be ready + WAIT_SERVER_TIMEOUT="${WAIT_SERVER_TIMEOUT:-2500}" + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting for all servers to be up (timeout=${WAIT_SERVER_TIMEOUT}s)..." + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait for prefill/decode /health endpoints" + else + _deadline=$(( $(date +%s) + WAIT_SERVER_TIMEOUT )) + for _ip in "${PREFILL_IPS[@]}"; do + echo "[wait] prefill http://${_ip}:${PREFILL_PORT}/health" + while ! curl -sf --max-time 10 "http://${_ip}:${PREFILL_PORT}/health" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_deadline ]]; then + echo "[wait][FAIL] prefill ${_ip}:${PREFILL_PORT} not ready after ${WAIT_SERVER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] prefill ${_ip}:${PREFILL_PORT} ready" + done + for _ip in "${DECODE_IPS[@]}"; do + echo "[wait] decode http://${_ip}:${DECODE_PORT}/health" + while ! curl -sf --max-time 10 "http://${_ip}:${DECODE_PORT}/health" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_deadline ]]; then + echo "[wait][FAIL] decode ${_ip}:${DECODE_PORT} not ready after ${WAIT_SERVER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] decode ${_ip}:${DECODE_PORT} ready" + done + fi + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "All servers up. Starting atomesh router..." + + ROUTER_CMD="/usr/local/bin/atomesh launch \ + --host 0.0.0.0 --port ${ROUTER_PORT} \ + --pd-disaggregation \ + ${PREFILL_ARGS} \ + ${DECODE_ARGS} \ + --policy random \ + --backend atom \ + --log-level info \ + --disable-health-check \ + --disable-circuit-breaker \ + --prometheus-port 29100" + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $ROUTER_CMD" + else + ROUTER_LOG_FILE="/tmp/slurm_job-${SLURM_JOB_ID}_router_${host_name}.log" + set -x + eval "$ROUTER_CMD" 2>&1 | tee "$ROUTER_LOG_FILE" & + set +x + proxy_pid=$! + + # Wait for router to accept connections + WAIT_ROUTER_TIMEOUT="${WAIT_ROUTER_TIMEOUT:-300}" + echo "[wait] router http://0.0.0.0:${ROUTER_PORT}/v1/models (timeout=${WAIT_ROUTER_TIMEOUT}s)" + _router_deadline=$(( $(date +%s) + WAIT_ROUTER_TIMEOUT )) + while ! curl -sf --max-time 10 "http://0.0.0.0:${ROUTER_PORT}/v1/models" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_router_deadline ]]; then + echo "[wait][FAIL] router ${ROUTER_PORT}/v1/models not ready after ${WAIT_ROUTER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] router /v1/models ready" + + echo "Router is ready for benchmarking" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Ready for benchmarking on ${host_name}:${host_ip}" + + cd $ATOM_WS_PATH + + export IS_MTP="false" + if [ "$SPEC_DECODING" = "mtp" ]; then + export IS_MTP="true" + fi + + BENCH_CMD="bash $ATOM_WS_PATH/bench.sh ${xP} ${yD} $((PREFILL_TP_SIZE*xP)) $((DECODE_TP_SIZE*yD)) \ + $MODEL_DIR $MODEL_NAME /run_logs/slurm_job-${SLURM_JOB_ID} ${BENCH_INPUT_LEN} \ + ${BENCH_OUTPUT_LEN} \"${BENCH_MAX_CONCURRENCY}\" ${BENCH_REQUEST_RATE} \ + ${BENCH_RANDOM_RANGE_RATIO} ${BENCH_NUM_PROMPTS_MULTIPLIER}" + + if [[ "${EVAL_ONLY:-false}" == "true" ]]; then + echo "EVAL_ONLY mode: skipping throughput benchmark" + elif [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $BENCH_CMD" + else + set -x + eval "$BENCH_CMD" + set +x + fi + + # Run evaluation if requested (before killing router) + if [[ "${RUN_EVAL:-false}" == "true" ]]; then + echo "Running lm-eval evaluation on Node 0..." + + # Health check: verify the router is still serving before running eval. + EVAL_HEALTH_OK=false + for _attempt in 1 2 3; do + if curl -sf --max-time 10 "http://0.0.0.0:${ROUTER_PORT}/health" >/dev/null 2>&1; then + EVAL_HEALTH_OK=true + break + fi + echo "Eval health check attempt $_attempt failed, retrying in 10s..." + sleep 10 + done + + if [[ "$EVAL_HEALTH_OK" != "true" ]]; then + echo "WARNING: Router health check failed after 3 attempts. Skipping eval." + else + pushd /workspace + + source /workspace/benchmarks/benchmark_lib.sh + + if [[ -n "${EVAL_CONC:-}" ]]; then + export EVAL_CONCURRENT_REQUESTS="${EVAL_CONC}" + else + export EVAL_CONCURRENT_REQUESTS=$(echo "$BENCH_MAX_CONCURRENCY" | tr 'x' '\n' | sort -n | tail -1) + fi + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: run_eval --framework lm-eval --port ${ROUTER_PORT} (conc=${EVAL_CONCURRENT_REQUESTS})" + else + MODEL_NAME="${MODEL_DIR}/${MODEL_NAME}" run_eval --framework lm-eval --port ${ROUTER_PORT} + eval_rc=$? + + if [[ $eval_rc -ne 0 ]]; then + echo "ERROR: run_eval exited rc=$eval_rc; skipping metadata write and eval artifact staging" >&2 + EVAL_FAILED=1 + else + export TP="${PREFILL_TP_SIZE}" + export CONC="${EVAL_CONCURRENT_REQUESTS}" + export PREFILL_TP="${PREFILL_TP_SIZE}" + export PREFILL_EP=1 + export PREFILL_NUM_WORKERS="${xP}" + export DECODE_TP="${DECODE_TP_SIZE}" + export DECODE_EP=1 + export DECODE_NUM_WORKERS="${yD}" + export ISL="${BENCH_INPUT_LEN}" + export OSL="${BENCH_OUTPUT_LEN}" + + MODEL_NAME="${MODEL_DIR}/${MODEL_NAME}" append_lm_eval_summary + + EVAL_COPY_DIR="/run_logs/slurm_job-${SLURM_JOB_ID}/eval_results" + mkdir -p "$EVAL_COPY_DIR" + for f in meta_env.json; do + [ -e "/workspace/$f" ] && cp -f "/workspace/$f" "$EVAL_COPY_DIR/" + done + find /workspace -maxdepth 1 -name 'results*.json' -exec cp -f {} "$EVAL_COPY_DIR/" \; + find /workspace -maxdepth 1 -name 'sample*.jsonl' -exec cp -f {} "$EVAL_COPY_DIR/" \; + + echo "Eval completed. Artifacts staged in $EVAL_COPY_DIR" + fi + fi + + popd + fi + fi + + # Copy results + LOGS_OUTPUT="${BENCHMARK_LOGS_DIR:-/run_logs}/logs" + mkdir -p "$LOGS_OUTPUT" + if [[ "$DRY_RUN" -eq 0 ]]; then + cp -r /run_logs/slurm_job-${SLURM_JOB_ID} "$LOGS_OUTPUT/" + echo "Copied results to $LOGS_OUTPUT/slurm_job-${SLURM_JOB_ID}" + fi + + echo "Waiting 60s before killing router and prefill server..." + sleep 60 + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Killing router and prefill server" + if [[ "$DRY_RUN" -eq 0 ]]; then + kill $proxy_pid + kill $prefill0_pid + fi + + if [[ "${EVAL_FAILED:-0}" -eq 1 ]]; then + echo "ERROR: eval failed; exiting node-0 with rc=1" + exit 1 + fi + +elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -lt "$NODE_OFFSET" ]; then + # ────────────────────────────────────────────────────────────────────────── + # Prefill nodes 1..N (kv_producer) + # ────────────────────────────────────────────────────────────────────────── + echo "${host_name}:${host_ip} is Prefill Node (rank ${NODE_RANK})" + + # Determine which prefill worker this node belongs to, and its headnode IP + prefill_worker_idx=$((NODE_RANK / PREFILL_NODES_PER_WORKER)) + PREFILL_HEADNODE_IP="${PREFILL_IPS[$prefill_worker_idx]}" + + PREFILL_CMD="python3 -m atom.entrypoints.openai_server \ + --model ${MODEL_DIR}/${MODEL_NAME} \ + --host 0.0.0.0 --server-port ${PREFILL_PORT} \ + --trust-remote-code \ + ${PREFILL_PARALLEL_ARGS[*]} \ + ${SPEC_ARGS[*]} \ + ${KV_CACHE_ARG} \ + --block-size ${BLOCK_SIZE} \ + --gpu-memory-utilization ${MEM_FRAC_STATIC} \ + --max-num-seqs ${MAX_NUM_SEQS} \ + ${MODEL_LEN_ARGS} \ + --no-enable_prefix_caching \ + ${HF_OVERRIDES_ARG} \ + --kv-transfer-config '{\"kv_role\":\"kv_producer\",\"kv_connector\":\"mooncake\",\"proxy_ip\":\"${host_ip}\",\"handshake_port\":${HANDSHAKE_PORT}}' \ + ${EXTRA_SERVER_ARGS}" + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $PREFILL_CMD" + else + set -x + eval "$PREFILL_CMD" \ + 2>&1 | tee /run_logs/slurm_job-${SLURM_JOB_ID}/prefill_${host_name}.log & + set +x + prefill_pid=$! + trap 'echo "Caught signal, killing prefill (pid=$prefill_pid)"; kill $prefill_pid 2>/dev/null; exit 0' SIGTERM SIGINT + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting for router to be up..." + WAIT_ROUTER_TIMEOUT="${WAIT_ROUTER_TIMEOUT:-2800}" + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait for router ${NODE0_ADDR}:${ROUTER_PORT}/health" + else + _router_deadline=$(( $(date +%s) + WAIT_ROUTER_TIMEOUT )) + while ! curl -sf --max-time 10 "http://${NODE0_ADDR}:${ROUTER_PORT}/health" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_router_deadline ]]; then + echo "[wait][FAIL] router ${NODE0_ADDR}:${ROUTER_PORT} not ready after ${WAIT_ROUTER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] router ${NODE0_ADDR}:${ROUTER_PORT} ready" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting until router closes..." + trap 'echo "Caught signal, killing prefill (pid=$prefill_pid)"; kill $prefill_pid 2>/dev/null; exit 0' SIGTERM SIGINT + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait until router ${NODE0_ADDR}:${ROUTER_PORT} closes" + else + while curl -sf --max-time 10 "http://${NODE0_ADDR}:${ROUTER_PORT}/health" >/dev/null 2>&1; do + sleep 10 & + wait $! + done + echo "[wait] router ${NODE0_ADDR}:${ROUTER_PORT} closed" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Killing prefill server (rank ${NODE_RANK})" + if [[ "$DRY_RUN" -eq 0 ]]; then kill $prefill_pid 2>/dev/null; fi + +else + # ────────────────────────────────────────────────────────────────────────── + # Decode nodes (kv_consumer) + # ────────────────────────────────────────────────────────────────────────── + RANK=$((NODE_RANK - NODE_OFFSET)) + echo "${host_name}:${host_ip} is Decode Node (rank ${RANK})" + + _MAX_CONC=$(echo "$BENCH_MAX_CONCURRENCY" | tr 'x' '\n' | sort -n | tail -1) + if [[ "$_MAX_CONC" -gt 2048 ]]; then + CUDAGRAPH_SIZES='[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124,128,132,136,140,144,148,152,156,160,164,168,172,176,180,184,188,192,196,200,204,208,212,216,220,224,228,232,236,240,244,248,252,256,512,1024,2048,4096]' + elif [[ "$_MAX_CONC" -gt 1024 ]]; then + CUDAGRAPH_SIZES='[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124,128,132,136,140,144,148,152,156,160,164,168,172,176,180,184,188,192,196,200,204,208,212,216,220,224,228,232,236,240,244,248,252,256,512,1024,2048]' + elif [[ "$_MAX_CONC" -gt 512 ]]; then + CUDAGRAPH_SIZES='[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124,128,132,136,140,144,148,152,156,160,164,168,172,176,180,184,188,192,196,200,204,208,212,216,220,224,228,232,236,240,244,248,252,256,512,768,1024]' + else + CUDAGRAPH_SIZES='[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124,128,132,136,140,144,148,152,156,160,164,168,172,176,180,184,188,192,196,200,204,208,212,216,220,224,228,232,236,240,244,248,252,256,512]' + fi + + if [[ "$BENCH_INPUT_LEN" == "1024" && "$BENCH_OUTPUT_LEN" == "1024" ]]; then + DECODE_MAX_NUM_SEQS="${_MAX_CONC}" + else + DECODE_MAX_NUM_SEQS="${MAX_NUM_SEQS}" + fi + + DECODE_CMD="python3 -m atom.entrypoints.openai_server \ + --model ${MODEL_DIR}/${MODEL_NAME} \ + --host 0.0.0.0 --server-port ${DECODE_PORT} \ + --trust-remote-code \ + ${DECODE_PARALLEL_ARGS[*]} \ + ${SPEC_ARGS[*]} \ + ${KV_CACHE_ARG} \ + --block-size ${BLOCK_SIZE} \ + --gpu-memory-utilization ${MEM_FRAC_STATIC} \ + --max-num-seqs ${DECODE_MAX_NUM_SEQS} \ + ${MODEL_LEN_ARGS} \ + --no-enable_prefix_caching \ + ${HF_OVERRIDES_ARG} \ + --kv-transfer-config '{\"kv_role\":\"kv_consumer\",\"kv_connector\":\"mooncake\",\"proxy_ip\":\"${host_ip}\",\"handshake_port\":${HANDSHAKE_PORT}}' \ + --cudagraph-capture-sizes "${CUDAGRAPH_SIZES}" \ + ${EXTRA_SERVER_ARGS}" + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $DECODE_CMD" + else + set -x + eval "$DECODE_CMD" \ + 2>&1 | tee /run_logs/slurm_job-${SLURM_JOB_ID}/decode_${host_name}.log & + set +x + decode_pid=$! + trap 'echo "Caught signal, killing decode (pid=$decode_pid)"; kill $decode_pid 2>/dev/null; exit 0' SIGTERM SIGINT + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting for router to be up..." + WAIT_ROUTER_TIMEOUT="${WAIT_ROUTER_TIMEOUT:-2800}" + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait for router ${NODE0_ADDR}:${ROUTER_PORT}/health" + else + _router_deadline=$(( $(date +%s) + WAIT_ROUTER_TIMEOUT )) + while ! curl -sf --max-time 10 "http://${NODE0_ADDR}:${ROUTER_PORT}/health" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_router_deadline ]]; then + echo "[wait][FAIL] router ${NODE0_ADDR}:${ROUTER_PORT} not ready after ${WAIT_ROUTER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] router ${NODE0_ADDR}:${ROUTER_PORT} ready" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting until router closes..." + trap 'echo "Caught signal, killing decode (pid=$decode_pid)"; kill $decode_pid 2>/dev/null; exit 0' SIGTERM SIGINT + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait until router ${NODE0_ADDR}:${ROUTER_PORT} closes" + else + while curl -sf --max-time 10 "http://${NODE0_ADDR}:${ROUTER_PORT}/health" >/dev/null 2>&1; do + sleep 10 & + wait $! + done + echo "[wait] router ${NODE0_ADDR}:${ROUTER_PORT} closed" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Killing decode server (rank ${RANK})" + if [[ "$DRY_RUN" -eq 0 ]]; then kill $decode_pid 2>/dev/null; fi +fi + +echo "Script completed successfully" +exit 0 diff --git a/scripts/DeepSeekV4/atom_disagg/setup_deps.sh b/scripts/DeepSeekV4/atom_disagg/setup_deps.sh new file mode 100755 index 00000000..d6eea0d9 --- /dev/null +++ b/scripts/DeepSeekV4/atom_disagg/setup_deps.sh @@ -0,0 +1,11 @@ +#!/bin/bash +# setup_deps.sh — sourced by server_atom.sh at startup. +# +# In the InferenceX CI this installs/patches runtime deps into the container. Our image +# (rocm/atom-dev:nightly_202606101403) already ships everything needed for DSv4 disagg — +# verified present: atomesh (/usr/local/bin/atomesh), mooncake (py3.12 site-packages + +# /usr/local/bin/mooncake_{master,client}), atom.entrypoints.openai_server. +# +# So this is intentionally a NO-OP stub. If a future image is missing a dep, add the +# install here. Kept as a separate file so server_atom.sh runs unmodified (verbatim vendor). +: diff --git a/scripts/DeepSeekV4/cluster.yaml b/scripts/DeepSeekV4/cluster.yaml new file mode 100644 index 00000000..b218c934 --- /dev/null +++ b/scripts/DeepSeekV4/cluster.yaml @@ -0,0 +1,28 @@ +# cluster.yaml — DSv4 disagg cluster definition. SCALE LEVER for multi-node PD. +# Unlike Qwen (1 node, independent replicas), DSv4 disagg spans MULTIPLE nodes per +# serving instance: xP prefill nodes + yD decode nodes + 1 router (on prefill node 0). + +reservation: +gpus_per_node: 8 # gfx950 +hardware: MI350 + +# Model staging root (same path on every node; local NVMe recommended). +models_root: /path/to/models + +results_root: ./results + +# Disagg ports (match InferenceX server_atom.sh defaults) +router_port: 8000 +prefill_port: 8010 +decode_port: 8020 +handshake_port: 6301 + +# Node pool. The launcher picks the first +# (xP*pf_nodes_per_worker + yD*dec_nodes_per_worker) nodes for a given topology. +nodes: + - + - + - + - + - + - diff --git a/scripts/DeepSeekV4/docs/proxy_and_disagg.md b/scripts/DeepSeekV4/docs/proxy_and_disagg.md new file mode 100644 index 00000000..0f82829e --- /dev/null +++ b/scripts/DeepSeekV4/docs/proxy_and_disagg.md @@ -0,0 +1,58 @@ +# DSv4-Pro disaggregated serving — proxy/router per engine (from InferenceX) + +DeepSeek-V4-Pro is a large MoE (61 layers, 384 routed experts, hidden 7168, `moe_intermediate_size=3072`). +Checkpoint is **FP8 block-quant (128×128, scale ue8m0)** with FP4 MoE expert weights applied at +runtime — i.e. "FP4+FP8 mixed". It does NOT fit/serve well aggregated on one node at scale, so +serving is **prefill/decode disaggregated (PD-disagg)** across nodes. + +Note: unlike Qwen3-Next FP8 (`moe_intermediate_size=512` → pure TP8 invalid), DSv4 has 3072/8=384 ≥ 128, +so **TP8 is valid** for DSv4 FP8. InferenceX runs TP8 throughout. + +## The proxy/router each engine uses + +| Engine | Router/proxy | KV-transfer backend | Ports | InferenceX source | +|--------|--------------|---------------------|-------|-------------------| +| **ATOM** | **atomesh** (`/usr/local/bin/atomesh launch --pd-disaggregation`) | **mooncake** (RDMA, `--kv-transfer-config` JSON) | router 8000, prefill 8010, decode 8020, handshake 6301 | `server_atom.sh` (DSv4 disagg config exists) | +| **SGLang** | **sglang_router** / mini-lb (`ENGINE=sglang-disagg`) | **MoRI** (MoRI-IO RDMA) | router 8000 (all servers 8000) | `server_sglang.sh` + `env.sh` (only DeepSeek-**R1** disagg config exists; adapt to DSv4) | +| **vLLM** | vLLM disagg proxy | **MoRI-IO** / Nixl (`MoRIIOConnector`) | — | `server_vllm.sh` (referenced by dispatcher `server.sh`; not in our snapshot) | + +Dispatcher `server.sh` selects by `ENGINE`: `atom-disagg`→server_atom.sh, `vllm-disagg`→server_vllm.sh, +else server_sglang.sh. + +## ATOM disagg mechanics (the one DSv4-validated path) + +- **Topology**: `xP` prefill workers + `yD` decode workers, each TP8 (1 worker = 1 node at TP8). + - 1P1D = 2 nodes; 2P1D + DP-attention = 3 nodes (high concurrency). +- **Prefill server** (kv_producer): + ``` + python3 -m atom.entrypoints.openai_server --model --host 0.0.0.0 --server-port 8010 \ + -tp 8 [--enable-dp-attention --enable-tbo] --kv-cache-dtype fp8 --block-size 16 \ + --gpu-memory-utilization 0.85 --max-num-seqs N --no-enable_prefix_caching \ + --kv-transfer-config '{"kv_role":"kv_producer","kv_connector":"mooncake","proxy_ip":"","handshake_port":6301}' + ``` +- **Decode server** (kv_consumer): same, `--server-port 8020`, `kv_role":"kv_consumer"`. +- **Router** (on prefill node 0, after all servers healthy): + ``` + atomesh launch --host 0.0.0.0 --port 8000 --pd-disaggregation \ + --prefill http://:8010 [...] --decode http://:8020 [...] \ + --policy random --backend atom --disable-health-check --disable-circuit-breaker + ``` +- **Bench**: drive `http://:8000` (OpenAI-compatible) at aggregate concurrency. +- **DSv4-specific env** (from env_atom.sh): `ATOM_MOE_GU_ITLV=1`, `AITER_BF16_FP8_MOE_BOUND=0`, + and for the TBO path `GPU_MAX_HW_QUEUES=5`, `ATOM_CPU_AFFINITY=1`. +- KV transfer uses mooncake's own RDMA device selection (IBDEVICES not passed as a server arg). + +## SGLang disagg mechanics (adapt from DeepSeek-R1 recipe — no DSv4 config yet) + +- `ENGINE=sglang-disagg`; MoRI-IO env in `env.sh` (`MORI_IO_*`, `MORI_RDMA_TC/SL` QoS by NIC). +- Prefill: `--disaggregation-mode prefill`; Decode: `--disaggregation-mode decode` + DP-attention/EP + (R1 decode uses `ep:8 dp-attn:true`, "DEP" mode). All servers on port 8000; sglang_router fronts. +- DSv4 caveat: needs an sglang image carrying the DSv4 model branch (the `-DSv4` rocm/sgl-dev tags). + R1 disagg uses `rocm/sgl-dev:sglang-0.5.9-rocm720-mi35x-mori-*` — DSv4 support must be confirmed. + +## Topologies to benchmark (from dsv4-fp4-mi355x-atom-disagg) +| shape | conc | topo | nodes | +|-------|------|------|-------| +| 1024/1024 | 4–1024 | 1P1D TP8 | 2 | +| 8192/1024 | 4–128 | 1P1D TP8 | 2 | +| 8192/1024 | 256–2048 | 2P1D + DP-attn TP8 | 3 | diff --git a/scripts/DeepSeekV4/lib/cfg.py b/scripts/DeepSeekV4/lib/cfg.py new file mode 100755 index 00000000..9ced83da --- /dev/null +++ b/scripts/DeepSeekV4/lib/cfg.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +"""cfg.py — tiny yaml field reader for the shell scripts. +Usage: + cfg.py cluster -> cluster.yaml top-level key + cfg.py model -> models.. + cfg.py engine -> engines.. +Lists print space-separated; scalars print as-is. +""" +import sys, yaml + +def emit(v): + if isinstance(v, list): print(" ".join(str(x) for x in v)) + elif isinstance(v, dict): print(" ".join(f"{k}={v}" for k, v in v.items())) + else: print("" if v is None else v) + +def main(): + f = sys.argv[1]; kind = sys.argv[2] + d = yaml.safe_load(open(f)) + if kind == "cluster": + emit(d.get(sys.argv[3])) + elif kind == "model": + emit(d["models"][sys.argv[3]].get(sys.argv[4])) + elif kind == "engine": + emit(d["engines"][sys.argv[3]].get(sys.argv[4])) + elif kind == "default": + # default : per-model override wins over global defaults{} + model = sys.argv[3]; key = sys.argv[4] + mv = d.get("models", {}).get(model, {}).get(key) + emit(mv if mv is not None else d.get("defaults", {}).get(key)) + else: + sys.exit(f"unknown kind {kind}") + +if __name__ == "__main__": + main() diff --git a/scripts/DeepSeekV4/lib/check_accuracy.py b/scripts/DeepSeekV4/lib/check_accuracy.py new file mode 100755 index 00000000..e89571ce --- /dev/null +++ b/scripts/DeepSeekV4/lib/check_accuracy.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +# Sanity/accuracy check: send 5 real prompts to a running OpenAI-compatible server, +# verify coherent + expected output. Catches corrupt-weight / bad-quant failures that +# throughput benchmarks (which use --ignore-eos on random tokens) miss. +# +# Usage: PORT, MODEL, OUT(json) via env. Exits 0 if all coherent, 1 otherwise. +import json, os, sys, time, urllib.request + +PORT = os.environ.get("PORT", "8000") +MODEL = os.environ["MODEL"] +OUT = os.environ.get("OUT", "/results/sanity.json") +HERE = os.path.dirname(os.path.abspath(__file__)) +prompts = json.load(open(os.path.join(HERE, "prompts.json"))) + +def chat(prompt, max_tokens): + body = json.dumps({ + "model": MODEL, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, + "temperature": 0.0, + }).encode() + req = urllib.request.Request(f"http://0.0.0.0:{PORT}/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=120) as r: + d = json.load(r) + return d["choices"][0]["message"]["content"], time.time() - t0 + +results = [] +allpass = True +for p in prompts: + try: + text, dt = chat(p["prompt"], p["max_tokens"]) + low = text.lower() + hits = [k for k in p["expect_contains"] if k.lower() in low] + # coherent = non-empty, not pure repetition/garbage + coherent = len(text.strip()) > 0 and len(set(text.split())) > 1 or len(text.strip()) <= 8 + ok = len(hits) > 0 and coherent + results.append({"id": p["id"], "ok": ok, "hits": hits, + "expected": p["expect_contains"], "latency_s": round(dt, 2), + "output": text.strip()[:200]}) + if not ok: + allpass = False + except Exception as e: + results.append({"id": p["id"], "ok": False, "error": str(e)[:150]}) + allpass = False + +summary = {"model": MODEL, "all_pass": allpass, + "passed": sum(1 for r in results if r.get("ok")), + "total": len(results), "results": results} +os.makedirs(os.path.dirname(OUT), exist_ok=True) +json.dump(summary, open(OUT, "w"), indent=2) + +print(f"=== SANITY: {summary['passed']}/{summary['total']} prompts passed ===") +for r in results: + mark = "PASS" if r.get("ok") else "FAIL" + out = r.get("output", r.get("error", "")) + print(f" [{mark}] {r['id']:12s} -> {out[:90]!r}") +sys.exit(0 if allpass else 1) diff --git a/scripts/DeepSeekV4/lib/clean_node.sh b/scripts/DeepSeekV4/lib/clean_node.sh new file mode 100755 index 00000000..6d1aea5a --- /dev/null +++ b/scripts/DeepSeekV4/lib/clean_node.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# clean_node.sh — ensure a node's GPUs are in a CLEAN state before deploying an engine. +# Run via srun on the target node (exclusive). Kills leftover serving containers/processes +# that squat VRAM (the #1 cause of OOM-at-init), then verifies VRAM is actually free. +# +# Usage (inside srun on the node): clean_node.sh [max_used_gb_threshold] +# default threshold = 5 GB; exits non-zero if any GPU still above it after cleanup. +set -uo pipefail +THRESH_GB="${1:-5}" + +echo "[clean] $(hostname): removing ALL containers + stray GPU procs (reservation is exclusive to us)" +# The reservation is slurm-exclusive to us, but Docker containers are NOT slurm-tracked, +# so orphaned/stale containers from prior runs can linger and squat VRAM. Remove them ALL. +cids=$(docker ps -aq 2>/dev/null) +[ -n "$cids" ] && { docker kill $cids >/dev/null 2>&1; docker rm -f $cids >/dev/null 2>&1; } || true +# belt-and-suspenders: kill stray engine procs that may hold VRAM +pkill -9 -f "vllm|atom.entrypoints|sglang.launch|EngineCore|sglang.launch_server" 2>/dev/null || true +sleep 5 + +# 3) verify VRAM free on every GPU. +# rocm-smi prints one line per GPU: "GPU[N] : VRAM Total Used Memory (B): ". +# Take the LAST integer on each line (the byte count), not the GPU index. +echo "[clean] verifying VRAM free (threshold ${THRESH_GB}GB)" +busy=0; i=0 +while IFS= read -r line; do + b=$(echo "$line" | grep -oE "[0-9]+" | tail -1) + [ -z "$b" ] && continue + gb=$(( b / 1000000000 )) + if [ "$gb" -gt "$THRESH_GB" ]; then echo " GPU$i: ${gb}GB USED (>threshold)"; busy=1; else echo " GPU$i: ${gb}GB ok"; fi + i=$((i+1)) +done < <(rocm-smi --showmeminfo vram 2>/dev/null | grep -i "Used Memory") +if [ "$busy" -ne 0 ]; then + echo "[clean] WARNING: some GPUs still hold VRAM after cleanup — listing holders:" + rocm-smi --showpids 2>/dev/null | grep -iE "vllm|atom|sglang|EngineCore" | head + exit 1 +fi +echo "[clean] $(hostname): all GPUs clean" +exit 0 diff --git a/scripts/DeepSeekV4/lib/lib_inferencex.sh b/scripts/DeepSeekV4/lib/lib_inferencex.sh new file mode 100755 index 00000000..95e063a3 --- /dev/null +++ b/scripts/DeepSeekV4/lib/lib_inferencex.sh @@ -0,0 +1,1281 @@ +#!/usr/bin/env bash + +# Shared benchmarking utilities for InferenceX + +# Keep Python bytecode out of the mounted workspace. Benchmark jobs often run as +# root inside containers, and root-owned cache directories break future checkout +# cleanup on self-hosted runners. +export PYTHONDONTWRITEBYTECODE=1 +export PYTHONPYCACHEPREFIX="${PYTHONPYCACHEPREFIX:-/tmp/inferencex-pycache}" +mkdir -p "$PYTHONPYCACHEPREFIX" 2>/dev/null || true + +# Inference server port shared by every benchmark recipe. Launchers that need +# a non-default value (e.g. launch_mi355x-amds.sh derives PORT from RUNNER_NAME +# to avoid collisions across concurrent gh-runners on a shared host) set PORT +# themselves before sourcing this file; the `:-` fallback only kicks in when +# nothing upstream set it. +export PORT="${PORT:-8888}" + +# -------------------------------- +# GPU monitoring helpers +# -------------------------------- + +GPU_MONITOR_PID="" +GPU_METRICS_CSV="/workspace/gpu_metrics.csv" +export GPU_METRICS_CSV + +# Start background GPU monitoring that logs metrics every second to CSV. +# Auto-detects NVIDIA (nvidia-smi) or AMD (amd-smi) GPUs. +# Usage: start_gpu_monitor [--output /path/to/output.csv] [--interval 1] +start_gpu_monitor() { + local output="$GPU_METRICS_CSV" + local interval=1 + + while [[ $# -gt 0 ]]; do + case $1 in + --output) output="$2"; shift 2 ;; + --interval) interval="$2"; shift 2 ;; + *) shift ;; + esac + done + + GPU_METRICS_CSV="$output" + export GPU_METRICS_CSV + + if command -v nvidia-smi &>/dev/null; then + nvidia-smi --query-gpu=timestamp,index,power.draw,temperature.gpu,clocks.current.sm,clocks.current.memory,utilization.gpu,utilization.memory \ + --format=csv -l "$interval" > "$output" 2>/dev/null & + GPU_MONITOR_PID=$! + echo "[GPU Monitor] Started NVIDIA (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" + elif command -v amd-smi &>/dev/null; then + # Use amd-smi native watch mode (-w) which includes timestamps automatically. + # Pipe through awk to: skip preamble lines, keep first CSV header, skip repeated headers. + amd-smi metric -p -c -t -u -w "$interval" --csv 2>/dev/null \ + | awk '/^timestamp,/{if(!h){print;h=1};next} h{print}' > "$output" & + GPU_MONITOR_PID=$! + echo "[GPU Monitor] Started AMD (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" + else + echo "[GPU Monitor] No GPU monitoring tool found (nvidia-smi or amd-smi), skipping" + return 0 + fi +} + +# Stop the background GPU monitor and report file size. +stop_gpu_monitor() { + if [[ -n "$GPU_MONITOR_PID" ]] && kill -0 "$GPU_MONITOR_PID" 2>/dev/null; then + kill "$GPU_MONITOR_PID" 2>/dev/null + wait "$GPU_MONITOR_PID" 2>/dev/null || true + echo "[GPU Monitor] Stopped (PID=$GPU_MONITOR_PID)" + if [[ -f "$GPU_METRICS_CSV" ]]; then + local lines + lines=$(wc -l < "$GPU_METRICS_CSV") + echo "[GPU Monitor] Collected $lines rows -> $GPU_METRICS_CSV" + fi + fi + GPU_MONITOR_PID="" +} + +# Check if required environment variables are set +# Usage: check_env_vars VAR1 VAR2 VAR3 ... +# Exits with code 1 if any variable is not set +check_env_vars() { + local missing_vars=() + + for var_name in "$@"; do + if [[ -z "${!var_name:-}" ]]; then + missing_vars+=("$var_name") + fi + done + + if [[ ${#missing_vars[@]} -gt 0 ]]; then + echo "Error: The following required environment variables are not set:" + for var in "${missing_vars[@]}"; do + echo " - $var" + done + exit 1 + fi +} + +# Wait for server to be ready by polling the health endpoint +# All parameters are required +# Parameters: +# --port: Server port +# --server-log: Path to server log file +# --server-pid: Server process ID (required) +# --sleep-interval: Sleep interval between health checks (optional, default: 5) +wait_for_server_ready() { + set +x + local port="" + local server_log="" + local server_pid="" + local sleep_interval=5 + + # Parse arguments + while [[ $# -gt 0 ]]; do + case $1 in + --port) + port="$2" + shift 2 + ;; + --server-log) + server_log="$2" + shift 2 + ;; + --server-pid) + server_pid="$2" + shift 2 + ;; + --sleep-interval) + sleep_interval="$2" + shift 2 + ;; + *) + echo "Unknown parameter: $1" + return 1 + ;; + esac + done + + # Validate required parameters + if [[ -z "$port" ]]; then + echo "Error: --port is required" + return 1 + fi + if [[ -z "$server_log" ]]; then + echo "Error: --server-log is required" + return 1 + fi + if [[ -z "$server_pid" ]]; then + echo "Error: --server-pid is required" + return 1 + fi + + # Wait for server log file to be created (container startup may delay this) + while [ ! -f "$server_log" ]; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "Server died before creating log file. Exiting." + exit 1 + fi + sleep 1 + done + + # Show logs until server is ready + tail -f -n +1 "$server_log" & + local TAIL_PID=$! + until curl --output /dev/null --silent --fail http://0.0.0.0:$port/health; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "Server died before becoming healthy. Exiting." + kill $TAIL_PID + exit 1 + fi + sleep "$sleep_interval" + done + kill $TAIL_PID +} + +# Run benchmark serving with standardized parameters +# All parameters are required except --endpoint, --use-chat-template, --dsv4, and --trust-remote-code +# Parameters: +# --model: Model name +# --port: Server port +# --backend: Backend type - e.g., 'vllm' or 'openai' +# --endpoint: Optional API endpoint override +# --input-len: Random input sequence length +# --output-len: Random output sequence length +# --random-range-ratio: Random range ratio +# --num-prompts: Number of prompts +# --max-concurrency: Max concurrency +# --result-filename: Result filename without extension +# --result-dir: Result directory +# --use-chat-template: Optional flag to enable chat template +# --dsv4: Optional flag to use the DeepSeek-V4 chat template +# (encoding_dsv4.py) instead of the tokenizer's built-in jinja +# template. Implies --use-chat-template. +# --trust-remote-code: Optional flag to trust remote code from HuggingFace +# --server-pid: Optional server process ID to monitor during benchmark +run_benchmark_serving() { + # In eval-only mode, skip the throughput benchmark entirely. + if [ "${EVAL_ONLY}" = "true" ]; then + echo "EVAL_ONLY mode: skipping throughput benchmark" + return 0 + fi + + set +x + local model="" + local port="" + local backend="" + local endpoint="" + local input_len="" + local output_len="" + local random_range_ratio="" + local num_prompts="" + local max_concurrency="" + local result_filename="" + local result_dir="" + local workspace_dir="" + local use_chat_template=false + local dsv4=false + local trust_remote_code=false + local server_pid="" + local tokenizer="" + + while [[ $# -gt 0 ]]; do + case $1 in + --model) + model="$2" + shift 2 + ;; + --port) + port="$2" + shift 2 + ;; + --backend) + backend="$2" + shift 2 + ;; + --endpoint) + endpoint="$2" + shift 2 + ;; + --input-len) + input_len="$2" + shift 2 + ;; + --output-len) + output_len="$2" + shift 2 + ;; + --random-range-ratio) + random_range_ratio="$2" + shift 2 + ;; + --num-prompts) + num_prompts="$2" + shift 2 + ;; + --max-concurrency) + max_concurrency="$2" + shift 2 + ;; + --result-filename) + result_filename="$2" + shift 2 + ;; + --result-dir) + result_dir="$2" + shift 2 + ;; + --bench-serving-dir) + workspace_dir="$2" + shift 2 + ;; + --use-chat-template) + use_chat_template=true + shift + ;; + --dsv4) + dsv4=true + use_chat_template=true + shift + ;; + --trust-remote-code) + trust_remote_code=true + shift + ;; + --server-pid) + server_pid="$2" + shift 2 + ;; + --tokenizer) + tokenizer="$2" + shift 2 + ;; + *) + echo "Unknown parameter: $1" + return 1 + ;; + esac + done + + # Validate all required parameters + if [[ -z "$model" ]]; then + echo "Error: --model is required" + return 1 + fi + if [[ -z "$port" ]]; then + echo "Error: --port is required" + return 1 + fi + if [[ -z "$backend" ]]; then + echo "Error: --backend is required" + return 1 + fi + if [[ -z "$input_len" ]]; then + echo "Error: --input-len is required" + return 1 + fi + if [[ -z "$output_len" ]]; then + echo "Error: --output-len is required" + return 1 + fi + if [[ -z "$random_range_ratio" ]]; then + echo "Error: --random-range-ratio is required" + return 1 + fi + if [[ -z "$num_prompts" ]]; then + echo "Error: --num-prompts is required" + return 1 + fi + if [[ -z "$max_concurrency" ]]; then + echo "Error: --max-concurrency is required" + return 1 + fi + if [[ -z "$result_filename" ]]; then + echo "Error: --result-filename is required" + return 1 + fi + if [[ -z "$result_dir" ]]; then + echo "Error: --result-dir is required" + return 1 + fi + + if [[ -z "$workspace_dir" ]]; then + workspace_dir=$(pwd) + fi + + # Profiling support: when PROFILE=1, ensure profiler dir exists, add --profile flag, + # and cap num_prompts to keep traces small. + local profile_flag=() + if [[ "${PROFILE:-}" == "1" ]]; then + local _prof_dir="${SGLANG_TORCH_PROFILER_DIR:-${VLLM_TORCH_PROFILER_DIR:-}}" + if [[ -n "$_prof_dir" ]]; then + mkdir -p "$_prof_dir" + fi + profile_flag+=(--profile) + num_prompts="$max_concurrency" + fi + + # Build benchmark command + local benchmark_cmd=( + python3 "$workspace_dir/utils/bench_serving/benchmark_serving.py" + --model "$model" + --backend "$backend" + --base-url "http://0.0.0.0:$port" + --dataset-name random + --random-input-len "$input_len" + --random-output-len "$output_len" + --random-range-ratio "$random_range_ratio" + --num-prompts "$num_prompts" + --max-concurrency "$max_concurrency" + --request-rate inf + --ignore-eos + "${profile_flag[@]}" + --save-result + --num-warmups "$((2 * max_concurrency))" \ + --percentile-metrics 'ttft,tpot,itl,e2el' + --result-dir "$result_dir" + --result-filename "$result_filename.json" + ) + + if [[ -n "$endpoint" ]]; then + benchmark_cmd+=(--endpoint "$endpoint") + fi + + # Add --use-chat-template if requested + if [[ "$use_chat_template" == true ]]; then + benchmark_cmd+=(--use-chat-template) + fi + + # Add --dsv4 if requested (requires --use-chat-template, which we + # auto-enable when --dsv4 is passed in). + if [[ "$dsv4" == true ]]; then + benchmark_cmd+=(--dsv4) + fi + + # Add --trust-remote-code if requested + if [[ "$trust_remote_code" == true ]]; then + benchmark_cmd+=(--trust-remote-code) + fi + + if [[ -n "$tokenizer" ]]; then + benchmark_cmd+=(--tokenizer "$tokenizer") + fi + + # Run benchmark with optional server monitoring + set -x + if [[ -n "$server_pid" ]]; then + # Run benchmark in background and monitor server health + "${benchmark_cmd[@]}" & + local benchmark_pid=$! + + # Monitor loop: check both benchmark and server status + while kill -0 "$benchmark_pid" 2>/dev/null; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "ERROR: Server process $server_pid died during benchmark" + kill "$benchmark_pid" 2>/dev/null + wait "$benchmark_pid" 2>/dev/null + set +x + return 1 + fi + sleep 2 + done + + # Benchmark finished, get its exit code + wait "$benchmark_pid" + local benchmark_exit_code=$? + else + # No server monitoring, run benchmark directly + "${benchmark_cmd[@]}" + local benchmark_exit_code=$? + fi + set +x + + # If profiling, move trace to relay-upload location + if [[ "${PROFILE:-}" == "1" ]]; then + move_profile_trace_for_relay + fi + + return $benchmark_exit_code +} + + +# -------------------------------- +# Profiling trace helpers +# -------------------------------- + +_find_latest_profile_trace() { + local latest="" + local dir="" candidate="" base="" + local -a search_roots=() + + for dir in "$@"; do + search_roots=() + if [[ -d "$dir" ]]; then + search_roots+=("$dir") + fi + if [[ -d "$dir/profiles" ]]; then + search_roots+=("$dir/profiles") + fi + if [[ ${#search_roots[@]} -eq 0 ]]; then + continue + fi + + while IFS= read -r -d '' candidate; do + base="$(basename "$candidate")" + if [[ "$base" == profile_*.trace.json.gz ]]; then + continue + fi + if [[ -z "$latest" || "$candidate" -nt "$latest" ]]; then + latest="$candidate" + fi + done < <( + find "${search_roots[@]}" -maxdepth 1 -type f \ + \( -name "*.trace.json" -o -name "*.trace.json.gz" -o -name "*trace*.json" -o -name "*trace*.json.gz" -o -name "*profile*.json" -o -name "*profile*.json.gz" \) \ + -print0 2>/dev/null + ) + done + + printf '%s' "$latest" +} + +# Move profiler trace into a stable workspace path for workflow relay/upload. +move_profile_trace_for_relay() { + if [[ "${PROFILE:-}" != "1" ]]; then + return 0 + fi + + if [[ -z "${RESULT_FILENAME:-}" ]]; then + echo "[PROFILE] RESULT_FILENAME is not set; skipping relay trace staging." >&2 + return 0 + fi + + local sglang_dir="${SGLANG_TORCH_PROFILER_DIR:-/workspace}" + local vllm_dir="${VLLM_TORCH_PROFILER_DIR:-/workspace}" + local -a search_dirs=() + local dir="" existing="" + local seen=0 + + for dir in "$sglang_dir" "$vllm_dir" "/workspace"; do + if [[ -z "$dir" ]]; then + continue + fi + seen=0 + for existing in "${search_dirs[@]}"; do + if [[ "$existing" == "$dir" ]]; then + seen=1 + break + fi + done + if [[ "$seen" -eq 0 ]]; then + search_dirs+=("$dir") + fi + done + + local trace_file="" + local wait_attempts=10 + for (( i=1; i<=wait_attempts; i++ )); do + trace_file="$(_find_latest_profile_trace "${search_dirs[@]}")" + if [[ -n "$trace_file" ]]; then + break + fi + sleep 10 + done + + if [[ -z "$trace_file" ]]; then + echo "[PROFILE] No trace found for relay under: ${search_dirs[*]}" >&2 + return 0 + fi + + local dest_trace="/workspace/profile_${RESULT_FILENAME}.trace.json.gz" + if [[ "$trace_file" == *.gz ]]; then + cp -f "$trace_file" "$dest_trace" + else + gzip -c "$trace_file" > "$dest_trace" + fi + + echo "[PROFILE] Relay trace prepared: $dest_trace (source: $trace_file)" +} + + +# ------------------------------ +# Eval (lm-eval-harness) helpers +# ------------------------------ + +_install_lm_eval_deps() { + # torchvision causes circular imports in ATOM; TRT-LLM/SGLang need it at module level. + if [[ "${IMAGE:-}" == *atom* ]]; then + python3 -m pip uninstall -y torchvision 2>/dev/null || true + fi + python3 -m pip install -q --no-cache-dir --break-system-packages "lm-eval[api]" || true + local lm_eval_ref="b315ef3b05176acc9732bb7fdec116abe1ecc476" + if command -v git >/dev/null 2>&1; then + if ! python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "git+https://github.com/EleutherAI/lm-evaluation-harness.git@${lm_eval_ref}"; then + python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true + fi + else + python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true + fi +} + +# Patch lm-eval filters to be robust to empty strings via sitecustomize +_patch_lm_eval() { + local patch_dir + patch_dir="$(mktemp -d)" + cat > "$patch_dir/sitecustomize.py" <<'PY' +# --- Patch LocalChatCompletion.parse_generations to handle empty content with reasoning_content --- +import re, sys, unicodedata, json +from lm_eval.filters import extraction as ex +from lm_eval.models.openai_completions import LocalChatCompletion as _LCC + +def _le_parse_generations(outputs, **kwargs): + res = [] + if not isinstance(outputs, list): + outputs = [outputs] + for out in (outputs or []): + try: + choices = out.get("choices", []) + tmp = ["" for _ in choices] + for choice in choices: + idx = choice.get("index", 0) + msg = (choice.get("message") or {}) + content = msg.get("content") + if content in (None, "", []): + content = msg.get("reasoning_content") or "" + tmp[idx] = content + except Exception: + tmp = [""] + res.extend(tmp) + return res + +# Keep staticmethod semantics +_LCC.parse_generations = staticmethod(_le_parse_generations) + +# --- Patch TemplateAPI.apply_chat_template to avoid injecting "type": "text" for TRT --- +try: + from lm_eval.models import api_models as _api_models + _TemplateAPI = _api_models.TemplateAPI + _JsonChatStr = _api_models.JsonChatStr +except Exception: + _TemplateAPI = None + _JsonChatStr = None + +if _TemplateAPI is not None and _JsonChatStr is not None: + _orig_apply_chat_template = _TemplateAPI.apply_chat_template + + def _patched_apply_chat_template( + self, + chat_history, + add_generation_prompt: bool = True, + ): + """Applies a chat template to a list of chat history between user and model.""" + if self.tokenizer_backend == "huggingface" and self.tokenized_requests: + return self.tokenizer.apply_chat_template( + chat_history, + tokenize=False, + add_generation_prompt=add_generation_prompt, + continue_final_message=not add_generation_prompt, + ) + elif self.tokenizer_backend == "remote" and self.tokenized_requests: + return chat_history + else: + # NOTE: we no longer inject `"type": "text"` when tokenizer is None / non-HF + return _JsonChatStr( + json.dumps( + [{**item} for item in chat_history], + ensure_ascii=False, + ) + ) + + _TemplateAPI.apply_chat_template = _patched_apply_chat_template +PY + export PYTHONPATH="${patch_dir}:${PYTHONPATH:-}" +} + +get_native_max_context_length() { + local model_path="$1" + # Prefer MODEL_PATH (local model directory) when available, since the + # argument may be a served-model name that is neither a valid HF repo + # ID nor a local path (e.g. "deepseek-r1-fp4" on the B300 cluster). + if [ -n "${MODEL_PATH:-}" ] && [ -d "${MODEL_PATH}" ]; then + model_path="${MODEL_PATH}" + fi + python3 -c " +try: + from transformers import AutoConfig + config = AutoConfig.from_pretrained('${model_path}', trust_remote_code=True) + for attr in ['max_position_embeddings', 'max_sequence_length', 'seq_length', 'n_positions']: + if hasattr(config, attr): + print(getattr(config, attr)) + break + else: + print(0) +except Exception: + print(0) +" +} + +# Compute the context length for eval-only mode. +# Uses the requested benchmark context capped at the model's native max. +# Sets EVAL_MAX_MODEL_LEN (needed by run_lm_eval). +# Echoes the computed value for scripts to capture. +# +# Usage: local ctx=$(compute_eval_context_length "$MODEL" "${current_ctx}") +compute_eval_context_length() { + local model="$1" + local benchmark_ctx="${2:-0}" + local native_max + native_max=$(get_native_max_context_length "$model") + native_max="${native_max:-0}" + + if [ "$benchmark_ctx" -eq 0 ] 2>/dev/null; then + benchmark_ctx="${native_max:-0}" + fi + local eval_ctx=$(( benchmark_ctx * 1 )) + if [ "$native_max" -gt 0 ] 2>/dev/null && [ "$eval_ctx" -gt "$native_max" ]; then + eval_ctx="$native_max" + fi + # If eval_ctx is still 0 (both benchmark_ctx and native_max were 0), fall back + if [ "$eval_ctx" -le 0 ] 2>/dev/null; then + echo "WARN: compute_eval_context_length could not determine context length for $model" >&2 + eval_ctx="${MAX_MODEL_LEN:-16384}" + fi + EVAL_MAX_MODEL_LEN="$eval_ctx" + echo "$eval_ctx" +} + +# Convenience wrapper: compute eval context from ISL/OSL and export EVAL_MAX_MODEL_LEN. +# Call directly (not in a subshell) so the export persists. +# Scripts then wire $EVAL_MAX_MODEL_LEN into whichever server variable they need. +setup_eval_context() { + EVAL_MAX_MODEL_LEN=$(compute_eval_context_length "$MODEL" "$((ISL + OSL + 256))") + export EVAL_MAX_MODEL_LEN +} + +run_lm_eval() { + local port="${PORT:-8888}" + local tasks_dir="${EVAL_TASKS_DIR:-utils/evals/gsm8k.yaml}" + local results_dir="${EVAL_RESULT_DIR:-$(mktemp -d /tmp/eval_out-XXXXXX)}" + local eval_context_len="${EVAL_MAX_MODEL_LEN:-16384}" + local temperature=0 + local top_p=1 + local concurrent_requests="${EVAL_CONCURRENT_REQUESTS:-${CONC:-64}}" + + while [[ $# -gt 0 ]]; do + case $1 in + --port) port="$2"; shift 2 ;; + --task) tasks_dir="$2"; shift 2 ;; + --results-dir) results_dir="$2"; shift 2 ;; + --gen-max-tokens) eval_context_len="$2"; shift 2 ;; + --temperature) temperature="$2"; shift 2 ;; + --top-p) top_p="$2"; shift 2 ;; + *) echo "Unknown parameter: $1"; return 1 ;; + esac + done + + if [ "${INFERENCEX_LM_EVAL_RUNTIME_READY:-false}" != "true" ]; then + _install_lm_eval_deps + _patch_lm_eval + export INFERENCEX_LM_EVAL_RUNTIME_READY=true + fi + + local openai_server_base="http://0.0.0.0:${port}" + local openai_chat_base="${openai_server_base}/v1/chat/completions" + export OPENAI_API_KEY=${OPENAI_API_KEY:-EMPTY} + MODEL_NAME=${MODEL_NAME:-$MODEL} # Prefer MODEL_NAME, else MODEL + + # Cap output tokens: must fit within context window (leave room for input), + # and avoid excessive KV cache reservation per request on TRT. + local max_output_tokens=$(( eval_context_len > 4096 ? eval_context_len - 4096 : eval_context_len / 2 )) + if [ "$max_output_tokens" -gt 16384 ]; then + max_output_tokens=16384 + fi + echo "Eval budget: eval_context_len=${eval_context_len}, max_output_tokens=${max_output_tokens}" + + # Export for append_lm_eval_summary to pick up + export EVAL_RESULT_DIR="$results_dir" + set -x + python3 -m lm_eval --model local-chat-completions --apply_chat_template \ + --tasks "${tasks_dir}" \ + --output_path "${results_dir}" \ + --log_samples \ + --model_args "model=${MODEL_NAME},base_url=${openai_chat_base},api_key=${OPENAI_API_KEY},eos_string=,max_retries=5,num_concurrent=${concurrent_requests},timeout=1800,tokenized_requests=False,max_length=${eval_context_len}" \ + --gen_kwargs "max_tokens=${max_output_tokens},temperature=${temperature},top_p=${top_p}" + local eval_exit=$? + set +x + return $eval_exit +} + +_stage_lm_eval_artifacts() { + local results_dir="$1" + local eval_conc="$2" + local moved=0 + local failed=0 + local jf base stem extension target suffix + + if [ ! -d "${results_dir}" ]; then + echo "WARN: eval result directory '${results_dir}' does not exist" >&2 + return 1 + fi + + while IFS= read -r -d '' jf; do + base=$(basename "$jf") + case "$base" in + meta_env.json) + continue + ;; + *.jsonl) + stem="${base%.jsonl}" + extension=".jsonl" + ;; + *.json) + stem="${base%.json}" + extension=".json" + ;; + *) + continue + ;; + esac + + target="./${stem}_conc${eval_conc}${extension}" + suffix=2 + while [ -e "$target" ]; do + target="./${stem}_conc${eval_conc}_${suffix}${extension}" + suffix=$((suffix + 1)) + done + + if mv -f "$jf" "$target"; then + moved=1 + else + echo "WARN: failed to stage eval artifact ${jf}" >&2 + failed=1 + fi + done < <( + find "${results_dir}" -type f \ + \( -name "*.json" -o -name "*.jsonl" \) -print0 2>/dev/null + ) + + rm -rf --one-file-system "${results_dir}" 2>/dev/null \ + || rm -rf "${results_dir}" \ + || true + + if [ "$moved" -eq 0 ]; then + echo "WARN: no eval artifacts were produced for concurrency ${eval_conc}" >&2 + return 1 + fi + return "$failed" +} + +_eval_concs_to_json() { + local values="$1" + local value + local joined="" + + for value in $values; do + if [[ ! "$value" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: invalid eval concurrency '${value}'" >&2 + return 1 + fi + if [ -n "$joined" ]; then + joined="${joined}, " + fi + joined="${joined}${value}" + done + + printf '[%s]' "$joined" +} + +append_lm_eval_summary() { + local batch_concs="${EVAL_BATCHED_CONCS:-}" + local results_dir="${EVAL_RESULT_DIR:-}" + local out_dir="${results_dir}" + local meta_json + local metadata_conc="${CONC:-1}" + local batch_metadata="" + + if [ -n "$batch_concs" ]; then + meta_json="./meta_env.json" + metadata_conc="${batch_concs%% *}" + + local eval_concs_json completed_concs_json failed_concs_json + eval_concs_json=$(_eval_concs_to_json "$batch_concs") || return 1 + completed_concs_json=$( + _eval_concs_to_json "${EVAL_BATCHED_COMPLETED_CONCS:-}" + ) || return 1 + failed_concs_json=$( + _eval_concs_to_json "${EVAL_BATCHED_FAILED_CONCS:-}" + ) || return 1 + printf -v batch_metadata \ + ' "eval_concs": %s,\n "completed_eval_concs": %s,\n "failed_eval_concs": %s,\n' \ + "$eval_concs_json" \ + "$completed_concs_json" \ + "$failed_concs_json" + else + if [ -z "${results_dir}" ]; then + echo "WARN: EVAL_RESULT_DIR is empty; skipping artifact collection" >&2 + return 1 + fi + if [ ! -d "${out_dir}" ]; then + echo "WARN: EVAL_RESULT_DIR='${out_dir}' does not exist; skipping artifact collection" >&2 + return 1 + fi + meta_json="${out_dir}/meta_env.json" + fi + + # Write minimal meta for collectors that expect it + local model_name="${MODEL_NAME:-$MODEL}" + local is_multinode_json="false" + if [ "${IS_MULTINODE:-false}" = "true" ]; then + is_multinode_json="true" + fi + + local prefill_tp="${PREFILL_TP:-${TP:-1}}" + local prefill_ep="${PREFILL_EP:-${EP_SIZE:-1}}" + local prefill_num_workers="${PREFILL_NUM_WORKERS:-1}" + local decode_tp="${DECODE_TP:-${TP:-1}}" + local decode_ep="${DECODE_EP:-${EP_SIZE:-1}}" + local decode_num_workers="${DECODE_NUM_WORKERS:-1}" + + local dp_json="false" + if [ "${DP_ATTENTION:-false}" = "true" ]; then dp_json="true"; fi + local prefill_dp_json="$dp_json" + if [ "${PREFILL_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then + prefill_dp_json="true" + else + prefill_dp_json="false" + fi + local decode_dp_json="$dp_json" + if [ "${DECODE_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then + decode_dp_json="true" + else + decode_dp_json="false" + fi + + # Derive framework/precision from env, fallback to parsing RESULT_FILENAME + # RESULT_FILENAME format (from workflow): + # ___tp<...>_ep<...>_dpa_<...>_conc<...>_ + local fw="${FRAMEWORK:-}" + local prec="${PRECISION:-}" + if [[ -z "$fw" || -z "$prec" ]]; then + if [[ -n "${RESULT_FILENAME}" ]]; then + # Extract the two fields immediately before "_tp" + # Handles arbitrary underscores in exp_name by matching from the end + local parsed + parsed=$(echo "${RESULT_FILENAME}" | sed -n 's/.*_\([^_][^_]*\)_\([^_][^_]*\)_tp.*/\1 \2/p') + local p1="${parsed%% *}" + local p2="${parsed#* }" + if [[ -z "$prec" && -n "$p1" && "$p1" != "$parsed" ]]; then + prec="$p1" + fi + if [[ -z "$fw" && -n "$p2" && "$p2" != "$parsed" ]]; then + fw="$p2" + fi + fi + fi + cat > "${meta_json}" <&2 + fi + if [ -d "${out_dir}" ]; then + while IFS= read -r -d '' jf; do + base=$(basename "$jf") + if [ "$base" != "meta_env.json" ]; then + mv -f "$jf" ./ || echo "WARN: failed to move ${jf}" >&2 + fi + done < <(find "${out_dir}" -type f -name "*.json*" -print0 2>/dev/null) + fi + + # Best-effort cleanup of the temp directory + if [ -n "${out_dir}" ] && [ -d "${out_dir}" ]; then + rm -rf --one-file-system "${out_dir}" || rm -rf "${out_dir}" || true + fi + + echo "Moved eval artifacts to: $(pwd)" +} + +# ------------------------------ +# Unified eval entrypoint +# ------------------------------ + +run_eval() { + local framework="${EVAL_FRAMEWORK:-lm-eval}" + local forwarded=() + + while [[ $# -gt 0 ]]; do + case "$1" in + --framework) framework="$2"; shift 2 ;; + *) forwarded+=("$1"); shift ;; + esac + done + + # Compute EVAL_MAX_MODEL_LEN if not already set by the calling script + if [ -z "${EVAL_MAX_MODEL_LEN:-}" ]; then + compute_eval_context_length "$MODEL" "${MAX_MODEL_LEN:-0}" > /dev/null + fi + + unset EVAL_BATCHED_CONCS + unset EVAL_BATCHED_COMPLETED_CONCS + unset EVAL_BATCHED_FAILED_CONCS + + local requested_concs="${EVAL_CONCURRENT_REQUESTS:-}" + local eval_concs=() + if [ -n "$requested_concs" ]; then + read -r -a eval_concs <<< "$requested_concs" + fi + + if [ "${#eval_concs[@]}" -gt 1 ]; then + if [[ "$framework" != "lm-eval" && "$framework" != "lm_eval" ]]; then + echo "ERROR: batched eval concurrency is only supported for lm-eval" >&2 + return 1 + fi + + local eval_conc results_dir eval_rc stage_rc + local completed_concs=() + local failed_concs=() + + for eval_conc in "${eval_concs[@]}"; do + if [[ ! "$eval_conc" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: invalid eval concurrency '${eval_conc}'" >&2 + return 1 + fi + + if ! results_dir=$(mktemp -d /tmp/eval_out-conc"${eval_conc}"-XXXXXX); then + echo "ERROR: failed to create eval output directory for concurrency ${eval_conc}" >&2 + failed_concs+=("$eval_conc") + continue + fi + + echo "Running lm-eval at concurrency ${eval_conc} using the existing engine" + export EVAL_CONCURRENT_REQUESTS="$eval_conc" + export CONC="$eval_conc" + eval_rc=0 + stage_rc=0 + run_lm_eval "${forwarded[@]}" --results-dir "$results_dir" \ + || eval_rc=$? + _stage_lm_eval_artifacts "$results_dir" "$eval_conc" \ + || stage_rc=$? + + if [ "$eval_rc" -eq 0 ] && [ "$stage_rc" -eq 0 ]; then + completed_concs+=("$eval_conc") + else + echo "ERROR: lm-eval failed at concurrency ${eval_conc} (eval_rc=${eval_rc}, stage_rc=${stage_rc})" >&2 + failed_concs+=("$eval_conc") + fi + done + + export EVAL_CONCURRENT_REQUESTS="$requested_concs" + export EVAL_RESULT_DIR="" + export EVAL_BATCHED_CONCS="${eval_concs[*]}" + export EVAL_BATCHED_COMPLETED_CONCS="${completed_concs[*]}" + export EVAL_BATCHED_FAILED_CONCS="${failed_concs[*]}" + + if [ "${#failed_concs[@]}" -gt 0 ]; then + echo "ERROR: batched eval failed for concurrency: ${failed_concs[*]}" >&2 + echo "Deferring failure until post-upload score validation preserves all artifacts" >&2 + fi + return 0 + fi + + local eval_rc=0 + case "$framework" in + lm-eval|lm_eval) run_lm_eval "${forwarded[@]}" || eval_rc=$? ;; + *) echo "Unknown framework '${framework}'"; eval_rc=1 ;; + esac + + if [ "$eval_rc" -ne 0 ]; then + echo "ERROR: run_eval failed with exit code $eval_rc" >&2 + if [ "${EVAL_ONLY}" = "true" ]; then + echo "Eval-only mode: failing after artifact collection" >&2 + return "$eval_rc" + fi + fi + return $eval_rc +} + + +# -------------------------------- +# Agentic trace replay helpers (aiperf driver) +# -------------------------------- + +INFMAX_CONTAINER_WORKSPACE="${INFMAX_CONTAINER_WORKSPACE:-/workspace}" +AGENTIC_DIR="${AGENTIC_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/agentic-benchmark}" +AIPERF_DIR="${AIPERF_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/aiperf}" + +agentic_pip_install() { + local pip_install=(python3 -m pip install) + if python3 -m pip install --help 2>/dev/null | grep -q -- "--break-system-packages"; then + pip_install+=(--break-system-packages) + fi + + "${pip_install[@]}" "$@" +} + +ensure_hf_cli() { + if command -v hf >/dev/null 2>&1; then + return 0 + fi + + # Some lean runtime images used by multinode SGLang include Python but not + # the Hugging Face CLI. Install just the hub CLI before prefetching traces. + agentic_pip_install --quiet "huggingface_hub[cli]>=0.25.0" +} + +resolve_trace_source() { + # Per-recipe override: set WEKA_LOADER_OVERRIDE to one of the aiperf + # public-dataset loader names allowed by the inferencex-agentx-mvp + # scenario. Used by recipes whose servers have non-default context + # caps (e.g. minimaxm2.5 at max_model_len ~256k can't replay the + # unfiltered 052726 corpus and switches to the 256k-capped variant). + local loader="${WEKA_LOADER_OVERRIDE:-semianalysis_cc_traces_weka_with_subagents}" + local dataset + case "$loader" in + semianalysis_cc_traces_weka_with_subagents) + dataset="semianalysisai/cc-traces-weka-with-subagents-052726" + ;; + semianalysis_cc_traces_weka_with_subagents_256k) + dataset="semianalysisai/cc-traces-weka-with-subagents-052726-256k" + ;; + *) + echo "Error: unknown WEKA_LOADER_OVERRIDE='$loader'. Allowed: semianalysis_cc_traces_weka_with_subagents, semianalysis_cc_traces_weka_with_subagents_256k" >&2 + exit 1 + ;; + esac + TRACE_SOURCE_FLAG="--public-dataset $loader" + echo "Loading traces via aiperf public-dataset: $loader ($dataset)" + # Pre-download the dataset into the shared HF_HUB_CACHE (same mount used + # for model weights) so subsequent runs read from cache instead of + # re-downloading every job. + ensure_hf_cli + hf download --repo-type dataset "$dataset" +} + +install_agentic_deps() { + # vllm/vllm-openai container ships without git. pip needs git to + # introspect the aiperf source tree on install. Install on demand; + # no-op when git is already present (e.g. AMD images that ship it). + if ! command -v git >/dev/null 2>&1; then + apt-get update && apt-get install -y git + fi + agentic_pip_install --quiet urllib3 requests 2>/dev/null || true + agentic_pip_install -q -r "$AGENTIC_DIR/requirements.txt" + # Editable install of aiperf from the submodule — gives us the + # `aiperf` CLI plus the inferencex-agentx-mvp scenario plugin. + # + # `--ignore-installed` sidesteps the distutils-uninstall error that + # vLLM containers hit on apt-managed system packages (blinker, etc.) + # when pip's resolver tries to upgrade one of aiperf's transitive + # deps. Installing fresh into the user/site location is safe — the + # system package stays in place and pip's import order picks up our + # newer copy first. + agentic_pip_install -q --ignore-installed -e "$AIPERF_DIR" + # Force-upgrade datasets: containers often ship an older version without + # the `Json` feature type used by the HF traces dataset. `Json` was added + # in datasets 4.7.0 (March 2025). Unpinned installs won't upgrade an + # already-present package. + agentic_pip_install --upgrade "datasets>=4.7.0" +} + +build_replay_cmd() { + # aiperf invocation for the inferencex-agentx-mvp scenario. + # + # Pre-canned assistant replay is the default: recorded assistant responses + # are used for future prompt construction, and live server responses are + # discarded. Set AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES=1 explicitly + # to use live-assistant mode, where the loader emits user-only deltas and + # the worker threads the server's live assistant response back into the + # session. + # + # The scenario plugin locks: --cache-bust first_turn_prefix and + # --trace-idle-gap-cap-seconds 60 (per-trace idle-gap compression + # against parent + subagent request-start timestamps; supersedes the + # legacy --use-think-time-only / --inter-turn-delay-cap-seconds path), + # and auto-injects them — so we do not pass them. See + # utils/aiperf/docs/tutorials/agentx-mvp.md. + local result_dir="$1" + local duration="$DURATION" + + export AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES="${AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES:-0}" + # Dataset configuration (load + reconstruct + inputs.json + mmap) + # routinely takes 4-5 min for the Weka corpus on fast /tmp + # (B300) but can stretch to 14 min on slower /tmp + parallel contention + # (observed on H200 where all 14 R3 jobs hit aiperf's 900s Configure + # Profiling timeout simultaneously). Bump to 1800s to absorb 3x + # worst-case slowdown — the post-setup measurement window is unaffected. + export AIPERF_DATASET_CONFIGURATION_TIMEOUT=1800 + # aiperf validates that SERVICE_PROFILE_CONFIGURE_TIMEOUT >= + # DATASET_CONFIGURATION_TIMEOUT at startup. Bump it in lockstep. + export AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT=1800 + REPLAY_CMD="aiperf profile --scenario inferencex-agentx-mvp" + REPLAY_CMD+=" --url http://localhost:$PORT" + REPLAY_CMD+=" --endpoint /v1/chat/completions" + REPLAY_CMD+=" --endpoint-type chat" + REPLAY_CMD+=" --streaming" + REPLAY_CMD+=" --model $MODEL" + REPLAY_CMD+=" --concurrency $CONC" + REPLAY_CMD+=" --benchmark-duration $duration" + REPLAY_CMD+=" --random-seed 42" + # Fail runs once more than 10% of requests error. This keeps known + # transient low-rate failures from killing long sweeps while still + # catching malformed payloads or server crashes before they get aggregated + # as benchmarkable data. + REPLAY_CMD+=" --failed-request-threshold 0.10" + # Sample each trajectory's warmup start position uniformly from + # [25%, 75%] of the trace's turn count (was hardcoded 0%-70% upstream). + # Avoids starting trajectories right at turn 0 where the KV cache is + # cold and skews early steady-state samples. + REPLAY_CMD+=" --trajectory-start-min-ratio 0.25" + REPLAY_CMD+=" --trajectory-start-max-ratio 0.75" + # Use server-reported usage fields (prompt_tokens / completion_tokens) for + # ISL/OSL instead of client-side tokenizer.encode(). Auto-enables + # stream_options.include_usage on the OpenAI chat endpoint. Skips the + # heavy per-record tokenization in the records pipeline that was pinning + # CPU on minimax-m2.5 at high concurrency. Lossless for vLLM (server + # usage is authoritative). + REPLAY_CMD+=" --use-server-token-count" + # aiperf's dataset manager (separate from the inference parser) loads + # the model's tokenizer for trace-prompt tokenization regardless of + # --use-server-token-count. Models like kimi (amd/Kimi-K2.5-MXFP4, + # moonshotai/Kimi-K2.5) ship a custom tokenizer in their HF repo and + # need trust_remote_code=True to load. Benign for models without + # custom tokenizer code, so we set it unconditionally. + REPLAY_CMD+=" --tokenizer-trust-remote-code" + # Keep replay inputs inside the same context window used to launch the + # server. The WEKA corpus contains a few very long parent/subagent traces; + # if we mmap and replay them against a smaller-context server they become + # deterministic 4xxs and can still pressure the engine while queued. + if [ -n "${MAX_MODEL_LEN:-}" ] && [ "$MAX_MODEL_LEN" != "0" ]; then + REPLAY_CMD+=" --max-context-length $MAX_MODEL_LEN" + fi + # Default --num-dataset-entries is 100; the with-subagents Weka corpus + # has 472. Cap at 472 so all unique traces are loaded (the loader treats + # this as a ``min(cap, available)`` ceiling, not a target — see + # semianalysis_cc_traces_weka.py). + REPLAY_CMD+=" --num-dataset-entries 472" + # 1-second timeslices on the server-metrics scrape so the post-run + # plotter has per-window time series (KV usage, cache hit rate, + # throughput, etc.). Matches kv-cache-tester's poll_interval=1.0 + # snapshot cadence so metrics_plots.png is visually comparable. + # Without this, aiperf only emits aggregate stats and the 6x2 panels + # collapse to flat lines. + REPLAY_CMD+=" --slice-duration 1.0" + REPLAY_CMD+=" --output-artifact-dir $result_dir/aiperf_artifacts" + # The inferencex-agentx-mvp scenario enforces a 900s minimum + # benchmark duration. For smoke tests with shorter durations, opt + # into --unsafe-override (the run's submission_valid will be flagged + # false; that's expected for non-canonical runs). + if [ "$duration" -lt 900 ] || [ "${AIPERF_UNSAFE_OVERRIDE:-false}" = "true" ]; then + REPLAY_CMD+=" --unsafe-override" + fi + REPLAY_CMD+=" $TRACE_SOURCE_FLAG" +} + +write_agentic_result_json() { + # Aggregate aiperf's profile_export.{json,jsonl} + server_metrics_export.json + # into $AGENTIC_OUTPUT_DIR/$RESULT_FILENAME.json. The workflow's existing + # retry-based existence check is the single success gate. + local result_dir="$1" + RESULT_DIR="$result_dir" AGENTIC_OUTPUT_DIR="${AGENTIC_OUTPUT_DIR:-$INFMAX_CONTAINER_WORKSPACE}" \ + python3 "$INFMAX_CONTAINER_WORKSPACE/utils/process_agentic_result.py" + + # Generate metrics_plots.png from the same aiperf artifacts. Best-effort: + # don't fail the launcher if plot generation has trouble (e.g. matplotlib + # missing in a stripped-down image). The agg JSON is the success gate. + python3 "$INFMAX_CONTAINER_WORKSPACE/utils/generate_aiperf_plots.py" "$result_dir" 2>&1 || true +} + +run_agentic_replay_and_write_outputs() { + local result_dir="$1" + local replay_rc + + echo "$REPLAY_CMD" > "$result_dir/benchmark_command.txt" + + set +e + set -x + $REPLAY_CMD 2>&1 | tee "$result_dir/benchmark.log" + replay_rc=${PIPESTATUS[0]} + set +x + set -e + + write_agentic_result_json "$result_dir" + + python3 "$AGENTIC_DIR/scripts/analyze_benchmark_distributions.py" \ + "$result_dir/aiperf_artifacts" -o "$result_dir" 2>&1 || true + + if [ "$replay_rc" -ne 0 ]; then + echo "ERROR: agentic trace replay exited with code $replay_rc after writing available results" >&2 + return "$replay_rc" + fi +} diff --git a/scripts/DeepSeekV4/lib/prompts.json b/scripts/DeepSeekV4/lib/prompts.json new file mode 100644 index 00000000..1b0d3755 --- /dev/null +++ b/scripts/DeepSeekV4/lib/prompts.json @@ -0,0 +1,32 @@ +[ + { + "id": "factual", + "prompt": "What is the capital of France? Answer in one word.", + "expect_contains": ["Paris"], + "max_tokens": 16 + }, + { + "id": "arithmetic", + "prompt": "Compute 17 multiplied by 23. Give only the number.", + "expect_contains": ["391"], + "max_tokens": 24 + }, + { + "id": "code", + "prompt": "Write a Python one-liner that returns the sum of a list named nums. Output only the code.", + "expect_contains": ["sum(", "nums"], + "max_tokens": 48 + }, + { + "id": "reasoning", + "prompt": "If a train travels 60 km in 1.5 hours, what is its average speed in km/h? Give only the number.", + "expect_contains": ["40"], + "max_tokens": 32 + }, + { + "id": "instruction", + "prompt": "List exactly three primary colors, comma-separated, lowercase.", + "expect_contains": ["red", "blue"], + "max_tokens": 24 + } +] diff --git a/scripts/DeepSeekV4/lib/run_disagg.sh b/scripts/DeepSeekV4/lib/run_disagg.sh new file mode 100755 index 00000000..e4eb5ddf --- /dev/null +++ b/scripts/DeepSeekV4/lib/run_disagg.sh @@ -0,0 +1,200 @@ +#!/bin/bash +# run_disagg.sh — DSv4-Pro PD-disaggregated serving + benchmark across multiple nodes. +# Faithful port of InferenceX server_atom.sh (atomesh + mooncake) and the sglang-disagg path. +# +# Env in: ENGINE=atom-disagg|sglang-disagg MODEL=dsv4-pro TOPO=1p1d|2p1d_dpa [ACTION=bench|serve] +# +# Layout per topology (topo.py resolves nodes): +# prefill worker(s) = kv_producer servers (port 8010) [worker 0's node also hosts the router] +# decode worker(s) = kv_consumer servers (port 8020) +# router (atomesh / sglang_router) on prefill-0, port 8000 <- bench drives THIS endpoint +# +# Each server runs in its own docker container via a per-node `srun --exclusive`. KV flows +# node->node over RDMA (mooncake for ATOM, MoRI for SGLang). Bench runs from a clean container +# against the router (aggregate throughput; no x N extrapolation — this is a real disagg instance). +set -uo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" # scripts_dsv4 +CLUSTER="$HERE/cluster.yaml"; MODELS="$HERE/model.yaml" +ENGINE="${ENGINE:?set ENGINE=atom-disagg|sglang-disagg}" +MODEL="${MODEL:-dsv4-pro}"; TOPO="${TOPO:?set TOPO=1p1d|2p1d_dpa}"; ACTION="${ACTION:-bench}" +cfg(){ python3 "$HERE/lib/cfg.py" "$@"; } +topo(){ python3 "$HERE/lib/topo.py" "$@"; } + +RESV=$(cfg "$CLUSTER" cluster reservation) +MROOT=$(cfg "$CLUSTER" cluster models_root) +RROOT=$(cfg "$CLUSTER" cluster results_root) +RPORT=$(cfg "$CLUSTER" cluster router_port) +PPORT=$(cfg "$CLUSTER" cluster prefill_port) +DPORT=$(cfg "$CLUSTER" cluster decode_port) +HPORT=$(cfg "$CLUSTER" cluster handshake_port) +MDIR=$(cfg "$MODELS" model "$MODEL" dir) +IMG=$(cfg "$MODELS" engine "$ENGINE" image) +ROUTER=$(cfg "$MODELS" engine "$ENGINE" router) +KVCONN=$(cfg "$MODELS" engine "$ENGINE" kv_connector) +GMU=$(cfg "$MODELS" default "$MODEL" gpu_memory_util) +MNS=$(cfg "$MODELS" default "$MODEL" max_num_seqs) +KVDT=$(cfg "$MODELS" default "$MODEL" kv_cache_dtype) +BS=$(cfg "$MODELS" default "$MODEL" block_size) +TP=$(cfg "$MODELS" default "$MODEL" tp) +SHAPES=$(topo "$MODELS" "$TOPO" field shapes) +CONC=$(topo "$MODELS" "$TOPO" field conc) +DPATTN=$(python3 -c "import yaml;t=yaml.safe_load(open('$MODELS'))['topologies']['$TOPO'];print('1' if t['decode'].get('dp_attn') else '0')") + +ENV_ARGS=""; for kv in $(cfg "$MODELS" engine "$ENGINE" env); do ENV_ARGS+=" -e $kv"; done + +STAMP=$(python3 -c "import time;print(time.strftime('%Y%m%d_%H%M%S'))") +RUN="$RROOT/${MODEL}_${ENGINE}_${TOPO}_${STAMP}"; mkdir -p "$RUN"; chmod 777 "$RUN" +topo "$CLUSTER" "$MODELS" "$TOPO" roles > "$RUN/roles.tsv" +echo "[disagg] ENGINE=$ENGINE TOPO=$TOPO router=$ROUTER kv=$KVCONN image=$IMG" +echo "[disagg] roles:"; sed 's/^/ /' "$RUN/roles.tsv" + +PREFILL_NODE0=$(awk -F'\t' '$1=="prefill" && $2=="0"{print $3}' "$RUN/roles.tsv") +PREFILL_IPS=""; DECODE_IPS="" + +ip_of(){ srun --reservation="$RESV" --nodelist="$1" --nodes=1 --ntasks=1 --overlap --mem=1G \ + bash -c "ip route get 1.1.1.1 2>/dev/null | awk '/src/{print \$7}'" 2>/dev/null | grep -oE '([0-9]+\.){3}[0-9]+' | head -1; } + +# --- PRE-FLIGHT: clean every node in this topology ------------------------------------------ +NODES=$(cut -f3 "$RUN/roles.tsv" | sort -u) +echo "[disagg] cleaning nodes: $(echo $NODES | tr '\n' ' ')" +for n in $NODES; do + srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=0 \ + bash "$HERE/lib/clean_node.sh" > "$RUN/clean_${n}.log" 2>&1 && echo " [clean] $n OK" || echo " [clean] $n WARN" +done + +# --- resolve router IP (prefill node 0) ----------------------------------------------------- +ROUTER_IP=$(ip_of "$PREFILL_NODE0") +echo "[disagg] router will be on $PREFILL_NODE0 ($ROUTER_IP):$RPORT" + +# parallel args per role from topology (ATOM flags shown; sglang path uses its own flags below) +dpa_arg(){ [[ "$1" == "1" ]] && echo "--enable-dp-attention --enable-tbo" || true; } + +# --- launch one server (prefill|decode) on a node in its own srun+docker -------------------- +launch_server() { + local role="$1" node="$2" port="$3" kvrole="$4" rdir="$RUN/${role}_${node}" + mkdir -p "$rdir"; chmod 777 "$rdir" + local pflags dflags extra="" + if [[ "$ENGINE" == "atom-disagg" ]]; then + [[ "$role" == "prefill" ]] && extra="$(cfg "$MODELS" engine "$ENGINE" prefill_flags)" || extra="$(cfg "$MODELS" engine "$ENGINE" decode_flags)" + [[ "$DPATTN" == "1" ]] && extra="$extra $(dpa_arg 1)" + local SRV="python3 -m atom.entrypoints.openai_server --model $MROOT/$MDIR \ + --host 0.0.0.0 --server-port $port -tp $TP \ + --kv-cache-dtype $KVDT --block-size $BS --gpu-memory-utilization $GMU --max-num-seqs $MNS \ + $extra \ + --kv-transfer-config '{\"kv_role\":\"$kvrole\",\"kv_connector\":\"$KVCONN\",\"proxy_ip\":\"$ROUTER_IP\",\"handshake_port\":$HPORT}'" + else # sglang-disagg + [[ "$role" == "prefill" ]] && extra="$(cfg "$MODELS" engine "$ENGINE" prefill_flags)" || extra="$(cfg "$MODELS" engine "$ENGINE" decode_flags)" + local SRV="python3 -m sglang.launch_server --model-path $MROOT/$MDIR \ + --host 0.0.0.0 --port $port --tp-size $TP --mem-fraction-static $GMU \ + --max-running-requests $MNS --kv-cache-dtype $KVDT $extra" + fi + srun --reservation="$RESV" --nodelist="$node" --nodes=1 --ntasks=1 --exclusive bash -c " + docker rm -f dsv4_${role} >/dev/null 2>&1 || true; sleep 3 + cid=\$(docker run -d --name dsv4_${role} --entrypoint bash \ + --device=/dev/kfd --device=/dev/dri --group-add video --ipc host --shm-size 128G --network host \ + --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --privileged \ + -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $ENV_ARGS \ + -v $MROOT:$MROOT -v $HERE:/scripts -v $rdir:/out \ + $IMG -c '$SRV') + echo \"[$role] container \$cid on $node:$port\" + for i in \$(seq 1 300); do + docker logs dsv4_${role} > $rdir/server.log 2>&1 || true + curl -sf http://localhost:$port/health >/dev/null 2>&1 && { echo READY > $rdir/READY; break; } + st=\$(docker inspect -f '{{.State.Status}}' dsv4_${role} 2>/dev/null || echo missing) + { [ \"\$st\" = exited ] || [ \"\$st\" = missing ]; } && { echo \"[$role] container \$st\"; break; } + sleep 5 + done + docker logs dsv4_${role} > $rdir/server.log 2>&1 || true + # hold the container alive for the duration (router + bench run elsewhere); wait on it + [ -f $rdir/READY ] && docker wait dsv4_${role} || { echo '[FAIL] tail:'; tail -25 $rdir/server.log; exit 1; } + " > "$rdir/srun.log" 2>&1 & + echo "$!" +} + +# Build router prefill/decode arg lists (need IPs) +declare -a SRV_PIDS=() +while IFS=$'\t' read -r role widx node port; do + ipn=$(ip_of "$node") + if [[ "$role" == "prefill" ]]; then PREFILL_IPS+="$ipn "; else DECODE_IPS+="$ipn "; fi +done < "$RUN/roles.tsv" + +echo "[disagg] launching servers..." +while IFS=$'\t' read -r role widx node port; do + kvrole="kv_producer"; [[ "$role" == "decode" ]] && kvrole="kv_consumer" + pid=$(launch_server "$role" "$node" "$port" "$kvrole"); SRV_PIDS+=("$pid") + echo " [$role $widx] $node:$port ($kvrole) srun-pid $pid" +done < "$RUN/roles.tsv" + +# --- wait for all servers READY ------------------------------------------------------------- +echo "[disagg] waiting for all servers healthy (timeout ~25min)..." +ALLREADY=0 +for i in $(seq 1 300); do + r=$(find "$RUN" -name READY 2>/dev/null | wc -l); n=$(wc -l < "$RUN/roles.tsv") + [[ "$r" -ge "$n" ]] && { ALLREADY=1; echo " all $n servers READY"; break; } + # bail early if any srun died + sleep 5 +done +if [[ "$ALLREADY" != "1" ]]; then echo "[disagg] NOT all servers came up — see $RUN/*/server.log"; fi + +# --- start the router on prefill node 0 ----------------------------------------------------- +# (atomesh for ATOM; sglang_router for SGLang). Built from PREFILL_IPS / DECODE_IPS. +RDIR="$RUN/router"; mkdir -p "$RDIR"; chmod 777 "$RDIR" +if [[ "$ENGINE" == "atom-disagg" ]]; then + PF_ARGS=""; for ip in $PREFILL_IPS; do PF_ARGS+=" --prefill http://$ip:$PPORT"; done + DC_ARGS=""; for ip in $DECODE_IPS; do DC_ARGS+=" --decode http://$ip:$DPORT"; done + RFLAGS=$(cfg "$MODELS" engine "$ENGINE" router_flags) + ROUTER_CMD="/usr/local/bin/atomesh launch --host 0.0.0.0 --port $RPORT --pd-disaggregation $PF_ARGS $DC_ARGS $RFLAGS" +else + # sglang_router: register prefill/decode workers (mini-lb style). Flags TBD per image. + PF_ARGS=""; for ip in $PREFILL_IPS; do PF_ARGS+=" --prefill http://$ip:$PPORT"; done + DC_ARGS=""; for ip in $DECODE_IPS; do DC_ARGS+=" --decode http://$ip:$DPORT"; done + ROUTER_CMD="python3 -m sglang_router.launch_router --host 0.0.0.0 --port $RPORT --pd-disaggregation $PF_ARGS $DC_ARGS" +fi +echo "[disagg] router cmd: $ROUTER_CMD" +srun --reservation="$RESV" --nodelist="$PREFILL_NODE0" --nodes=1 --ntasks=1 --overlap --mem=8G bash -c " + docker rm -f dsv4_router >/dev/null 2>&1 || true + cid=\$(docker run -d --name dsv4_router --entrypoint bash --network host \ + -v $HERE:/scripts $IMG -c '$ROUTER_CMD') + for i in \$(seq 1 60); do + docker logs dsv4_router > $RDIR/router.log 2>&1 || true + curl -sf http://localhost:$RPORT/v1/models >/dev/null 2>&1 && { echo READY > $RDIR/READY; break; } + sleep 5 + done + docker logs dsv4_router > $RDIR/router.log 2>&1 || true +" > "$RDIR/srun.log" 2>&1 & + +for i in $(seq 1 60); do [[ -f "$RDIR/READY" ]] && { echo "[disagg] router READY on $ROUTER_IP:$RPORT"; break; }; sleep 5; done +[[ -f "$RDIR/READY" ]] || echo "[disagg] router did NOT become ready — see $RDIR/router.log" + +if [[ "$ACTION" == "serve" ]]; then + echo "[disagg] ACTION=serve — endpoint http://$ROUTER_IP:$RPORT (router). Servers held; Ctrl-C / scancel to stop." + wait; exit 0 +fi + +# --- ACTION=bench: drive the router at each (shape,conc) ------------------------------------- +BENCH_IMG=$(cfg "$MODELS" engine "$ENGINE" bench_image) +BENCH_IMG="${BENCH_IMG:-vllm/vllm-openai-rocm:nightly-9037498c22891e55b594f567fb91d9b4efbf3e99}" +echo "[disagg] benchmarking router http://$ROUTER_IP:$RPORT shapes='$SHAPES' conc='$CONC'" +for shape in $SHAPES; do + ISL=${shape%,*}; OSL=${shape#*,} + for c in $CONC; do + echo " [bench] isl=$ISL osl=$OSL conc=$c" + srun --reservation="$RESV" --nodelist="$PREFILL_NODE0" --nodes=1 --ntasks=1 --overlap --mem=8G bash -c " + docker run --rm --entrypoint bash --network host -v $HERE:/scripts -v $RUN:/out \ + $BENCH_IMG -c 'python3 /scripts/utils/bench_serving/benchmark_serving.py \ + --model $MROOT/$MDIR --backend vllm --base-url http://$ROUTER_IP:$RPORT --dataset-name random \ + --random-input-len $ISL --random-output-len $OSL --random-range-ratio 1.0 \ + --num-prompts \$(( $c * 10 )) --num-warmups \$(( $c * 2 )) --max-concurrency $c \ + --request-rate inf --ignore-eos --save-result --percentile-metrics ttft,tpot,itl,e2el \ + --result-dir /out --result-filename disagg_isl${ISL}_osl${OSL}_c${c}.json' + " > "$RUN/bench_isl${ISL}_osl${OSL}_c${c}.log" 2>&1 || echo " [bench] FAILED isl=$ISL conc=$c" + done +done + +# --- teardown ------------------------------------------------------------------------------- +echo "[disagg] teardown: stopping servers + router" +for n in $NODES; do + srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=1G \ + bash -c "docker rm -f dsv4_prefill dsv4_decode dsv4_router >/dev/null 2>&1" 2>/dev/null || true +done +echo "[disagg] done -> $RUN" diff --git a/scripts/DeepSeekV4/lib/topo.py b/scripts/DeepSeekV4/lib/topo.py new file mode 100644 index 00000000..21cdc6a0 --- /dev/null +++ b/scripts/DeepSeekV4/lib/topo.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +"""topo.py — resolve a DSv4 disagg topology into a node-role placement. + +Usage: + topo.py roles + -> TSV lines: roleworker_idxnodeport + role in {prefill, decode}; node 0 of prefill also hosts the router. + topo.py field + -> scalar/list field of the topology (xP, yD, shapes, conc, ...) + +Topology drives node count: nodes_needed = xP*pf_nodes_per_worker + yD*dec_nodes_per_worker, +with nodes_per_worker = ceil(TP / gpus_per_node) (TP8 on 8-GPU node => 1 node/worker). +""" +import sys, yaml, math + +def load(f): return yaml.safe_load(open(f)) + +def roles(cluster_f, model_f, topo_name): + c = load(cluster_f); m = load(model_f) + gpn = c["gpus_per_node"]; nodes = c["nodes"] + t = m["topologies"][topo_name] + xP, yD = t["xP"], t["yD"] + pf_tp = t["prefill"]["tp"]; dec_tp = t["decode"]["tp"] + pf_npw = math.ceil(pf_tp / gpn); dec_npw = math.ceil(dec_tp / gpn) + need = xP * pf_npw + yD * dec_npw + if need > len(nodes): + sys.exit(f"ERROR: topo {topo_name} needs {need} nodes, cluster has {len(nodes)}") + pf_port = c["prefill_port"]; dec_port = c["decode_port"] + out = []; i = 0 + for w in range(xP): + out.append(f"prefill\t{w}\t{nodes[i]}\t{pf_port}"); i += pf_npw + for w in range(yD): + out.append(f"decode\t{w}\t{nodes[i]}\t{dec_port}"); i += dec_npw + return out + +def main(): + a = sys.argv + if a[2:3] == ["roles"] or (len(a) >= 4 and a[3] == "roles"): + # form: topo.py cluster model topo roles -> but support both orderings + pass + if len(a) >= 5 and a[4] == "roles": + print("\n".join(roles(a[1], a[2], a[3]))); return + if len(a) >= 5 and a[3] == "field": + m = load(a[1]); t = m["topologies"][a[2]]; v = t.get(a[4]) + if isinstance(v, list): print(" ".join(str(x) for x in v)) + elif isinstance(v, dict): print(" ".join(f"{k}={vv}" for k, vv in v.items())) + else: print("" if v is None else v) + return + sys.exit("usage: topo.py roles | topo.py field ") + +if __name__ == "__main__": + main() diff --git a/scripts/DeepSeekV4/model.yaml b/scripts/DeepSeekV4/model.yaml new file mode 100644 index 00000000..96d29b46 --- /dev/null +++ b/scripts/DeepSeekV4/model.yaml @@ -0,0 +1,72 @@ +# model.yaml — DSv4-Pro model def + per-engine disagg images/flags/router. +# +# DeepSeek-V4-Pro: deepseek_v4, 61 layers, 384 routed experts (6 active), hidden 7168, +# moe_intermediate_size 3072. FP8 block-quant (128x128, ue8m0) + FP4 MoE experts. +# TP8 VALID (3072/8=384 >= 128) — InferenceX runs TP8 throughout. PD-disaggregated. + +defaults: + gpu_memory_util: 0.85 # ATOM --gpu-memory-utilization / SGLang --mem-fraction-static + max_num_seqs: 256 + kv_cache_dtype: fp8 + block_size: 16 + tp: 8 # per prefill/decode worker + +models: + dsv4-pro: + dir: DeepSeek-V4-Pro # under cluster.models_root (staged to scratch) + arch: deepseek_v4 + status: WEIGHTS_VALID # 64/64 shards integrity-checked on NFS; staging to scratch + # DSv4-specific env applied by the launcher (from InferenceX env_atom.sh): + # ATOM_MOE_GU_ITLV=1, AITER_BF16_FP8_MOE_BOUND=0; TBO path adds GPU_MAX_HW_QUEUES=5, + # ATOM_CPU_AFFINITY=1. + +# ---- engines: disagg image + router + KV-transfer backend ---------------------------------- +engines: + atom-disagg: + # The InferenceX-VALIDATED DSv4 disagg path. atomesh router + mooncake RDMA KV. + image: rocm/atom-dev:nightly_202606101403 + # bench client runs from a clean vLLM image (ATOM's own image crashes the bench client) + bench_image: vllm/vllm-openai-rocm:nightly-9037498c22891e55b594f567fb91d9b4efbf3e99 + router: atomesh # /usr/local/bin/atomesh launch --pd-disaggregation + kv_connector: mooncake # --kv-transfer-config '{"kv_connector":"mooncake",...}' + entrypoint: "python3 -m atom.entrypoints.openai_server" + env: + ATOM_MOE_GU_ITLV: "1" + AITER_BF16_FP8_MOE_BOUND: "0" + SAFETENSORS_FAST_GPU: "1" + # engine-unique server flags (shared knobs come from defaults{}) + prefill_flags: "--trust-remote-code --no-enable_prefix_caching" + decode_flags: "--trust-remote-code --no-enable_prefix_caching" + router_flags: "--policy random --backend atom --disable-health-check --disable-circuit-breaker" + status: PRIMARY # validated DSv4 atom-disagg config + + sglang-disagg: + # Adapted from DeepSeek-R1 sglang-disagg (NO DSv4-sglang-disagg config exists upstream yet). + # MoRI-IO KV transfer + sglang_router. Needs an sglang image carrying DSv4 model support. + image: rocm/sgl-dev:rocm720-mi35x-0363e6c-20260509-DSv4 # DSv4-capable; disagg support TBC + router: sglang_router + kv_connector: mori # MoRI-IO (env: MORI_IO_*, MORI_RDMA_TC/SL) + entrypoint: "python3 -m sglang.launch_server" + env: + SGLANG_USE_AITER: "1" + SGLANG_USE_AITER_UNIFIED_ATTN: "1" + prefill_flags: "--trust-remote-code --disaggregation-mode prefill --disable-radix-cache" + decode_flags: "--trust-remote-code --disaggregation-mode decode --disable-radix-cache --enable-dp-attention --ep-size 8" + status: EXPERIMENTAL # R1 template; DSv4 disagg must be validated + +# ---- disagg topologies (from dsv4-fp4-mi355x-atom-disagg search-space) ---------------------- +topologies: + 1p1d: # 2 nodes: 1 prefill (TP8) + 1 decode (TP8) + xP: 1 + yD: 1 + prefill: { tp: 8, ep: 1, dp_attn: false } + decode: { tp: 8, ep: 1, dp_attn: false } + shapes: ["1024,1024", "8192,1024"] + conc: [4, 8, 16, 32, 64, 128] + 2p1d_dpa: # 3 nodes: 2 prefill (TP8, DP-attn) + 1 decode (TP8, DP-attn) + xP: 2 + yD: 1 + prefill: { tp: 8, ep: 1, dp_attn: true } + decode: { tp: 8, ep: 1, dp_attn: true } + shapes: ["8192,1024"] + conc: [256, 512, 768, 1024, 2048] diff --git a/scripts/DeepSeekV4/run_atom_disagg.sh b/scripts/DeepSeekV4/run_atom_disagg.sh new file mode 100755 index 00000000..1e68f193 --- /dev/null +++ b/scripts/DeepSeekV4/run_atom_disagg.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# run_atom_disagg.sh — DSv4 ATOM PD-disagg (atomesh + mooncake). Thin wrapper. +# Usage: MODEL=dsv4-pro TOPO=1p1d|2p1d_dpa [ACTION=bench|serve] ./run_atom_disagg.sh +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENGINE=atom-disagg exec bash "$HERE/lib/run_disagg.sh" diff --git a/scripts/DeepSeekV4/run_sglang_disagg.sh b/scripts/DeepSeekV4/run_sglang_disagg.sh new file mode 100755 index 00000000..183b5995 --- /dev/null +++ b/scripts/DeepSeekV4/run_sglang_disagg.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# run_sglang_disagg.sh — DSv4 SGLang PD-disagg (sglang_router + MoRI). EXPERIMENTAL (R1 template). +# Usage: MODEL=dsv4-pro TOPO=1p1d|2p1d_dpa [ACTION=bench|serve] ./run_sglang_disagg.sh +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENGINE=sglang-disagg exec bash "$HERE/lib/run_disagg.sh" diff --git a/scripts/DeepSeekV4/utils/bench_serving/KNOWN_LIMITATION.md b/scripts/DeepSeekV4/utils/bench_serving/KNOWN_LIMITATION.md new file mode 100644 index 00000000..342de691 --- /dev/null +++ b/scripts/DeepSeekV4/utils/bench_serving/KNOWN_LIMITATION.md @@ -0,0 +1,7 @@ +# KNOWN_LIMITATION + +When benchmarking small models (e.g., Gemma 1B or Llama 8B) and at ultra-high querys per second (QPS) and with short input and output lengths, the current InferenceX bench_serving client becomes client-bound. + +InferenceX does not currently benchmark in this regime and has no plans to. Our roadmap skews the opposite direction, larger models, longer ISL/OSL (i.e., agentic workloads), and interactive TTFT & tok/s/user scenarios rather than ultra-high QPS. + +Should the need arise, we shall fix this known limitation & migrate to a multi-process benchmark client. If your needs are benchmarking small models at high QPS at small ISL/OSL lengths, we recommend xyz benchmark instead. diff --git a/scripts/DeepSeekV4/utils/bench_serving/backend_request_func.py b/scripts/DeepSeekV4/utils/bench_serving/backend_request_func.py new file mode 100644 index 00000000..4c8820f8 --- /dev/null +++ b/scripts/DeepSeekV4/utils/bench_serving/backend_request_func.py @@ -0,0 +1,574 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +import sys +import time +import traceback +from dataclasses import dataclass, field +from typing import List, Optional, Union + +import aiohttp +import huggingface_hub.constants +from tqdm.asyncio import tqdm +from transformers import (AutoTokenizer, PreTrainedTokenizer, + PreTrainedTokenizerFast) + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) + + +@dataclass +class RequestFuncInput: + prompt: str + api_url: str + prompt_len: int + output_len: int + model: str + model_name: Optional[str] = None + best_of: int = 1 + logprobs: Optional[int] = None + extra_body: Optional[dict] = None + multi_modal_content: Optional[dict] = None + ignore_eos: bool = False + + +@dataclass +class RequestFuncOutput: + generated_text: str = "" + success: bool = False + latency: float = 0.0 + output_tokens: int = 0 + ttft: float = 0.0 # Time to first token + itl: List[float] = field( + default_factory=list) # List of inter-token latencies + tpot: float = 0.0 # avg next-token latencies + prompt_len: int = 0 + error: str = "" + + +async def async_request_tgi( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith("generate_stream") + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + params = { + "best_of": request_func_input.best_of, + "max_new_tokens": request_func_input.output_len, + "do_sample": True, + "temperature": 0.01, # TGI does not accept 0.0 temperature. + "top_p": 0.99, # TGI does not accept 1.0 top_p. + "truncate": request_func_input.prompt_len, + # TGI does not accept ignore_eos flag. + } + payload = { + "inputs": request_func_input.prompt, + "parameters": params, + } + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + chunk_bytes = chunk_bytes.decode("utf-8") + + # NOTE: Sometimes TGI returns a ping response without + # any data, we should skip it. + if chunk_bytes.startswith(":"): + continue + chunk = chunk_bytes.removeprefix("data:") + + data = json.loads(chunk) + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + most_recent_timestamp = timestamp + + output.latency = most_recent_timestamp - st + output.success = True + output.generated_text = data["generated_text"] + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_trt_llm( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith("generate_stream") + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + assert request_func_input.best_of == 1 + payload = { + "accumulate_tokens": True, + "text_input": request_func_input.prompt, + "temperature": 0.0, + "top_p": 1.0, + "max_tokens": request_func_input.output_len, + "stream": True, + } + if request_func_input.ignore_eos: + payload["min_length"] = request_func_input.output_len + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = chunk_bytes.decode("utf-8").removeprefix( + "data:") + + data = json.loads(chunk) + output.generated_text += data["text_output"] + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + most_recent_timestamp = timestamp + + output.latency = most_recent_timestamp - st + output.success = True + + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_deepspeed_mii( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + assert request_func_input.best_of == 1 + + payload = { + "prompt": request_func_input.prompt, + "max_tokens": request_func_input.output_len, + "temperature": 0.01, # deepspeed-mii does not accept 0.0 temp. + "top_p": 1.0, + } + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + # NOTE: DeepSpeed-MII doesn't support streaming as of Jan 28 2024, + # will use 0 as placeholder. + # See https://github.com/microsoft/DeepSpeed-MII/pull/311 + output.ttft = 0 + + st = time.perf_counter() + try: + async with session.post(url=request_func_input.api_url, + json=payload) as response: + if response.status == 200: + parsed_resp = await response.json() + output.latency = time.perf_counter() - st + output.generated_text = parsed_resp["text"][0] + output.success = True + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + ("completions", "profile") + ), "OpenAI Completions API URL must end with 'completions' or 'profile'." + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "model": request_func_input.model_name \ + if request_func_input.model_name else request_func_input.model, + "prompt": request_func_input.prompt, + "temperature": 0.0, + "best_of": request_func_input.best_of, + "max_tokens": request_func_input.output_len, + "logprobs": request_func_input.logprobs, + "stream": True, + "stream_options": { + "include_usage": True, + }, + } + if request_func_input.ignore_eos: + payload["ignore_eos"] = request_func_input.ignore_eos + if request_func_input.extra_body: + payload.update(request_func_input.extra_body) + headers = { + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}" + } + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload, + headers=headers) as response: + if response.status == 200: + first_chunk_received = False + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = chunk_bytes.decode("utf-8").removeprefix( + "data: ") + if chunk != "[DONE]": + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if choices := data.get("choices"): + # Note that text could be empty here + # e.g. for special tokens + text = choices[0].get("text") + timestamp = time.perf_counter() + # First token + if not first_chunk_received: + first_chunk_received = True + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += text or "" + elif usage := data.get("usage"): + output.output_tokens = usage.get( + "completion_tokens") + if first_chunk_received: + output.success = True + else: + output.success = False + output.error = ( + "Never received a valid chunk to calculate TTFT." + "This response will be marked as failed!") + output.generated_text = generated_text + output.latency = most_recent_timestamp - st + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_openai_chat_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + "chat/completions" + ), "OpenAI Chat Completions API URL must end with 'chat/completions'." + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + content = request_func_input.prompt + if request_func_input.multi_modal_content: + content = [{"type": "text", "text": request_func_input.prompt}] + content.append(request_func_input.multi_modal_content) + payload = { + "model": request_func_input.model_name \ + if request_func_input.model_name else request_func_input.model, + "messages": [ + { + "role": "user", + "content": content + }, + ], + "temperature": 0.0, + "max_completion_tokens": request_func_input.output_len, + "stream": True, + "stream_options": { + "include_usage": True, + }, + } + if request_func_input.ignore_eos: + payload["ignore_eos"] = request_func_input.ignore_eos + if request_func_input.extra_body: + payload.update(request_func_input.extra_body) + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + } + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload, + headers=headers) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = chunk_bytes.decode("utf-8").removeprefix( + "data: ") + if chunk != "[DONE]": + timestamp = time.perf_counter() + data = json.loads(chunk) + + if choices := data.get("choices"): + content = choices[0]["delta"].get("content") + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + generated_text += content or "" + elif usage := data.get("usage"): + output.output_tokens = usage.get( + "completion_tokens") + + most_recent_timestamp = timestamp + + output.generated_text = generated_text + output.success = True + output.latency = most_recent_timestamp - st + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +def get_model(pretrained_model_name_or_path: str) -> str: + if os.getenv('VLLM_USE_MODELSCOPE', 'False').lower() == 'true': + from modelscope import snapshot_download + + model_path = snapshot_download( + model_id=pretrained_model_name_or_path, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"]) + + return model_path + return pretrained_model_name_or_path + + +def _fix_tokenizer_for_sglang(tokenizer, model_path): + """Fix transformers v5 tokenizer to match sglang server-side behavior. + + Root cause: transformers v5 (>= 5.0) changed how tokenizers are loaded. + Specifically, LlamaTokenizerFast.__init__ in v5 rebuilds the pre_tokenizer + and decoder from scratch using class-specific components, discarding the + originals from tokenizer.json. For models like DeepSeek-R1 that declare + LlamaTokenizerFast but actually use a ByteLevel/Sequence tokenizer + architecture, v5 incorrectly replaces the original Sequence pre_tokenizer + with Metaspace, and the original ByteLevel decoder with Sequence. + See: https://github.com/sgl-project/sglang/blob/9238bd08a2895fa3b7ec79ea567e5c27ac951343/python/sglang/srt/utils/hf_transformers_utils.py#L836 + + The sglang server applies fixes for this in hf_transformers_utils.py + (_fix_v5_tokenizer_components and _fix_v5_add_bos_eos_token), but the + benchmark client loads the tokenizer directly via AutoTokenizer without + these fixes. This mismatch causes the client to encode text differently + from the server -- e.g. a 7000-token prompt on the client becomes ~35000 + tokens on the server, leading to ~5x TTFT inflation and false performance + regressions in benchmarks. + + This function replicates the same fixes so the benchmark client tokenizes + identically to the sglang server. It is a no-op on transformers v4. + """ + import json + from pathlib import Path + + def _resolve(filename): + """Return a filesystem path for `filename`, whether `model_path` is a + local directory or an HF Hub repo id. Returns None and logs a warning + on failure so we don't silently fail to apply the v5 fix.""" + local = Path(model_path) / filename + if local.is_file(): + return str(local) + try: + from huggingface_hub import hf_hub_download + return hf_hub_download(repo_id=model_path, filename=filename) + except Exception as e: + print( + f"v5 tokenizer fix: cannot resolve {filename} for {model_path!r} " + f"({type(e).__name__}: {e}); fix will not apply.", + flush=True, + ) + return None + + backend = getattr(tokenizer, "_tokenizer", None) + if backend is not None: + tok_file = _resolve("tokenizer.json") + if tok_file is not None: + from tokenizers import Tokenizer as RawTokenizer + raw = RawTokenizer.from_file(tok_file) + raw_pre = type(raw.pre_tokenizer).__name__ if raw.pre_tokenizer else None + loaded_pre = type(backend.pre_tokenizer).__name__ if backend.pre_tokenizer else None + if raw_pre and loaded_pre and raw_pre != loaded_pre: + print( + f"v5 tokenizer fix: {model_path} pre_tokenizer {loaded_pre} -> {raw_pre}, " + f"decoder {type(backend.decoder).__name__ if backend.decoder else None} -> " + f"{type(raw.decoder).__name__ if raw.decoder else None}", + flush=True, + ) + backend.pre_tokenizer = raw.pre_tokenizer + backend.decoder = raw.decoder + + config_file = _resolve("tokenizer_config.json") + if config_file is not None: + with open(config_file) as f: + config = json.load(f) + tok_class = config.get("tokenizer_class", "") + bos_eos_classes = { + "LlamaTokenizer", "LlamaTokenizerFast", + "CodeLlamaTokenizer", "CodeLlamaTokenizerFast", + "GemmaTokenizer", "GemmaTokenizerFast", "CohereTokenizerFast", + } + if tok_class in bos_eos_classes: + defaults = {"add_bos_token": True, "add_eos_token": False} + changed = False + for attr in ("add_bos_token", "add_eos_token"): + val = config.get(attr) + if val is None: + val = defaults.get(attr, False) + if getattr(tokenizer, attr, None) != val: + setattr(tokenizer, f"_{attr}", val) + changed = True + if changed and hasattr(tokenizer, "update_post_processor"): + tokenizer.update_post_processor() + + return tokenizer + + +def get_tokenizer( + pretrained_model_name_or_path: str, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + **kwargs, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + if pretrained_model_name_or_path is not None and not os.path.exists( + pretrained_model_name_or_path): + pretrained_model_name_or_path = get_model( + pretrained_model_name_or_path) + if tokenizer_mode == "slow": + if kwargs.get("use_fast", False): + raise ValueError( + "Cannot use the fast tokenizer in slow tokenizer mode.") + kwargs["use_fast"] = False + if tokenizer_mode == "mistral": + try: + from vllm.transformers_utils.tokenizer import MistralTokenizer + except ImportError as e: + raise ImportError("MistralTokenizer requires vllm package.\n" + "Please install it with `pip install vllm` " + "to use mistral tokenizer mode.") from e + return MistralTokenizer.from_pretrained( + str(pretrained_model_name_or_path)) + else: + tokenizer = AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, + trust_remote_code=trust_remote_code, + **kwargs, + ) + return _fix_tokenizer_for_sglang(tokenizer, pretrained_model_name_or_path) + + +ASYNC_REQUEST_FUNCS = { + "tgi": async_request_tgi, + "vllm": async_request_openai_completions, + "lmdeploy": async_request_openai_completions, + "deepspeed-mii": async_request_deepspeed_mii, + "openai": async_request_openai_completions, + "openai-chat": async_request_openai_chat_completions, + "tensorrt-llm": async_request_trt_llm, + "scalellm": async_request_openai_completions, + "sglang": async_request_openai_completions, +} diff --git a/scripts/DeepSeekV4/utils/bench_serving/benchmark_serving.py b/scripts/DeepSeekV4/utils/bench_serving/benchmark_serving.py new file mode 100644 index 00000000..7cc806ec --- /dev/null +++ b/scripts/DeepSeekV4/utils/bench_serving/benchmark_serving.py @@ -0,0 +1,1276 @@ +# SPDX-License-Identifier: Apache-2.0 +r"""Benchmark online serving throughput. + +On the server side, run one of the following commands: + vLLM OpenAI API server + vllm serve \ + --swap-space 16 \ + --disable-log-requests + + (TGI backend) + ./launch_tgi_server.sh + +On the client side, run: + python benchmarks/benchmark_serving.py \ + --backend \ + --model \ + --dataset-name sharegpt \ + --dataset-path \ + --request-rate \ # By default is inf + --num-prompts # By default is 1000 + + when using tgi backend, add + --endpoint /generate_stream + to the end of the command above. +""" +import argparse +import asyncio +import base64 +import contextlib +import gc +import io +import json +import os +import random +import time +import warnings +from dataclasses import dataclass +from datetime import datetime +from multiprocessing import Pool, cpu_count +from typing import Any, AsyncGenerator, Collection, Dict, List, Optional, Tuple + +import numpy as np +from backend_request_func import (ASYNC_REQUEST_FUNCS, RequestFuncInput, + RequestFuncOutput) +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +try: + from backend_request_func import get_tokenizer +except ImportError: + from vllm.transformers_utils.tokenizer import get_tokenizer + +try: + from vllm.utils import FlexibleArgumentParser +except ImportError: + from argparse import ArgumentParser as FlexibleArgumentParser + +from benchmark_utils import convert_to_pytorch_benchmark_format +from encoding_dsv4 import encode_messages as dsv4_encode_messages + +MILLISECONDS_TO_SECONDS_CONVERSION = 1000 + + +@dataclass +class BenchmarkMetrics: + completed: int + total_input: int + total_output: int + request_throughput: float + request_goodput: float + output_throughput: float + total_token_throughput: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + percentiles_ttft_ms: List[Tuple[float, float]] + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + percentiles_tpot_ms: List[Tuple[float, float]] + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + percentiles_itl_ms: List[Tuple[float, float]] + # E2EL stands for end-to-end latency per request. + # It is the time taken on the client side from sending + # a request to receiving a complete response. + mean_e2el_ms: float + median_e2el_ms: float + std_e2el_ms: float + percentiles_e2el_ms: List[Tuple[float, float]] + + +# --- Multiprocessing helpers for sample_random_requests --- +_worker_tokenizer = None + + +def _load_tokenizer(tokenizer_id, tokenizer_mode, trust_remote_code): + """Load tokenizer for random-prompt generation. + + vLLM's get_tokenizer can raise AttributeError when transformers removes + LlamaTokenizer.all_special_tokens_extended (e.g. Qwen3.5 with newer + transformers). Prefer backend_request_func.get_tokenizer on fallback so + client tokenization stays aligned with the sglang server (#1381, #1428). + """ + try: + return get_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) + except AttributeError as exc: + if "all_special_tokens_extended" not in str(exc): + raise + try: + from backend_request_func import get_tokenizer as _backend_get_tokenizer + return _backend_get_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) + except ImportError: + from transformers import AutoTokenizer + use_fast = tokenizer_mode != "slow" + return AutoTokenizer.from_pretrained( + tokenizer_id, + trust_remote_code=trust_remote_code, + use_fast=use_fast, + ) + + +def _init_tokenizer_worker(tokenizer_id, tokenizer_mode, trust_remote_code): + """Initialize tokenizer once per worker process.""" + global _worker_tokenizer + _worker_tokenizer = _load_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) + + +def _apply_chat_template(prompt, tokenizer, dsv4): + """Render a single user message into the appropriate chat-template prompt. + + When `dsv4` is True we use the self-contained DeepSeek-V4 encoder + (encoding_dsv4.encode_messages) which emits the + ... framing the model expects. Otherwise we + fall back to the tokenizer's built-in jinja chat template. + """ + if dsv4: + return dsv4_encode_messages( + [{"role": "user", "content": prompt}], + thinking_mode="thinking", + ) + return tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + ) + + +def _process_prompt_chunk(chunk_args): + """Generate a chunk of random prompts in a worker process.""" + (indices, prefix_token_ids, input_lens, output_lens, offsets, + prefix_len, vocab_size, use_chat_template, dsv4, seed) = chunk_args + + rng = np.random.RandomState(seed) + tokenizer = _worker_tokenizer + + results = [] + for local_idx, global_idx in enumerate(indices): + tgt_prompt_len = prefix_len + input_lens[local_idx] + prompt_token_ids = prefix_token_ids + [ + (offsets[local_idx] + global_idx + j) % vocab_size + for j in range(input_lens[local_idx]) + ] + prompt = tokenizer.decode(prompt_token_ids) + + max_retries = 10 + for _ in range(max_retries): + prompt_token_ids = tokenizer.encode(prompt, add_special_tokens=False) + if len(prompt_token_ids) < tgt_prompt_len: + num_extras = tgt_prompt_len - len(prompt_token_ids) + prompt_token_ids.extend( + rng.randint(0, vocab_size, size=num_extras).tolist()) + elif len(prompt_token_ids) > tgt_prompt_len: + prompt_token_ids = prompt_token_ids[:tgt_prompt_len] + else: + break + prompt = tokenizer.decode(prompt_token_ids) + + if use_chat_template: + prompt = _apply_chat_template(prompt, tokenizer, dsv4) + + prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False)) + mismatch = prompt_len - tgt_prompt_len + results.append((prompt, prompt_len, output_lens[local_idx], None, mismatch)) + + return results + + +def sample_random_requests( + prefix_len: int, + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + use_chat_template: bool = False, + dsv4: bool = False, + tokenizer_id: Optional[str] = None, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + num_workers: int = 0, +) -> List[Tuple[str, int, int]]: + vocab_size = tokenizer.vocab_size + prefix_token_ids = np.random.randint(0, vocab_size, size=prefix_len).tolist() + + if dsv4 and not use_chat_template: + raise ValueError("--dsv4 requires --use-chat-template to be set.") + + if use_chat_template: + chat_template_dummy = _apply_chat_template("a", tokenizer, dsv4) + tokenized_chat_template_dummy = tokenizer.encode(chat_template_dummy, add_special_tokens=False) + chat_template_len = len(tokenized_chat_template_dummy) - 1 + input_len = input_len - chat_template_len + + def sample_uniform(seq_len): + lower = int(seq_len * range_ratio) + upper = seq_len + seq_lens = np.random.randint(lower, upper+1, size=num_prompts).tolist() + return seq_lens + + input_lens = sample_uniform(input_len) + output_lens = sample_uniform(output_len) + offsets = np.random.randint(0, vocab_size, size=num_prompts) + + # Create a local RNG for retry-loop padding so that neither serial nor + # parallel path consumes global np.random draws beyond this point. + # This ensures downstream code (e.g. gamma draws for inter-arrival times) + # sees identical global RNG state regardless of num_workers. + local_rng = np.random.RandomState( + np.random.get_state()[1][:4].tolist() # derive seed from current state without advancing it + ) + + # Decide whether to use multiprocessing + if num_workers <= 0: + num_workers = min(cpu_count() or 1, 8) + use_parallel = num_workers > 1 and tokenizer_id is not None + + if use_parallel: + # Split work into chunks, one per worker + chunk_size = (num_prompts + num_workers - 1) // num_workers + chunk_args_list = [] + for w in range(num_workers): + start = w * chunk_size + end = min(start + chunk_size, num_prompts) + if start >= num_prompts: + break + chunk_args_list.append(( + list(range(start, end)), + prefix_token_ids, + input_lens[start:end], + output_lens[start:end], + offsets[start:end].tolist(), + prefix_len, + vocab_size, + use_chat_template, + dsv4, + int(local_rng.randint(0, 2**31)), + )) + + actual_workers = len(chunk_args_list) + print(f"Generating {num_prompts} prompts using {actual_workers} worker processes...") + t0 = time.perf_counter() + with Pool( + processes=actual_workers, + initializer=_init_tokenizer_worker, + initargs=(tokenizer_id, tokenizer_mode, trust_remote_code), + ) as pool: + chunk_results = pool.map(_process_prompt_chunk, chunk_args_list) + + input_requests = [] + mismatches = [] + for chunk in chunk_results: + for prompt, prompt_len, out_len, mm_content, mismatch in chunk: + input_requests.append((prompt, prompt_len, out_len, mm_content)) + mismatches.append(mismatch) + elapsed = time.perf_counter() - t0 + print(f"Prompt generation completed in {elapsed:.1f}s") + else: + # Original serial path — also uses local_rng for retry-loop padding + # to keep global RNG consumption identical to the parallel path. + if tokenizer_id is None and num_workers > 1: + print("Warning: tokenizer_id not provided, falling back to serial prompt generation.") + input_requests = [] + mismatches = [] + for i in range(num_prompts): + tgt_prompt_len = prefix_len + input_lens[i] + prompt_token_ids = prefix_token_ids + [(offsets[i] + i + j) % vocab_size for j in range(input_lens[i])] + prompt = tokenizer.decode(prompt_token_ids) + + max_retries = 10 + for _ in range(max_retries): + prompt_token_ids = tokenizer.encode(prompt, add_special_tokens=False) + if len(prompt_token_ids) < tgt_prompt_len: + num_extras = tgt_prompt_len - len(prompt_token_ids) + prompt_token_ids.extend(local_rng.randint(0, vocab_size, size=num_extras).tolist()) + elif len(prompt_token_ids) > tgt_prompt_len: + prompt_token_ids = prompt_token_ids[:tgt_prompt_len] + else: + break + prompt = tokenizer.decode(prompt_token_ids) + + if use_chat_template: + prompt = _apply_chat_template(prompt, tokenizer, dsv4) + + prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False)) + mismatches.append(prompt_len - tgt_prompt_len) + input_requests.append((prompt, prompt_len, output_lens[i], None)) + + header_str = f'{"-"*19} Input/Output Length Statistics {"-"*19}' + print(header_str) + print( + f' input_lens : ' + f'min={min(r[1] for r in input_requests):<4d} ' + f'max={max(r[1] for r in input_requests):<4d} ' + f'mean={np.mean([r[1] for r in input_requests]):<7.2f} ' + f'avg_token_mismatch={np.mean(mismatches):<5.2f} ' + ) + print( + f' output_lens: ' + f'min={min(r[2] for r in input_requests):<4d} ' + f'max={max(r[2] for r in input_requests):<4d} ' + f'mean={np.mean([r[2] for r in input_requests]):<7.2f} ' + ) + print('-' * len(header_str), '\n') + + return input_requests + + +async def get_request( + input_requests: List[Tuple[str, int, int]], + request_rate: float, + burstiness: float = 1.0, +) -> AsyncGenerator[Tuple[str, int, int], None]: + """ + Asynchronously generates requests at a specified rate + with OPTIONAL burstiness. + + Args: + input_requests: + A list of input requests, each represented as a tuple. + request_rate: + The rate at which requests are generated (requests/s). + burstiness (optional): + The burstiness factor of the request generation. + Only takes effect when request_rate is not inf. + Default value is 1, which follows a Poisson process. + Otherwise, the request intervals follow a gamma distribution. + A lower burstiness value (0 < burstiness < 1) results + in more bursty requests, while a higher burstiness value + (burstiness > 1) results in a more uniform arrival of requests. + """ + input_requests = iter(input_requests) + + # Calculate scale parameter theta to maintain the desired request_rate. + assert burstiness > 0, ( + f"A positive burstiness factor is expected, but given {burstiness}.") + theta = 1.0 / (request_rate * burstiness) + + for request in input_requests: + yield request + + if request_rate == float("inf"): + # If the request rate is infinity, then we don't need to wait. + continue + + # Sample the request interval from the gamma distribution. + # If burstiness is 1, it follows exponential distribution. + interval = np.random.gamma(shape=burstiness, scale=theta) + # The next request will be sent after the interval. + await asyncio.sleep(interval) + + +def calculate_metrics( + input_requests: List[Tuple[str, int, int]], + outputs: List[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase, + selected_percentile_metrics: List[str], + selected_percentiles: List[float], + goodput_config_dict: Dict[str, float], +) -> Tuple[BenchmarkMetrics, List[int]]: + actual_output_lens: List[int] = [] + total_input = 0 + completed = 0 + good_completed = 0 + itls: List[float] = [] + tpots: List[float] = [] + all_tpots: List[float] = [] + ttfts: List[float] = [] + e2els: List[float] = [] + for i in range(len(outputs)): + if outputs[i].success: + output_len = outputs[i].output_tokens + + if output_len is None: + # We use the tokenizer to count the number of output tokens + # for some serving backends instead of looking at + # len(outputs[i].itl) since multiple output tokens may be + # bundled together + # Note : this may inflate the output token count slightly + output_len = len( + tokenizer(outputs[i].generated_text, + add_special_tokens=False).input_ids) + actual_output_lens.append(output_len) + total_input += input_requests[i][1] + tpot = 0 + if output_len > 1: + latency_minus_ttft = outputs[i].latency - outputs[i].ttft + tpot = latency_minus_ttft / (output_len - 1) + tpots.append(tpot) + # Note: if output_len <= 1, we regard tpot as 0 for goodput + all_tpots.append(tpot) + itls += outputs[i].itl + ttfts.append(outputs[i].ttft) + e2els.append(outputs[i].latency) + completed += 1 + else: + actual_output_lens.append(0) + + if goodput_config_dict: + valid_metrics = [] + slo_values = [] + + if "ttft" in goodput_config_dict: + valid_metrics.append(ttfts) + slo_values.append(goodput_config_dict["ttft"] / + MILLISECONDS_TO_SECONDS_CONVERSION) + if "tpot" in goodput_config_dict: + valid_metrics.append(all_tpots) + slo_values.append(goodput_config_dict["tpot"] / + MILLISECONDS_TO_SECONDS_CONVERSION) + if "e2el" in goodput_config_dict: + valid_metrics.append(e2els) + slo_values.append(goodput_config_dict["e2el"] / + MILLISECONDS_TO_SECONDS_CONVERSION) + + for req_metric in zip(*valid_metrics): + is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)]) + if is_good_req: + good_completed += 1 + + if completed == 0: + warnings.warn( + "All requests failed. This is likely due to a misconfiguration " + "on the benchmark arguments.", + stacklevel=2) + metrics = BenchmarkMetrics( + completed=completed, + total_input=total_input, + total_output=sum(actual_output_lens), + request_throughput=completed / dur_s, + request_goodput=good_completed / dur_s, + output_throughput=sum(actual_output_lens) / dur_s, + total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s, + mean_ttft_ms=np.mean(ttfts or 0) * + 1000, # ttfts is empty if streaming is not supported by backend + std_ttft_ms=np.std(ttfts or 0) * 1000, + median_ttft_ms=np.median(ttfts or 0) * 1000, + percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) + for p in selected_percentiles], + mean_tpot_ms=np.mean(tpots or 0) * 1000, + std_tpot_ms=np.std(tpots or 0) * 1000, + median_tpot_ms=np.median(tpots or 0) * 1000, + percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) + for p in selected_percentiles], + mean_itl_ms=np.mean(itls or 0) * 1000, + std_itl_ms=np.std(itls or 0) * 1000, + median_itl_ms=np.median(itls or 0) * 1000, + percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) + for p in selected_percentiles], + mean_e2el_ms=np.mean(e2els or 0) * 1000, + std_e2el_ms=np.std(e2els or 0) * 1000, + median_e2el_ms=np.median(e2els or 0) * 1000, + percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000) + for p in selected_percentiles], + ) + + return metrics, actual_output_lens + + +async def benchmark( + backend: str, + api_url: str, + base_url: str, + model_id: str, + model_name: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: List[Tuple[str, int, int]], + logprobs: Optional[int], + best_of: int, + request_rate: float, + burstiness: float, + disable_tqdm: bool, + num_warmups: int, + profile: bool, + selected_percentile_metrics: List[str], + selected_percentiles: List[str], + ignore_eos: bool, + goodput_config_dict: Dict[str, float], + max_concurrency: Optional[int], + lora_modules: Optional[List[str]], +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + print("Starting initial single prompt test run...") + test_prompt, test_prompt_len, test_output_len, test_mm_content = ( + input_requests[0]) + if backend != "openai-chat" and test_mm_content is not None: + # multi-modal benchmark is only available on OpenAI Chat backend. + raise ValueError( + "Multi-modal content is only supported on 'openai-chat' backend.") + test_input = RequestFuncInput( + model=model_id, + model_name=model_name, + prompt=test_prompt, + api_url=api_url, + prompt_len=test_prompt_len, + output_len=test_output_len, + logprobs=logprobs, + best_of=best_of, + multi_modal_content=test_mm_content, + ignore_eos=ignore_eos, + ) + + if num_warmups > 0: + print(f"Warming up with {num_warmups} requests...") + warmup_pbar = None if disable_tqdm else tqdm(total=num_warmups) + warmup_semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None + + async def warmup_limited_req_fn(): + if warmup_semaphore is None: + return await request_func(request_func_input=test_input, pbar=warmup_pbar) + async with warmup_semaphore: + return await request_func(request_func_input=test_input, pbar=warmup_pbar) + + warmup_tasks = [] + for _ in range(num_warmups): + task = asyncio.create_task(warmup_limited_req_fn()) + warmup_tasks.append(task) + _ = await asyncio.gather(*warmup_tasks) + + if warmup_pbar is not None: + warmup_pbar.close() + print("Warmup completed.") + + if lora_modules: + # For each input request, choose a LoRA module at random. + lora_modules = iter( + [random.choice(lora_modules) for _ in range(len(input_requests))]) + + if profile: + print("Starting profiler...") + profile_input = RequestFuncInput(model=model_id, + model_name=model_name, + prompt=test_prompt, + api_url=base_url + "/start_profile", + prompt_len=test_prompt_len, + output_len=test_output_len, + extra_body={"num_steps": 1, "merge_profiles": True, "profile_by_stage": True}, + logprobs=logprobs, + best_of=best_of, + multi_modal_content=test_mm_content, + ignore_eos=ignore_eos) + profile_output = await request_func(request_func_input=profile_input) + if profile_output.success: + print("Profiler started") + + if burstiness == 1.0: + distribution = "Poisson process" + else: + distribution = "Gamma distribution" + + print(f"Traffic request rate: {request_rate}") + print(f"Burstiness factor: {burstiness} ({distribution})") + print(f"Maximum request concurrency: {max_concurrency}") + + pbar = None if disable_tqdm else tqdm(total=len(input_requests)) + + semaphore = (asyncio.Semaphore(max_concurrency) + if max_concurrency else None) + + async def limited_request_func(request_func_input, pbar): + if semaphore is None: + return await request_func(request_func_input=request_func_input, + pbar=pbar) + async with semaphore: + return await request_func(request_func_input=request_func_input, + pbar=pbar) + + print("Starting main benchmark run...") + + benchmark_start_time = time.perf_counter() + benchmark_start_time_unix = time.time() + tasks: List[asyncio.Task] = [] + async for request in get_request(input_requests, request_rate, burstiness): + prompt, prompt_len, output_len, mm_content = request + req_model_id, req_model_name = model_id, model_name + if lora_modules: + req_lora_module = next(lora_modules) + req_model_id, req_model_name = req_lora_module, req_lora_module + + request_func_input = RequestFuncInput(model=req_model_id, + model_name=req_model_name, + prompt=prompt, + api_url=api_url, + prompt_len=prompt_len, + output_len=output_len, + logprobs=logprobs, + best_of=best_of, + multi_modal_content=mm_content, + ignore_eos=ignore_eos) + tasks.append( + asyncio.create_task( + limited_request_func(request_func_input=request_func_input, + pbar=pbar))) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + if profile: + print("Stopping profiler...") + profile_input = RequestFuncInput( + model=model_id, + prompt=test_prompt, + api_url=base_url + "/stop_profile", + prompt_len=test_prompt_len, + output_len=test_output_len, + logprobs=logprobs, + best_of=best_of, + ) + profile_output = await request_func(request_func_input=profile_input) + if profile_output.success: + print("Profiler stopped") + + if pbar is not None: + pbar.close() + + benchmark_duration = time.perf_counter() - benchmark_start_time + benchmark_end_time_unix = time.time() + + metrics, actual_output_lens = calculate_metrics( + input_requests=input_requests, + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + selected_percentile_metrics=selected_percentile_metrics, + selected_percentiles=selected_percentiles, + goodput_config_dict=goodput_config_dict, + ) + + print("{s:{c}^{n}}".format(s=' Serving Benchmark Result ', n=50, c='=')) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", + benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", + metrics.total_output)) + print("{:<40} {:<10.2f}".format("Request throughput (req/s):", + metrics.request_throughput)) + if goodput_config_dict: + print("{:<40} {:<10.2f}".format("Request goodput (req/s):", + metrics.request_goodput)) + print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):", + metrics.output_throughput)) + print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):", + metrics.total_token_throughput)) + + result = { + "duration": benchmark_duration, + "benchmark_start_time_unix": benchmark_start_time_unix, + "benchmark_end_time_unix": benchmark_end_time_unix, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "request_throughput": metrics.request_throughput, + "request_goodput:": + metrics.request_goodput if goodput_config_dict else None, + "output_throughput": metrics.output_throughput, + "total_token_throughput": metrics.total_token_throughput, + "input_lens": [output.prompt_len for output in outputs], + "output_lens": actual_output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + } + + def process_one_metric( + # E.g., "ttft" + metric_attribute_name: str, + # E.g., "TTFT" + metric_name: str, + # E.g., "Time to First Token" + metric_header: str, + ): + # This function prints and adds statistics of the specified + # metric. + if metric_attribute_name not in selected_percentile_metrics: + return + print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-')) + print("{:<40} {:<10.2f}".format( + f"Mean {metric_name} (ms):", + getattr(metrics, f"mean_{metric_attribute_name}_ms"))) + print("{:<40} {:<10.2f}".format( + f"Median {metric_name} (ms):", + getattr(metrics, f"median_{metric_attribute_name}_ms"))) + result[f"mean_{metric_attribute_name}_ms"] = getattr( + metrics, f"mean_{metric_attribute_name}_ms") + result[f"median_{metric_attribute_name}_ms"] = getattr( + metrics, f"median_{metric_attribute_name}_ms") + result[f"std_{metric_attribute_name}_ms"] = getattr( + metrics, f"std_{metric_attribute_name}_ms") + for p, value in getattr(metrics, + f"percentiles_{metric_attribute_name}_ms"): + p_word = str(int(p)) if int(p) == p else str(p) + print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):", + value)) + result[f"p{p_word}_{metric_attribute_name}_ms"] = value + + process_one_metric("ttft", "TTFT", "Time to First Token") + process_one_metric("tpot", "TPOT", + "Time per Output Token (excl. 1st token)") + process_one_metric("itl", "ITL", "Inter-token Latency") + process_one_metric("e2el", "E2EL", "End-to-end Latency") + + print("=" * 50) + + return result + + +def check_goodput_args(args): + # Check and parse goodput arguments + goodput_config_dict = {} + VALID_NAMES = ["ttft", "tpot", "e2el"] + if args.goodput: + goodput_config_dict = parse_goodput(args.goodput) + for slo_name, slo_val in goodput_config_dict.items(): + if slo_name not in VALID_NAMES: + raise ValueError( + f"Invalid metric name found, {slo_name}: {slo_val}. " + "The service level objective name should be one of " + f"{str(VALID_NAMES)}. ") + if slo_val < 0: + raise ValueError( + f"Invalid value found, {slo_name}: {slo_val}. " + "The service level objective value should be " + "non-negative.") + return goodput_config_dict + + +def parse_goodput(slo_pairs): + goodput_config_dict = {} + try: + for slo_pair in slo_pairs: + slo_name, slo_val = slo_pair.split(":") + goodput_config_dict[slo_name] = float(slo_val) + except ValueError as err: + raise argparse.ArgumentTypeError( + "Invalid format found for service level objectives. " + "Specify service level objectives for goodput as \"KEY:VALUE\" " + "pairs, where the key is a metric name, and the value is a " + "number in milliseconds.") from err + return goodput_config_dict + + +def save_to_pytorch_benchmark_format(args: argparse.Namespace, + results: Dict[str, Any], + file_name: str) -> None: + metrics = [ + "median_ttft_ms", "mean_ttft_ms", "std_ttft_ms", "p99_ttft_ms", + "mean_tpot_ms", "median_tpot_ms", "std_tpot_ms", "p99_tpot_ms", + "median_itl_ms", "mean_itl_ms", "std_itl_ms", "p99_itl_ms" + ] + # These raw data might be useful, but they are rather big. They can be added + # later if needed + ignored_metrics = ["ttfts", "itls", "generated_texts", "errors"] + pt_records = convert_to_pytorch_benchmark_format( + args=args, + metrics={k: [results[k]] + for k in metrics}, + extra_info={ + k: results[k] + for k in results if k not in metrics and k not in ignored_metrics + }) + if pt_records: + # Don't use json suffix here as we don't want CI to pick it up + pt_file = f"{os.path.splitext(file_name)[0]}.pytorch.json" + with open(pt_file, "w") as f: + json.dump(pt_records, f) + + +def main(args: argparse.Namespace): + print(args) + random.seed(args.seed) + np.random.seed(args.seed) + + backend = args.backend + model_id = args.model + model_name = args.served_model_name + tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model + tokenizer_mode = args.tokenizer_mode + + if args.base_url is not None: + api_url = f"{args.base_url}{args.endpoint}" + base_url = f"{args.base_url}" + else: + api_url = f"http://{args.host}:{args.port}{args.endpoint}" + base_url = f"http://{args.host}:{args.port}" + + tokenizer = _load_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=args.trust_remote_code, + ) + + + if args.dataset_name == "random": + input_requests = sample_random_requests( + prefix_len=args.random_prefix_len, + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + use_chat_template=args.use_chat_template, + dsv4=args.dsv4, + tokenizer_id=tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=args.trust_remote_code, + num_workers=args.random_num_workers, + ) + + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + + goodput_config_dict = check_goodput_args(args) + + # Avoid GC processing "static" data - reduce pause times. + gc.collect() + gc.freeze() + + benchmark_result = asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + base_url=base_url, + model_id=model_id, + model_name=model_name, + tokenizer=tokenizer, + input_requests=input_requests, + logprobs=args.logprobs, + best_of=args.best_of, + request_rate=args.request_rate, + burstiness=args.burstiness, + disable_tqdm=args.disable_tqdm, + num_warmups=args.num_warmups, + profile=args.profile, + selected_percentile_metrics=args.percentile_metrics.split(","), + selected_percentiles=[ + float(p) for p in args.metric_percentiles.split(",") + ], + ignore_eos=args.ignore_eos, + goodput_config_dict=goodput_config_dict, + max_concurrency=args.max_concurrency, + lora_modules=args.lora_modules, + )) + + # Save config and results to json + if args.save_result: + result_json: Dict[str, Any] = {} + + # Setup + current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") + result_json["date"] = current_dt + result_json["backend"] = backend + result_json["model_id"] = model_id + result_json["tokenizer_id"] = tokenizer_id + result_json["best_of"] = args.best_of + result_json["num_prompts"] = args.num_prompts + + # Metadata + if args.metadata: + for item in args.metadata: + if "=" in item: + kvstring = item.split("=") + result_json[kvstring[0].strip()] = kvstring[1].strip() + else: + raise ValueError( + "Invalid metadata format. Please use KEY=VALUE format." + ) + + # Traffic + result_json["request_rate"] = (args.request_rate if args.request_rate + < float("inf") else "inf") + result_json["burstiness"] = args.burstiness + result_json["max_concurrency"] = args.max_concurrency + + # Merge with benchmark result + result_json = {**result_json, **benchmark_result} + + if not args.save_detailed: + # Remove fields with too many data points + for field in [ + "ttfts", + "itls", + "generated_texts", + "errors", + ]: + if field in result_json: + del result_json[field] + if field in benchmark_result: + del benchmark_result[field] + + # Save to file + base_model_id = model_id.split("/")[-1] + max_concurrency_str = (f"-concurrency{args.max_concurrency}" + if args.max_concurrency is not None else "") + file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json" #noqa + if args.result_filename: + file_name = args.result_filename + if args.result_dir: + file_name = os.path.join(args.result_dir, file_name) + with open(file_name, "w", encoding='utf-8') as outfile: + json.dump(result_json, outfile) + save_to_pytorch_benchmark_format(args, result_json, file_name) + + max_failure_rate = 0.05 + completed = benchmark_result["completed"] + failure_rate = 1 - completed / args.num_prompts + if failure_rate > max_failure_rate: + raise SystemExit( + f"FAIL: request failure rate {failure_rate:.1%} exceeds " + f"{max_failure_rate:.0%} threshold " + f"({completed}/{args.num_prompts} completed)" + ) + + +if __name__ == "__main__": + parser = FlexibleArgumentParser( + description="Benchmark the online serving throughput.") + parser.add_argument( + "--backend", + type=str, + default="vllm", + choices=list(ASYNC_REQUEST_FUNCS.keys()), + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + # Use 127.0.0.1 here instead of localhost to force the use of ipv4 + parser.add_argument("--host", type=str, default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument( + "--endpoint", + type=str, + default="/v1/completions", + help="API endpoint.", + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=["random"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument("--dataset-path", + type=str, + default=None, + help="Path to the sharegpt/sonnet dataset. " + "Or the huggingface dataset ID if using HF dataset.") + parser.add_argument( + "--max-concurrency", + type=int, + default=None, + help="Maximum number of concurrent requests. This can be used " + "to help simulate an environment where a higher level component " + "is enforcing a maximum number of concurrent requests. While the " + "--request-rate argument controls the rate at which requests are " + "initiated, this argument will control how many are actually allowed " + "to execute at a time. This means that when used in combination, the " + "actual request rate may be lower than specified with --request-rate, " + "if the server is not processing requests fast enough to keep up.") + + parser.add_argument( + "--model", + type=str, + required=True, + help="Name of the model.", + ) + parser.add_argument( + "--tokenizer", + type=str, + help= + "Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501 + ) + parser.add_argument( + "--best-of", + type=int, + default=1, + help="Generates `best_of` sequences per prompt and " + "returns the best one.", + ) + parser.add_argument("--use-beam-search", action="store_true") + parser.add_argument( + "--num-prompts", + type=int, + default=1000, + help="Number of prompts to process.", + ) + parser.add_argument( + "--logprobs", + type=int, + default=None, + help=("Number of logprobs-per-token to compute & return as part of " + "the request. If unspecified, then either (1) if beam search " + "is disabled, no logprobs are computed & a single dummy " + "logprob is returned for each token; or (2) if beam search " + "is enabled 1 logprob per token is computed"), + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, " + "then all the requests are sent at time 0. " + "Otherwise, we use Poisson process or gamma distribution " + "to synthesize the request arrival times.", + ) + parser.add_argument( + "--burstiness", + type=float, + default=1.0, + help="Burstiness factor of the request generation. " + "Only take effect when request_rate is not inf. " + "Default value is 1, which follows Poisson process. " + "Otherwise, the request intervals follow a gamma distribution. " + "A lower burstiness value (0 < burstiness < 1) results in more " + "bursty requests. A higher burstiness value (burstiness > 1) " + "results in a more uniform arrival of requests.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust remote code from huggingface", + ) + parser.add_argument( + "--disable-tqdm", + action="store_true", + help="Specify to disable tqdm progress bar.", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "VLLM_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--save-result", + action="store_true", + help="Specify to save benchmark results to a json file", + ) + parser.add_argument( + "--save-detailed", + action="store_true", + default=False, + help="When saving results, include detailed per-request data " + "(input_lens, output_lens, ttfts, itls, generated_texts, errors). " + "By default, only aggregated metrics are saved to reduce file size.", + ) + parser.add_argument( + "--metadata", + metavar="KEY=VALUE", + nargs="*", + help="Key-value pairs (e.g, --metadata version=0.3.3 tp=1) " + "for metadata of this run to be saved in the result JSON file " + "for record keeping purposes.", + ) + parser.add_argument( + "--result-dir", + type=str, + default=None, + help="Specify directory to save benchmark json results." + "If not specified, results are saved in the current directory.", + ) + parser.add_argument( + "--result-filename", + type=str, + default=None, + help="Specify the filename to save benchmark json results." + "If not specified, results will be saved in " + "{backend}-{args.request_rate}qps-{base_model_id}-{current_dt}.json" + " format.", + ) + parser.add_argument( + "--ignore-eos", + action="store_true", + help="Set ignore_eos flag when sending the benchmark request." + "Warning: ignore_eos is not supported in deepspeed_mii and tgi.") + parser.add_argument( + "--percentile-metrics", + type=str, + default="ttft,tpot,itl,e2el", + help="Comma-seperated list of selected metrics to report percentils. " + "This argument specifies the metrics to report percentiles. " + "Allowed metric names are \"ttft\", \"tpot\", \"itl\", \"e2el\". " + "Default value is \"ttft,tpot,itl,e2el\".") + parser.add_argument( + "--metric-percentiles", + type=str, + default="90,99,99.9", + help="Comma-seperated list of percentiles for selected metrics. " + "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". " + "Default value is \"90,99,99.9\". " + "Use \"--percentile-metrics\" to select metrics.", + ) + parser.add_argument( + "--goodput", + nargs="+", + required=False, + help="Specify service level objectives for goodput as \"KEY:VALUE\" " + "pairs, where the key is a metric name, and the value is in " + "milliseconds. Multiple \"KEY:VALUE\" pairs can be provided, " + "separated by spaces. Allowed request level metric names are " + "\"ttft\", \"tpot\", \"e2el\". For more context on the definition of " + "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 " + "and the blog: https://hao-ai-lab.github.io/blogs/distserve") + + # group for dataset specific arguments + sonnet_group = parser.add_argument_group("sonnet dataset options") + sonnet_group.add_argument( + "--sonnet-input-len", + type=int, + default=550, + help= + "Number of input tokens per request, used only for sonnet dataset.", + ) + sonnet_group.add_argument( + "--sonnet-output-len", + type=int, + default=150, + help= + "Number of output tokens per request, used only for sonnet dataset.", + ) + sonnet_group.add_argument( + "--sonnet-prefix-len", + type=int, + default=200, + help= + "Number of prefix tokens per request, used only for sonnet dataset.", + ) + + sharegpt_group = parser.add_argument_group("sharegpt dataset options") + sharegpt_group.add_argument( + "--sharegpt-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output length " + "from the ShareGPT dataset.") + + random_group = parser.add_argument_group("random dataset options") + random_group.add_argument( + "--random-input-len", + type=int, + default=1024, + help= + "Number of input tokens per request, used only for random sampling.", + ) + random_group.add_argument( + "--random-output-len", + type=int, + default=128, + help= + "Number of output tokens per request, used only for random sampling.", + ) + random_group.add_argument( + "--random-range-ratio", + type=float, + default=1.0, + help="Range of sampled ratio of input/output length, " + "used only for random sampling.", + ) + random_group.add_argument( + "--random-prefix-len", + type=int, + default=0, + help="Number of fixed prefix tokens before random " + " context. The length range of context in a random " + " request is [random-prefix-len, " + " random-prefix-len + random-prefix-len * random-range-ratio).") + random_group.add_argument( + "--use-chat-template", + action="store_true", + help="Use chat template to format the prompt.", + ) + random_group.add_argument( + '--random-num-workers', + type=int, + default=0, + help="Number of worker processes for parallel random prompt generation. " + "Only used with --dataset-name random. " + "0 (default) = auto (min(cpu_count, 8)). 1 = serial (no multiprocessing).", + ) + + dsv4_group = parser.add_argument_group("DeepSeek-V4 chat template options") + dsv4_group.add_argument( + "--dsv4", + action="store_true", + help="Use the DeepSeek-V4 chat template (encoding_dsv4.py) instead of " + "the tokenizer's built-in jinja chat template. Requires " + "--use-chat-template to also be set. Applies to the random dataset.", + ) + + hf_group = parser.add_argument_group("hf dataset options") + hf_group.add_argument("--hf-subset", + type=str, + default=None, + help="Subset of the HF dataset.") + hf_group.add_argument("--hf-split", + type=str, + default=None, + help="Split of the HF dataset.") + hf_group.add_argument( + "--hf-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output lengths " + "from the sampled HF dataset.", + ) + + parser.add_argument( + '--tokenizer-mode', + type=str, + default="auto", + choices=['auto', 'slow', 'mistral', 'custom'], + help='The tokenizer mode.\n\n* "auto" will use the ' + 'fast tokenizer if available.\n* "slow" will ' + 'always use the slow tokenizer. \n* ' + '"mistral" will always use the `mistral_common` tokenizer. \n*' + '"custom" will use --tokenizer to select the preregistered tokenizer.') + + parser.add_argument("--served-model-name", + type=str, + default=None, + help="The model name used in the API. " + "If not specified, the model name will be the " + "same as the ``--model`` argument. ") + + parser.add_argument("--lora-modules", + nargs='+', + default=None, + help="A subset of LoRA module names passed in when " + "launching the server. For each request, the " + "script chooses a LoRA module at random.") + + parser.add_argument('--num-warmups', type=int, default=0) + + args = parser.parse_args() + main(args) diff --git a/scripts/DeepSeekV4/utils/bench_serving/benchmark_utils.py b/scripts/DeepSeekV4/utils/bench_serving/benchmark_utils.py new file mode 100644 index 00000000..dc6d31f6 --- /dev/null +++ b/scripts/DeepSeekV4/utils/bench_serving/benchmark_utils.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +from typing import Any, Dict, List + + +def convert_to_pytorch_benchmark_format(args: argparse.Namespace, + metrics: Dict[str, List], + extra_info: Dict[str, Any]) -> List: + """ + Save the benchmark results in the format used by PyTorch OSS benchmark with + on metric per record + https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database + """ + records = [] + if not os.environ.get("SAVE_TO_PYTORCH_BENCHMARK_FORMAT", False): + return records + + for name, benchmark_values in metrics.items(): + record = { + "benchmark": { + "name": "vLLM benchmark", + "extra_info": { + "args": vars(args), + }, + }, + "model": { + "name": args.model, + }, + "metric": { + "name": name, + "benchmark_values": benchmark_values, + "extra_info": extra_info, + }, + } + records.append(record) + + return records \ No newline at end of file diff --git a/scripts/DeepSeekV4/utils/bench_serving/encoding_dsv4.py b/scripts/DeepSeekV4/utils/bench_serving/encoding_dsv4.py new file mode 100644 index 00000000..cdb68f23 --- /dev/null +++ b/scripts/DeepSeekV4/utils/bench_serving/encoding_dsv4.py @@ -0,0 +1,744 @@ +""" +DeepSeek-V4 Encoding + +A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages +with tool calling, thinking mode, and quick instruction task support. +""" + +from typing import Any, Dict, List, Union, Optional, Tuple +import copy +import json +import re + +# ============================================================ +# Special Tokens +# ============================================================ + +bos_token: str = "<｜begin▁of▁sentence｜>" +eos_token: str = "<｜end▁of▁sentence｜>" +thinking_start_token: str = "" +thinking_end_token: str = "" +dsml_token: str = "｜DSML｜" + +USER_SP_TOKEN = "<｜User｜>" +ASSISTANT_SP_TOKEN = "<｜Assistant｜>" +LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>" + +# Task special tokens for internal classification tasks +DS_TASK_SP_TOKENS = { + "action": "<｜action｜>", + "query": "<｜query｜>", + "authority": "<｜authority｜>", + "domain": "<｜domain｜>", + "title": "<｜title｜>", + "read_url": "<｜read_url｜>", +} +VALID_TASKS = set(DS_TASK_SP_TOKENS.keys()) + +# ============================================================ +# Templates +# ============================================================ + +system_msg_template: str = "{content}" +user_msg_template: str = "{content}" +latest_reminder_msg_template: str = "{content}" +assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token +assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}" +thinking_template: str = "{reasoning_content}" + +response_format_template: str = ( + "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" +) +tool_call_template: str = ( + "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n" +) +tool_calls_template = ( + "<{dsml_token}{tc_block_name}>\n{tool_calls}\n" +) +tool_calls_block_name: str = "tool_calls" + +tool_output_template: str = ( + "{content}" +) + +REASONING_EFFORT_MAX = ( + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" + "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n" + "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n" +) + +TOOLS_TEMPLATE = """## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following: + +<{dsml_token}tool_calls> +<{dsml_token}invoke name="$TOOL_NAME"> +<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<{dsml_token}invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response. + +Otherwise, output directly after {thinking_end_token} with tool calls or final response. + +### Available Tool Schemas + +{tool_schemas} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +""" + +# ============================================================ +# Utility Functions +# ============================================================ + +def to_json(value: Any) -> str: + """Serialize a value to JSON string.""" + try: + return json.dumps(value, ensure_ascii=False) + except: + return json.dumps(value, ensure_ascii=True) + + +def tools_from_openai_format(tools): + """Extract function definitions from OpenAI-format tool list.""" + return [tool["function"] for tool in tools] + + +def tool_calls_from_openai_format(tool_calls): + """Convert OpenAI-format tool calls to internal format.""" + return [ + { + "name": tool_call["function"]["name"], + "arguments": tool_call["function"]["arguments"], + } + for tool_call in tool_calls + ] + + +def tool_calls_to_openai_format(tool_calls): + """Convert internal tool calls to OpenAI format.""" + return [ + { + "type": "function", + "function": { + "name": tool_call["name"], + "arguments": tool_call["arguments"], + } + } + for tool_call in tool_calls + ] + + +def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str: + """ + Encode tool call arguments into DSML parameter format. + + Args: + tool_call: Dict with "name" and "arguments" (JSON string) keys. + + Returns: + DSML-formatted parameter string. + """ + p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}' + P_dsml_strs = [] + + try: + arguments = json.loads(tool_call["arguments"]) + except Exception as err: + arguments = {"arguments": tool_call["arguments"]} + + for k, v in arguments.items(): + p_dsml_str = p_dsml_template.format( + dsml_token=dsml_token, + key=k, + is_str="true" if isinstance(v, str) else "false", + value=v if isinstance(v, str) else to_json(v), + ) + P_dsml_strs.append(p_dsml_str) + + return "\n".join(P_dsml_strs) + + +def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]: + """ + Decode DSML parameters back to a tool call dict. + + Args: + tool_name: Name of the tool. + tool_args: Dict mapping param_name -> (value, is_string_flag). + + Returns: + Dict with "name" and "arguments" (JSON string) keys. + """ + def _decode_value(key: str, value: str, string: str): + if string == "true": + value = to_json(value) + return f"{to_json(key)}: {value}" + + tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}" + return dict(name=tool_name, arguments=tool_args_json) + + +def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str: + """ + Render tool schemas into the system prompt format. + + Args: + tools: List of tool schema dicts (each with name, description, parameters). + + Returns: + Formatted tools section string. + """ + tools_json = [to_json(t) for t in tools] + + return TOOLS_TEMPLATE.format( + tool_schemas="\n".join(tools_json), + dsml_token=dsml_token, + thinking_start_token=thinking_start_token, + thinking_end_token=thinking_end_token, + ) + + +def find_last_user_index(messages: List[Dict[str, Any]]) -> int: + """Find the index of the last user/developer message.""" + last_user_index = -1 + for idx in range(len(messages) - 1, -1, -1): + if messages[idx].get("role") in ["user", "developer"]: + last_user_index = idx + break + return last_user_index + + +# ============================================================ +# Message Rendering +# ============================================================ + +def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str: + """ + Render a single message at the given index into its encoded string form. + + This is the core function that converts each message in the conversation + into the DeepSeek-V4 format. + + Args: + index: Index of the message to render. + messages: Full list of messages in the conversation. + thinking_mode: Either "chat" or "thinking". + drop_thinking: Whether to drop reasoning content from earlier turns. + reasoning_effort: Optional reasoning effort level ("max", "high", or None). + + Returns: + Encoded string for this message. + """ + assert 0 <= index < len(messages) + assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`" + + prompt = "" + msg = messages[index] + last_user_idx = find_last_user_index(messages) + + role = msg.get("role") + content = msg.get("content") + tools = msg.get("tools") + response_format = msg.get("response_format") + tool_calls = msg.get("tool_calls") + reasoning_content = msg.get("reasoning_content") + wo_eos = msg.get("wo_eos", False) + + if tools: + tools = tools_from_openai_format(tools) + if tool_calls: + tool_calls = tool_calls_from_openai_format(tool_calls) + + # Reasoning effort prefix (only at index 0 in thinking mode with max effort) + assert reasoning_effort in ['max', None, 'high'], f"Invalid reasoning effort: {reasoning_effort}" + if index == 0 and thinking_mode == "thinking" and reasoning_effort == 'max': + prompt += REASONING_EFFORT_MAX + + if role == "system": + prompt += system_msg_template.format(content=content or "") + if tools: + prompt += "\n\n" + render_tools(tools) + if response_format: + prompt += "\n\n" + response_format_template.format(schema=to_json(response_format)) + + elif role == "developer": + assert content, f"Invalid message for role `{role}`: {msg}" + + content_developer = USER_SP_TOKEN + content_developer += content + + if tools: + content_developer += "\n\n" + render_tools(tools) + if response_format: + content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format)) + + prompt += user_msg_template.format(content=content_developer) + + elif role == "user": + prompt += USER_SP_TOKEN + + # Handle content blocks (tool results mixed with text) + content_blocks = msg.get("content_blocks") + if content_blocks: + parts = [] + for block in content_blocks: + block_type = block.get("type") + if block_type == "text": + parts.append(block.get("text", "")) + elif block_type == "tool_result": + tool_content = block.get("content", "") + if isinstance(tool_content, list): + text_parts = [] + for b in tool_content: + if b.get("type") == "text": + text_parts.append(b.get("text", "")) + else: + text_parts.append(f"[Unsupported {b.get('type')}]") + tool_content = "\n\n".join(text_parts) + parts.append(tool_output_template.format(content=tool_content)) + else: + parts.append(f"[Unsupported {block_type}]") + prompt += "\n\n".join(parts) + else: + prompt += content or "" + + elif role == "latest_reminder": + prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content) + + elif role == "tool": + raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()") + + elif role == "assistant": + thinking_part = "" + tc_content = "" + + if tool_calls: + tc_list = [ + tool_call_template.format( + dsml_token=dsml_token, + name=tc.get("name"), + arguments=encode_arguments_to_dsml(tc) + ) + for tc in tool_calls + ] + tc_content += '\n\n' + tool_calls_template.format( + dsml_token=dsml_token, + tool_calls="\n".join(tc_list), + tc_block_name=tool_calls_block_name, + ) + + summary_content = content or "" + rc = reasoning_content or "" + + # Check if previous message has a task - if so, this is a task output (no thinking) + prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None + + if thinking_mode == "thinking" and not prev_has_task: + if not drop_thinking or index > last_user_idx: + thinking_part = thinking_template.format(reasoning_content=rc) + thinking_end_token + else: + thinking_part = "" + + if wo_eos: + prompt += assistant_msg_wo_eos_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) + else: + prompt += assistant_msg_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) + else: + raise NotImplementedError(f"Unknown role: {role}") + + # Append transition tokens based on what follows + if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]: + return prompt + + task = messages[index].get("task") + if task is not None: + # Task special token for internal classification tasks + assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}" + task_sp_token = DS_TASK_SP_TOKENS[task] + + if task != "action": + # Non-action tasks: append task sp token directly after the message + prompt += task_sp_token + else: + # Action task: append Assistant + thinking token + action sp token + prompt += ASSISTANT_SP_TOKEN + prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token + prompt += task_sp_token + + elif messages[index].get("role") in ["user", "developer"]: + # Normal generation: append Assistant + thinking token + prompt += ASSISTANT_SP_TOKEN + if not drop_thinking and thinking_mode == "thinking": + prompt += thinking_start_token + elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx: + prompt += thinking_start_token + else: + prompt += thinking_end_token + + return prompt + + +# ============================================================ +# Preprocessing +# ============================================================ + +def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Merge tool messages into the preceding user message using content_blocks format. + + DeepSeek-V4 does not have a standalone "tool" role; instead, tool results + are encoded as blocks within user messages. + + This function converts a standard OpenAI-format conversation (with separate + "tool" role messages) into V4 format where tool results are merged into + user messages. + + Args: + messages: List of message dicts in OpenAI format. + + Returns: + Processed message list with tool messages merged into user messages. + """ + merged: List[Dict[str, Any]] = [] + + for msg in messages: + msg = copy.deepcopy(msg) + role = msg.get("role") + + if role == "tool": + # Convert tool message to a user message with tool_result block + tool_block = { + "type": "tool_result", + "tool_use_id": msg.get("tool_call_id", ""), + "content": msg.get("content", ""), + } + # Merge into previous message if it's already a user (merged tool) + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]: + merged[-1]["content_blocks"].append(tool_block) + else: + merged.append({ + "role": "user", + "content_blocks": [tool_block], + }) + elif role == "user": + text_block = {"type": "text", "text": msg.get("content", "")} + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None: + merged[-1]["content_blocks"].append(text_block) + else: + new_msg = { + "role": "user", + "content": msg.get("content", ""), + "content_blocks": [text_block], + } + # Preserve extra fields (task, wo_eos, mask, etc.) + for key in ("task", "wo_eos", "mask"): + if key in msg: + new_msg[key] = msg[key] + merged.append(new_msg) + else: + merged.append(msg) + + return merged + + +def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Sort tool_result blocks within user messages by the order of tool_calls + in the preceding assistant message. + + Args: + messages: Preprocessed message list (after merge_tool_messages). + + Returns: + Message list with sorted tool result blocks. + """ + last_tool_call_order: Dict[str, int] = {} + + for msg in messages: + role = msg.get("role") + if role == "assistant" and msg.get("tool_calls"): + last_tool_call_order = {} + for idx, tc in enumerate(msg["tool_calls"]): + tc_id = tc.get("id") or tc.get("function", {}).get("id", "") + if tc_id: + last_tool_call_order[tc_id] = idx + + elif role == "user" and msg.get("content_blocks"): + tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"] + if len(tool_blocks) > 1 and last_tool_call_order: + sorted_blocks = sorted( + tool_blocks, + key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0) + ) + sorted_idx = 0 + new_blocks = [] + for block in msg["content_blocks"]: + if block.get("type") == "tool_result": + new_blocks.append(sorted_blocks[sorted_idx]) + sorted_idx += 1 + else: + new_blocks.append(block) + msg["content_blocks"] = new_blocks + + return messages + + +# ============================================================ +# Main Encoding Function +# ============================================================ + +def encode_messages( + messages: List[Dict[str, Any]], + thinking_mode: str, + context: Optional[List[Dict[str, Any]]] = None, + drop_thinking: bool = True, + add_default_bos_token: bool = True, + reasoning_effort: Optional[str] = None, +) -> str: + """ + Encode a list of messages into the DeepSeek-V4 prompt format. + + This is the main entry point for encoding conversations. It handles: + - BOS token insertion + - Thinking mode with optional reasoning content dropping + - Tool message merging into user messages + - Multi-turn conversation context + + Args: + messages: List of message dicts to encode. + thinking_mode: Either "chat" or "thinking". + context: Optional preceding context messages (already encoded prefix). + drop_thinking: If True, drop reasoning_content from earlier assistant turns + (only keep reasoning for messages after the last user message). + add_default_bos_token: Whether to prepend BOS token at conversation start. + reasoning_effort: Optional reasoning effort level ("max", "high", or None). + + Returns: + The encoded prompt string. + """ + context = context if context else [] + + # Preprocess: merge tool messages and sort tool results + messages = merge_tool_messages(messages) + messages = sort_tool_results_by_call_order(context + messages)[len(context):] + if context: + context = merge_tool_messages(context) + context = sort_tool_results_by_call_order(context) + + full_messages = context + messages + + prompt = bos_token if add_default_bos_token and len(context) == 0 else "" + + # Resolve drop_thinking: if any message has tools defined, don't drop thinking + effective_drop_thinking = drop_thinking + if any(m.get("tools") for m in full_messages): + effective_drop_thinking = False + + if thinking_mode == "thinking" and effective_drop_thinking: + full_messages = _drop_thinking_messages(full_messages) + # After dropping, recalculate how many messages to render + # (context may have shrunk too) + num_to_render = len(full_messages) - len(_drop_thinking_messages(context)) + context_len = len(full_messages) - num_to_render + else: + num_to_render = len(messages) + context_len = len(context) + + for idx in range(num_to_render): + prompt += render_message( + idx + context_len, + full_messages, + thinking_mode=thinking_mode, + drop_thinking=effective_drop_thinking, + reasoning_effort=reasoning_effort, + ) + + return prompt + + +def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Drop reasoning_content and non-essential messages before the last user message. + + Behavior: + - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept. + - Messages at or after the last user index are always kept. + - Assistant messages before the last user get reasoning_content removed. + - Developer messages before the last user are dropped entirely. + """ + last_user_idx = find_last_user_index(messages) + result = [] + keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"} + + for idx, msg in enumerate(messages): + role = msg.get("role") + if role in keep_roles or idx >= last_user_idx: + result.append(msg) + elif role == "assistant": + msg = copy.copy(msg) + msg.pop("reasoning_content", None) + result.append(msg) + # developer and other roles before last_user_idx are dropped + + return result + + +# ============================================================ +# Parsing (Decoding model output) +# ============================================================ + +def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]: + """ + Read text from index until one of the stop strings is found. + + Returns: + Tuple of (new_index, content_before_stop, matched_stop_string_or_None). + """ + min_pos = len(text) + matched_stop = None + + for s in stop: + pos = text.find(s, index) + if pos != -1 and pos < min_pos: + min_pos = pos + matched_stop = s + + if matched_stop: + content = text[index:min_pos] + return min_pos + len(matched_stop), content, matched_stop + else: + content = text[index:] + return len(text), content, None + + +def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]: + """ + Parse DSML tool calls from text starting at the given index. + + Args: + index: Starting position in text. + text: The full text to parse. + + Returns: + Tuple of (new_index, last_stop_token, list_of_tool_call_dicts). + Each tool call dict has "name" and "arguments" keys. + """ + tool_calls: List[Dict[str, Any]] = [] + stop_token = None + tool_calls_end_token = f"" + + while index < len(text): + index, _, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token]) + if _ != ">\n": + raise ValueError(f"Tool call format error: expected '>\\n' but got '{_}'") + + if stop_token == tool_calls_end_token: + break + + if stop_token is None: + raise ValueError("Missing special token in tool calls") + + index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n$', tool_name_content, flags=re.DOTALL) + if len(p_tool_name) != 1: + raise ValueError(f"Tool name format error: '{tool_name_content}'") + tool_name = p_tool_name[0] + + tool_args: Dict[str, Tuple[str, str]] = {} + while stop_token == f"<{dsml_token}parameter": + index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"]) + + param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL) + if len(param_kv) != 1: + raise ValueError(f"Parameter format error: '{param_content}'") + param_name, string, param_value = param_kv[0] + + if param_name in tool_args: + raise ValueError(f"Duplicate parameter name: '{param_name}'") + tool_args[param_name] = (param_value, string) + + index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n": + raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'") + + tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args) + tool_calls.append(tool_call) + + return index, stop_token, tool_calls + + +def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]: + """ + Parse a model completion text into a structured assistant message. + + This function takes the raw text output from the model (a single assistant turn) + and extracts: + - reasoning_content (thinking block) + - content (summary/response) + - tool_calls (if any) + + NOTE: This function is designed to parse only correctly formatted strings and + will raise ValueError for malformed output. + + Args: + text: The raw completion text (including EOS token). + thinking_mode: Either "chat" or "thinking". + + Returns: + Dict with keys: "role", "content", "reasoning_content", "tool_calls". + tool_calls are in OpenAI format. + """ + summary_content, reasoning_content, tool_calls = "", "", [] + index, stop_token = 0, None + tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}" + + is_thinking = thinking_mode == "thinking" + is_tool_calling = False + + if is_thinking: + index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token]) + reasoning_content = content_delta + assert stop_token == thinking_end_token, "Invalid thinking format: missing " + + index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token]) + summary_content = content_delta + if stop_token == tool_calls_start_token: + is_tool_calling = True + else: + assert stop_token == eos_token, "Invalid format: missing EOS token" + + if is_tool_calling: + index, stop_token, tool_calls = parse_tool_calls(index, text) + + index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token]) + assert not tool_ends_text, "Unexpected content after tool calls" + + assert len(text) == index and stop_token in [eos_token, None], "Unexpected content at end" + + for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]: + assert sp_token not in summary_content and sp_token not in reasoning_content, \ + f"Unexpected special token '{sp_token}' in content" + + return { + "role": "assistant", + "content": summary_content, + "reasoning_content": reasoning_content, + "tool_calls": tool_calls_to_openai_format(tool_calls) + } From 41070cb7ae9394faebad39eaf0e4ece2f98436d8 Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Wed, 8 Jul 2026 18:43:37 -0500 Subject: [PATCH 02/12] Add Qwen3.5-35B-A3B aggregated inference recipe (scripts/Qwen3.5) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Single-node aggregated serving/benchmark harness for Qwen3.5-35B-A3B (3B-active hybrid-GDN MoE) on AMD Instinct MI350 (gfx950), engines SGLang + vLLM. Recommended default: SGLang + MXFP4 + TP1 (8 single-GPU instances/node) — a 3B-active MoE fits one GPU, so fill the node with independent instances rather than tensor-parallel. Findings: TP1 > TP2 > TP4; SGLang > vLLM; MXFP4 > FP8 (and ~half the HBM). Both engines serve MXFP4 correctly (no garbling). ATOM excluded — crashes at ISL>=8192. Code/scripts + README only (no result data). Engine flags baked into model.yaml per AMD Day-0 recipe (SGLang --disable-radix-cache required for the hybrid-GDN arch; vLLM --enable-expert-parallel --reasoning-parser qwen3). cluster.yaml uses placeholders. Co-Authored-By: Claude --- README.md | 1 + scripts/Qwen3.5/README.md | 113 ++ scripts/Qwen3.5/cluster.yaml | 29 + scripts/Qwen3.5/docs/TP8_invalid_for_fp8.md | 60 + scripts/Qwen3.5/docs/native_DP_broken_fp8.md | 28 + scripts/Qwen3.5/model.yaml | 79 + scripts/Qwen3.5/patches/instantiator.py | 145 ++ scripts/Qwen3.5/run_atom.sh | 5 + scripts/Qwen3.5/run_sglang.sh | 5 + scripts/Qwen3.5/run_vllm.sh | 5 + scripts/Qwen3.5/sweep/sweep_config.sh | 100 ++ .../utils/bench_serving/KNOWN_LIMITATION.md | 7 + .../bench_serving/backend_request_func.py | 574 ++++++++ .../utils/bench_serving/benchmark_serving.py | 1276 +++++++++++++++++ .../utils/bench_serving/benchmark_utils.py | 39 + .../utils/bench_serving/encoding_dsv4.py | 744 ++++++++++ 16 files changed, 3210 insertions(+) create mode 100644 scripts/Qwen3.5/README.md create mode 100644 scripts/Qwen3.5/cluster.yaml create mode 100644 scripts/Qwen3.5/docs/TP8_invalid_for_fp8.md create mode 100644 scripts/Qwen3.5/docs/native_DP_broken_fp8.md create mode 100644 scripts/Qwen3.5/model.yaml create mode 100644 scripts/Qwen3.5/patches/instantiator.py create mode 100755 scripts/Qwen3.5/run_atom.sh create mode 100755 scripts/Qwen3.5/run_sglang.sh create mode 100755 scripts/Qwen3.5/run_vllm.sh create mode 100755 scripts/Qwen3.5/sweep/sweep_config.sh create mode 100644 scripts/Qwen3.5/utils/bench_serving/KNOWN_LIMITATION.md create mode 100644 scripts/Qwen3.5/utils/bench_serving/backend_request_func.py create mode 100644 scripts/Qwen3.5/utils/bench_serving/benchmark_serving.py create mode 100644 scripts/Qwen3.5/utils/bench_serving/benchmark_utils.py create mode 100644 scripts/Qwen3.5/utils/bench_serving/encoding_dsv4.py diff --git a/README.md b/README.md index e99868b3..da9004f1 100644 --- a/README.md +++ b/README.md @@ -25,6 +25,7 @@ Below are blueprints of supported models along with their documentation. | [**SGLang disaggregated P/D inference**](scripts/sglang_disagg/README.MD) | Distributed Inference P/D disaggregation with SGLang (MoRI IO, Mooncake) | Llama-3.1-8B, Qwen3-32B, Llama-3.3-70B-FP8, Llama-3.1-405B-FP8, Mixtral-8x7B, DeepSeek-V3, DeepSeek-R1 | | [**SGLang disaggregated P/D inference with WideEP/LargeEP**](scripts/sglang_disagg/README.MD) | Distributed Inference P/D disaggregation with SGLang with WideEP/LargeEP | DeepSeek-V3, DeepSeek-R1 | | [**DeepSeek-V4-Pro disaggregated P/D inference**](scripts/DeepSeekV4/README.md) | Distributed Inference P/D disaggregation for DeepSeek-V4-Pro (ATOM+mooncake / SGLang+MoRI) | DeepSeek-V4-Pro | +| [**Qwen3.5-35B-A3B aggregated inference**](scripts/Qwen3.5/README.md) | Single-node aggregated serving/benchmark (SGLang/vLLM); recommended default SGLang+MXFP4+TP1 (8 instances/node) | Qwen3.5-35B-A3B (MXFP4, FP8) | | [**KVCache Transfer Bench**](scripts/kvcache_transfer_bench/README.md) | Inter-node Transfer Benchmark | no specific models | ## Table of Contents diff --git a/scripts/Qwen3.5/README.md b/scripts/Qwen3.5/README.md new file mode 100644 index 00000000..f15020ba --- /dev/null +++ b/scripts/Qwen3.5/README.md @@ -0,0 +1,113 @@ +# Qwen3.5-35B-A3B serving/benchmark harness — AMD Instinct MI350 (gfx950) + +Aggregated serving (3B-active MoE fits one GPU — **no disagg**). Scales 1→N nodes by editing one +file. Core knob: **(TP, DP)** replicas per node. + +## ⭐ RECOMMENDED DEFAULT — SGLang · MXFP4 · TP1 (8 instances/node) + +> **Serve Qwen3.5-35B-A3B as 8 independent single-GPU MXFP4 instances per node, on SGLang.** +> Benchmarking on MI350 (gfx950) found this the highest-throughput setup: +> **SGLang > vLLM · MXFP4 > FP8 (and ~½ the HBM) · TP1 > TP2 > TP4**. + +```bash +cd Qwen3.5 +# production serving — 8 single-GPU MXFP4 instances fill the node: +MODEL=qwen35-moe-mxfp4 TP=1 DP=8 ACTION=serve ./run_sglang.sh +# reproduce the benchmark sweep: +MODEL=qwen35-moe-mxfp4 TP=1 ACTION=sweep ./run_sglang.sh +``` +Why TP1 and not TP-shard: a 3B-active model fits one GPU, so tensor-parallel only adds cross-GPU +comm cost. Fill the node with independent instances instead. (ATOM excluded — crashes at ISL≥8192.) + +## Design in one picture +``` +Each node (8 GPUs): DP replicas, each TP-sharded. DP × TP ≤ 8. Full util ⇔ DP×TP = 8. + TP1→DP8 (⭐ best) TP2→DP4 TP4→DP2 +N nodes = N × (8//TP) independent replicas. No inter-node TP. No disagg. +Scale 1→N nodes = add node names to cluster.yaml. Nothing else changes. +``` + +## Config layering (set a knob ONCE, applied to all engines) +- `defaults:` block in `model.yaml` holds shared logical knobs (`max_model_len`, `gpu_memory_util`, + `max_num_seqs`). `lib/replica_entry.sh` translates each to the per-engine flag name + (vLLM/ATOM `--gpu-memory-utilization`, SGLang `--mem-fraction-static`; etc.). +- engine `serve_flags` hold ONLY engine-unique flags. Per-model override > global defaults. + +## Files +| File | Role | +|------|------| +| `cluster.yaml` | **SCALE LEVER** — node list, gpus_per_node, reservation, paths, port base | +| `model.yaml` | model defs + shared `defaults{}` + per-engine image/flags/fixes | +| `run_vllm.sh` / `run_sglang.sh` / `run_atom.sh` | thin wrappers → `lib/run_engine.sh` | +| `lib/run_engine.sh` | **single source of truth** launcher: yaml → clean GPUs → place → launch. Two bench paths (see below) | +| `lib/clean_node.sh` | pre-flight GPU clean-state guard (kill ALL containers incl. zombie VRAM, verify free) | +| `lib/placement.py` | (TP,DP,nodes) → one line per replica (node, gpus, port, idx); `--nodes` override | +| `lib/replica_entry.sh` | runs INSIDE each container: start server (clean env) → health-wait → sanity/serve/sweep | +| `lib/check_accuracy.py` + `lib/prompts.json` | 5-prompt accuracy gate | +| `lib/cfg.py` | yaml reader (cluster/model/engine/default lookups) | +| `lib/lib_inferencex.sh` | vendored InferenceX helpers (health-wait, gpu monitor, bench) | +| `utils/bench_serving/` | vendored InferenceX benchmark client (Apache-2.0, see NOTICE) | +| `sweep/sweep_config.sh` | **SUPERSEDED** — one-node ATOM reproducer only; logic folded into `run_engine.sh` | +| `patches/instantiator.py` | torch patch (kept for reference; NOT needed — see ATOM note) | + +## Two benchmark paths in `run_engine.sh` (engine-dependent) +- **In-container** (vLLM, SGLang — `bench_external: false`): the bench client runs inside the engine + container via `replica_entry.sh`. One backgrounded `--overlap` srun per replica; replicas run concurrently. +- **External / sibling** (ATOM — `bench_external: true`): ATOM's own image crashes the bench client, so the + bench runs from a clean vLLM container. Server + bench are **sibling containers inside ONE `--exclusive` + srun** per node (`launch_external_bench_node`). This avoids a deadlock the old design hit: a *second* + `--overlap` bench srun could never schedule because the server srun (`--mem=0`) already claimed all node + memory. One srun = one allocation = no contention. Server launched detached (`-d`) to avoid an empty-log + startup race on the 75 GB image (poll real `docker inspect` state, not a `docker ps --filter` first-tick). + +## Usage +Models (in `model.yaml`): `qwen35-moe-mxfp4` (⭐ recommended), `qwen35-moe-fp8` (baseline). +```bash +# RECOMMENDED — SGLang + MXFP4 + TP1 (8 instances fill the node) +MODEL=qwen35-moe-mxfp4 TP=1 DP=8 ACTION=serve ./run_sglang.sh +MODEL=qwen35-moe-mxfp4 TP=1 ACTION=sweep ./run_sglang.sh # benchmark + +# other engines / precisions for comparison +MODEL=qwen35-moe-fp8 TP=1 ACTION=sweep ./run_vllm.sh +MODEL=qwen35-moe-mxfp4 TP=2 ACTION=sweep ./run_sglang.sh + +# scale: edit cluster.yaml `nodes`, re-run same command. Target nodes: NODES=node-a,node-b ... +``` +`ACTION=sweep`: clean GPUs → boot → accuracy gate → full perf sweep (shapes × conc). +`ACTION=serve`: boot + stay up. `ACTION=sanity`: accuracy + one smoke point. + +## Findings (MI350 / gfx950) + +Benchmarked across 2 precisions × {SGLang, vLLM} × TP{1,2,4} × concurrency. Three consistent +findings drive the recommended default: +- **TP1 > TP2 > TP4** — a 3B-active MoE fits one GPU; fill the node with independent instances + rather than tensor-parallel (TP only adds cross-GPU comm cost). +- **SGLang > vLLM** — higher throughput at matched config. +- **MXFP4 > FP8** — higher throughput AND ~half the HBM (MI355X has native MXFP4 hardware). + +→ default = **SGLang · MXFP4 · TP1** (8 single-GPU instances per node). + +**Correctness:** both engines serve MXFP4 correctly out-of-box (no garbling). Qwen3.5 is a +reasoning model — allow generous `max_tokens`, or use vLLM `--reasoning-parser qwen3`. + +**Scope / caveats:** +- **ATOM excluded** — serves 35B MXFP4 at short context but crashes at ISL≥8192 + (`could not broadcast (513,) into (512,)`). Use SGLang or vLLM. +- Set `max_model_len ≥ ISL+OSL` for your workload (default 17408 covers 16384/1024). + +## Discovered FIXES (baked into model.yaml / scripts) +- **Reasoning model**: Qwen3.5 emits a "Thinking Process" before the answer. Use generous + `max_tokens` (accuracy gate uses 512), or vLLM `--reasoning-parser qwen3` to split reasoning/answer. +- **SGLang**: `--disable-radix-cache` REQUIRED (qwen3_5_moe is hybrid-GDN), else + `AssertionError: extra_buffer needs CUDA/MUSA/NPU (FLA)`. Plus `--enable-aiter-allreduce-fusion + --page-size 16` (AMD Day-0 recipe). +- **vLLM**: image ENTRYPOINT is `vllm` → override to `bash`; rolling tag `vllm/vllm-openai-rocm:nightly` + (pinned nightlies get GC'd). Flags `--enable-expert-parallel --reasoning-parser qwen3`. +- **max_model_len ≥ ISL+OSL**: set to 17408 to cover 16384/1024 (else long-context requests are + rejected as over-context → 0 throughput). +- **Weights on local scratch**: cache NFS→`/local_datasets` first (NFS cold-load adds ~7 min/boot). +- **All engines**: pre-flight `clean_node.sh` kills zombie VRAM from prior runs (the #1 OOM-at-init cause). + +## NOT in scope (by design) +- **Disagg / KV-transfer**: irrelevant — 35B fits one GPU (see the `dsv4/` harness for disagg). +- **TP-sharding as the fill mechanism**: proven suboptimal here — use TP1 × N instances instead. diff --git a/scripts/Qwen3.5/cluster.yaml b/scripts/Qwen3.5/cluster.yaml new file mode 100644 index 00000000..1da3ee03 --- /dev/null +++ b/scripts/Qwen3.5/cluster.yaml @@ -0,0 +1,29 @@ +# cluster.yaml — the SCALE LEVER. Edit `nodes` to scale 1 -> N. +# Each node is self-contained: it runs (gpus_per_node // TP) replicas, all node-local. +# Scaling is linear: total_replicas = len(nodes) * (gpus_per_node // TP). No inter-node TP. + +# Slurm allocation flag prepended to every srun. Set ONE of: +# - "--reservation=" if you hold a reservation +# - "--jobid=" to run inside an existing (e.g. salloc/interactive) allocation +# - "" to srun with no allocation (nodes must be free) +srun_alloc: "--reservation=" +reservation: # legacy fallback if srun_alloc is empty +gpus_per_node: 8 # gfx950, 8 GPUs/node +hardware: MI350 + +# Local model staging root (same path on every node; local NVMe recommended). +# Cache weights here first (NFS cold-load adds ~7 min/boot). If local scratch is not +# provisioned, point this at your shared model path instead. +models_root: /path/to/models + +# Where result JSON / logs land (shared filesystem, visible from all nodes). +results_root: ./results + +# Port base: replica i on a node binds (port_base + i). +port_base: 8000 + +# The node pool. START SMALL (1 node) to settle single-node configs, then add nodes. +# 35B-A3B fits ONE GPU: a single node runs the full TP1 x8 sweep. +nodes: + - +# - diff --git a/scripts/Qwen3.5/docs/TP8_invalid_for_fp8.md b/scripts/Qwen3.5/docs/TP8_invalid_for_fp8.md new file mode 100644 index 00000000..ce8ce12e --- /dev/null +++ b/scripts/Qwen3.5/docs/TP8_invalid_for_fp8.md @@ -0,0 +1,60 @@ +# Why pure TP8 is INVALID for Qwen3-Next-80B-A3B-Instruct-FP8 (and how EP fixes it) + +**Short version:** the FP8 checkpoint uses **block quantization with a 128×128 block size**. +Pure tensor-parallel sharding must split weight dimensions into chunks that are still divisible +by the quant block size. At TP8, the MoE expert intermediate dimension (512) splits to 64 per GPU, +which is **< 128** and **not divisible by 128** → the server fails to load. + +Valid **pure-TP** for this FP8 model: **1, 2, 4** (set as `valid_tp: [1,2,4]` in `model.yaml`). + +> ⚠️ **This limit is specific to TENSOR-sharding the experts.** With **Expert Parallelism +> (`--ep-size 8` / `--enable-expert-parallel`)** each GPU holds **whole experts** — the expert +> weight is never tensor-split, so the 128-block constraint is never violated and **TP8+EP8 is +> valid**. This is exactly the shape InferenceX and the customer run in production (`tp:8, ep:8`). +> So "TP8 invalid" means *pure* TP8, NOT TP8+EP. See `docs/native_DP_broken_fp8.md` and the +> README "EP next-step" note. EP boot on this model is still to be validated empirically here. + +--- + +## The exact error +`vllm serve ... --tensor-parallel-size 8` (or ATOM `-tp 8`) fails at weight creation with: +``` +ValueError: Weight input_size_per_partition = 64 is not divisible by weight quantization +block_k = 128. +``` +(raised in the FP8 block-quant `create_weights` path). + +## The arithmetic +From the model `config.json`: +- `moe_intermediate_size = 512` +- FP8 block quant: `weight_block_size = [128, 128]` + +Per-GPU expert intermediate size after TP sharding = `512 / TP`: + +| TP | 512 / TP | ≥128 and divisible by 128? | Valid? | +|----|----------|----------------------------|--------| +| 1 | 512 | yes (512 = 4×128) | ✅ | +| 2 | 256 | yes (256 = 2×128) | ✅ | +| 4 | 128 | yes (128 = 1×128) | ✅ | +| 8 | **64** | **no (64 < 128)** | ❌ | + +At TP8 each shard would hold a 64-wide slice of a tensor whose quant scales are defined per +128-element block — the block can't be split across two GPUs, so the loader rejects it. + +## Why this is fundamental, not a flag/config issue +- It is a property of **how the checkpoint was quantized** (128-block FP8), combined with the + model's small `moe_intermediate_size` (512). No serving flag changes the block size. +- Confirmed empirically: vLLM TP8 FP8 fails at `create_weights` with the error above + (WorkerProc failed to start). Not an OOM, not an env issue — a hard divisibility constraint. + +## What WOULD allow TP8 +- A different quantization with a **smaller block size** (e.g. MXFP4 uses `group_size = 32`: + 512/8 = 64, and 64 is divisible by 32 → TP8 *may* be valid for MXFP4). This is why + `model.yaml` lists `valid_tp: [1,2,4,8]` for the MXFP4 model (note: MXFP4 is currently PARKED + for a separate engine-correctness reason — see README). +- A model with a larger `moe_intermediate_size` (so 512/8=64 wouldn't be the limiter). + +## Practical implication for the sweep +For FP8 full-node (8-GPU) packing, the valid (TP, DP) combos that use all 8 GPUs are: +- **TP1 × DP8**, **TP2 × DP4**, **TP4 × DP2**. +- **TP8 × DP1 is NOT runnable for FP8** — omit it from FP8 sweeps. diff --git a/scripts/Qwen3.5/docs/native_DP_broken_fp8.md b/scripts/Qwen3.5/docs/native_DP_broken_fp8.md new file mode 100644 index 00000000..7c53670f --- /dev/null +++ b/scripts/Qwen3.5/docs/native_DP_broken_fp8.md @@ -0,0 +1,28 @@ +# vLLM native data-parallel (--data-parallel-size) is BROKEN for Qwen3-Next-80B FP8 + +## Finding (2026-06-28) +vLLM's native DP fails to start for the FP8 model, at ANY TP including TP1: +``` +ValueError: The output_size of gate's and up's weight = 64 is not divisible by +weight quantization block_n = 128. +``` +- `--data-parallel-size 8 --tensor-parallel-size 1` → FAILS (Worker_DP1/DP4 ...) +- `--data-parallel-size 2 --tensor-parallel-size 4` → FAILS (Worker_DP0_TP3 ...) +- Plain `--tensor-parallel-size 4` (DP1, single replica) → WORKS (validated Phase B). + +## Why +vLLM's DP path appears to apply an extra split to the MoE gate/up projection that produces a +64-wide shard, violating the FP8 128-block quant (same class as the TP8 issue, but triggered by +the DP machinery even at TP1). It is NOT just a TP-sharding effect — DP8/TP1 has no TP split yet +still hits 64. + +## Decision +Do NOT use engine-native DP for the FP8 sweep. Instead use the original v2 approach: +**N independent single-server replicas** (each a plain TP-k server on its own GPU group), launched +by run_engine.sh placement. Full-node = (8/TP) independent replicas. Aggregate is measured by +fanning the benchmark across the replica ports (or a router), NOT by engine DP. + +## Implication for "inbuilt router" +The engines DO expose DP flags, but vLLM's DP is unusable for this FP8 checkpoint. So the +manual-replica + external aggregation path is required for FP8. (SGLang/ATOM DP untested for FP8 +given vLLM's failure; assume manual replicas for uniformity.) diff --git a/scripts/Qwen3.5/model.yaml b/scripts/Qwen3.5/model.yaml new file mode 100644 index 00000000..027890a9 --- /dev/null +++ b/scripts/Qwen3.5/model.yaml @@ -0,0 +1,79 @@ +# model.yaml — model defs, shared runtime defaults, and per-engine images + UNIQUE flags. +# +# CONFIG LAYERING (so a knob is set ONCE and applied identically to all engines): +# defaults: logical knobs shared by all engines (max_model_len, gpu_memory_util, max_num_seqs) +# -> replica_entry.sh translates each to the engine's specific flag name. +# models{}: per-model identity (dir/arch/valid_tp) + optional per-model overrides. +# engines{}.serve_flags: ONLY engine-UNIQUE flags (attention backend, radix-cache, etc.). +# Shared knobs are NO LONGER here — they come from `defaults` to avoid drift. + +# ---- shared logical knobs (identical intent across vLLM / SGLang / ATOM) ------------------- +defaults: + max_model_len: 17408 # covers 16384+1024 sweep shape. (Was 8192 -> rejected all 8192/16384 + # ISL requests as over-context -> 0 throughput. Fixed 2026-07-08.) + # Context cap; same for all engines. Override per-run via MAX_MODEL_LEN + # or per-model below (e.g. long-context sweeps need up to 262144). + gpu_memory_util: 0.8 # fraction of VRAM per replica. vLLM/ATOM: --gpu-memory-utilization, + # SGLang: --mem-fraction-static (translated in replica_entry.sh). + max_num_seqs: 256 # max concurrent sequences per replica. + +models: + # ACTIVE small model: Qwen3.5-35B-A3B (arch qwen3_5_moe) — a 3B-active hybrid-GDN MoE that + # fits ONE MI350 GPU. Replaces the prior Qwen3-Next-80B. Deploy = single-GPU instance x8 per + # node (TP1 multi-replica), NOT tensor-parallel: a 3B-active model gains nothing from TP and + # pays cross-GPU comm cost. Full-node aggregate = per-replica x8. (Proven on the 80B: TP1 > + # TP2 > TP4 in every cell.) valid_tp lists 1,2,4 so the sweep confirms TP1 wins. + qwen35-moe-mxfp4: # ⭐ RECOMMENDED DEFAULT — run on SGLang, TP1, 8/node + dir: Qwen3.5-35B-A3B-MXFP4 # amd/Qwen3.5-35B-A3B-MXFP4 (AMD Quark v0.12) + arch: qwen3_5_moe # hybrid GDN + MoE (gated-delta-rule kernels) + valid_tp: [1, 2, 4] # TP1 is best (see REPORT_qwen35.md); fill node w/ 8 replicas + status: WORKING # BEST config: SGLang+MXFP4+TP1 = ~96k tok/s/node @1024/256. MXFP4 > FP8 ~12-15%, + # ~half the HBM, native MI355X FP4. Serves correctly on SGLang+vLLM (no garbling). 1024 validated; + # long-context (8192/16384) re-run pending after max_model_len fix. + + qwen35-moe-fp8: # baseline for comparison (MXFP4 beats it) + dir: Qwen3.5-35B-A3B-FP8 # Qwen/Qwen3.5-35B-A3B-FP8 (official FP8) + arch: qwen3_5_moe + valid_tp: [1, 2, 4] + status: WORKING # ~12-15% slower than MXFP4, 2x the HBM. Use only if MXFP4 unavailable. + +# ---- engines: image + env + ONLY engine-unique serve flags -------------------------------- +engines: + vllm: + # vLLM ROCm — Jun22 nightly (per user choice). Multi-arch build incl. gfx950. + image: vllm/vllm-openai-rocm:nightly + entrypoint_override: bash # image ENTRYPOINT is `vllm`; override to run scripts + env: + VLLM_ROCM_USE_AITER: "1" + SAFETENSORS_FAST_GPU: "1" + # unique only: attention backend, fp8 kv, prefix-caching off. (mem-util/max-seqs/ctx from defaults) + serve_flags: >- + --attention-backend ROCM_AITER_UNIFIED_ATTN --kv-cache-dtype fp8 --no-enable-prefix-caching + --enable-expert-parallel --reasoning-parser qwen3 + # TP1 -> --no-async-scheduling ; TP>1 -> --async-scheduling (handled in replica_entry.sh) + + sglang: + image: lmsysorg/sglang:v0.5.14-rocm720-mi35x + env: + SGLANG_USE_AITER_UNIFIED_ATTN: "1" + SGLANG_USE_AITER: "1" + # unique only: aiter attn, trust-remote-code, the qwen3_next radix-cache FIX, chunked prefill. + # FIX: --disable-radix-cache REQUIRED for qwen3_next (hybrid GDN/mamba) — else crashes + # AssertionError: extra_buffer needs CUDA/MUSA/NPU (FLA) in _validate_mamba_extra_buffer. + serve_flags: >- + --attention-backend aiter --trust-remote-code --disable-radix-cache + --enable-aiter-allreduce-fusion --page-size 16 --chunked-prefill-size 32768 + + atom: + image: rocm/atom-dev:vllm-v0.22.0-nightly_20260617 + entrypoint_override: bash + # ATOM works out of the box in a CLEAN env on a CLEAN GPU. Launch server before sourcing any + # helper lib (replica_entry handles this). Shared mem-util/max-seqs/max-model-len now come from + # `defaults` (previously ATOM was missing --max-model-len -> fell back to native 262k -> OOM). + env: {} + serve_flags: "--no-enable_prefix_caching" + # ATOM image's Python env crashes the vendored bench client (importlib cache_from_source + # IndexError). Run the perf benchmark from a CLEAN container (bench_image) against ATOM's port. + bench_external: true + bench_image: vllm/vllm-openai-rocm:nightly + status: WORKING # serves + accuracy verified; perf via external bench diff --git a/scripts/Qwen3.5/patches/instantiator.py b/scripts/Qwen3.5/patches/instantiator.py new file mode 100644 index 00000000..27a6f5e5 --- /dev/null +++ b/scripts/Qwen3.5/patches/instantiator.py @@ -0,0 +1,145 @@ +#!/usr/bin/python3 +# mypy: allow-untyped-defs +import importlib.abc +import importlib.util +import sys + +import torch +from torch.distributed.nn.jit.templates.remote_module_template import ( + get_remote_module_template, +) + + +_FILE_PREFIX = "_remote_module_" + + +def get_arg_return_types_from_interface(module_interface): + assert getattr(module_interface, "__torch_script_interface__", False), ( + "Expect a TorchScript class interface decorated by @torch.jit.interface." + ) + qualified_name = torch._jit_internal._qualified_name(module_interface) + cu = torch.jit._state._python_cu + module_interface_c = cu.get_interface(qualified_name) + assert "forward" in module_interface_c.getMethodNames(), ( + f"Expect forward in interface methods, while it has {module_interface_c.getMethodNames()}" + ) + method_schema = module_interface_c.getMethod("forward") + + arg_str_list = [] + arg_type_str_list = [] + assert method_schema is not None + for argument in method_schema.arguments: + arg_str_list.append(argument.name) + + if argument.has_default_value(): + default_value_str = f" = {argument.default_value}" + else: + default_value_str = "" + arg_type_str = f"{argument.name}: {argument.type}{default_value_str}" + arg_type_str_list.append(arg_type_str) + + arg_str_list = arg_str_list[1:] # Remove "self". + args_str = ", ".join(arg_str_list) + + arg_type_str_list = arg_type_str_list[1:] # Remove "self". + arg_types_str = ", ".join(arg_type_str_list) + + assert len(method_schema.returns) == 1 + argument = method_schema.returns[0] + return_type_str = str(argument.type) + + return args_str, arg_types_str, return_type_str + + +class _StringLoader(importlib.abc.SourceLoader): + def __init__(self, data): + self.data = data + + def get_source(self, fullname): + return self.data + + def get_data(self, path): + return self.data.encode("utf-8") + + def get_filename(self, fullname): + # PATCH: return a path WITH a .py suffix. Upstream returns bare `fullname` + # (no extension), which makes importlib.cache_from_source() do string slicing + # that throws "IndexError: string index out of range" on this CPython/torch + # combo. A proper .py-suffixed path makes cache path computation well-formed. + return fullname + ".py" + + +def _do_instantiate_remote_module_template( + generated_module_name, str_dict, enable_moving_cpu_tensors_to_cuda +): + if generated_module_name in sys.modules: + return sys.modules[generated_module_name] + + loader = _StringLoader( + get_remote_module_template(enable_moving_cpu_tensors_to_cuda).format(**str_dict) + ) + spec = importlib.util.spec_from_loader( + generated_module_name, loader, origin="torch-git" + ) + assert spec is not None + module = importlib.util.module_from_spec(spec) + sys.modules[generated_module_name] = module + loader.exec_module(module) + return module + + +def instantiate_scriptable_remote_module_template( + module_interface_cls, enable_moving_cpu_tensors_to_cuda=True +): + if not getattr(module_interface_cls, "__torch_script_interface__", False): + raise ValueError( + f"module_interface_cls {module_interface_cls} must be a type object decorated by " + "@torch.jit.interface" + ) + + # Generate the template instance name. + module_interface_cls_name = torch._jit_internal._qualified_name( + module_interface_cls + ).replace(".", "_") + generated_module_name = f"{_FILE_PREFIX}{module_interface_cls_name}" + + # Generate type annotation strs. + assign_module_interface_cls_str = ( + f"from {module_interface_cls.__module__} import " + f"{module_interface_cls.__name__} as module_interface_cls" + ) + args_str, arg_types_str, return_type_str = get_arg_return_types_from_interface( + module_interface_cls + ) + kwargs_str = "" + arrow_and_return_type_str = f" -> {return_type_str}" + arrow_and_future_return_type_str = f" -> Future[{return_type_str}]" + + str_dict = dict( + assign_module_interface_cls=assign_module_interface_cls_str, + arg_types=arg_types_str, + arrow_and_return_type=arrow_and_return_type_str, + arrow_and_future_return_type=arrow_and_future_return_type_str, + args=args_str, + kwargs=kwargs_str, + jit_script_decorator="@torch.jit.script", + ) + return _do_instantiate_remote_module_template( + generated_module_name, str_dict, enable_moving_cpu_tensors_to_cuda + ) + + +def instantiate_non_scriptable_remote_module_template(): + generated_module_name = f"{_FILE_PREFIX}non_scriptable" + str_dict = dict( + assign_module_interface_cls="module_interface_cls = None", + args="*args", + kwargs="**kwargs", + arg_types="*args, **kwargs", + arrow_and_return_type="", + arrow_and_future_return_type="", + jit_script_decorator="", + ) + # For a non-scriptable template, always enable moving CPU tensors to a cuda device, + # because there is no syntax limitation on the extra handling caused by the script. + return _do_instantiate_remote_module_template(generated_module_name, str_dict, True) diff --git a/scripts/Qwen3.5/run_atom.sh b/scripts/Qwen3.5/run_atom.sh new file mode 100755 index 00000000..4e88e56a --- /dev/null +++ b/scripts/Qwen3.5/run_atom.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# run_atom.sh — launch atom replicas across the cluster. Thin wrapper over lib/run_engine.sh. +# Usage: MODEL=qwen3-next-fp8 TP=4 [DP=2] [ACTION=sanity|serve] ./run_atom.sh +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENGINE=atom exec bash "$HERE/lib/run_engine.sh" diff --git a/scripts/Qwen3.5/run_sglang.sh b/scripts/Qwen3.5/run_sglang.sh new file mode 100755 index 00000000..426ce1a5 --- /dev/null +++ b/scripts/Qwen3.5/run_sglang.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# run_sglang.sh — launch sglang replicas across the cluster. Thin wrapper over lib/run_engine.sh. +# Usage: MODEL=qwen3-next-fp8 TP=4 [DP=2] [ACTION=sanity|serve] ./run_sglang.sh +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENGINE=sglang exec bash "$HERE/lib/run_engine.sh" diff --git a/scripts/Qwen3.5/run_vllm.sh b/scripts/Qwen3.5/run_vllm.sh new file mode 100755 index 00000000..23f222d3 --- /dev/null +++ b/scripts/Qwen3.5/run_vllm.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# run_vllm.sh — launch vllm replicas across the cluster. Thin wrapper over lib/run_engine.sh. +# Usage: MODEL=qwen3-next-fp8 TP=4 [DP=2] [ACTION=sanity|serve] ./run_vllm.sh +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENGINE=vllm exec bash "$HERE/lib/run_engine.sh" diff --git a/scripts/Qwen3.5/sweep/sweep_config.sh b/scripts/Qwen3.5/sweep/sweep_config.sh new file mode 100755 index 00000000..02233028 --- /dev/null +++ b/scripts/Qwen3.5/sweep/sweep_config.sh @@ -0,0 +1,100 @@ +#!/bin/bash +# ============================================================================================ +# SUPERSEDED (2026-06-28). The proven per-node pattern (host pre-clean via clean_node.sh, +# detached -d launch to dodge the empty-log race, and server+bench as SIBLING containers in +# ONE --exclusive srun to avoid the --overlap deadlock) is now folded into the single source +# of truth: lib/run_engine.sh (launch_external_bench_node). Prefer: +# ENGINE=atom MODEL=qwen3-next-fp8 TP=4 ACTION=sweep bash lib/run_engine.sh +# This standalone driver is kept only as a minimal one-node reproducer for the ATOM team. +# ============================================================================================ +# sweep_config.sh — run ONE (engine, TP) FP8 config on ONE node, full-node via MANUAL REPLICAS. +# Engine-native DP is BROKEN for this FP8 model (see docs/native_DP_broken_fp8.md), so we launch +# N = 8/TP INDEPENDENT single-server replicas, each plain TP-k on its own GPU group + own port. +# Replicas are identical & independent -> per-replica metrics; full-node aggregate = per-replica +# throughput x N. For the sweep we benchmark ONE representative replica (replica 0) per point +# (all replicas identical); node aggregate is computed offline as (metric x N). +# +# Args (env): ENGINE=vllm|sglang|atom TP=1|2|4 NODE= +set -uo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" # scripts_qwen_v2 +RESV= +MROOT=/path/to/models +MDIR=Qwen3-Next-80B-A3B-Instruct-FP8 +ENGINE="${ENGINE:?}"; TP="${TP:?}"; NODE="${NODE:?}" +DP=$(( 8 / TP )) # number of independent replicas to fill the node +MML=17408 # covers 16384+1024 +SHAPES="${SHAPES:-1024,1024 8192,1024 16384,1024}" +CONC="${CONC:-4 8 16 32 64 128 256}" +STAMP=$(python3 -c "import time;print(time.strftime('%Y%m%d_%H%M%S'))") +RES="$HERE/../results_v2/sweep_fp8_${ENGINE}_tp${TP}_x${DP}_${STAMP}"; mkdir -p "$RES"; chmod 777 "$RES" + +VLLM_IMG=vllm/vllm-openai-rocm:nightly-9037498c22891e55b594f567fb91d9b4efbf3e99 +SGL_IMG=lmsysorg/sglang:v0.5.14-rocm720-mi35x +ATOM_IMG=rocm/atom-dev:vllm-v0.22.0-nightly_20260617 +case "$ENGINE" in vllm) IMG=$VLLM_IMG;; sglang) IMG=$SGL_IMG;; atom) IMG=$ATOM_IMG;; esac +BENCH_IMG=$VLLM_IMG # clean bench client for ALL engines (uniform; sidesteps ATOM bench crash) + +# GPUs for replica 0 = first TP GPUs; we benchmark replica 0 (all replicas identical). +GPUS0=$(seq -s, 0 $((TP-1))) +PORT=8000 +sched="--no-async-scheduling"; [[ $TP -gt 1 ]] && sched="--async-scheduling" +# Knobs aligned with the run_engine.sh vLLM/SGLang sweep for apples-to-apples: +# gpu_mem_util=0.8, max_num_seqs=256, max_model_len=17408 (=$MML). +case "$ENGINE" in + vllm) SRVCMD="vllm serve $MROOT/$MDIR --host 0.0.0.0 --port $PORT --tensor-parallel-size $TP --max-model-len $MML --gpu-memory-utilization 0.8 --max-num-seqs 256 --no-enable-prefix-caching $sched --attention-backend ROCM_AITER_UNIFIED_ATTN --kv-cache-dtype fp8" ;; + sglang) SRVCMD="python3 -m sglang.launch_server --model-path $MROOT/$MDIR --host 0.0.0.0 --port $PORT --tp-size $TP --context-length $MML --mem-fraction-static 0.8 --max-running-requests 256 --trust-remote-code --disable-radix-cache --attention-backend aiter" ;; + atom) SRVCMD="python -m atom.entrypoints.openai_server --model $MROOT/$MDIR -tp $TP --port $PORT --max-model-len $MML --gpu-memory-utilization 0.8 --max-num-seqs 256 --no-enable_prefix_caching" ;; +esac + +echo "[sweep] $ENGINE TP=$TP (rep0 of $DP) node=$NODE -> $RES" +echo "[sweep] NOTE: benchmarking replica 0 (1 TP-$TP server on GPUs $GPUS0); node aggregate = metric x $DP" + +# Pre-flight clean on the host (frees VRAM zombies from prior --rm containers; the #1 cause of +# ATOM OOM-at-init). Weak `docker kill` alone misses detached zombies — use the full cleaner. +echo "[sweep] pre-flight clean $NODE" +srun --reservation="$RESV" --nodelist="$NODE" --nodes=1 --ntasks=1 --overlap --mem=0 \ + bash "$HERE/lib/clean_node.sh" > "$RES/clean.log" 2>&1 || echo "[sweep] WARN clean (see clean.log)" + +srun --reservation="$RESV" --nodelist="$NODE" --nodes=1 --ntasks=1 --exclusive bash -c " + set -uo pipefail + docker rm -f sweep_${ENGINE}_tp${TP} >/dev/null 2>&1 || true # remove any lingering named container (name conflict guard) + sleep 4 + echo '[node] start $ENGINE TP=$TP server on GPUs $GPUS0' + # Launch DETACHED (-d) so we have a real container id. Logs via 'docker logs' to server.log. + # (Old approach raced: 'docker run ... &' + immediate 'docker ps --filter name' on a 75GB image + # saw no container on the first loop tick -> false 'exited' with an empty log.) + cid=\$(docker run -d --name sweep_${ENGINE}_tp${TP} --entrypoint bash \ + --device=/dev/kfd --device=/dev/dri --group-add video --ipc host --shm-size 128G --network host \ + --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e HIP_VISIBLE_DEVICES=$GPUS0 -e VLLM_ROCM_USE_AITER=1 -e SAFETENSORS_FAST_GPU=1 -e SGLANG_USE_AITER=1 \ + -v $MROOT:$MROOT \ + $IMG -c '$SRVCMD') + echo \"[node] container \$cid\" + up=0; for i in \$(seq 1 180); do + docker logs sweep_${ENGINE}_tp${TP} > $RES/server.log 2>&1 || true + curl -sf http://localhost:$PORT/health >/dev/null 2>&1 && { up=1; break; } + # only declare exit if the container is actually gone (not just slow to register) + st=\$(docker inspect -f '{{.State.Status}}' sweep_${ENGINE}_tp${TP} 2>/dev/null || echo missing) + [ \"\$st\" = exited ] || [ \"\$st\" = missing ] && { echo \"[node] server container \$st\"; break; } + sleep 5 + done + docker logs sweep_${ENGINE}_tp${TP} > $RES/server.log 2>&1 || true + echo \"[node] SERVER_UP=\$up\" + if [ \$up -ne 1 ]; then echo '[node] FAILED tail:'; tail -25 $RES/server.log; docker kill sweep_${ENGINE}_tp${TP} >/dev/null 2>&1; exit 1; fi + for shape in $SHAPES; do + ISL=\${shape%,*}; OSL=\${shape#*,} + for c in $CONC; do + echo \"[bench] isl=\$ISL osl=\$OSL conc=\$c\" + docker run --rm --entrypoint bash --network host -v $HERE:/scripts -v $MROOT:$MROOT -v $RES:/out \ + $BENCH_IMG -c \"python3 /scripts/utils/bench_serving/benchmark_serving.py \ + --model $MROOT/$MDIR --backend vllm --base-url http://0.0.0.0:$PORT --dataset-name random \ + --random-input-len \$ISL --random-output-len \$OSL --random-range-ratio 1.0 \ + --num-prompts \$(( c * 10 )) --num-warmups \$(( c * 2 )) --max-concurrency \$c \ + --request-rate inf --ignore-eos --save-result --percentile-metrics ttft,tpot,itl,e2el \ + --result-dir /out --result-filename sweep_isl\${ISL}_osl\${OSL}_c\${c}.json\" \ + > $RES/bench_isl\${ISL}_osl\${OSL}_c\${c}.log 2>&1 || echo \" [bench] FAILED isl=\$ISL conc=\$c\" + done + done + echo '[node] standdown'; docker kill sweep_${ENGINE}_tp${TP} >/dev/null 2>&1; sleep 3 +" > "$RES/run.log" 2>&1 +echo "[sweep] $ENGINE TP=$TP done -> $RES" diff --git a/scripts/Qwen3.5/utils/bench_serving/KNOWN_LIMITATION.md b/scripts/Qwen3.5/utils/bench_serving/KNOWN_LIMITATION.md new file mode 100644 index 00000000..342de691 --- /dev/null +++ b/scripts/Qwen3.5/utils/bench_serving/KNOWN_LIMITATION.md @@ -0,0 +1,7 @@ +# KNOWN_LIMITATION + +When benchmarking small models (e.g., Gemma 1B or Llama 8B) and at ultra-high querys per second (QPS) and with short input and output lengths, the current InferenceX bench_serving client becomes client-bound. + +InferenceX does not currently benchmark in this regime and has no plans to. Our roadmap skews the opposite direction, larger models, longer ISL/OSL (i.e., agentic workloads), and interactive TTFT & tok/s/user scenarios rather than ultra-high QPS. + +Should the need arise, we shall fix this known limitation & migrate to a multi-process benchmark client. If your needs are benchmarking small models at high QPS at small ISL/OSL lengths, we recommend xyz benchmark instead. diff --git a/scripts/Qwen3.5/utils/bench_serving/backend_request_func.py b/scripts/Qwen3.5/utils/bench_serving/backend_request_func.py new file mode 100644 index 00000000..4c8820f8 --- /dev/null +++ b/scripts/Qwen3.5/utils/bench_serving/backend_request_func.py @@ -0,0 +1,574 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +import sys +import time +import traceback +from dataclasses import dataclass, field +from typing import List, Optional, Union + +import aiohttp +import huggingface_hub.constants +from tqdm.asyncio import tqdm +from transformers import (AutoTokenizer, PreTrainedTokenizer, + PreTrainedTokenizerFast) + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) + + +@dataclass +class RequestFuncInput: + prompt: str + api_url: str + prompt_len: int + output_len: int + model: str + model_name: Optional[str] = None + best_of: int = 1 + logprobs: Optional[int] = None + extra_body: Optional[dict] = None + multi_modal_content: Optional[dict] = None + ignore_eos: bool = False + + +@dataclass +class RequestFuncOutput: + generated_text: str = "" + success: bool = False + latency: float = 0.0 + output_tokens: int = 0 + ttft: float = 0.0 # Time to first token + itl: List[float] = field( + default_factory=list) # List of inter-token latencies + tpot: float = 0.0 # avg next-token latencies + prompt_len: int = 0 + error: str = "" + + +async def async_request_tgi( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith("generate_stream") + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + params = { + "best_of": request_func_input.best_of, + "max_new_tokens": request_func_input.output_len, + "do_sample": True, + "temperature": 0.01, # TGI does not accept 0.0 temperature. + "top_p": 0.99, # TGI does not accept 1.0 top_p. + "truncate": request_func_input.prompt_len, + # TGI does not accept ignore_eos flag. + } + payload = { + "inputs": request_func_input.prompt, + "parameters": params, + } + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + chunk_bytes = chunk_bytes.decode("utf-8") + + # NOTE: Sometimes TGI returns a ping response without + # any data, we should skip it. + if chunk_bytes.startswith(":"): + continue + chunk = chunk_bytes.removeprefix("data:") + + data = json.loads(chunk) + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + most_recent_timestamp = timestamp + + output.latency = most_recent_timestamp - st + output.success = True + output.generated_text = data["generated_text"] + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_trt_llm( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith("generate_stream") + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + assert request_func_input.best_of == 1 + payload = { + "accumulate_tokens": True, + "text_input": request_func_input.prompt, + "temperature": 0.0, + "top_p": 1.0, + "max_tokens": request_func_input.output_len, + "stream": True, + } + if request_func_input.ignore_eos: + payload["min_length"] = request_func_input.output_len + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = chunk_bytes.decode("utf-8").removeprefix( + "data:") + + data = json.loads(chunk) + output.generated_text += data["text_output"] + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + most_recent_timestamp = timestamp + + output.latency = most_recent_timestamp - st + output.success = True + + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_deepspeed_mii( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + assert request_func_input.best_of == 1 + + payload = { + "prompt": request_func_input.prompt, + "max_tokens": request_func_input.output_len, + "temperature": 0.01, # deepspeed-mii does not accept 0.0 temp. + "top_p": 1.0, + } + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + # NOTE: DeepSpeed-MII doesn't support streaming as of Jan 28 2024, + # will use 0 as placeholder. + # See https://github.com/microsoft/DeepSpeed-MII/pull/311 + output.ttft = 0 + + st = time.perf_counter() + try: + async with session.post(url=request_func_input.api_url, + json=payload) as response: + if response.status == 200: + parsed_resp = await response.json() + output.latency = time.perf_counter() - st + output.generated_text = parsed_resp["text"][0] + output.success = True + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + ("completions", "profile") + ), "OpenAI Completions API URL must end with 'completions' or 'profile'." + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "model": request_func_input.model_name \ + if request_func_input.model_name else request_func_input.model, + "prompt": request_func_input.prompt, + "temperature": 0.0, + "best_of": request_func_input.best_of, + "max_tokens": request_func_input.output_len, + "logprobs": request_func_input.logprobs, + "stream": True, + "stream_options": { + "include_usage": True, + }, + } + if request_func_input.ignore_eos: + payload["ignore_eos"] = request_func_input.ignore_eos + if request_func_input.extra_body: + payload.update(request_func_input.extra_body) + headers = { + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}" + } + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload, + headers=headers) as response: + if response.status == 200: + first_chunk_received = False + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = chunk_bytes.decode("utf-8").removeprefix( + "data: ") + if chunk != "[DONE]": + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if choices := data.get("choices"): + # Note that text could be empty here + # e.g. for special tokens + text = choices[0].get("text") + timestamp = time.perf_counter() + # First token + if not first_chunk_received: + first_chunk_received = True + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += text or "" + elif usage := data.get("usage"): + output.output_tokens = usage.get( + "completion_tokens") + if first_chunk_received: + output.success = True + else: + output.success = False + output.error = ( + "Never received a valid chunk to calculate TTFT." + "This response will be marked as failed!") + output.generated_text = generated_text + output.latency = most_recent_timestamp - st + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_openai_chat_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + "chat/completions" + ), "OpenAI Chat Completions API URL must end with 'chat/completions'." + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + content = request_func_input.prompt + if request_func_input.multi_modal_content: + content = [{"type": "text", "text": request_func_input.prompt}] + content.append(request_func_input.multi_modal_content) + payload = { + "model": request_func_input.model_name \ + if request_func_input.model_name else request_func_input.model, + "messages": [ + { + "role": "user", + "content": content + }, + ], + "temperature": 0.0, + "max_completion_tokens": request_func_input.output_len, + "stream": True, + "stream_options": { + "include_usage": True, + }, + } + if request_func_input.ignore_eos: + payload["ignore_eos"] = request_func_input.ignore_eos + if request_func_input.extra_body: + payload.update(request_func_input.extra_body) + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + } + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload, + headers=headers) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = chunk_bytes.decode("utf-8").removeprefix( + "data: ") + if chunk != "[DONE]": + timestamp = time.perf_counter() + data = json.loads(chunk) + + if choices := data.get("choices"): + content = choices[0]["delta"].get("content") + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + generated_text += content or "" + elif usage := data.get("usage"): + output.output_tokens = usage.get( + "completion_tokens") + + most_recent_timestamp = timestamp + + output.generated_text = generated_text + output.success = True + output.latency = most_recent_timestamp - st + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +def get_model(pretrained_model_name_or_path: str) -> str: + if os.getenv('VLLM_USE_MODELSCOPE', 'False').lower() == 'true': + from modelscope import snapshot_download + + model_path = snapshot_download( + model_id=pretrained_model_name_or_path, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"]) + + return model_path + return pretrained_model_name_or_path + + +def _fix_tokenizer_for_sglang(tokenizer, model_path): + """Fix transformers v5 tokenizer to match sglang server-side behavior. + + Root cause: transformers v5 (>= 5.0) changed how tokenizers are loaded. + Specifically, LlamaTokenizerFast.__init__ in v5 rebuilds the pre_tokenizer + and decoder from scratch using class-specific components, discarding the + originals from tokenizer.json. For models like DeepSeek-R1 that declare + LlamaTokenizerFast but actually use a ByteLevel/Sequence tokenizer + architecture, v5 incorrectly replaces the original Sequence pre_tokenizer + with Metaspace, and the original ByteLevel decoder with Sequence. + See: https://github.com/sgl-project/sglang/blob/9238bd08a2895fa3b7ec79ea567e5c27ac951343/python/sglang/srt/utils/hf_transformers_utils.py#L836 + + The sglang server applies fixes for this in hf_transformers_utils.py + (_fix_v5_tokenizer_components and _fix_v5_add_bos_eos_token), but the + benchmark client loads the tokenizer directly via AutoTokenizer without + these fixes. This mismatch causes the client to encode text differently + from the server -- e.g. a 7000-token prompt on the client becomes ~35000 + tokens on the server, leading to ~5x TTFT inflation and false performance + regressions in benchmarks. + + This function replicates the same fixes so the benchmark client tokenizes + identically to the sglang server. It is a no-op on transformers v4. + """ + import json + from pathlib import Path + + def _resolve(filename): + """Return a filesystem path for `filename`, whether `model_path` is a + local directory or an HF Hub repo id. Returns None and logs a warning + on failure so we don't silently fail to apply the v5 fix.""" + local = Path(model_path) / filename + if local.is_file(): + return str(local) + try: + from huggingface_hub import hf_hub_download + return hf_hub_download(repo_id=model_path, filename=filename) + except Exception as e: + print( + f"v5 tokenizer fix: cannot resolve {filename} for {model_path!r} " + f"({type(e).__name__}: {e}); fix will not apply.", + flush=True, + ) + return None + + backend = getattr(tokenizer, "_tokenizer", None) + if backend is not None: + tok_file = _resolve("tokenizer.json") + if tok_file is not None: + from tokenizers import Tokenizer as RawTokenizer + raw = RawTokenizer.from_file(tok_file) + raw_pre = type(raw.pre_tokenizer).__name__ if raw.pre_tokenizer else None + loaded_pre = type(backend.pre_tokenizer).__name__ if backend.pre_tokenizer else None + if raw_pre and loaded_pre and raw_pre != loaded_pre: + print( + f"v5 tokenizer fix: {model_path} pre_tokenizer {loaded_pre} -> {raw_pre}, " + f"decoder {type(backend.decoder).__name__ if backend.decoder else None} -> " + f"{type(raw.decoder).__name__ if raw.decoder else None}", + flush=True, + ) + backend.pre_tokenizer = raw.pre_tokenizer + backend.decoder = raw.decoder + + config_file = _resolve("tokenizer_config.json") + if config_file is not None: + with open(config_file) as f: + config = json.load(f) + tok_class = config.get("tokenizer_class", "") + bos_eos_classes = { + "LlamaTokenizer", "LlamaTokenizerFast", + "CodeLlamaTokenizer", "CodeLlamaTokenizerFast", + "GemmaTokenizer", "GemmaTokenizerFast", "CohereTokenizerFast", + } + if tok_class in bos_eos_classes: + defaults = {"add_bos_token": True, "add_eos_token": False} + changed = False + for attr in ("add_bos_token", "add_eos_token"): + val = config.get(attr) + if val is None: + val = defaults.get(attr, False) + if getattr(tokenizer, attr, None) != val: + setattr(tokenizer, f"_{attr}", val) + changed = True + if changed and hasattr(tokenizer, "update_post_processor"): + tokenizer.update_post_processor() + + return tokenizer + + +def get_tokenizer( + pretrained_model_name_or_path: str, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + **kwargs, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + if pretrained_model_name_or_path is not None and not os.path.exists( + pretrained_model_name_or_path): + pretrained_model_name_or_path = get_model( + pretrained_model_name_or_path) + if tokenizer_mode == "slow": + if kwargs.get("use_fast", False): + raise ValueError( + "Cannot use the fast tokenizer in slow tokenizer mode.") + kwargs["use_fast"] = False + if tokenizer_mode == "mistral": + try: + from vllm.transformers_utils.tokenizer import MistralTokenizer + except ImportError as e: + raise ImportError("MistralTokenizer requires vllm package.\n" + "Please install it with `pip install vllm` " + "to use mistral tokenizer mode.") from e + return MistralTokenizer.from_pretrained( + str(pretrained_model_name_or_path)) + else: + tokenizer = AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, + trust_remote_code=trust_remote_code, + **kwargs, + ) + return _fix_tokenizer_for_sglang(tokenizer, pretrained_model_name_or_path) + + +ASYNC_REQUEST_FUNCS = { + "tgi": async_request_tgi, + "vllm": async_request_openai_completions, + "lmdeploy": async_request_openai_completions, + "deepspeed-mii": async_request_deepspeed_mii, + "openai": async_request_openai_completions, + "openai-chat": async_request_openai_chat_completions, + "tensorrt-llm": async_request_trt_llm, + "scalellm": async_request_openai_completions, + "sglang": async_request_openai_completions, +} diff --git a/scripts/Qwen3.5/utils/bench_serving/benchmark_serving.py b/scripts/Qwen3.5/utils/bench_serving/benchmark_serving.py new file mode 100644 index 00000000..7cc806ec --- /dev/null +++ b/scripts/Qwen3.5/utils/bench_serving/benchmark_serving.py @@ -0,0 +1,1276 @@ +# SPDX-License-Identifier: Apache-2.0 +r"""Benchmark online serving throughput. + +On the server side, run one of the following commands: + vLLM OpenAI API server + vllm serve \ + --swap-space 16 \ + --disable-log-requests + + (TGI backend) + ./launch_tgi_server.sh + +On the client side, run: + python benchmarks/benchmark_serving.py \ + --backend \ + --model \ + --dataset-name sharegpt \ + --dataset-path \ + --request-rate \ # By default is inf + --num-prompts # By default is 1000 + + when using tgi backend, add + --endpoint /generate_stream + to the end of the command above. +""" +import argparse +import asyncio +import base64 +import contextlib +import gc +import io +import json +import os +import random +import time +import warnings +from dataclasses import dataclass +from datetime import datetime +from multiprocessing import Pool, cpu_count +from typing import Any, AsyncGenerator, Collection, Dict, List, Optional, Tuple + +import numpy as np +from backend_request_func import (ASYNC_REQUEST_FUNCS, RequestFuncInput, + RequestFuncOutput) +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +try: + from backend_request_func import get_tokenizer +except ImportError: + from vllm.transformers_utils.tokenizer import get_tokenizer + +try: + from vllm.utils import FlexibleArgumentParser +except ImportError: + from argparse import ArgumentParser as FlexibleArgumentParser + +from benchmark_utils import convert_to_pytorch_benchmark_format +from encoding_dsv4 import encode_messages as dsv4_encode_messages + +MILLISECONDS_TO_SECONDS_CONVERSION = 1000 + + +@dataclass +class BenchmarkMetrics: + completed: int + total_input: int + total_output: int + request_throughput: float + request_goodput: float + output_throughput: float + total_token_throughput: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + percentiles_ttft_ms: List[Tuple[float, float]] + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + percentiles_tpot_ms: List[Tuple[float, float]] + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + percentiles_itl_ms: List[Tuple[float, float]] + # E2EL stands for end-to-end latency per request. + # It is the time taken on the client side from sending + # a request to receiving a complete response. + mean_e2el_ms: float + median_e2el_ms: float + std_e2el_ms: float + percentiles_e2el_ms: List[Tuple[float, float]] + + +# --- Multiprocessing helpers for sample_random_requests --- +_worker_tokenizer = None + + +def _load_tokenizer(tokenizer_id, tokenizer_mode, trust_remote_code): + """Load tokenizer for random-prompt generation. + + vLLM's get_tokenizer can raise AttributeError when transformers removes + LlamaTokenizer.all_special_tokens_extended (e.g. Qwen3.5 with newer + transformers). Prefer backend_request_func.get_tokenizer on fallback so + client tokenization stays aligned with the sglang server (#1381, #1428). + """ + try: + return get_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) + except AttributeError as exc: + if "all_special_tokens_extended" not in str(exc): + raise + try: + from backend_request_func import get_tokenizer as _backend_get_tokenizer + return _backend_get_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) + except ImportError: + from transformers import AutoTokenizer + use_fast = tokenizer_mode != "slow" + return AutoTokenizer.from_pretrained( + tokenizer_id, + trust_remote_code=trust_remote_code, + use_fast=use_fast, + ) + + +def _init_tokenizer_worker(tokenizer_id, tokenizer_mode, trust_remote_code): + """Initialize tokenizer once per worker process.""" + global _worker_tokenizer + _worker_tokenizer = _load_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) + + +def _apply_chat_template(prompt, tokenizer, dsv4): + """Render a single user message into the appropriate chat-template prompt. + + When `dsv4` is True we use the self-contained DeepSeek-V4 encoder + (encoding_dsv4.encode_messages) which emits the + ... framing the model expects. Otherwise we + fall back to the tokenizer's built-in jinja chat template. + """ + if dsv4: + return dsv4_encode_messages( + [{"role": "user", "content": prompt}], + thinking_mode="thinking", + ) + return tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + ) + + +def _process_prompt_chunk(chunk_args): + """Generate a chunk of random prompts in a worker process.""" + (indices, prefix_token_ids, input_lens, output_lens, offsets, + prefix_len, vocab_size, use_chat_template, dsv4, seed) = chunk_args + + rng = np.random.RandomState(seed) + tokenizer = _worker_tokenizer + + results = [] + for local_idx, global_idx in enumerate(indices): + tgt_prompt_len = prefix_len + input_lens[local_idx] + prompt_token_ids = prefix_token_ids + [ + (offsets[local_idx] + global_idx + j) % vocab_size + for j in range(input_lens[local_idx]) + ] + prompt = tokenizer.decode(prompt_token_ids) + + max_retries = 10 + for _ in range(max_retries): + prompt_token_ids = tokenizer.encode(prompt, add_special_tokens=False) + if len(prompt_token_ids) < tgt_prompt_len: + num_extras = tgt_prompt_len - len(prompt_token_ids) + prompt_token_ids.extend( + rng.randint(0, vocab_size, size=num_extras).tolist()) + elif len(prompt_token_ids) > tgt_prompt_len: + prompt_token_ids = prompt_token_ids[:tgt_prompt_len] + else: + break + prompt = tokenizer.decode(prompt_token_ids) + + if use_chat_template: + prompt = _apply_chat_template(prompt, tokenizer, dsv4) + + prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False)) + mismatch = prompt_len - tgt_prompt_len + results.append((prompt, prompt_len, output_lens[local_idx], None, mismatch)) + + return results + + +def sample_random_requests( + prefix_len: int, + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + use_chat_template: bool = False, + dsv4: bool = False, + tokenizer_id: Optional[str] = None, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + num_workers: int = 0, +) -> List[Tuple[str, int, int]]: + vocab_size = tokenizer.vocab_size + prefix_token_ids = np.random.randint(0, vocab_size, size=prefix_len).tolist() + + if dsv4 and not use_chat_template: + raise ValueError("--dsv4 requires --use-chat-template to be set.") + + if use_chat_template: + chat_template_dummy = _apply_chat_template("a", tokenizer, dsv4) + tokenized_chat_template_dummy = tokenizer.encode(chat_template_dummy, add_special_tokens=False) + chat_template_len = len(tokenized_chat_template_dummy) - 1 + input_len = input_len - chat_template_len + + def sample_uniform(seq_len): + lower = int(seq_len * range_ratio) + upper = seq_len + seq_lens = np.random.randint(lower, upper+1, size=num_prompts).tolist() + return seq_lens + + input_lens = sample_uniform(input_len) + output_lens = sample_uniform(output_len) + offsets = np.random.randint(0, vocab_size, size=num_prompts) + + # Create a local RNG for retry-loop padding so that neither serial nor + # parallel path consumes global np.random draws beyond this point. + # This ensures downstream code (e.g. gamma draws for inter-arrival times) + # sees identical global RNG state regardless of num_workers. + local_rng = np.random.RandomState( + np.random.get_state()[1][:4].tolist() # derive seed from current state without advancing it + ) + + # Decide whether to use multiprocessing + if num_workers <= 0: + num_workers = min(cpu_count() or 1, 8) + use_parallel = num_workers > 1 and tokenizer_id is not None + + if use_parallel: + # Split work into chunks, one per worker + chunk_size = (num_prompts + num_workers - 1) // num_workers + chunk_args_list = [] + for w in range(num_workers): + start = w * chunk_size + end = min(start + chunk_size, num_prompts) + if start >= num_prompts: + break + chunk_args_list.append(( + list(range(start, end)), + prefix_token_ids, + input_lens[start:end], + output_lens[start:end], + offsets[start:end].tolist(), + prefix_len, + vocab_size, + use_chat_template, + dsv4, + int(local_rng.randint(0, 2**31)), + )) + + actual_workers = len(chunk_args_list) + print(f"Generating {num_prompts} prompts using {actual_workers} worker processes...") + t0 = time.perf_counter() + with Pool( + processes=actual_workers, + initializer=_init_tokenizer_worker, + initargs=(tokenizer_id, tokenizer_mode, trust_remote_code), + ) as pool: + chunk_results = pool.map(_process_prompt_chunk, chunk_args_list) + + input_requests = [] + mismatches = [] + for chunk in chunk_results: + for prompt, prompt_len, out_len, mm_content, mismatch in chunk: + input_requests.append((prompt, prompt_len, out_len, mm_content)) + mismatches.append(mismatch) + elapsed = time.perf_counter() - t0 + print(f"Prompt generation completed in {elapsed:.1f}s") + else: + # Original serial path — also uses local_rng for retry-loop padding + # to keep global RNG consumption identical to the parallel path. + if tokenizer_id is None and num_workers > 1: + print("Warning: tokenizer_id not provided, falling back to serial prompt generation.") + input_requests = [] + mismatches = [] + for i in range(num_prompts): + tgt_prompt_len = prefix_len + input_lens[i] + prompt_token_ids = prefix_token_ids + [(offsets[i] + i + j) % vocab_size for j in range(input_lens[i])] + prompt = tokenizer.decode(prompt_token_ids) + + max_retries = 10 + for _ in range(max_retries): + prompt_token_ids = tokenizer.encode(prompt, add_special_tokens=False) + if len(prompt_token_ids) < tgt_prompt_len: + num_extras = tgt_prompt_len - len(prompt_token_ids) + prompt_token_ids.extend(local_rng.randint(0, vocab_size, size=num_extras).tolist()) + elif len(prompt_token_ids) > tgt_prompt_len: + prompt_token_ids = prompt_token_ids[:tgt_prompt_len] + else: + break + prompt = tokenizer.decode(prompt_token_ids) + + if use_chat_template: + prompt = _apply_chat_template(prompt, tokenizer, dsv4) + + prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False)) + mismatches.append(prompt_len - tgt_prompt_len) + input_requests.append((prompt, prompt_len, output_lens[i], None)) + + header_str = f'{"-"*19} Input/Output Length Statistics {"-"*19}' + print(header_str) + print( + f' input_lens : ' + f'min={min(r[1] for r in input_requests):<4d} ' + f'max={max(r[1] for r in input_requests):<4d} ' + f'mean={np.mean([r[1] for r in input_requests]):<7.2f} ' + f'avg_token_mismatch={np.mean(mismatches):<5.2f} ' + ) + print( + f' output_lens: ' + f'min={min(r[2] for r in input_requests):<4d} ' + f'max={max(r[2] for r in input_requests):<4d} ' + f'mean={np.mean([r[2] for r in input_requests]):<7.2f} ' + ) + print('-' * len(header_str), '\n') + + return input_requests + + +async def get_request( + input_requests: List[Tuple[str, int, int]], + request_rate: float, + burstiness: float = 1.0, +) -> AsyncGenerator[Tuple[str, int, int], None]: + """ + Asynchronously generates requests at a specified rate + with OPTIONAL burstiness. + + Args: + input_requests: + A list of input requests, each represented as a tuple. + request_rate: + The rate at which requests are generated (requests/s). + burstiness (optional): + The burstiness factor of the request generation. + Only takes effect when request_rate is not inf. + Default value is 1, which follows a Poisson process. + Otherwise, the request intervals follow a gamma distribution. + A lower burstiness value (0 < burstiness < 1) results + in more bursty requests, while a higher burstiness value + (burstiness > 1) results in a more uniform arrival of requests. + """ + input_requests = iter(input_requests) + + # Calculate scale parameter theta to maintain the desired request_rate. + assert burstiness > 0, ( + f"A positive burstiness factor is expected, but given {burstiness}.") + theta = 1.0 / (request_rate * burstiness) + + for request in input_requests: + yield request + + if request_rate == float("inf"): + # If the request rate is infinity, then we don't need to wait. + continue + + # Sample the request interval from the gamma distribution. + # If burstiness is 1, it follows exponential distribution. + interval = np.random.gamma(shape=burstiness, scale=theta) + # The next request will be sent after the interval. + await asyncio.sleep(interval) + + +def calculate_metrics( + input_requests: List[Tuple[str, int, int]], + outputs: List[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase, + selected_percentile_metrics: List[str], + selected_percentiles: List[float], + goodput_config_dict: Dict[str, float], +) -> Tuple[BenchmarkMetrics, List[int]]: + actual_output_lens: List[int] = [] + total_input = 0 + completed = 0 + good_completed = 0 + itls: List[float] = [] + tpots: List[float] = [] + all_tpots: List[float] = [] + ttfts: List[float] = [] + e2els: List[float] = [] + for i in range(len(outputs)): + if outputs[i].success: + output_len = outputs[i].output_tokens + + if output_len is None: + # We use the tokenizer to count the number of output tokens + # for some serving backends instead of looking at + # len(outputs[i].itl) since multiple output tokens may be + # bundled together + # Note : this may inflate the output token count slightly + output_len = len( + tokenizer(outputs[i].generated_text, + add_special_tokens=False).input_ids) + actual_output_lens.append(output_len) + total_input += input_requests[i][1] + tpot = 0 + if output_len > 1: + latency_minus_ttft = outputs[i].latency - outputs[i].ttft + tpot = latency_minus_ttft / (output_len - 1) + tpots.append(tpot) + # Note: if output_len <= 1, we regard tpot as 0 for goodput + all_tpots.append(tpot) + itls += outputs[i].itl + ttfts.append(outputs[i].ttft) + e2els.append(outputs[i].latency) + completed += 1 + else: + actual_output_lens.append(0) + + if goodput_config_dict: + valid_metrics = [] + slo_values = [] + + if "ttft" in goodput_config_dict: + valid_metrics.append(ttfts) + slo_values.append(goodput_config_dict["ttft"] / + MILLISECONDS_TO_SECONDS_CONVERSION) + if "tpot" in goodput_config_dict: + valid_metrics.append(all_tpots) + slo_values.append(goodput_config_dict["tpot"] / + MILLISECONDS_TO_SECONDS_CONVERSION) + if "e2el" in goodput_config_dict: + valid_metrics.append(e2els) + slo_values.append(goodput_config_dict["e2el"] / + MILLISECONDS_TO_SECONDS_CONVERSION) + + for req_metric in zip(*valid_metrics): + is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)]) + if is_good_req: + good_completed += 1 + + if completed == 0: + warnings.warn( + "All requests failed. This is likely due to a misconfiguration " + "on the benchmark arguments.", + stacklevel=2) + metrics = BenchmarkMetrics( + completed=completed, + total_input=total_input, + total_output=sum(actual_output_lens), + request_throughput=completed / dur_s, + request_goodput=good_completed / dur_s, + output_throughput=sum(actual_output_lens) / dur_s, + total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s, + mean_ttft_ms=np.mean(ttfts or 0) * + 1000, # ttfts is empty if streaming is not supported by backend + std_ttft_ms=np.std(ttfts or 0) * 1000, + median_ttft_ms=np.median(ttfts or 0) * 1000, + percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) + for p in selected_percentiles], + mean_tpot_ms=np.mean(tpots or 0) * 1000, + std_tpot_ms=np.std(tpots or 0) * 1000, + median_tpot_ms=np.median(tpots or 0) * 1000, + percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) + for p in selected_percentiles], + mean_itl_ms=np.mean(itls or 0) * 1000, + std_itl_ms=np.std(itls or 0) * 1000, + median_itl_ms=np.median(itls or 0) * 1000, + percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) + for p in selected_percentiles], + mean_e2el_ms=np.mean(e2els or 0) * 1000, + std_e2el_ms=np.std(e2els or 0) * 1000, + median_e2el_ms=np.median(e2els or 0) * 1000, + percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000) + for p in selected_percentiles], + ) + + return metrics, actual_output_lens + + +async def benchmark( + backend: str, + api_url: str, + base_url: str, + model_id: str, + model_name: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: List[Tuple[str, int, int]], + logprobs: Optional[int], + best_of: int, + request_rate: float, + burstiness: float, + disable_tqdm: bool, + num_warmups: int, + profile: bool, + selected_percentile_metrics: List[str], + selected_percentiles: List[str], + ignore_eos: bool, + goodput_config_dict: Dict[str, float], + max_concurrency: Optional[int], + lora_modules: Optional[List[str]], +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + print("Starting initial single prompt test run...") + test_prompt, test_prompt_len, test_output_len, test_mm_content = ( + input_requests[0]) + if backend != "openai-chat" and test_mm_content is not None: + # multi-modal benchmark is only available on OpenAI Chat backend. + raise ValueError( + "Multi-modal content is only supported on 'openai-chat' backend.") + test_input = RequestFuncInput( + model=model_id, + model_name=model_name, + prompt=test_prompt, + api_url=api_url, + prompt_len=test_prompt_len, + output_len=test_output_len, + logprobs=logprobs, + best_of=best_of, + multi_modal_content=test_mm_content, + ignore_eos=ignore_eos, + ) + + if num_warmups > 0: + print(f"Warming up with {num_warmups} requests...") + warmup_pbar = None if disable_tqdm else tqdm(total=num_warmups) + warmup_semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None + + async def warmup_limited_req_fn(): + if warmup_semaphore is None: + return await request_func(request_func_input=test_input, pbar=warmup_pbar) + async with warmup_semaphore: + return await request_func(request_func_input=test_input, pbar=warmup_pbar) + + warmup_tasks = [] + for _ in range(num_warmups): + task = asyncio.create_task(warmup_limited_req_fn()) + warmup_tasks.append(task) + _ = await asyncio.gather(*warmup_tasks) + + if warmup_pbar is not None: + warmup_pbar.close() + print("Warmup completed.") + + if lora_modules: + # For each input request, choose a LoRA module at random. + lora_modules = iter( + [random.choice(lora_modules) for _ in range(len(input_requests))]) + + if profile: + print("Starting profiler...") + profile_input = RequestFuncInput(model=model_id, + model_name=model_name, + prompt=test_prompt, + api_url=base_url + "/start_profile", + prompt_len=test_prompt_len, + output_len=test_output_len, + extra_body={"num_steps": 1, "merge_profiles": True, "profile_by_stage": True}, + logprobs=logprobs, + best_of=best_of, + multi_modal_content=test_mm_content, + ignore_eos=ignore_eos) + profile_output = await request_func(request_func_input=profile_input) + if profile_output.success: + print("Profiler started") + + if burstiness == 1.0: + distribution = "Poisson process" + else: + distribution = "Gamma distribution" + + print(f"Traffic request rate: {request_rate}") + print(f"Burstiness factor: {burstiness} ({distribution})") + print(f"Maximum request concurrency: {max_concurrency}") + + pbar = None if disable_tqdm else tqdm(total=len(input_requests)) + + semaphore = (asyncio.Semaphore(max_concurrency) + if max_concurrency else None) + + async def limited_request_func(request_func_input, pbar): + if semaphore is None: + return await request_func(request_func_input=request_func_input, + pbar=pbar) + async with semaphore: + return await request_func(request_func_input=request_func_input, + pbar=pbar) + + print("Starting main benchmark run...") + + benchmark_start_time = time.perf_counter() + benchmark_start_time_unix = time.time() + tasks: List[asyncio.Task] = [] + async for request in get_request(input_requests, request_rate, burstiness): + prompt, prompt_len, output_len, mm_content = request + req_model_id, req_model_name = model_id, model_name + if lora_modules: + req_lora_module = next(lora_modules) + req_model_id, req_model_name = req_lora_module, req_lora_module + + request_func_input = RequestFuncInput(model=req_model_id, + model_name=req_model_name, + prompt=prompt, + api_url=api_url, + prompt_len=prompt_len, + output_len=output_len, + logprobs=logprobs, + best_of=best_of, + multi_modal_content=mm_content, + ignore_eos=ignore_eos) + tasks.append( + asyncio.create_task( + limited_request_func(request_func_input=request_func_input, + pbar=pbar))) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + if profile: + print("Stopping profiler...") + profile_input = RequestFuncInput( + model=model_id, + prompt=test_prompt, + api_url=base_url + "/stop_profile", + prompt_len=test_prompt_len, + output_len=test_output_len, + logprobs=logprobs, + best_of=best_of, + ) + profile_output = await request_func(request_func_input=profile_input) + if profile_output.success: + print("Profiler stopped") + + if pbar is not None: + pbar.close() + + benchmark_duration = time.perf_counter() - benchmark_start_time + benchmark_end_time_unix = time.time() + + metrics, actual_output_lens = calculate_metrics( + input_requests=input_requests, + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + selected_percentile_metrics=selected_percentile_metrics, + selected_percentiles=selected_percentiles, + goodput_config_dict=goodput_config_dict, + ) + + print("{s:{c}^{n}}".format(s=' Serving Benchmark Result ', n=50, c='=')) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", + benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", + metrics.total_output)) + print("{:<40} {:<10.2f}".format("Request throughput (req/s):", + metrics.request_throughput)) + if goodput_config_dict: + print("{:<40} {:<10.2f}".format("Request goodput (req/s):", + metrics.request_goodput)) + print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):", + metrics.output_throughput)) + print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):", + metrics.total_token_throughput)) + + result = { + "duration": benchmark_duration, + "benchmark_start_time_unix": benchmark_start_time_unix, + "benchmark_end_time_unix": benchmark_end_time_unix, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "request_throughput": metrics.request_throughput, + "request_goodput:": + metrics.request_goodput if goodput_config_dict else None, + "output_throughput": metrics.output_throughput, + "total_token_throughput": metrics.total_token_throughput, + "input_lens": [output.prompt_len for output in outputs], + "output_lens": actual_output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + } + + def process_one_metric( + # E.g., "ttft" + metric_attribute_name: str, + # E.g., "TTFT" + metric_name: str, + # E.g., "Time to First Token" + metric_header: str, + ): + # This function prints and adds statistics of the specified + # metric. + if metric_attribute_name not in selected_percentile_metrics: + return + print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-')) + print("{:<40} {:<10.2f}".format( + f"Mean {metric_name} (ms):", + getattr(metrics, f"mean_{metric_attribute_name}_ms"))) + print("{:<40} {:<10.2f}".format( + f"Median {metric_name} (ms):", + getattr(metrics, f"median_{metric_attribute_name}_ms"))) + result[f"mean_{metric_attribute_name}_ms"] = getattr( + metrics, f"mean_{metric_attribute_name}_ms") + result[f"median_{metric_attribute_name}_ms"] = getattr( + metrics, f"median_{metric_attribute_name}_ms") + result[f"std_{metric_attribute_name}_ms"] = getattr( + metrics, f"std_{metric_attribute_name}_ms") + for p, value in getattr(metrics, + f"percentiles_{metric_attribute_name}_ms"): + p_word = str(int(p)) if int(p) == p else str(p) + print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):", + value)) + result[f"p{p_word}_{metric_attribute_name}_ms"] = value + + process_one_metric("ttft", "TTFT", "Time to First Token") + process_one_metric("tpot", "TPOT", + "Time per Output Token (excl. 1st token)") + process_one_metric("itl", "ITL", "Inter-token Latency") + process_one_metric("e2el", "E2EL", "End-to-end Latency") + + print("=" * 50) + + return result + + +def check_goodput_args(args): + # Check and parse goodput arguments + goodput_config_dict = {} + VALID_NAMES = ["ttft", "tpot", "e2el"] + if args.goodput: + goodput_config_dict = parse_goodput(args.goodput) + for slo_name, slo_val in goodput_config_dict.items(): + if slo_name not in VALID_NAMES: + raise ValueError( + f"Invalid metric name found, {slo_name}: {slo_val}. " + "The service level objective name should be one of " + f"{str(VALID_NAMES)}. ") + if slo_val < 0: + raise ValueError( + f"Invalid value found, {slo_name}: {slo_val}. " + "The service level objective value should be " + "non-negative.") + return goodput_config_dict + + +def parse_goodput(slo_pairs): + goodput_config_dict = {} + try: + for slo_pair in slo_pairs: + slo_name, slo_val = slo_pair.split(":") + goodput_config_dict[slo_name] = float(slo_val) + except ValueError as err: + raise argparse.ArgumentTypeError( + "Invalid format found for service level objectives. " + "Specify service level objectives for goodput as \"KEY:VALUE\" " + "pairs, where the key is a metric name, and the value is a " + "number in milliseconds.") from err + return goodput_config_dict + + +def save_to_pytorch_benchmark_format(args: argparse.Namespace, + results: Dict[str, Any], + file_name: str) -> None: + metrics = [ + "median_ttft_ms", "mean_ttft_ms", "std_ttft_ms", "p99_ttft_ms", + "mean_tpot_ms", "median_tpot_ms", "std_tpot_ms", "p99_tpot_ms", + "median_itl_ms", "mean_itl_ms", "std_itl_ms", "p99_itl_ms" + ] + # These raw data might be useful, but they are rather big. They can be added + # later if needed + ignored_metrics = ["ttfts", "itls", "generated_texts", "errors"] + pt_records = convert_to_pytorch_benchmark_format( + args=args, + metrics={k: [results[k]] + for k in metrics}, + extra_info={ + k: results[k] + for k in results if k not in metrics and k not in ignored_metrics + }) + if pt_records: + # Don't use json suffix here as we don't want CI to pick it up + pt_file = f"{os.path.splitext(file_name)[0]}.pytorch.json" + with open(pt_file, "w") as f: + json.dump(pt_records, f) + + +def main(args: argparse.Namespace): + print(args) + random.seed(args.seed) + np.random.seed(args.seed) + + backend = args.backend + model_id = args.model + model_name = args.served_model_name + tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model + tokenizer_mode = args.tokenizer_mode + + if args.base_url is not None: + api_url = f"{args.base_url}{args.endpoint}" + base_url = f"{args.base_url}" + else: + api_url = f"http://{args.host}:{args.port}{args.endpoint}" + base_url = f"http://{args.host}:{args.port}" + + tokenizer = _load_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=args.trust_remote_code, + ) + + + if args.dataset_name == "random": + input_requests = sample_random_requests( + prefix_len=args.random_prefix_len, + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + use_chat_template=args.use_chat_template, + dsv4=args.dsv4, + tokenizer_id=tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=args.trust_remote_code, + num_workers=args.random_num_workers, + ) + + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + + goodput_config_dict = check_goodput_args(args) + + # Avoid GC processing "static" data - reduce pause times. + gc.collect() + gc.freeze() + + benchmark_result = asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + base_url=base_url, + model_id=model_id, + model_name=model_name, + tokenizer=tokenizer, + input_requests=input_requests, + logprobs=args.logprobs, + best_of=args.best_of, + request_rate=args.request_rate, + burstiness=args.burstiness, + disable_tqdm=args.disable_tqdm, + num_warmups=args.num_warmups, + profile=args.profile, + selected_percentile_metrics=args.percentile_metrics.split(","), + selected_percentiles=[ + float(p) for p in args.metric_percentiles.split(",") + ], + ignore_eos=args.ignore_eos, + goodput_config_dict=goodput_config_dict, + max_concurrency=args.max_concurrency, + lora_modules=args.lora_modules, + )) + + # Save config and results to json + if args.save_result: + result_json: Dict[str, Any] = {} + + # Setup + current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") + result_json["date"] = current_dt + result_json["backend"] = backend + result_json["model_id"] = model_id + result_json["tokenizer_id"] = tokenizer_id + result_json["best_of"] = args.best_of + result_json["num_prompts"] = args.num_prompts + + # Metadata + if args.metadata: + for item in args.metadata: + if "=" in item: + kvstring = item.split("=") + result_json[kvstring[0].strip()] = kvstring[1].strip() + else: + raise ValueError( + "Invalid metadata format. Please use KEY=VALUE format." + ) + + # Traffic + result_json["request_rate"] = (args.request_rate if args.request_rate + < float("inf") else "inf") + result_json["burstiness"] = args.burstiness + result_json["max_concurrency"] = args.max_concurrency + + # Merge with benchmark result + result_json = {**result_json, **benchmark_result} + + if not args.save_detailed: + # Remove fields with too many data points + for field in [ + "ttfts", + "itls", + "generated_texts", + "errors", + ]: + if field in result_json: + del result_json[field] + if field in benchmark_result: + del benchmark_result[field] + + # Save to file + base_model_id = model_id.split("/")[-1] + max_concurrency_str = (f"-concurrency{args.max_concurrency}" + if args.max_concurrency is not None else "") + file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json" #noqa + if args.result_filename: + file_name = args.result_filename + if args.result_dir: + file_name = os.path.join(args.result_dir, file_name) + with open(file_name, "w", encoding='utf-8') as outfile: + json.dump(result_json, outfile) + save_to_pytorch_benchmark_format(args, result_json, file_name) + + max_failure_rate = 0.05 + completed = benchmark_result["completed"] + failure_rate = 1 - completed / args.num_prompts + if failure_rate > max_failure_rate: + raise SystemExit( + f"FAIL: request failure rate {failure_rate:.1%} exceeds " + f"{max_failure_rate:.0%} threshold " + f"({completed}/{args.num_prompts} completed)" + ) + + +if __name__ == "__main__": + parser = FlexibleArgumentParser( + description="Benchmark the online serving throughput.") + parser.add_argument( + "--backend", + type=str, + default="vllm", + choices=list(ASYNC_REQUEST_FUNCS.keys()), + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + # Use 127.0.0.1 here instead of localhost to force the use of ipv4 + parser.add_argument("--host", type=str, default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument( + "--endpoint", + type=str, + default="/v1/completions", + help="API endpoint.", + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=["random"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument("--dataset-path", + type=str, + default=None, + help="Path to the sharegpt/sonnet dataset. " + "Or the huggingface dataset ID if using HF dataset.") + parser.add_argument( + "--max-concurrency", + type=int, + default=None, + help="Maximum number of concurrent requests. This can be used " + "to help simulate an environment where a higher level component " + "is enforcing a maximum number of concurrent requests. While the " + "--request-rate argument controls the rate at which requests are " + "initiated, this argument will control how many are actually allowed " + "to execute at a time. This means that when used in combination, the " + "actual request rate may be lower than specified with --request-rate, " + "if the server is not processing requests fast enough to keep up.") + + parser.add_argument( + "--model", + type=str, + required=True, + help="Name of the model.", + ) + parser.add_argument( + "--tokenizer", + type=str, + help= + "Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501 + ) + parser.add_argument( + "--best-of", + type=int, + default=1, + help="Generates `best_of` sequences per prompt and " + "returns the best one.", + ) + parser.add_argument("--use-beam-search", action="store_true") + parser.add_argument( + "--num-prompts", + type=int, + default=1000, + help="Number of prompts to process.", + ) + parser.add_argument( + "--logprobs", + type=int, + default=None, + help=("Number of logprobs-per-token to compute & return as part of " + "the request. If unspecified, then either (1) if beam search " + "is disabled, no logprobs are computed & a single dummy " + "logprob is returned for each token; or (2) if beam search " + "is enabled 1 logprob per token is computed"), + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, " + "then all the requests are sent at time 0. " + "Otherwise, we use Poisson process or gamma distribution " + "to synthesize the request arrival times.", + ) + parser.add_argument( + "--burstiness", + type=float, + default=1.0, + help="Burstiness factor of the request generation. " + "Only take effect when request_rate is not inf. " + "Default value is 1, which follows Poisson process. " + "Otherwise, the request intervals follow a gamma distribution. " + "A lower burstiness value (0 < burstiness < 1) results in more " + "bursty requests. A higher burstiness value (burstiness > 1) " + "results in a more uniform arrival of requests.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust remote code from huggingface", + ) + parser.add_argument( + "--disable-tqdm", + action="store_true", + help="Specify to disable tqdm progress bar.", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "VLLM_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--save-result", + action="store_true", + help="Specify to save benchmark results to a json file", + ) + parser.add_argument( + "--save-detailed", + action="store_true", + default=False, + help="When saving results, include detailed per-request data " + "(input_lens, output_lens, ttfts, itls, generated_texts, errors). " + "By default, only aggregated metrics are saved to reduce file size.", + ) + parser.add_argument( + "--metadata", + metavar="KEY=VALUE", + nargs="*", + help="Key-value pairs (e.g, --metadata version=0.3.3 tp=1) " + "for metadata of this run to be saved in the result JSON file " + "for record keeping purposes.", + ) + parser.add_argument( + "--result-dir", + type=str, + default=None, + help="Specify directory to save benchmark json results." + "If not specified, results are saved in the current directory.", + ) + parser.add_argument( + "--result-filename", + type=str, + default=None, + help="Specify the filename to save benchmark json results." + "If not specified, results will be saved in " + "{backend}-{args.request_rate}qps-{base_model_id}-{current_dt}.json" + " format.", + ) + parser.add_argument( + "--ignore-eos", + action="store_true", + help="Set ignore_eos flag when sending the benchmark request." + "Warning: ignore_eos is not supported in deepspeed_mii and tgi.") + parser.add_argument( + "--percentile-metrics", + type=str, + default="ttft,tpot,itl,e2el", + help="Comma-seperated list of selected metrics to report percentils. " + "This argument specifies the metrics to report percentiles. " + "Allowed metric names are \"ttft\", \"tpot\", \"itl\", \"e2el\". " + "Default value is \"ttft,tpot,itl,e2el\".") + parser.add_argument( + "--metric-percentiles", + type=str, + default="90,99,99.9", + help="Comma-seperated list of percentiles for selected metrics. " + "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". " + "Default value is \"90,99,99.9\". " + "Use \"--percentile-metrics\" to select metrics.", + ) + parser.add_argument( + "--goodput", + nargs="+", + required=False, + help="Specify service level objectives for goodput as \"KEY:VALUE\" " + "pairs, where the key is a metric name, and the value is in " + "milliseconds. Multiple \"KEY:VALUE\" pairs can be provided, " + "separated by spaces. Allowed request level metric names are " + "\"ttft\", \"tpot\", \"e2el\". For more context on the definition of " + "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 " + "and the blog: https://hao-ai-lab.github.io/blogs/distserve") + + # group for dataset specific arguments + sonnet_group = parser.add_argument_group("sonnet dataset options") + sonnet_group.add_argument( + "--sonnet-input-len", + type=int, + default=550, + help= + "Number of input tokens per request, used only for sonnet dataset.", + ) + sonnet_group.add_argument( + "--sonnet-output-len", + type=int, + default=150, + help= + "Number of output tokens per request, used only for sonnet dataset.", + ) + sonnet_group.add_argument( + "--sonnet-prefix-len", + type=int, + default=200, + help= + "Number of prefix tokens per request, used only for sonnet dataset.", + ) + + sharegpt_group = parser.add_argument_group("sharegpt dataset options") + sharegpt_group.add_argument( + "--sharegpt-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output length " + "from the ShareGPT dataset.") + + random_group = parser.add_argument_group("random dataset options") + random_group.add_argument( + "--random-input-len", + type=int, + default=1024, + help= + "Number of input tokens per request, used only for random sampling.", + ) + random_group.add_argument( + "--random-output-len", + type=int, + default=128, + help= + "Number of output tokens per request, used only for random sampling.", + ) + random_group.add_argument( + "--random-range-ratio", + type=float, + default=1.0, + help="Range of sampled ratio of input/output length, " + "used only for random sampling.", + ) + random_group.add_argument( + "--random-prefix-len", + type=int, + default=0, + help="Number of fixed prefix tokens before random " + " context. The length range of context in a random " + " request is [random-prefix-len, " + " random-prefix-len + random-prefix-len * random-range-ratio).") + random_group.add_argument( + "--use-chat-template", + action="store_true", + help="Use chat template to format the prompt.", + ) + random_group.add_argument( + '--random-num-workers', + type=int, + default=0, + help="Number of worker processes for parallel random prompt generation. " + "Only used with --dataset-name random. " + "0 (default) = auto (min(cpu_count, 8)). 1 = serial (no multiprocessing).", + ) + + dsv4_group = parser.add_argument_group("DeepSeek-V4 chat template options") + dsv4_group.add_argument( + "--dsv4", + action="store_true", + help="Use the DeepSeek-V4 chat template (encoding_dsv4.py) instead of " + "the tokenizer's built-in jinja chat template. Requires " + "--use-chat-template to also be set. Applies to the random dataset.", + ) + + hf_group = parser.add_argument_group("hf dataset options") + hf_group.add_argument("--hf-subset", + type=str, + default=None, + help="Subset of the HF dataset.") + hf_group.add_argument("--hf-split", + type=str, + default=None, + help="Split of the HF dataset.") + hf_group.add_argument( + "--hf-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output lengths " + "from the sampled HF dataset.", + ) + + parser.add_argument( + '--tokenizer-mode', + type=str, + default="auto", + choices=['auto', 'slow', 'mistral', 'custom'], + help='The tokenizer mode.\n\n* "auto" will use the ' + 'fast tokenizer if available.\n* "slow" will ' + 'always use the slow tokenizer. \n* ' + '"mistral" will always use the `mistral_common` tokenizer. \n*' + '"custom" will use --tokenizer to select the preregistered tokenizer.') + + parser.add_argument("--served-model-name", + type=str, + default=None, + help="The model name used in the API. " + "If not specified, the model name will be the " + "same as the ``--model`` argument. ") + + parser.add_argument("--lora-modules", + nargs='+', + default=None, + help="A subset of LoRA module names passed in when " + "launching the server. For each request, the " + "script chooses a LoRA module at random.") + + parser.add_argument('--num-warmups', type=int, default=0) + + args = parser.parse_args() + main(args) diff --git a/scripts/Qwen3.5/utils/bench_serving/benchmark_utils.py b/scripts/Qwen3.5/utils/bench_serving/benchmark_utils.py new file mode 100644 index 00000000..dc6d31f6 --- /dev/null +++ b/scripts/Qwen3.5/utils/bench_serving/benchmark_utils.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +from typing import Any, Dict, List + + +def convert_to_pytorch_benchmark_format(args: argparse.Namespace, + metrics: Dict[str, List], + extra_info: Dict[str, Any]) -> List: + """ + Save the benchmark results in the format used by PyTorch OSS benchmark with + on metric per record + https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database + """ + records = [] + if not os.environ.get("SAVE_TO_PYTORCH_BENCHMARK_FORMAT", False): + return records + + for name, benchmark_values in metrics.items(): + record = { + "benchmark": { + "name": "vLLM benchmark", + "extra_info": { + "args": vars(args), + }, + }, + "model": { + "name": args.model, + }, + "metric": { + "name": name, + "benchmark_values": benchmark_values, + "extra_info": extra_info, + }, + } + records.append(record) + + return records \ No newline at end of file diff --git a/scripts/Qwen3.5/utils/bench_serving/encoding_dsv4.py b/scripts/Qwen3.5/utils/bench_serving/encoding_dsv4.py new file mode 100644 index 00000000..cdb68f23 --- /dev/null +++ b/scripts/Qwen3.5/utils/bench_serving/encoding_dsv4.py @@ -0,0 +1,744 @@ +""" +DeepSeek-V4 Encoding + +A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages +with tool calling, thinking mode, and quick instruction task support. +""" + +from typing import Any, Dict, List, Union, Optional, Tuple +import copy +import json +import re + +# ============================================================ +# Special Tokens +# ============================================================ + +bos_token: str = "<｜begin▁of▁sentence｜>" +eos_token: str = "<｜end▁of▁sentence｜>" +thinking_start_token: str = "" +thinking_end_token: str = "" +dsml_token: str = "｜DSML｜" + +USER_SP_TOKEN = "<｜User｜>" +ASSISTANT_SP_TOKEN = "<｜Assistant｜>" +LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>" + +# Task special tokens for internal classification tasks +DS_TASK_SP_TOKENS = { + "action": "<｜action｜>", + "query": "<｜query｜>", + "authority": "<｜authority｜>", + "domain": "<｜domain｜>", + "title": "<｜title｜>", + "read_url": "<｜read_url｜>", +} +VALID_TASKS = set(DS_TASK_SP_TOKENS.keys()) + +# ============================================================ +# Templates +# ============================================================ + +system_msg_template: str = "{content}" +user_msg_template: str = "{content}" +latest_reminder_msg_template: str = "{content}" +assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token +assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}" +thinking_template: str = "{reasoning_content}" + +response_format_template: str = ( + "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" +) +tool_call_template: str = ( + "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n" +) +tool_calls_template = ( + "<{dsml_token}{tc_block_name}>\n{tool_calls}\n" +) +tool_calls_block_name: str = "tool_calls" + +tool_output_template: str = ( + "{content}" +) + +REASONING_EFFORT_MAX = ( + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" + "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n" + "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n" +) + +TOOLS_TEMPLATE = """## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following: + +<{dsml_token}tool_calls> +<{dsml_token}invoke name="$TOOL_NAME"> +<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<{dsml_token}invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response. + +Otherwise, output directly after {thinking_end_token} with tool calls or final response. + +### Available Tool Schemas + +{tool_schemas} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +""" + +# ============================================================ +# Utility Functions +# ============================================================ + +def to_json(value: Any) -> str: + """Serialize a value to JSON string.""" + try: + return json.dumps(value, ensure_ascii=False) + except: + return json.dumps(value, ensure_ascii=True) + + +def tools_from_openai_format(tools): + """Extract function definitions from OpenAI-format tool list.""" + return [tool["function"] for tool in tools] + + +def tool_calls_from_openai_format(tool_calls): + """Convert OpenAI-format tool calls to internal format.""" + return [ + { + "name": tool_call["function"]["name"], + "arguments": tool_call["function"]["arguments"], + } + for tool_call in tool_calls + ] + + +def tool_calls_to_openai_format(tool_calls): + """Convert internal tool calls to OpenAI format.""" + return [ + { + "type": "function", + "function": { + "name": tool_call["name"], + "arguments": tool_call["arguments"], + } + } + for tool_call in tool_calls + ] + + +def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str: + """ + Encode tool call arguments into DSML parameter format. + + Args: + tool_call: Dict with "name" and "arguments" (JSON string) keys. + + Returns: + DSML-formatted parameter string. + """ + p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}' + P_dsml_strs = [] + + try: + arguments = json.loads(tool_call["arguments"]) + except Exception as err: + arguments = {"arguments": tool_call["arguments"]} + + for k, v in arguments.items(): + p_dsml_str = p_dsml_template.format( + dsml_token=dsml_token, + key=k, + is_str="true" if isinstance(v, str) else "false", + value=v if isinstance(v, str) else to_json(v), + ) + P_dsml_strs.append(p_dsml_str) + + return "\n".join(P_dsml_strs) + + +def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]: + """ + Decode DSML parameters back to a tool call dict. + + Args: + tool_name: Name of the tool. + tool_args: Dict mapping param_name -> (value, is_string_flag). + + Returns: + Dict with "name" and "arguments" (JSON string) keys. + """ + def _decode_value(key: str, value: str, string: str): + if string == "true": + value = to_json(value) + return f"{to_json(key)}: {value}" + + tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}" + return dict(name=tool_name, arguments=tool_args_json) + + +def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str: + """ + Render tool schemas into the system prompt format. + + Args: + tools: List of tool schema dicts (each with name, description, parameters). + + Returns: + Formatted tools section string. + """ + tools_json = [to_json(t) for t in tools] + + return TOOLS_TEMPLATE.format( + tool_schemas="\n".join(tools_json), + dsml_token=dsml_token, + thinking_start_token=thinking_start_token, + thinking_end_token=thinking_end_token, + ) + + +def find_last_user_index(messages: List[Dict[str, Any]]) -> int: + """Find the index of the last user/developer message.""" + last_user_index = -1 + for idx in range(len(messages) - 1, -1, -1): + if messages[idx].get("role") in ["user", "developer"]: + last_user_index = idx + break + return last_user_index + + +# ============================================================ +# Message Rendering +# ============================================================ + +def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str: + """ + Render a single message at the given index into its encoded string form. + + This is the core function that converts each message in the conversation + into the DeepSeek-V4 format. + + Args: + index: Index of the message to render. + messages: Full list of messages in the conversation. + thinking_mode: Either "chat" or "thinking". + drop_thinking: Whether to drop reasoning content from earlier turns. + reasoning_effort: Optional reasoning effort level ("max", "high", or None). + + Returns: + Encoded string for this message. + """ + assert 0 <= index < len(messages) + assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`" + + prompt = "" + msg = messages[index] + last_user_idx = find_last_user_index(messages) + + role = msg.get("role") + content = msg.get("content") + tools = msg.get("tools") + response_format = msg.get("response_format") + tool_calls = msg.get("tool_calls") + reasoning_content = msg.get("reasoning_content") + wo_eos = msg.get("wo_eos", False) + + if tools: + tools = tools_from_openai_format(tools) + if tool_calls: + tool_calls = tool_calls_from_openai_format(tool_calls) + + # Reasoning effort prefix (only at index 0 in thinking mode with max effort) + assert reasoning_effort in ['max', None, 'high'], f"Invalid reasoning effort: {reasoning_effort}" + if index == 0 and thinking_mode == "thinking" and reasoning_effort == 'max': + prompt += REASONING_EFFORT_MAX + + if role == "system": + prompt += system_msg_template.format(content=content or "") + if tools: + prompt += "\n\n" + render_tools(tools) + if response_format: + prompt += "\n\n" + response_format_template.format(schema=to_json(response_format)) + + elif role == "developer": + assert content, f"Invalid message for role `{role}`: {msg}" + + content_developer = USER_SP_TOKEN + content_developer += content + + if tools: + content_developer += "\n\n" + render_tools(tools) + if response_format: + content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format)) + + prompt += user_msg_template.format(content=content_developer) + + elif role == "user": + prompt += USER_SP_TOKEN + + # Handle content blocks (tool results mixed with text) + content_blocks = msg.get("content_blocks") + if content_blocks: + parts = [] + for block in content_blocks: + block_type = block.get("type") + if block_type == "text": + parts.append(block.get("text", "")) + elif block_type == "tool_result": + tool_content = block.get("content", "") + if isinstance(tool_content, list): + text_parts = [] + for b in tool_content: + if b.get("type") == "text": + text_parts.append(b.get("text", "")) + else: + text_parts.append(f"[Unsupported {b.get('type')}]") + tool_content = "\n\n".join(text_parts) + parts.append(tool_output_template.format(content=tool_content)) + else: + parts.append(f"[Unsupported {block_type}]") + prompt += "\n\n".join(parts) + else: + prompt += content or "" + + elif role == "latest_reminder": + prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content) + + elif role == "tool": + raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()") + + elif role == "assistant": + thinking_part = "" + tc_content = "" + + if tool_calls: + tc_list = [ + tool_call_template.format( + dsml_token=dsml_token, + name=tc.get("name"), + arguments=encode_arguments_to_dsml(tc) + ) + for tc in tool_calls + ] + tc_content += '\n\n' + tool_calls_template.format( + dsml_token=dsml_token, + tool_calls="\n".join(tc_list), + tc_block_name=tool_calls_block_name, + ) + + summary_content = content or "" + rc = reasoning_content or "" + + # Check if previous message has a task - if so, this is a task output (no thinking) + prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None + + if thinking_mode == "thinking" and not prev_has_task: + if not drop_thinking or index > last_user_idx: + thinking_part = thinking_template.format(reasoning_content=rc) + thinking_end_token + else: + thinking_part = "" + + if wo_eos: + prompt += assistant_msg_wo_eos_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) + else: + prompt += assistant_msg_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) + else: + raise NotImplementedError(f"Unknown role: {role}") + + # Append transition tokens based on what follows + if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]: + return prompt + + task = messages[index].get("task") + if task is not None: + # Task special token for internal classification tasks + assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}" + task_sp_token = DS_TASK_SP_TOKENS[task] + + if task != "action": + # Non-action tasks: append task sp token directly after the message + prompt += task_sp_token + else: + # Action task: append Assistant + thinking token + action sp token + prompt += ASSISTANT_SP_TOKEN + prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token + prompt += task_sp_token + + elif messages[index].get("role") in ["user", "developer"]: + # Normal generation: append Assistant + thinking token + prompt += ASSISTANT_SP_TOKEN + if not drop_thinking and thinking_mode == "thinking": + prompt += thinking_start_token + elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx: + prompt += thinking_start_token + else: + prompt += thinking_end_token + + return prompt + + +# ============================================================ +# Preprocessing +# ============================================================ + +def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Merge tool messages into the preceding user message using content_blocks format. + + DeepSeek-V4 does not have a standalone "tool" role; instead, tool results + are encoded as blocks within user messages. + + This function converts a standard OpenAI-format conversation (with separate + "tool" role messages) into V4 format where tool results are merged into + user messages. + + Args: + messages: List of message dicts in OpenAI format. + + Returns: + Processed message list with tool messages merged into user messages. + """ + merged: List[Dict[str, Any]] = [] + + for msg in messages: + msg = copy.deepcopy(msg) + role = msg.get("role") + + if role == "tool": + # Convert tool message to a user message with tool_result block + tool_block = { + "type": "tool_result", + "tool_use_id": msg.get("tool_call_id", ""), + "content": msg.get("content", ""), + } + # Merge into previous message if it's already a user (merged tool) + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]: + merged[-1]["content_blocks"].append(tool_block) + else: + merged.append({ + "role": "user", + "content_blocks": [tool_block], + }) + elif role == "user": + text_block = {"type": "text", "text": msg.get("content", "")} + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None: + merged[-1]["content_blocks"].append(text_block) + else: + new_msg = { + "role": "user", + "content": msg.get("content", ""), + "content_blocks": [text_block], + } + # Preserve extra fields (task, wo_eos, mask, etc.) + for key in ("task", "wo_eos", "mask"): + if key in msg: + new_msg[key] = msg[key] + merged.append(new_msg) + else: + merged.append(msg) + + return merged + + +def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Sort tool_result blocks within user messages by the order of tool_calls + in the preceding assistant message. + + Args: + messages: Preprocessed message list (after merge_tool_messages). + + Returns: + Message list with sorted tool result blocks. + """ + last_tool_call_order: Dict[str, int] = {} + + for msg in messages: + role = msg.get("role") + if role == "assistant" and msg.get("tool_calls"): + last_tool_call_order = {} + for idx, tc in enumerate(msg["tool_calls"]): + tc_id = tc.get("id") or tc.get("function", {}).get("id", "") + if tc_id: + last_tool_call_order[tc_id] = idx + + elif role == "user" and msg.get("content_blocks"): + tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"] + if len(tool_blocks) > 1 and last_tool_call_order: + sorted_blocks = sorted( + tool_blocks, + key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0) + ) + sorted_idx = 0 + new_blocks = [] + for block in msg["content_blocks"]: + if block.get("type") == "tool_result": + new_blocks.append(sorted_blocks[sorted_idx]) + sorted_idx += 1 + else: + new_blocks.append(block) + msg["content_blocks"] = new_blocks + + return messages + + +# ============================================================ +# Main Encoding Function +# ============================================================ + +def encode_messages( + messages: List[Dict[str, Any]], + thinking_mode: str, + context: Optional[List[Dict[str, Any]]] = None, + drop_thinking: bool = True, + add_default_bos_token: bool = True, + reasoning_effort: Optional[str] = None, +) -> str: + """ + Encode a list of messages into the DeepSeek-V4 prompt format. + + This is the main entry point for encoding conversations. It handles: + - BOS token insertion + - Thinking mode with optional reasoning content dropping + - Tool message merging into user messages + - Multi-turn conversation context + + Args: + messages: List of message dicts to encode. + thinking_mode: Either "chat" or "thinking". + context: Optional preceding context messages (already encoded prefix). + drop_thinking: If True, drop reasoning_content from earlier assistant turns + (only keep reasoning for messages after the last user message). + add_default_bos_token: Whether to prepend BOS token at conversation start. + reasoning_effort: Optional reasoning effort level ("max", "high", or None). + + Returns: + The encoded prompt string. + """ + context = context if context else [] + + # Preprocess: merge tool messages and sort tool results + messages = merge_tool_messages(messages) + messages = sort_tool_results_by_call_order(context + messages)[len(context):] + if context: + context = merge_tool_messages(context) + context = sort_tool_results_by_call_order(context) + + full_messages = context + messages + + prompt = bos_token if add_default_bos_token and len(context) == 0 else "" + + # Resolve drop_thinking: if any message has tools defined, don't drop thinking + effective_drop_thinking = drop_thinking + if any(m.get("tools") for m in full_messages): + effective_drop_thinking = False + + if thinking_mode == "thinking" and effective_drop_thinking: + full_messages = _drop_thinking_messages(full_messages) + # After dropping, recalculate how many messages to render + # (context may have shrunk too) + num_to_render = len(full_messages) - len(_drop_thinking_messages(context)) + context_len = len(full_messages) - num_to_render + else: + num_to_render = len(messages) + context_len = len(context) + + for idx in range(num_to_render): + prompt += render_message( + idx + context_len, + full_messages, + thinking_mode=thinking_mode, + drop_thinking=effective_drop_thinking, + reasoning_effort=reasoning_effort, + ) + + return prompt + + +def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Drop reasoning_content and non-essential messages before the last user message. + + Behavior: + - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept. + - Messages at or after the last user index are always kept. + - Assistant messages before the last user get reasoning_content removed. + - Developer messages before the last user are dropped entirely. + """ + last_user_idx = find_last_user_index(messages) + result = [] + keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"} + + for idx, msg in enumerate(messages): + role = msg.get("role") + if role in keep_roles or idx >= last_user_idx: + result.append(msg) + elif role == "assistant": + msg = copy.copy(msg) + msg.pop("reasoning_content", None) + result.append(msg) + # developer and other roles before last_user_idx are dropped + + return result + + +# ============================================================ +# Parsing (Decoding model output) +# ============================================================ + +def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]: + """ + Read text from index until one of the stop strings is found. + + Returns: + Tuple of (new_index, content_before_stop, matched_stop_string_or_None). + """ + min_pos = len(text) + matched_stop = None + + for s in stop: + pos = text.find(s, index) + if pos != -1 and pos < min_pos: + min_pos = pos + matched_stop = s + + if matched_stop: + content = text[index:min_pos] + return min_pos + len(matched_stop), content, matched_stop + else: + content = text[index:] + return len(text), content, None + + +def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]: + """ + Parse DSML tool calls from text starting at the given index. + + Args: + index: Starting position in text. + text: The full text to parse. + + Returns: + Tuple of (new_index, last_stop_token, list_of_tool_call_dicts). + Each tool call dict has "name" and "arguments" keys. + """ + tool_calls: List[Dict[str, Any]] = [] + stop_token = None + tool_calls_end_token = f"" + + while index < len(text): + index, _, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token]) + if _ != ">\n": + raise ValueError(f"Tool call format error: expected '>\\n' but got '{_}'") + + if stop_token == tool_calls_end_token: + break + + if stop_token is None: + raise ValueError("Missing special token in tool calls") + + index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n$', tool_name_content, flags=re.DOTALL) + if len(p_tool_name) != 1: + raise ValueError(f"Tool name format error: '{tool_name_content}'") + tool_name = p_tool_name[0] + + tool_args: Dict[str, Tuple[str, str]] = {} + while stop_token == f"<{dsml_token}parameter": + index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"]) + + param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL) + if len(param_kv) != 1: + raise ValueError(f"Parameter format error: '{param_content}'") + param_name, string, param_value = param_kv[0] + + if param_name in tool_args: + raise ValueError(f"Duplicate parameter name: '{param_name}'") + tool_args[param_name] = (param_value, string) + + index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n": + raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'") + + tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args) + tool_calls.append(tool_call) + + return index, stop_token, tool_calls + + +def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]: + """ + Parse a model completion text into a structured assistant message. + + This function takes the raw text output from the model (a single assistant turn) + and extracts: + - reasoning_content (thinking block) + - content (summary/response) + - tool_calls (if any) + + NOTE: This function is designed to parse only correctly formatted strings and + will raise ValueError for malformed output. + + Args: + text: The raw completion text (including EOS token). + thinking_mode: Either "chat" or "thinking". + + Returns: + Dict with keys: "role", "content", "reasoning_content", "tool_calls". + tool_calls are in OpenAI format. + """ + summary_content, reasoning_content, tool_calls = "", "", [] + index, stop_token = 0, None + tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}" + + is_thinking = thinking_mode == "thinking" + is_tool_calling = False + + if is_thinking: + index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token]) + reasoning_content = content_delta + assert stop_token == thinking_end_token, "Invalid thinking format: missing " + + index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token]) + summary_content = content_delta + if stop_token == tool_calls_start_token: + is_tool_calling = True + else: + assert stop_token == eos_token, "Invalid format: missing EOS token" + + if is_tool_calling: + index, stop_token, tool_calls = parse_tool_calls(index, text) + + index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token]) + assert not tool_ends_text, "Unexpected content after tool calls" + + assert len(text) == index and stop_token in [eos_token, None], "Unexpected content at end" + + for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]: + assert sp_token not in summary_content and sp_token not in reasoning_content, \ + f"Unexpected special token '{sp_token}' in content" + + return { + "role": "assistant", + "content": summary_content, + "reasoning_content": reasoning_content, + "tool_calls": tool_calls_to_openai_format(tool_calls) + } From c44d29b90f86bc13c19f065697de42413ae2755c Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Wed, 8 Jul 2026 19:10:35 -0500 Subject: [PATCH 03/12] Qwen3.5: add lib/ files (force-add past .gitignore lib/ rule) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The repo .gitignore has `lib/` (Python build-dir convention), which silently excluded scripts/Qwen3.5/lib/* from the prior commit. Force-add the 8 harness lib files — run_engine.sh (core launcher), replica_entry.sh, cfg.py, clean_node.sh, placement.py, check_accuracy.py, lib_inferencex.sh, prompts.json. Without these the recipe cannot run. Co-Authored-By: Claude --- scripts/Qwen3.5/lib/cfg.py | 34 + scripts/Qwen3.5/lib/check_accuracy.py | 59 ++ scripts/Qwen3.5/lib/clean_node.sh | 38 + scripts/Qwen3.5/lib/lib_inferencex.sh | 1281 +++++++++++++++++++++++++ scripts/Qwen3.5/lib/placement.py | 50 + scripts/Qwen3.5/lib/prompts.json | 32 + scripts/Qwen3.5/lib/replica_entry.sh | 110 +++ scripts/Qwen3.5/lib/run_engine.sh | 151 +++ 8 files changed, 1755 insertions(+) create mode 100755 scripts/Qwen3.5/lib/cfg.py create mode 100755 scripts/Qwen3.5/lib/check_accuracy.py create mode 100755 scripts/Qwen3.5/lib/clean_node.sh create mode 100755 scripts/Qwen3.5/lib/lib_inferencex.sh create mode 100755 scripts/Qwen3.5/lib/placement.py create mode 100644 scripts/Qwen3.5/lib/prompts.json create mode 100755 scripts/Qwen3.5/lib/replica_entry.sh create mode 100755 scripts/Qwen3.5/lib/run_engine.sh diff --git a/scripts/Qwen3.5/lib/cfg.py b/scripts/Qwen3.5/lib/cfg.py new file mode 100755 index 00000000..9ced83da --- /dev/null +++ b/scripts/Qwen3.5/lib/cfg.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +"""cfg.py — tiny yaml field reader for the shell scripts. +Usage: + cfg.py cluster -> cluster.yaml top-level key + cfg.py model -> models.. + cfg.py engine -> engines.. +Lists print space-separated; scalars print as-is. +""" +import sys, yaml + +def emit(v): + if isinstance(v, list): print(" ".join(str(x) for x in v)) + elif isinstance(v, dict): print(" ".join(f"{k}={v}" for k, v in v.items())) + else: print("" if v is None else v) + +def main(): + f = sys.argv[1]; kind = sys.argv[2] + d = yaml.safe_load(open(f)) + if kind == "cluster": + emit(d.get(sys.argv[3])) + elif kind == "model": + emit(d["models"][sys.argv[3]].get(sys.argv[4])) + elif kind == "engine": + emit(d["engines"][sys.argv[3]].get(sys.argv[4])) + elif kind == "default": + # default : per-model override wins over global defaults{} + model = sys.argv[3]; key = sys.argv[4] + mv = d.get("models", {}).get(model, {}).get(key) + emit(mv if mv is not None else d.get("defaults", {}).get(key)) + else: + sys.exit(f"unknown kind {kind}") + +if __name__ == "__main__": + main() diff --git a/scripts/Qwen3.5/lib/check_accuracy.py b/scripts/Qwen3.5/lib/check_accuracy.py new file mode 100755 index 00000000..e89571ce --- /dev/null +++ b/scripts/Qwen3.5/lib/check_accuracy.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +# Sanity/accuracy check: send 5 real prompts to a running OpenAI-compatible server, +# verify coherent + expected output. Catches corrupt-weight / bad-quant failures that +# throughput benchmarks (which use --ignore-eos on random tokens) miss. +# +# Usage: PORT, MODEL, OUT(json) via env. Exits 0 if all coherent, 1 otherwise. +import json, os, sys, time, urllib.request + +PORT = os.environ.get("PORT", "8000") +MODEL = os.environ["MODEL"] +OUT = os.environ.get("OUT", "/results/sanity.json") +HERE = os.path.dirname(os.path.abspath(__file__)) +prompts = json.load(open(os.path.join(HERE, "prompts.json"))) + +def chat(prompt, max_tokens): + body = json.dumps({ + "model": MODEL, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, + "temperature": 0.0, + }).encode() + req = urllib.request.Request(f"http://0.0.0.0:{PORT}/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=120) as r: + d = json.load(r) + return d["choices"][0]["message"]["content"], time.time() - t0 + +results = [] +allpass = True +for p in prompts: + try: + text, dt = chat(p["prompt"], p["max_tokens"]) + low = text.lower() + hits = [k for k in p["expect_contains"] if k.lower() in low] + # coherent = non-empty, not pure repetition/garbage + coherent = len(text.strip()) > 0 and len(set(text.split())) > 1 or len(text.strip()) <= 8 + ok = len(hits) > 0 and coherent + results.append({"id": p["id"], "ok": ok, "hits": hits, + "expected": p["expect_contains"], "latency_s": round(dt, 2), + "output": text.strip()[:200]}) + if not ok: + allpass = False + except Exception as e: + results.append({"id": p["id"], "ok": False, "error": str(e)[:150]}) + allpass = False + +summary = {"model": MODEL, "all_pass": allpass, + "passed": sum(1 for r in results if r.get("ok")), + "total": len(results), "results": results} +os.makedirs(os.path.dirname(OUT), exist_ok=True) +json.dump(summary, open(OUT, "w"), indent=2) + +print(f"=== SANITY: {summary['passed']}/{summary['total']} prompts passed ===") +for r in results: + mark = "PASS" if r.get("ok") else "FAIL" + out = r.get("output", r.get("error", "")) + print(f" [{mark}] {r['id']:12s} -> {out[:90]!r}") +sys.exit(0 if allpass else 1) diff --git a/scripts/Qwen3.5/lib/clean_node.sh b/scripts/Qwen3.5/lib/clean_node.sh new file mode 100755 index 00000000..6d1aea5a --- /dev/null +++ b/scripts/Qwen3.5/lib/clean_node.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# clean_node.sh — ensure a node's GPUs are in a CLEAN state before deploying an engine. +# Run via srun on the target node (exclusive). Kills leftover serving containers/processes +# that squat VRAM (the #1 cause of OOM-at-init), then verifies VRAM is actually free. +# +# Usage (inside srun on the node): clean_node.sh [max_used_gb_threshold] +# default threshold = 5 GB; exits non-zero if any GPU still above it after cleanup. +set -uo pipefail +THRESH_GB="${1:-5}" + +echo "[clean] $(hostname): removing ALL containers + stray GPU procs (reservation is exclusive to us)" +# The reservation is slurm-exclusive to us, but Docker containers are NOT slurm-tracked, +# so orphaned/stale containers from prior runs can linger and squat VRAM. Remove them ALL. +cids=$(docker ps -aq 2>/dev/null) +[ -n "$cids" ] && { docker kill $cids >/dev/null 2>&1; docker rm -f $cids >/dev/null 2>&1; } || true +# belt-and-suspenders: kill stray engine procs that may hold VRAM +pkill -9 -f "vllm|atom.entrypoints|sglang.launch|EngineCore|sglang.launch_server" 2>/dev/null || true +sleep 5 + +# 3) verify VRAM free on every GPU. +# rocm-smi prints one line per GPU: "GPU[N] : VRAM Total Used Memory (B): ". +# Take the LAST integer on each line (the byte count), not the GPU index. +echo "[clean] verifying VRAM free (threshold ${THRESH_GB}GB)" +busy=0; i=0 +while IFS= read -r line; do + b=$(echo "$line" | grep -oE "[0-9]+" | tail -1) + [ -z "$b" ] && continue + gb=$(( b / 1000000000 )) + if [ "$gb" -gt "$THRESH_GB" ]; then echo " GPU$i: ${gb}GB USED (>threshold)"; busy=1; else echo " GPU$i: ${gb}GB ok"; fi + i=$((i+1)) +done < <(rocm-smi --showmeminfo vram 2>/dev/null | grep -i "Used Memory") +if [ "$busy" -ne 0 ]; then + echo "[clean] WARNING: some GPUs still hold VRAM after cleanup — listing holders:" + rocm-smi --showpids 2>/dev/null | grep -iE "vllm|atom|sglang|EngineCore" | head + exit 1 +fi +echo "[clean] $(hostname): all GPUs clean" +exit 0 diff --git a/scripts/Qwen3.5/lib/lib_inferencex.sh b/scripts/Qwen3.5/lib/lib_inferencex.sh new file mode 100755 index 00000000..95e063a3 --- /dev/null +++ b/scripts/Qwen3.5/lib/lib_inferencex.sh @@ -0,0 +1,1281 @@ +#!/usr/bin/env bash + +# Shared benchmarking utilities for InferenceX + +# Keep Python bytecode out of the mounted workspace. Benchmark jobs often run as +# root inside containers, and root-owned cache directories break future checkout +# cleanup on self-hosted runners. +export PYTHONDONTWRITEBYTECODE=1 +export PYTHONPYCACHEPREFIX="${PYTHONPYCACHEPREFIX:-/tmp/inferencex-pycache}" +mkdir -p "$PYTHONPYCACHEPREFIX" 2>/dev/null || true + +# Inference server port shared by every benchmark recipe. Launchers that need +# a non-default value (e.g. launch_mi355x-amds.sh derives PORT from RUNNER_NAME +# to avoid collisions across concurrent gh-runners on a shared host) set PORT +# themselves before sourcing this file; the `:-` fallback only kicks in when +# nothing upstream set it. +export PORT="${PORT:-8888}" + +# -------------------------------- +# GPU monitoring helpers +# -------------------------------- + +GPU_MONITOR_PID="" +GPU_METRICS_CSV="/workspace/gpu_metrics.csv" +export GPU_METRICS_CSV + +# Start background GPU monitoring that logs metrics every second to CSV. +# Auto-detects NVIDIA (nvidia-smi) or AMD (amd-smi) GPUs. +# Usage: start_gpu_monitor [--output /path/to/output.csv] [--interval 1] +start_gpu_monitor() { + local output="$GPU_METRICS_CSV" + local interval=1 + + while [[ $# -gt 0 ]]; do + case $1 in + --output) output="$2"; shift 2 ;; + --interval) interval="$2"; shift 2 ;; + *) shift ;; + esac + done + + GPU_METRICS_CSV="$output" + export GPU_METRICS_CSV + + if command -v nvidia-smi &>/dev/null; then + nvidia-smi --query-gpu=timestamp,index,power.draw,temperature.gpu,clocks.current.sm,clocks.current.memory,utilization.gpu,utilization.memory \ + --format=csv -l "$interval" > "$output" 2>/dev/null & + GPU_MONITOR_PID=$! + echo "[GPU Monitor] Started NVIDIA (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" + elif command -v amd-smi &>/dev/null; then + # Use amd-smi native watch mode (-w) which includes timestamps automatically. + # Pipe through awk to: skip preamble lines, keep first CSV header, skip repeated headers. + amd-smi metric -p -c -t -u -w "$interval" --csv 2>/dev/null \ + | awk '/^timestamp,/{if(!h){print;h=1};next} h{print}' > "$output" & + GPU_MONITOR_PID=$! + echo "[GPU Monitor] Started AMD (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" + else + echo "[GPU Monitor] No GPU monitoring tool found (nvidia-smi or amd-smi), skipping" + return 0 + fi +} + +# Stop the background GPU monitor and report file size. +stop_gpu_monitor() { + if [[ -n "$GPU_MONITOR_PID" ]] && kill -0 "$GPU_MONITOR_PID" 2>/dev/null; then + kill "$GPU_MONITOR_PID" 2>/dev/null + wait "$GPU_MONITOR_PID" 2>/dev/null || true + echo "[GPU Monitor] Stopped (PID=$GPU_MONITOR_PID)" + if [[ -f "$GPU_METRICS_CSV" ]]; then + local lines + lines=$(wc -l < "$GPU_METRICS_CSV") + echo "[GPU Monitor] Collected $lines rows -> $GPU_METRICS_CSV" + fi + fi + GPU_MONITOR_PID="" +} + +# Check if required environment variables are set +# Usage: check_env_vars VAR1 VAR2 VAR3 ... +# Exits with code 1 if any variable is not set +check_env_vars() { + local missing_vars=() + + for var_name in "$@"; do + if [[ -z "${!var_name:-}" ]]; then + missing_vars+=("$var_name") + fi + done + + if [[ ${#missing_vars[@]} -gt 0 ]]; then + echo "Error: The following required environment variables are not set:" + for var in "${missing_vars[@]}"; do + echo " - $var" + done + exit 1 + fi +} + +# Wait for server to be ready by polling the health endpoint +# All parameters are required +# Parameters: +# --port: Server port +# --server-log: Path to server log file +# --server-pid: Server process ID (required) +# --sleep-interval: Sleep interval between health checks (optional, default: 5) +wait_for_server_ready() { + set +x + local port="" + local server_log="" + local server_pid="" + local sleep_interval=5 + + # Parse arguments + while [[ $# -gt 0 ]]; do + case $1 in + --port) + port="$2" + shift 2 + ;; + --server-log) + server_log="$2" + shift 2 + ;; + --server-pid) + server_pid="$2" + shift 2 + ;; + --sleep-interval) + sleep_interval="$2" + shift 2 + ;; + *) + echo "Unknown parameter: $1" + return 1 + ;; + esac + done + + # Validate required parameters + if [[ -z "$port" ]]; then + echo "Error: --port is required" + return 1 + fi + if [[ -z "$server_log" ]]; then + echo "Error: --server-log is required" + return 1 + fi + if [[ -z "$server_pid" ]]; then + echo "Error: --server-pid is required" + return 1 + fi + + # Wait for server log file to be created (container startup may delay this) + while [ ! -f "$server_log" ]; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "Server died before creating log file. Exiting." + exit 1 + fi + sleep 1 + done + + # Show logs until server is ready + tail -f -n +1 "$server_log" & + local TAIL_PID=$! + until curl --output /dev/null --silent --fail http://0.0.0.0:$port/health; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "Server died before becoming healthy. Exiting." + kill $TAIL_PID + exit 1 + fi + sleep "$sleep_interval" + done + kill $TAIL_PID +} + +# Run benchmark serving with standardized parameters +# All parameters are required except --endpoint, --use-chat-template, --dsv4, and --trust-remote-code +# Parameters: +# --model: Model name +# --port: Server port +# --backend: Backend type - e.g., 'vllm' or 'openai' +# --endpoint: Optional API endpoint override +# --input-len: Random input sequence length +# --output-len: Random output sequence length +# --random-range-ratio: Random range ratio +# --num-prompts: Number of prompts +# --max-concurrency: Max concurrency +# --result-filename: Result filename without extension +# --result-dir: Result directory +# --use-chat-template: Optional flag to enable chat template +# --dsv4: Optional flag to use the DeepSeek-V4 chat template +# (encoding_dsv4.py) instead of the tokenizer's built-in jinja +# template. Implies --use-chat-template. +# --trust-remote-code: Optional flag to trust remote code from HuggingFace +# --server-pid: Optional server process ID to monitor during benchmark +run_benchmark_serving() { + # In eval-only mode, skip the throughput benchmark entirely. + if [ "${EVAL_ONLY}" = "true" ]; then + echo "EVAL_ONLY mode: skipping throughput benchmark" + return 0 + fi + + set +x + local model="" + local port="" + local backend="" + local endpoint="" + local input_len="" + local output_len="" + local random_range_ratio="" + local num_prompts="" + local max_concurrency="" + local result_filename="" + local result_dir="" + local workspace_dir="" + local use_chat_template=false + local dsv4=false + local trust_remote_code=false + local server_pid="" + local tokenizer="" + + while [[ $# -gt 0 ]]; do + case $1 in + --model) + model="$2" + shift 2 + ;; + --port) + port="$2" + shift 2 + ;; + --backend) + backend="$2" + shift 2 + ;; + --endpoint) + endpoint="$2" + shift 2 + ;; + --input-len) + input_len="$2" + shift 2 + ;; + --output-len) + output_len="$2" + shift 2 + ;; + --random-range-ratio) + random_range_ratio="$2" + shift 2 + ;; + --num-prompts) + num_prompts="$2" + shift 2 + ;; + --max-concurrency) + max_concurrency="$2" + shift 2 + ;; + --result-filename) + result_filename="$2" + shift 2 + ;; + --result-dir) + result_dir="$2" + shift 2 + ;; + --bench-serving-dir) + workspace_dir="$2" + shift 2 + ;; + --use-chat-template) + use_chat_template=true + shift + ;; + --dsv4) + dsv4=true + use_chat_template=true + shift + ;; + --trust-remote-code) + trust_remote_code=true + shift + ;; + --server-pid) + server_pid="$2" + shift 2 + ;; + --tokenizer) + tokenizer="$2" + shift 2 + ;; + *) + echo "Unknown parameter: $1" + return 1 + ;; + esac + done + + # Validate all required parameters + if [[ -z "$model" ]]; then + echo "Error: --model is required" + return 1 + fi + if [[ -z "$port" ]]; then + echo "Error: --port is required" + return 1 + fi + if [[ -z "$backend" ]]; then + echo "Error: --backend is required" + return 1 + fi + if [[ -z "$input_len" ]]; then + echo "Error: --input-len is required" + return 1 + fi + if [[ -z "$output_len" ]]; then + echo "Error: --output-len is required" + return 1 + fi + if [[ -z "$random_range_ratio" ]]; then + echo "Error: --random-range-ratio is required" + return 1 + fi + if [[ -z "$num_prompts" ]]; then + echo "Error: --num-prompts is required" + return 1 + fi + if [[ -z "$max_concurrency" ]]; then + echo "Error: --max-concurrency is required" + return 1 + fi + if [[ -z "$result_filename" ]]; then + echo "Error: --result-filename is required" + return 1 + fi + if [[ -z "$result_dir" ]]; then + echo "Error: --result-dir is required" + return 1 + fi + + if [[ -z "$workspace_dir" ]]; then + workspace_dir=$(pwd) + fi + + # Profiling support: when PROFILE=1, ensure profiler dir exists, add --profile flag, + # and cap num_prompts to keep traces small. + local profile_flag=() + if [[ "${PROFILE:-}" == "1" ]]; then + local _prof_dir="${SGLANG_TORCH_PROFILER_DIR:-${VLLM_TORCH_PROFILER_DIR:-}}" + if [[ -n "$_prof_dir" ]]; then + mkdir -p "$_prof_dir" + fi + profile_flag+=(--profile) + num_prompts="$max_concurrency" + fi + + # Build benchmark command + local benchmark_cmd=( + python3 "$workspace_dir/utils/bench_serving/benchmark_serving.py" + --model "$model" + --backend "$backend" + --base-url "http://0.0.0.0:$port" + --dataset-name random + --random-input-len "$input_len" + --random-output-len "$output_len" + --random-range-ratio "$random_range_ratio" + --num-prompts "$num_prompts" + --max-concurrency "$max_concurrency" + --request-rate inf + --ignore-eos + "${profile_flag[@]}" + --save-result + --num-warmups "$((2 * max_concurrency))" \ + --percentile-metrics 'ttft,tpot,itl,e2el' + --result-dir "$result_dir" + --result-filename "$result_filename.json" + ) + + if [[ -n "$endpoint" ]]; then + benchmark_cmd+=(--endpoint "$endpoint") + fi + + # Add --use-chat-template if requested + if [[ "$use_chat_template" == true ]]; then + benchmark_cmd+=(--use-chat-template) + fi + + # Add --dsv4 if requested (requires --use-chat-template, which we + # auto-enable when --dsv4 is passed in). + if [[ "$dsv4" == true ]]; then + benchmark_cmd+=(--dsv4) + fi + + # Add --trust-remote-code if requested + if [[ "$trust_remote_code" == true ]]; then + benchmark_cmd+=(--trust-remote-code) + fi + + if [[ -n "$tokenizer" ]]; then + benchmark_cmd+=(--tokenizer "$tokenizer") + fi + + # Run benchmark with optional server monitoring + set -x + if [[ -n "$server_pid" ]]; then + # Run benchmark in background and monitor server health + "${benchmark_cmd[@]}" & + local benchmark_pid=$! + + # Monitor loop: check both benchmark and server status + while kill -0 "$benchmark_pid" 2>/dev/null; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "ERROR: Server process $server_pid died during benchmark" + kill "$benchmark_pid" 2>/dev/null + wait "$benchmark_pid" 2>/dev/null + set +x + return 1 + fi + sleep 2 + done + + # Benchmark finished, get its exit code + wait "$benchmark_pid" + local benchmark_exit_code=$? + else + # No server monitoring, run benchmark directly + "${benchmark_cmd[@]}" + local benchmark_exit_code=$? + fi + set +x + + # If profiling, move trace to relay-upload location + if [[ "${PROFILE:-}" == "1" ]]; then + move_profile_trace_for_relay + fi + + return $benchmark_exit_code +} + + +# -------------------------------- +# Profiling trace helpers +# -------------------------------- + +_find_latest_profile_trace() { + local latest="" + local dir="" candidate="" base="" + local -a search_roots=() + + for dir in "$@"; do + search_roots=() + if [[ -d "$dir" ]]; then + search_roots+=("$dir") + fi + if [[ -d "$dir/profiles" ]]; then + search_roots+=("$dir/profiles") + fi + if [[ ${#search_roots[@]} -eq 0 ]]; then + continue + fi + + while IFS= read -r -d '' candidate; do + base="$(basename "$candidate")" + if [[ "$base" == profile_*.trace.json.gz ]]; then + continue + fi + if [[ -z "$latest" || "$candidate" -nt "$latest" ]]; then + latest="$candidate" + fi + done < <( + find "${search_roots[@]}" -maxdepth 1 -type f \ + \( -name "*.trace.json" -o -name "*.trace.json.gz" -o -name "*trace*.json" -o -name "*trace*.json.gz" -o -name "*profile*.json" -o -name "*profile*.json.gz" \) \ + -print0 2>/dev/null + ) + done + + printf '%s' "$latest" +} + +# Move profiler trace into a stable workspace path for workflow relay/upload. +move_profile_trace_for_relay() { + if [[ "${PROFILE:-}" != "1" ]]; then + return 0 + fi + + if [[ -z "${RESULT_FILENAME:-}" ]]; then + echo "[PROFILE] RESULT_FILENAME is not set; skipping relay trace staging." >&2 + return 0 + fi + + local sglang_dir="${SGLANG_TORCH_PROFILER_DIR:-/workspace}" + local vllm_dir="${VLLM_TORCH_PROFILER_DIR:-/workspace}" + local -a search_dirs=() + local dir="" existing="" + local seen=0 + + for dir in "$sglang_dir" "$vllm_dir" "/workspace"; do + if [[ -z "$dir" ]]; then + continue + fi + seen=0 + for existing in "${search_dirs[@]}"; do + if [[ "$existing" == "$dir" ]]; then + seen=1 + break + fi + done + if [[ "$seen" -eq 0 ]]; then + search_dirs+=("$dir") + fi + done + + local trace_file="" + local wait_attempts=10 + for (( i=1; i<=wait_attempts; i++ )); do + trace_file="$(_find_latest_profile_trace "${search_dirs[@]}")" + if [[ -n "$trace_file" ]]; then + break + fi + sleep 10 + done + + if [[ -z "$trace_file" ]]; then + echo "[PROFILE] No trace found for relay under: ${search_dirs[*]}" >&2 + return 0 + fi + + local dest_trace="/workspace/profile_${RESULT_FILENAME}.trace.json.gz" + if [[ "$trace_file" == *.gz ]]; then + cp -f "$trace_file" "$dest_trace" + else + gzip -c "$trace_file" > "$dest_trace" + fi + + echo "[PROFILE] Relay trace prepared: $dest_trace (source: $trace_file)" +} + + +# ------------------------------ +# Eval (lm-eval-harness) helpers +# ------------------------------ + +_install_lm_eval_deps() { + # torchvision causes circular imports in ATOM; TRT-LLM/SGLang need it at module level. + if [[ "${IMAGE:-}" == *atom* ]]; then + python3 -m pip uninstall -y torchvision 2>/dev/null || true + fi + python3 -m pip install -q --no-cache-dir --break-system-packages "lm-eval[api]" || true + local lm_eval_ref="b315ef3b05176acc9732bb7fdec116abe1ecc476" + if command -v git >/dev/null 2>&1; then + if ! python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "git+https://github.com/EleutherAI/lm-evaluation-harness.git@${lm_eval_ref}"; then + python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true + fi + else + python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true + fi +} + +# Patch lm-eval filters to be robust to empty strings via sitecustomize +_patch_lm_eval() { + local patch_dir + patch_dir="$(mktemp -d)" + cat > "$patch_dir/sitecustomize.py" <<'PY' +# --- Patch LocalChatCompletion.parse_generations to handle empty content with reasoning_content --- +import re, sys, unicodedata, json +from lm_eval.filters import extraction as ex +from lm_eval.models.openai_completions import LocalChatCompletion as _LCC + +def _le_parse_generations(outputs, **kwargs): + res = [] + if not isinstance(outputs, list): + outputs = [outputs] + for out in (outputs or []): + try: + choices = out.get("choices", []) + tmp = ["" for _ in choices] + for choice in choices: + idx = choice.get("index", 0) + msg = (choice.get("message") or {}) + content = msg.get("content") + if content in (None, "", []): + content = msg.get("reasoning_content") or "" + tmp[idx] = content + except Exception: + tmp = [""] + res.extend(tmp) + return res + +# Keep staticmethod semantics +_LCC.parse_generations = staticmethod(_le_parse_generations) + +# --- Patch TemplateAPI.apply_chat_template to avoid injecting "type": "text" for TRT --- +try: + from lm_eval.models import api_models as _api_models + _TemplateAPI = _api_models.TemplateAPI + _JsonChatStr = _api_models.JsonChatStr +except Exception: + _TemplateAPI = None + _JsonChatStr = None + +if _TemplateAPI is not None and _JsonChatStr is not None: + _orig_apply_chat_template = _TemplateAPI.apply_chat_template + + def _patched_apply_chat_template( + self, + chat_history, + add_generation_prompt: bool = True, + ): + """Applies a chat template to a list of chat history between user and model.""" + if self.tokenizer_backend == "huggingface" and self.tokenized_requests: + return self.tokenizer.apply_chat_template( + chat_history, + tokenize=False, + add_generation_prompt=add_generation_prompt, + continue_final_message=not add_generation_prompt, + ) + elif self.tokenizer_backend == "remote" and self.tokenized_requests: + return chat_history + else: + # NOTE: we no longer inject `"type": "text"` when tokenizer is None / non-HF + return _JsonChatStr( + json.dumps( + [{**item} for item in chat_history], + ensure_ascii=False, + ) + ) + + _TemplateAPI.apply_chat_template = _patched_apply_chat_template +PY + export PYTHONPATH="${patch_dir}:${PYTHONPATH:-}" +} + +get_native_max_context_length() { + local model_path="$1" + # Prefer MODEL_PATH (local model directory) when available, since the + # argument may be a served-model name that is neither a valid HF repo + # ID nor a local path (e.g. "deepseek-r1-fp4" on the B300 cluster). + if [ -n "${MODEL_PATH:-}" ] && [ -d "${MODEL_PATH}" ]; then + model_path="${MODEL_PATH}" + fi + python3 -c " +try: + from transformers import AutoConfig + config = AutoConfig.from_pretrained('${model_path}', trust_remote_code=True) + for attr in ['max_position_embeddings', 'max_sequence_length', 'seq_length', 'n_positions']: + if hasattr(config, attr): + print(getattr(config, attr)) + break + else: + print(0) +except Exception: + print(0) +" +} + +# Compute the context length for eval-only mode. +# Uses the requested benchmark context capped at the model's native max. +# Sets EVAL_MAX_MODEL_LEN (needed by run_lm_eval). +# Echoes the computed value for scripts to capture. +# +# Usage: local ctx=$(compute_eval_context_length "$MODEL" "${current_ctx}") +compute_eval_context_length() { + local model="$1" + local benchmark_ctx="${2:-0}" + local native_max + native_max=$(get_native_max_context_length "$model") + native_max="${native_max:-0}" + + if [ "$benchmark_ctx" -eq 0 ] 2>/dev/null; then + benchmark_ctx="${native_max:-0}" + fi + local eval_ctx=$(( benchmark_ctx * 1 )) + if [ "$native_max" -gt 0 ] 2>/dev/null && [ "$eval_ctx" -gt "$native_max" ]; then + eval_ctx="$native_max" + fi + # If eval_ctx is still 0 (both benchmark_ctx and native_max were 0), fall back + if [ "$eval_ctx" -le 0 ] 2>/dev/null; then + echo "WARN: compute_eval_context_length could not determine context length for $model" >&2 + eval_ctx="${MAX_MODEL_LEN:-16384}" + fi + EVAL_MAX_MODEL_LEN="$eval_ctx" + echo "$eval_ctx" +} + +# Convenience wrapper: compute eval context from ISL/OSL and export EVAL_MAX_MODEL_LEN. +# Call directly (not in a subshell) so the export persists. +# Scripts then wire $EVAL_MAX_MODEL_LEN into whichever server variable they need. +setup_eval_context() { + EVAL_MAX_MODEL_LEN=$(compute_eval_context_length "$MODEL" "$((ISL + OSL + 256))") + export EVAL_MAX_MODEL_LEN +} + +run_lm_eval() { + local port="${PORT:-8888}" + local tasks_dir="${EVAL_TASKS_DIR:-utils/evals/gsm8k.yaml}" + local results_dir="${EVAL_RESULT_DIR:-$(mktemp -d /tmp/eval_out-XXXXXX)}" + local eval_context_len="${EVAL_MAX_MODEL_LEN:-16384}" + local temperature=0 + local top_p=1 + local concurrent_requests="${EVAL_CONCURRENT_REQUESTS:-${CONC:-64}}" + + while [[ $# -gt 0 ]]; do + case $1 in + --port) port="$2"; shift 2 ;; + --task) tasks_dir="$2"; shift 2 ;; + --results-dir) results_dir="$2"; shift 2 ;; + --gen-max-tokens) eval_context_len="$2"; shift 2 ;; + --temperature) temperature="$2"; shift 2 ;; + --top-p) top_p="$2"; shift 2 ;; + *) echo "Unknown parameter: $1"; return 1 ;; + esac + done + + if [ "${INFERENCEX_LM_EVAL_RUNTIME_READY:-false}" != "true" ]; then + _install_lm_eval_deps + _patch_lm_eval + export INFERENCEX_LM_EVAL_RUNTIME_READY=true + fi + + local openai_server_base="http://0.0.0.0:${port}" + local openai_chat_base="${openai_server_base}/v1/chat/completions" + export OPENAI_API_KEY=${OPENAI_API_KEY:-EMPTY} + MODEL_NAME=${MODEL_NAME:-$MODEL} # Prefer MODEL_NAME, else MODEL + + # Cap output tokens: must fit within context window (leave room for input), + # and avoid excessive KV cache reservation per request on TRT. + local max_output_tokens=$(( eval_context_len > 4096 ? eval_context_len - 4096 : eval_context_len / 2 )) + if [ "$max_output_tokens" -gt 16384 ]; then + max_output_tokens=16384 + fi + echo "Eval budget: eval_context_len=${eval_context_len}, max_output_tokens=${max_output_tokens}" + + # Export for append_lm_eval_summary to pick up + export EVAL_RESULT_DIR="$results_dir" + set -x + python3 -m lm_eval --model local-chat-completions --apply_chat_template \ + --tasks "${tasks_dir}" \ + --output_path "${results_dir}" \ + --log_samples \ + --model_args "model=${MODEL_NAME},base_url=${openai_chat_base},api_key=${OPENAI_API_KEY},eos_string=,max_retries=5,num_concurrent=${concurrent_requests},timeout=1800,tokenized_requests=False,max_length=${eval_context_len}" \ + --gen_kwargs "max_tokens=${max_output_tokens},temperature=${temperature},top_p=${top_p}" + local eval_exit=$? + set +x + return $eval_exit +} + +_stage_lm_eval_artifacts() { + local results_dir="$1" + local eval_conc="$2" + local moved=0 + local failed=0 + local jf base stem extension target suffix + + if [ ! -d "${results_dir}" ]; then + echo "WARN: eval result directory '${results_dir}' does not exist" >&2 + return 1 + fi + + while IFS= read -r -d '' jf; do + base=$(basename "$jf") + case "$base" in + meta_env.json) + continue + ;; + *.jsonl) + stem="${base%.jsonl}" + extension=".jsonl" + ;; + *.json) + stem="${base%.json}" + extension=".json" + ;; + *) + continue + ;; + esac + + target="./${stem}_conc${eval_conc}${extension}" + suffix=2 + while [ -e "$target" ]; do + target="./${stem}_conc${eval_conc}_${suffix}${extension}" + suffix=$((suffix + 1)) + done + + if mv -f "$jf" "$target"; then + moved=1 + else + echo "WARN: failed to stage eval artifact ${jf}" >&2 + failed=1 + fi + done < <( + find "${results_dir}" -type f \ + \( -name "*.json" -o -name "*.jsonl" \) -print0 2>/dev/null + ) + + rm -rf --one-file-system "${results_dir}" 2>/dev/null \ + || rm -rf "${results_dir}" \ + || true + + if [ "$moved" -eq 0 ]; then + echo "WARN: no eval artifacts were produced for concurrency ${eval_conc}" >&2 + return 1 + fi + return "$failed" +} + +_eval_concs_to_json() { + local values="$1" + local value + local joined="" + + for value in $values; do + if [[ ! "$value" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: invalid eval concurrency '${value}'" >&2 + return 1 + fi + if [ -n "$joined" ]; then + joined="${joined}, " + fi + joined="${joined}${value}" + done + + printf '[%s]' "$joined" +} + +append_lm_eval_summary() { + local batch_concs="${EVAL_BATCHED_CONCS:-}" + local results_dir="${EVAL_RESULT_DIR:-}" + local out_dir="${results_dir}" + local meta_json + local metadata_conc="${CONC:-1}" + local batch_metadata="" + + if [ -n "$batch_concs" ]; then + meta_json="./meta_env.json" + metadata_conc="${batch_concs%% *}" + + local eval_concs_json completed_concs_json failed_concs_json + eval_concs_json=$(_eval_concs_to_json "$batch_concs") || return 1 + completed_concs_json=$( + _eval_concs_to_json "${EVAL_BATCHED_COMPLETED_CONCS:-}" + ) || return 1 + failed_concs_json=$( + _eval_concs_to_json "${EVAL_BATCHED_FAILED_CONCS:-}" + ) || return 1 + printf -v batch_metadata \ + ' "eval_concs": %s,\n "completed_eval_concs": %s,\n "failed_eval_concs": %s,\n' \ + "$eval_concs_json" \ + "$completed_concs_json" \ + "$failed_concs_json" + else + if [ -z "${results_dir}" ]; then + echo "WARN: EVAL_RESULT_DIR is empty; skipping artifact collection" >&2 + return 1 + fi + if [ ! -d "${out_dir}" ]; then + echo "WARN: EVAL_RESULT_DIR='${out_dir}' does not exist; skipping artifact collection" >&2 + return 1 + fi + meta_json="${out_dir}/meta_env.json" + fi + + # Write minimal meta for collectors that expect it + local model_name="${MODEL_NAME:-$MODEL}" + local is_multinode_json="false" + if [ "${IS_MULTINODE:-false}" = "true" ]; then + is_multinode_json="true" + fi + + local prefill_tp="${PREFILL_TP:-${TP:-1}}" + local prefill_ep="${PREFILL_EP:-${EP_SIZE:-1}}" + local prefill_num_workers="${PREFILL_NUM_WORKERS:-1}" + local decode_tp="${DECODE_TP:-${TP:-1}}" + local decode_ep="${DECODE_EP:-${EP_SIZE:-1}}" + local decode_num_workers="${DECODE_NUM_WORKERS:-1}" + + local dp_json="false" + if [ "${DP_ATTENTION:-false}" = "true" ]; then dp_json="true"; fi + local prefill_dp_json="$dp_json" + if [ "${PREFILL_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then + prefill_dp_json="true" + else + prefill_dp_json="false" + fi + local decode_dp_json="$dp_json" + if [ "${DECODE_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then + decode_dp_json="true" + else + decode_dp_json="false" + fi + + # Derive framework/precision from env, fallback to parsing RESULT_FILENAME + # RESULT_FILENAME format (from workflow): + # ___tp<...>_ep<...>_dpa_<...>_conc<...>_ + local fw="${FRAMEWORK:-}" + local prec="${PRECISION:-}" + if [[ -z "$fw" || -z "$prec" ]]; then + if [[ -n "${RESULT_FILENAME}" ]]; then + # Extract the two fields immediately before "_tp" + # Handles arbitrary underscores in exp_name by matching from the end + local parsed + parsed=$(echo "${RESULT_FILENAME}" | sed -n 's/.*_\([^_][^_]*\)_\([^_][^_]*\)_tp.*/\1 \2/p') + local p1="${parsed%% *}" + local p2="${parsed#* }" + if [[ -z "$prec" && -n "$p1" && "$p1" != "$parsed" ]]; then + prec="$p1" + fi + if [[ -z "$fw" && -n "$p2" && "$p2" != "$parsed" ]]; then + fw="$p2" + fi + fi + fi + cat > "${meta_json}" <&2 + fi + if [ -d "${out_dir}" ]; then + while IFS= read -r -d '' jf; do + base=$(basename "$jf") + if [ "$base" != "meta_env.json" ]; then + mv -f "$jf" ./ || echo "WARN: failed to move ${jf}" >&2 + fi + done < <(find "${out_dir}" -type f -name "*.json*" -print0 2>/dev/null) + fi + + # Best-effort cleanup of the temp directory + if [ -n "${out_dir}" ] && [ -d "${out_dir}" ]; then + rm -rf --one-file-system "${out_dir}" || rm -rf "${out_dir}" || true + fi + + echo "Moved eval artifacts to: $(pwd)" +} + +# ------------------------------ +# Unified eval entrypoint +# ------------------------------ + +run_eval() { + local framework="${EVAL_FRAMEWORK:-lm-eval}" + local forwarded=() + + while [[ $# -gt 0 ]]; do + case "$1" in + --framework) framework="$2"; shift 2 ;; + *) forwarded+=("$1"); shift ;; + esac + done + + # Compute EVAL_MAX_MODEL_LEN if not already set by the calling script + if [ -z "${EVAL_MAX_MODEL_LEN:-}" ]; then + compute_eval_context_length "$MODEL" "${MAX_MODEL_LEN:-0}" > /dev/null + fi + + unset EVAL_BATCHED_CONCS + unset EVAL_BATCHED_COMPLETED_CONCS + unset EVAL_BATCHED_FAILED_CONCS + + local requested_concs="${EVAL_CONCURRENT_REQUESTS:-}" + local eval_concs=() + if [ -n "$requested_concs" ]; then + read -r -a eval_concs <<< "$requested_concs" + fi + + if [ "${#eval_concs[@]}" -gt 1 ]; then + if [[ "$framework" != "lm-eval" && "$framework" != "lm_eval" ]]; then + echo "ERROR: batched eval concurrency is only supported for lm-eval" >&2 + return 1 + fi + + local eval_conc results_dir eval_rc stage_rc + local completed_concs=() + local failed_concs=() + + for eval_conc in "${eval_concs[@]}"; do + if [[ ! "$eval_conc" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: invalid eval concurrency '${eval_conc}'" >&2 + return 1 + fi + + if ! results_dir=$(mktemp -d /tmp/eval_out-conc"${eval_conc}"-XXXXXX); then + echo "ERROR: failed to create eval output directory for concurrency ${eval_conc}" >&2 + failed_concs+=("$eval_conc") + continue + fi + + echo "Running lm-eval at concurrency ${eval_conc} using the existing engine" + export EVAL_CONCURRENT_REQUESTS="$eval_conc" + export CONC="$eval_conc" + eval_rc=0 + stage_rc=0 + run_lm_eval "${forwarded[@]}" --results-dir "$results_dir" \ + || eval_rc=$? + _stage_lm_eval_artifacts "$results_dir" "$eval_conc" \ + || stage_rc=$? + + if [ "$eval_rc" -eq 0 ] && [ "$stage_rc" -eq 0 ]; then + completed_concs+=("$eval_conc") + else + echo "ERROR: lm-eval failed at concurrency ${eval_conc} (eval_rc=${eval_rc}, stage_rc=${stage_rc})" >&2 + failed_concs+=("$eval_conc") + fi + done + + export EVAL_CONCURRENT_REQUESTS="$requested_concs" + export EVAL_RESULT_DIR="" + export EVAL_BATCHED_CONCS="${eval_concs[*]}" + export EVAL_BATCHED_COMPLETED_CONCS="${completed_concs[*]}" + export EVAL_BATCHED_FAILED_CONCS="${failed_concs[*]}" + + if [ "${#failed_concs[@]}" -gt 0 ]; then + echo "ERROR: batched eval failed for concurrency: ${failed_concs[*]}" >&2 + echo "Deferring failure until post-upload score validation preserves all artifacts" >&2 + fi + return 0 + fi + + local eval_rc=0 + case "$framework" in + lm-eval|lm_eval) run_lm_eval "${forwarded[@]}" || eval_rc=$? ;; + *) echo "Unknown framework '${framework}'"; eval_rc=1 ;; + esac + + if [ "$eval_rc" -ne 0 ]; then + echo "ERROR: run_eval failed with exit code $eval_rc" >&2 + if [ "${EVAL_ONLY}" = "true" ]; then + echo "Eval-only mode: failing after artifact collection" >&2 + return "$eval_rc" + fi + fi + return $eval_rc +} + + +# -------------------------------- +# Agentic trace replay helpers (aiperf driver) +# -------------------------------- + +INFMAX_CONTAINER_WORKSPACE="${INFMAX_CONTAINER_WORKSPACE:-/workspace}" +AGENTIC_DIR="${AGENTIC_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/agentic-benchmark}" +AIPERF_DIR="${AIPERF_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/aiperf}" + +agentic_pip_install() { + local pip_install=(python3 -m pip install) + if python3 -m pip install --help 2>/dev/null | grep -q -- "--break-system-packages"; then + pip_install+=(--break-system-packages) + fi + + "${pip_install[@]}" "$@" +} + +ensure_hf_cli() { + if command -v hf >/dev/null 2>&1; then + return 0 + fi + + # Some lean runtime images used by multinode SGLang include Python but not + # the Hugging Face CLI. Install just the hub CLI before prefetching traces. + agentic_pip_install --quiet "huggingface_hub[cli]>=0.25.0" +} + +resolve_trace_source() { + # Per-recipe override: set WEKA_LOADER_OVERRIDE to one of the aiperf + # public-dataset loader names allowed by the inferencex-agentx-mvp + # scenario. Used by recipes whose servers have non-default context + # caps (e.g. minimaxm2.5 at max_model_len ~256k can't replay the + # unfiltered 052726 corpus and switches to the 256k-capped variant). + local loader="${WEKA_LOADER_OVERRIDE:-semianalysis_cc_traces_weka_with_subagents}" + local dataset + case "$loader" in + semianalysis_cc_traces_weka_with_subagents) + dataset="semianalysisai/cc-traces-weka-with-subagents-052726" + ;; + semianalysis_cc_traces_weka_with_subagents_256k) + dataset="semianalysisai/cc-traces-weka-with-subagents-052726-256k" + ;; + *) + echo "Error: unknown WEKA_LOADER_OVERRIDE='$loader'. Allowed: semianalysis_cc_traces_weka_with_subagents, semianalysis_cc_traces_weka_with_subagents_256k" >&2 + exit 1 + ;; + esac + TRACE_SOURCE_FLAG="--public-dataset $loader" + echo "Loading traces via aiperf public-dataset: $loader ($dataset)" + # Pre-download the dataset into the shared HF_HUB_CACHE (same mount used + # for model weights) so subsequent runs read from cache instead of + # re-downloading every job. + ensure_hf_cli + hf download --repo-type dataset "$dataset" +} + +install_agentic_deps() { + # vllm/vllm-openai container ships without git. pip needs git to + # introspect the aiperf source tree on install. Install on demand; + # no-op when git is already present (e.g. AMD images that ship it). + if ! command -v git >/dev/null 2>&1; then + apt-get update && apt-get install -y git + fi + agentic_pip_install --quiet urllib3 requests 2>/dev/null || true + agentic_pip_install -q -r "$AGENTIC_DIR/requirements.txt" + # Editable install of aiperf from the submodule — gives us the + # `aiperf` CLI plus the inferencex-agentx-mvp scenario plugin. + # + # `--ignore-installed` sidesteps the distutils-uninstall error that + # vLLM containers hit on apt-managed system packages (blinker, etc.) + # when pip's resolver tries to upgrade one of aiperf's transitive + # deps. Installing fresh into the user/site location is safe — the + # system package stays in place and pip's import order picks up our + # newer copy first. + agentic_pip_install -q --ignore-installed -e "$AIPERF_DIR" + # Force-upgrade datasets: containers often ship an older version without + # the `Json` feature type used by the HF traces dataset. `Json` was added + # in datasets 4.7.0 (March 2025). Unpinned installs won't upgrade an + # already-present package. + agentic_pip_install --upgrade "datasets>=4.7.0" +} + +build_replay_cmd() { + # aiperf invocation for the inferencex-agentx-mvp scenario. + # + # Pre-canned assistant replay is the default: recorded assistant responses + # are used for future prompt construction, and live server responses are + # discarded. Set AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES=1 explicitly + # to use live-assistant mode, where the loader emits user-only deltas and + # the worker threads the server's live assistant response back into the + # session. + # + # The scenario plugin locks: --cache-bust first_turn_prefix and + # --trace-idle-gap-cap-seconds 60 (per-trace idle-gap compression + # against parent + subagent request-start timestamps; supersedes the + # legacy --use-think-time-only / --inter-turn-delay-cap-seconds path), + # and auto-injects them — so we do not pass them. See + # utils/aiperf/docs/tutorials/agentx-mvp.md. + local result_dir="$1" + local duration="$DURATION" + + export AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES="${AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES:-0}" + # Dataset configuration (load + reconstruct + inputs.json + mmap) + # routinely takes 4-5 min for the Weka corpus on fast /tmp + # (B300) but can stretch to 14 min on slower /tmp + parallel contention + # (observed on H200 where all 14 R3 jobs hit aiperf's 900s Configure + # Profiling timeout simultaneously). Bump to 1800s to absorb 3x + # worst-case slowdown — the post-setup measurement window is unaffected. + export AIPERF_DATASET_CONFIGURATION_TIMEOUT=1800 + # aiperf validates that SERVICE_PROFILE_CONFIGURE_TIMEOUT >= + # DATASET_CONFIGURATION_TIMEOUT at startup. Bump it in lockstep. + export AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT=1800 + REPLAY_CMD="aiperf profile --scenario inferencex-agentx-mvp" + REPLAY_CMD+=" --url http://localhost:$PORT" + REPLAY_CMD+=" --endpoint /v1/chat/completions" + REPLAY_CMD+=" --endpoint-type chat" + REPLAY_CMD+=" --streaming" + REPLAY_CMD+=" --model $MODEL" + REPLAY_CMD+=" --concurrency $CONC" + REPLAY_CMD+=" --benchmark-duration $duration" + REPLAY_CMD+=" --random-seed 42" + # Fail runs once more than 10% of requests error. This keeps known + # transient low-rate failures from killing long sweeps while still + # catching malformed payloads or server crashes before they get aggregated + # as benchmarkable data. + REPLAY_CMD+=" --failed-request-threshold 0.10" + # Sample each trajectory's warmup start position uniformly from + # [25%, 75%] of the trace's turn count (was hardcoded 0%-70% upstream). + # Avoids starting trajectories right at turn 0 where the KV cache is + # cold and skews early steady-state samples. + REPLAY_CMD+=" --trajectory-start-min-ratio 0.25" + REPLAY_CMD+=" --trajectory-start-max-ratio 0.75" + # Use server-reported usage fields (prompt_tokens / completion_tokens) for + # ISL/OSL instead of client-side tokenizer.encode(). Auto-enables + # stream_options.include_usage on the OpenAI chat endpoint. Skips the + # heavy per-record tokenization in the records pipeline that was pinning + # CPU on minimax-m2.5 at high concurrency. Lossless for vLLM (server + # usage is authoritative). + REPLAY_CMD+=" --use-server-token-count" + # aiperf's dataset manager (separate from the inference parser) loads + # the model's tokenizer for trace-prompt tokenization regardless of + # --use-server-token-count. Models like kimi (amd/Kimi-K2.5-MXFP4, + # moonshotai/Kimi-K2.5) ship a custom tokenizer in their HF repo and + # need trust_remote_code=True to load. Benign for models without + # custom tokenizer code, so we set it unconditionally. + REPLAY_CMD+=" --tokenizer-trust-remote-code" + # Keep replay inputs inside the same context window used to launch the + # server. The WEKA corpus contains a few very long parent/subagent traces; + # if we mmap and replay them against a smaller-context server they become + # deterministic 4xxs and can still pressure the engine while queued. + if [ -n "${MAX_MODEL_LEN:-}" ] && [ "$MAX_MODEL_LEN" != "0" ]; then + REPLAY_CMD+=" --max-context-length $MAX_MODEL_LEN" + fi + # Default --num-dataset-entries is 100; the with-subagents Weka corpus + # has 472. Cap at 472 so all unique traces are loaded (the loader treats + # this as a ``min(cap, available)`` ceiling, not a target — see + # semianalysis_cc_traces_weka.py). + REPLAY_CMD+=" --num-dataset-entries 472" + # 1-second timeslices on the server-metrics scrape so the post-run + # plotter has per-window time series (KV usage, cache hit rate, + # throughput, etc.). Matches kv-cache-tester's poll_interval=1.0 + # snapshot cadence so metrics_plots.png is visually comparable. + # Without this, aiperf only emits aggregate stats and the 6x2 panels + # collapse to flat lines. + REPLAY_CMD+=" --slice-duration 1.0" + REPLAY_CMD+=" --output-artifact-dir $result_dir/aiperf_artifacts" + # The inferencex-agentx-mvp scenario enforces a 900s minimum + # benchmark duration. For smoke tests with shorter durations, opt + # into --unsafe-override (the run's submission_valid will be flagged + # false; that's expected for non-canonical runs). + if [ "$duration" -lt 900 ] || [ "${AIPERF_UNSAFE_OVERRIDE:-false}" = "true" ]; then + REPLAY_CMD+=" --unsafe-override" + fi + REPLAY_CMD+=" $TRACE_SOURCE_FLAG" +} + +write_agentic_result_json() { + # Aggregate aiperf's profile_export.{json,jsonl} + server_metrics_export.json + # into $AGENTIC_OUTPUT_DIR/$RESULT_FILENAME.json. The workflow's existing + # retry-based existence check is the single success gate. + local result_dir="$1" + RESULT_DIR="$result_dir" AGENTIC_OUTPUT_DIR="${AGENTIC_OUTPUT_DIR:-$INFMAX_CONTAINER_WORKSPACE}" \ + python3 "$INFMAX_CONTAINER_WORKSPACE/utils/process_agentic_result.py" + + # Generate metrics_plots.png from the same aiperf artifacts. Best-effort: + # don't fail the launcher if plot generation has trouble (e.g. matplotlib + # missing in a stripped-down image). The agg JSON is the success gate. + python3 "$INFMAX_CONTAINER_WORKSPACE/utils/generate_aiperf_plots.py" "$result_dir" 2>&1 || true +} + +run_agentic_replay_and_write_outputs() { + local result_dir="$1" + local replay_rc + + echo "$REPLAY_CMD" > "$result_dir/benchmark_command.txt" + + set +e + set -x + $REPLAY_CMD 2>&1 | tee "$result_dir/benchmark.log" + replay_rc=${PIPESTATUS[0]} + set +x + set -e + + write_agentic_result_json "$result_dir" + + python3 "$AGENTIC_DIR/scripts/analyze_benchmark_distributions.py" \ + "$result_dir/aiperf_artifacts" -o "$result_dir" 2>&1 || true + + if [ "$replay_rc" -ne 0 ]; then + echo "ERROR: agentic trace replay exited with code $replay_rc after writing available results" >&2 + return "$replay_rc" + fi +} diff --git a/scripts/Qwen3.5/lib/placement.py b/scripts/Qwen3.5/lib/placement.py new file mode 100755 index 00000000..47b7c1d1 --- /dev/null +++ b/scripts/Qwen3.5/lib/placement.py @@ -0,0 +1,50 @@ +#!/usr/bin/env python3 +"""placement.py — compute replica placement from (TP, DP) across cluster nodes. + +Core abstraction for scalable serving of a model that FITS IN ONE NODE: +each replica is node-local (TP <= gpus_per_node), and we place +(DP per node) replicas on each node, replicated across all nodes. + + replicas_per_node = DP (caller picks DP; DP*TP must be <= gpus_per_node) + total_replicas = len(nodes) * DP + full utilization <=> DP * TP == gpus_per_node + +Emits one line per replica: NODE GPUS PORT REPLICA_GLOBAL_IDX +e.g. node-01 0,1,2,3 8000 0 +""" +import sys, yaml, os + +def load(p): + with open(p) as f: return yaml.safe_load(f) + +def main(): + import argparse + ap = argparse.ArgumentParser() + ap.add_argument("--cluster", required=True) + ap.add_argument("--tp", type=int, required=True) + ap.add_argument("--dp", type=int, default=None, + help="replicas per node; default = gpus_per_node // tp (full utilization)") + ap.add_argument("--nodes", default=None, + help="comma-separated node override (default: cluster.yaml nodes)") + args = ap.parse_args() + + c = load(args.cluster) + G = int(c["gpus_per_node"]); port_base = int(c.get("port_base", 8000)) + nodes = args.nodes.split(",") if args.nodes else c["nodes"] + tp = args.tp + dp = args.dp if args.dp is not None else (G // tp) + + if tp > G: + sys.exit(f"ERROR: TP={tp} > gpus_per_node={G} (model must fit in one node; no multi-node TP)") + if dp * tp > G: + sys.exit(f"ERROR: DP*TP = {dp}*{tp} = {dp*tp} > gpus_per_node={G}") + + for ni, node in enumerate(nodes): + for r in range(dp): # r = replica index on this node + gpus = ",".join(str(g) for g in range(r*tp, (r+1)*tp)) + port = port_base + r + gidx = ni * dp + r # global replica index + print(f"{node}\t{gpus}\t{port}\t{gidx}") + +if __name__ == "__main__": + main() diff --git a/scripts/Qwen3.5/lib/prompts.json b/scripts/Qwen3.5/lib/prompts.json new file mode 100644 index 00000000..ba98e3ef --- /dev/null +++ b/scripts/Qwen3.5/lib/prompts.json @@ -0,0 +1,32 @@ +[ + { + "id": "factual", + "prompt": "What is the capital of France? Answer in one word.", + "expect_contains": ["Paris"], + "max_tokens": 512 + }, + { + "id": "arithmetic", + "prompt": "Compute 17 multiplied by 23. Give only the number.", + "expect_contains": ["391"], + "max_tokens": 512 + }, + { + "id": "code", + "prompt": "Write a Python one-liner that returns the sum of a list named nums. Output only the code.", + "expect_contains": ["sum(", "nums"], + "max_tokens": 512 + }, + { + "id": "reasoning", + "prompt": "If a train travels 60 km in 1.5 hours, what is its average speed in km/h? Give only the number.", + "expect_contains": ["40"], + "max_tokens": 512 + }, + { + "id": "instruction", + "prompt": "List exactly three primary colors, comma-separated, lowercase.", + "expect_contains": ["red", "blue"], + "max_tokens": 512 + } +] diff --git a/scripts/Qwen3.5/lib/replica_entry.sh b/scripts/Qwen3.5/lib/replica_entry.sh new file mode 100755 index 00000000..4baf1cf8 --- /dev/null +++ b/scripts/Qwen3.5/lib/replica_entry.sh @@ -0,0 +1,110 @@ +#!/bin/bash +# replica_entry.sh — runs INSIDE one replica's container. +# Starts the engine server on $PORT, waits healthy, then runs ACTION. +# Env (from launcher): ENGINE_KIND, MODEL_PATH, PORT, TP, SERVE_FLAGS, SCHED, ACTION +# ACTION=sanity -> 5-prompt accuracy + perf smoke (conc32 128/128) +# ACTION=serve -> boot + stay up (for router-fronted serving) +set -uo pipefail +SC=/scripts +# IMPORTANT: do NOT source lib_inferencex.sh before launching the server. +# Sourcing it perturbs the env (PYTHON*/cache state) and makes ATOM's torch.distributed.nn +# JIT remote_module instantiation crash with `IndexError: string index out of range` in +# importlib cache_from_source. ATOM works perfectly when launched in a clean env (== the +# bare runbook command). So: launch server FIRST, source the bench lib AFTER. + +KIND="${ENGINE_KIND:-vllm}" +MODEL_PATH="${MODEL_PATH:?}"; PORT="${PORT:?}"; TP="${TP:?}" +ACTION="${ACTION:-sanity}" +SERVER_LOG=/out/server.log +# Shared logical knobs (from model.yaml `defaults`, resolved+passed by run_engine.sh). +# Translated below to each engine's specific flag name so ALL engines get the SAME values. +MAX_MODEL_LEN="${MAX_MODEL_LEN:-8192}" +GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.8}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-256}" + +start_server() { + local sched="--no-async-scheduling"; [[ "$TP" -gt 1 ]] && sched="--async-scheduling" + case "$KIND" in + vllm) + vllm serve "$MODEL_PATH" --host 0.0.0.0 --port "$PORT" --tensor-parallel-size "$TP" \ + --max-model-len "$MAX_MODEL_LEN" --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --max-num-seqs "$MAX_NUM_SEQS" $sched ${SERVE_FLAGS:-} > "$SERVER_LOG" 2>&1 & ;; + sglang) + # SGLang flag names differ: --context-length (ctx), --mem-fraction-static (mem-util), --max-running-requests (seqs) + python3 -m sglang.launch_server --model-path "$MODEL_PATH" --host 0.0.0.0 --port "$PORT" \ + --tensor-parallel-size "$TP" --context-length "$MAX_MODEL_LEN" \ + --mem-fraction-static "$GPU_MEM_UTIL" --max-running-requests "$MAX_NUM_SEQS" \ + ${SERVE_FLAGS:-} > "$SERVER_LOG" 2>&1 & ;; + atom) + # bare runbook invocation, clean env. Same shared knobs as vLLM (matching flag names). + python -m atom.entrypoints.openai_server --model "$MODEL_PATH" -tp "$TP" --port "$PORT" \ + --max-model-len "$MAX_MODEL_LEN" --gpu-memory-utilization "$GPU_MEM_UTIL" \ + --max-num-seqs "$MAX_NUM_SEQS" ${SERVE_FLAGS:-} > "$SERVER_LOG" 2>&1 & ;; + *) echo "bad ENGINE_KIND=$KIND"; exit 2 ;; + esac + SERVER_PID=$! +} + +# 1) launch server in CLEAN env +start_server + +# 2) NOW source the InferenceX helpers (health-wait + bench) — after the server process exists +source "$SC/lib/lib_inferencex.sh" +export EVAL_ONLY=false RUN_EVAL=false PROFILE=0 +export SGLANG_TORCH_PROFILER_DIR="" VLLM_TORCH_PROFILER_DIR="" GPU_MONITOR_PID="" + +cleanup(){ [[ -n "${SERVER_PID:-}" ]] && kill "$SERVER_PID" 2>/dev/null || true; } +trap cleanup EXIT + +wait_for_server_ready --port "$PORT" --server-log "$SERVER_LOG" --server-pid "$SERVER_PID" +echo "[replica] server healthy on :$PORT" + +if [[ "$ACTION" == "serve" ]]; then + echo "[replica] ACTION=serve — staying up"; wait "$SERVER_PID"; exit 0 +fi + +# ACTION=sweep: accuracy gate, then loop SWEEP_SHAPES x SWEEP_CONC through the bench client. +# For BENCH_EXTERNAL engines (ATOM), signal ready + hold; run_engine.sh drives the external sweep. +if [[ "$ACTION" == "sweep" ]]; then + PORT="$PORT" MODEL="$MODEL_PATH" OUT=/out/sanity_accuracy.json python3 "$SC/lib/check_accuracy.py" || true + if [[ "${BENCH_EXTERNAL:-0}" == "1" ]]; then + echo "[replica] sweep + BENCH_EXTERNAL=1 — ready; external sweep driven by host"; touch /out/SERVER_READY + wait "$SERVER_PID"; exit 0 + fi + SWEEP_SHAPES="${SWEEP_SHAPES:-1024,1024 8192,1024 16384,1024}" + SWEEP_CONC="${SWEEP_CONC:-4 8 16 32 64 128 256}" + for shape in $SWEEP_SHAPES; do + ISL=${shape%,*}; OSL=${shape#*,} + for c in $SWEEP_CONC; do + echo "[sweep] isl=$ISL osl=$OSL conc=$c" + run_benchmark_serving --model "$MODEL_PATH" --port "$PORT" --backend vllm \ + --input-len "$ISL" --output-len "$OSL" --random-range-ratio 1.0 \ + --num-prompts "$((c * 10))" --max-concurrency "$c" \ + --result-filename "sweep_isl${ISL}_osl${OSL}_c${c}" --result-dir /out \ + --bench-serving-dir "$SC" || echo "[sweep] FAILED isl=$ISL osl=$OSL conc=$c" + done + done + echo "[replica] sweep done"; exit 0 +fi + +# ACTION=sanity: accuracy (always — pure urllib, safe in any image) then perf smoke. +SMOKE_CONC="${SMOKE_CONC:-8}"; SMOKE_ISL="${SMOKE_ISL:-128}"; SMOKE_OSL="${SMOKE_OSL:-128}" +PORT="$PORT" MODEL="$MODEL_PATH" OUT=/out/sanity_accuracy.json python3 "$SC/lib/check_accuracy.py" || true + +# BENCH_EXTERNAL=1 (e.g. ATOM): the vendored bench client crashes in this image's Python env +# (importlib cache_from_source IndexError). Skip in-container bench; the host (run_engine.sh) +# runs it from a clean container against this server's port, then this replica stays up until killed. +if [[ "${BENCH_EXTERNAL:-0}" == "1" ]]; then + echo "[replica] BENCH_EXTERNAL=1 — server ready for external bench; holding on :$PORT" + touch /out/SERVER_READY + wait "$SERVER_PID"; exit 0 +fi + +# NOTE: run_benchmark_serving (lib_inferencex.sh) already injects --num-warmups=2*conc internally; +# do NOT pass --num-warmups here (its arg parser rejects it -> "Unknown parameter"). +run_benchmark_serving --model "$MODEL_PATH" --port "$PORT" --backend vllm \ + --input-len "$SMOKE_ISL" --output-len "$SMOKE_OSL" --random-range-ratio 1.0 \ + --num-prompts "$((SMOKE_CONC * 10))" --max-concurrency "$SMOKE_CONC" \ + --result-filename "perf_c${SMOKE_CONC}_${SMOKE_ISL}x${SMOKE_OSL}" --result-dir /out \ + --bench-serving-dir "$SC" || true +echo "[replica] sanity done" diff --git a/scripts/Qwen3.5/lib/run_engine.sh b/scripts/Qwen3.5/lib/run_engine.sh new file mode 100755 index 00000000..10b2f4c4 --- /dev/null +++ b/scripts/Qwen3.5/lib/run_engine.sh @@ -0,0 +1,151 @@ +#!/bin/bash +# run_engine.sh — generic launcher for ANY engine across the cluster, given (ENGINE, MODEL, TP, DP). +# Called by run_vllm.sh / run_sglang.sh / run_atom.sh (which just set ENGINE). +# Aggregated serving; each replica node-local. Scales 1->6 nodes via cluster.yaml `nodes`. +# +# Env in: ENGINE, MODEL, TP, [DP], [ACTION=sanity|serve], [MAX_MODEL_LEN] +set -uo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" # scripts_qwen_v2 +CLUSTER="$HERE/cluster.yaml"; MODELS="$HERE/model.yaml" +ENGINE="${ENGINE:?}"; MODEL="${MODEL:?}"; TP="${TP:?}"; ACTION="${ACTION:-sanity}" +cfg(){ python3 "$HERE/lib/cfg.py" "$@"; } + +RESV=$(cfg "$CLUSTER" cluster reservation) +ALLOC=$(cfg "$CLUSTER" cluster srun_alloc); ALLOC="${ALLOC:---reservation=$RESV}" +GPN=$(cfg "$CLUSTER" cluster gpus_per_node) +MROOT=$(cfg "$CLUSTER" cluster models_root) +RROOT=$(cfg "$CLUSTER" cluster results_root) +DP="${DP:-$((GPN / TP))}" +MDIR=$(cfg "$MODELS" model "$MODEL" dir) +VALID_TP=$(cfg "$MODELS" model "$MODEL" valid_tp) +IMG=$(cfg "$MODELS" engine "$ENGINE" image) +SERVE_FLAGS=$(cfg "$MODELS" engine "$ENGINE" serve_flags) +ENTRY=$(cfg "$MODELS" engine "$ENGINE" entrypoint_override); ENTRY="${ENTRY:-bash}" +PATCH=$(cfg "$MODELS" engine "$ENGINE" patch_mount) +BENCH_EXTERNAL=$(cfg "$MODELS" engine "$ENGINE" bench_external) # "True"/"" +BENCH_IMAGE=$(cfg "$MODELS" engine "$ENGINE" bench_image) +[[ "$BENCH_EXTERNAL" == "True" ]] && BX=1 || BX=0 + +# shared logical knobs: env override > per-model override > global defaults{} +MAX_MODEL_LEN="${MAX_MODEL_LEN:-$(cfg "$MODELS" default "$MODEL" max_model_len)}" +GPU_MEM_UTIL="${GPU_MEM_UTIL:-$(cfg "$MODELS" default "$MODEL" gpu_memory_util)}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-$(cfg "$MODELS" default "$MODEL" max_num_seqs)}" + +case " $VALID_TP " in *" $TP "*) :;; *) echo "ERROR: TP=$TP not in valid_tp=[$VALID_TP] for $MODEL"; exit 1;; esac +[[ $((DP*TP)) -le $GPN ]] || { echo "ERROR: DP*TP=$((DP*TP)) > gpus_per_node=$GPN"; exit 1; } + +# engine env vars (space-separated k=v from yaml) -> docker -e args +ENV_ARGS="" +for kv in $(cfg "$MODELS" engine "$ENGINE" env); do ENV_ARGS+=" -e $kv"; done +PATCH_ARG=""; [[ -n "$PATCH" ]] && PATCH_ARG="-v $HERE/$PATCH" + +STAMP=$(python3 -c "import time;print(time.strftime('%Y%m%d_%H%M%S'))") +RUNROOT="$RROOT/${MODEL}_${ENGINE}_tp${TP}_dp${DP}_${STAMP}"; mkdir -p "$RUNROOT"; chmod 777 "$RUNROOT" +echo "[$ENGINE] MODEL=$MODEL TP=$TP DP=$DP action=$ACTION image=$IMG" +NODES_ARG=""; [[ -n "${NODES:-}" ]] && NODES_ARG="--nodes $NODES" +python3 "$HERE/lib/placement.py" --cluster "$CLUSTER" --tp "$TP" --dp "$DP" $NODES_ARG > "$RUNROOT/placement.tsv" +echo "[$ENGINE] placement ($(wc -l < "$RUNROOT/placement.tsv") replicas):"; sed 's/^/ /' "$RUNROOT/placement.tsv" +echo "[$ENGINE] config: max_model_len=$MAX_MODEL_LEN gpu_mem_util=$GPU_MEM_UTIL max_num_seqs=$MAX_NUM_SEQS" + +# --- PRE-FLIGHT: ensure each target node's GPUs are CLEAN before deploying (kills zombie VRAM) --- +UNIQ_NODES=$(cut -f1 "$RUNROOT/placement.tsv" | sort -u) +echo "[$ENGINE] cleaning GPUs on: $(echo $UNIQ_NODES | tr '\n' ' ')" +for node in $UNIQ_NODES; do + srun $ALLOC --nodelist="$node" --nodes=1 --ntasks=1 --overlap --mem=2G \ + bash "$HERE/lib/clean_node.sh" > "$RUNROOT/clean_${node}.log" 2>&1 \ + && echo " [clean] $node OK" || echo " [clean] $node WARN (see clean_${node}.log)" +done + +# --- IN-CONTAINER BENCH (vLLM/SGLang, BX=0) ------------------------------------------------- +# The bench client runs INSIDE the engine container via replica_entry.sh (ACTION drives it). +# One backgrounded --overlap srun per replica; replicas are independent and run concurrently. +launch_replica() { + local node="$1" gpus="$2" port="$3" idx="$4" + local rdir="$RUNROOT/replica_${idx}_${node}_p${port}"; mkdir -p "$rdir"; chmod 777 "$rdir" + srun $ALLOC --nodelist="$node" --nodes=1 --ntasks=1 --overlap --mem=0 bash -c " + docker run --rm --entrypoint $ENTRY \ + --device=/dev/kfd --device=/dev/dri --group-add video \ + --ipc host --shm-size 64G --network host \ + --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e HIP_VISIBLE_DEVICES=$gpus -e CUDA_VISIBLE_DEVICES=$gpus \ + $ENV_ARGS \ + -e ENGINE_KIND=$ENGINE -e ACTION=$ACTION -e PORT=$port -e TP=$TP -e BENCH_EXTERNAL=$BX \ + -e MAX_MODEL_LEN=$MAX_MODEL_LEN -e GPU_MEM_UTIL=$GPU_MEM_UTIL -e MAX_NUM_SEQS=$MAX_NUM_SEQS \ + -e SMOKE_CONC=${SMOKE_CONC:-8} -e SMOKE_ISL=${SMOKE_ISL:-128} -e SMOKE_OSL=${SMOKE_OSL:-128} \ + -e SWEEP_SHAPES='${SWEEP_SHAPES:-1024,1024 8192,1024 16384,1024}' -e SWEEP_CONC='${SWEEP_CONC:-4 8 16 32 64 128 256}' \ + -e MODEL_PATH=$MROOT/$MDIR -e SERVE_FLAGS='$SERVE_FLAGS' \ + -v $MROOT:$MROOT -v $HERE:/scripts -v $rdir:/out $PATCH_ARG \ + $IMG /scripts/lib/replica_entry.sh + " > "$rdir/replica.log" 2>&1 & +} + +# --- EXTERNAL BENCH (ATOM, BX=1) ------------------------------------------------------------ +# ATOM's vendored bench client crashes inside its own image, so the bench must run from a CLEAN +# (vLLM) container. CRITICAL: server + bench run as SIBLING containers inside ONE --exclusive +# srun on the node. The old design used a 2nd `--overlap` bench srun, which DEADLOCKED: the +# server's `--mem=0` srun claims all node memory, so the bench srun's `--mem=4G` request never +# schedules. One srun = one allocation = no contention. (Proven pattern; see git history.) +launch_external_bench_node() { + local node="$1" gpus="$2" port="$3" idx="$4" + local rdir="$RUNROOT/replica_${idx}_${node}_p${port}"; mkdir -p "$rdir"; chmod 777 "$rdir" + local cname="extbench_${ENGINE}_tp${TP}_${idx}" + local shapes="${SWEEP_SHAPES:-1024,1024 8192,1024 16384,1024}" + local concs="${SWEEP_CONC:-4 8 16 32 64 128 256}" + [[ "$ACTION" == "sanity" ]] && { shapes="${SMOKE_ISL:-128},${SMOKE_OSL:-128}"; concs="${SMOKE_CONC:-8}"; } + srun $ALLOC --nodelist="$node" --nodes=1 --ntasks=1 --exclusive bash -c " + set -uo pipefail + docker rm -f $cname >/dev/null 2>&1 || true + sleep 4 + echo '[node] start $ENGINE TP=$TP server on GPUs $gpus :$port' + # Detached (-d): get a real container id and stream logs; avoids the empty-log launch race. + cid=\$(docker run -d --name $cname --entrypoint $ENTRY \ + --device=/dev/kfd --device=/dev/dri --group-add video --ipc host --shm-size 64G --network host \ + --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e HIP_VISIBLE_DEVICES=$gpus -e CUDA_VISIBLE_DEVICES=$gpus $ENV_ARGS \ + -e ENGINE_KIND=$ENGINE -e ACTION=serve -e PORT=$port -e TP=$TP \ + -e MAX_MODEL_LEN=$MAX_MODEL_LEN -e GPU_MEM_UTIL=$GPU_MEM_UTIL -e MAX_NUM_SEQS=$MAX_NUM_SEQS \ + -e MODEL_PATH=$MROOT/$MDIR -e SERVE_FLAGS='$SERVE_FLAGS' \ + -v $MROOT:$MROOT -v $HERE:/scripts -v $rdir:/out $PATCH_ARG \ + $IMG /scripts/lib/replica_entry.sh) + echo \"[node] container \$cid\" + up=0; for i in \$(seq 1 180); do + docker logs $cname > $rdir/server.log 2>&1 || true + curl -sf http://localhost:$port/health >/dev/null 2>&1 && { up=1; break; } + st=\$(docker inspect -f '{{.State.Status}}' $cname 2>/dev/null || echo missing) + { [ \"\$st\" = exited ] || [ \"\$st\" = missing ]; } && { echo \"[node] server \$st\"; break; } + sleep 5 + done + docker logs $cname > $rdir/server.log 2>&1 || true + echo \"[node] SERVER_UP=\$up\" + [ \$up -ne 1 ] && { echo '[node] FAILED tail:'; tail -25 $rdir/server.log; docker rm -f $cname >/dev/null 2>&1; exit 1; } + for shape in $shapes; do + ISL=\${shape%,*}; OSL=\${shape#*,} + for c in $concs; do + echo \"[bench] isl=\$ISL osl=\$OSL conc=\$c\" + docker run --rm --entrypoint bash --network host -v $HERE:/scripts -v $MROOT:$MROOT -v $rdir:/out \ + $BENCH_IMAGE -c \"python3 /scripts/utils/bench_serving/benchmark_serving.py \ + --model $MROOT/$MDIR --backend vllm --base-url http://0.0.0.0:$port --dataset-name random \ + --random-input-len \$ISL --random-output-len \$OSL --random-range-ratio 1.0 \ + --num-prompts \$(( c * 10 )) --num-warmups \$(( c * 2 )) --max-concurrency \$c \ + --request-rate inf --ignore-eos --save-result --percentile-metrics ttft,tpot,itl,e2el \ + --result-dir /out --result-filename sweep_isl\${ISL}_osl\${OSL}_c\${c}.json\" \ + > $rdir/bench_isl\${ISL}_osl\${OSL}_c\${c}.log 2>&1 || echo \" [bench] FAILED isl=\$ISL conc=\$c\" + done + done + echo '[node] standdown'; docker rm -f $cname >/dev/null 2>&1; sleep 3 + " > "$rdir/replica.log" 2>&1 & +} + +n=0 +while IFS=$'\t' read -r node gpus port idx; do + if [[ "$BX" == "1" && ( "$ACTION" == "sanity" || "$ACTION" == "sweep" ) ]]; then + launch_external_bench_node "$node" "$gpus" "$port" "$idx" # ATOM: server+bench in one srun + else + launch_replica "$node" "$gpus" "$port" "$idx" # vLLM/SGLang: in-container bench + fi + echo "[$ENGINE] launched replica $idx on $node gpus=$gpus port=$port"; n=$((n+1)) +done < "$RUNROOT/placement.tsv" +echo "[$ENGINE] $n replicas launching -> $RUNROOT" + +wait +echo "[$ENGINE] done -> $RUNROOT (check replica_*/replica.log, server.log, sweep_*.json)" From 61beced2b6a1106def892e8a96a5dba5a2c369e0 Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Wed, 8 Jul 2026 19:40:17 -0500 Subject: [PATCH 04/12] Qwen3.5 README: add highlighted 4-step Quickstart for the default config Makes the recommended SGLang+MXFP4+TP1 default runnable cold: download weights -> edit cluster.yaml -> one run_sglang.sh command -> query. Spells out the prerequisites (weights path, cluster.yaml placeholders, Slurm alloc) that were previously implicit. Co-Authored-By: Claude --- scripts/Qwen3.5/README.md | 20 ++++++++++++++++++++ 1 file changed, 20 insertions(+) diff --git a/scripts/Qwen3.5/README.md b/scripts/Qwen3.5/README.md index f15020ba..4ff092c6 100644 --- a/scripts/Qwen3.5/README.md +++ b/scripts/Qwen3.5/README.md @@ -19,6 +19,26 @@ MODEL=qwen35-moe-mxfp4 TP=1 ACTION=sweep ./run_sglang.sh Why TP1 and not TP-shard: a 3B-active model fits one GPU, so tensor-parallel only adds cross-GPU comm cost. Fill the node with independent instances instead. (ATOM excluded — crashes at ISL≥8192.) +### Quickstart (run the default in 4 steps) +```bash +# 1. Get the weights (MXFP4, ~23 GB) to a path visible on every node: +hf download amd/Qwen3.5-35B-A3B-MXFP4 --local-dir /path/to/models/Qwen3.5-35B-A3B-MXFP4 + +# 2. Edit cluster.yaml: set `nodes:` to your node(s), `models_root: /path/to/models`, +# and `srun_alloc:` to your reservation/jobid (or "" if nodes are free). + +# 3. From inside a Slurm allocation on those nodes, run the default: +cd scripts/Qwen3.5 +MODEL=qwen35-moe-mxfp4 TP=1 DP=8 ACTION=serve ./run_sglang.sh # 8 MXFP4 instances/node + +# 4. Query any instance (OpenAI-compatible, ports 8000..8007): +curl http://localhost:8000/v1/chat/completions -H 'Content-Type: application/json' \ + -d '{"model":"/path/to/models/Qwen3.5-35B-A3B-MXFP4", + "messages":[{"role":"user","content":"Hello"}],"max_tokens":256}' +``` +`ACTION=sweep` instead of `serve` reproduces the benchmark. Notes: it's a **reasoning model** (allow +generous `max_tokens`); the SGLang image + all flags are preset in `model.yaml` — nothing else to set. + ## Design in one picture ``` Each node (8 GPUs): DP replicas, each TP-sharded. DP × TP ≤ 8. Full util ⇔ DP×TP = 8. From f2239a4ee03e7067896a2a61f2716d93268f319f Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Thu, 16 Jul 2026 07:02:03 +0000 Subject: [PATCH 05/12] Add MiniMax-M3 MXFP4 disaggregated P/D inference recipe (scripts/MiniMax-M3) Multi-node prefill/decode disaggregated serving/benchmark harness for MiniMax-M3 (MXFP4) on AMD MI355X (gfx950) via ATOM (atomesh + mooncake RDMA), ported from the InferenceX minimaxm3-fp4-mi355x-atom-disagg config (top-throughput MiniMax-M3 disagg config on MI355X: ~10,829 tok/s/GPU at 8k/1k, 2P1D+DPA). All workers TP4 (no EP; TP4 beats TP8/TP4-EP on gfx950); STP only (DECODE_MTP_SIZE=0). Topologies 1p1d (2 nodes) and 2p1d_dpa (3 nodes). Mirrors the DeepSeekV4 harness structure (model.yaml/cluster.yaml/lib/atom_disagg/utils). Scaffolded from InferenceX; needs a live gfx950 dry-run once amd/MiniMax-M3-MXFP4 weights are staged. Also adds the MiniMax-M3 row to the root README models table. Co-Authored-By: Claude --- README.md | 1 + scripts/MiniMax-M3/README.md | 72 + .../MiniMax-M3/atom_disagg/accuracy_ctx.py | 73 + scripts/MiniMax-M3/atom_disagg/bench.sh | 47 + scripts/MiniMax-M3/atom_disagg/env_atom.sh | 75 + scripts/MiniMax-M3/atom_disagg/launch.sh | 132 ++ scripts/MiniMax-M3/atom_disagg/server_atom.sh | 585 ++++++++ scripts/MiniMax-M3/atom_disagg/setup_deps.sh | 11 + scripts/MiniMax-M3/cluster.yaml | 28 + scripts/MiniMax-M3/docs/proxy_and_disagg.md | 52 + scripts/MiniMax-M3/lib/cfg.py | 34 + scripts/MiniMax-M3/lib/check_accuracy.py | 59 + scripts/MiniMax-M3/lib/clean_node.sh | 38 + scripts/MiniMax-M3/lib/lib_inferencex.sh | 1281 +++++++++++++++++ scripts/MiniMax-M3/lib/prompts.json | 32 + scripts/MiniMax-M3/lib/run_disagg.sh | 192 +++ scripts/MiniMax-M3/lib/topo.py | 52 + scripts/MiniMax-M3/model.yaml | 64 + scripts/MiniMax-M3/run_atom_disagg.sh | 5 + .../utils/bench_serving/KNOWN_LIMITATION.md | 7 + .../bench_serving/backend_request_func.py | 574 ++++++++ .../utils/bench_serving/benchmark_serving.py | 1276 ++++++++++++++++ .../utils/bench_serving/benchmark_utils.py | 39 + .../utils/bench_serving/encoding_dsv4.py | 744 ++++++++++ 24 files changed, 5473 insertions(+) create mode 100644 scripts/MiniMax-M3/README.md create mode 100644 scripts/MiniMax-M3/atom_disagg/accuracy_ctx.py create mode 100644 scripts/MiniMax-M3/atom_disagg/bench.sh create mode 100644 scripts/MiniMax-M3/atom_disagg/env_atom.sh create mode 100644 scripts/MiniMax-M3/atom_disagg/launch.sh create mode 100644 scripts/MiniMax-M3/atom_disagg/server_atom.sh create mode 100644 scripts/MiniMax-M3/atom_disagg/setup_deps.sh create mode 100644 scripts/MiniMax-M3/cluster.yaml create mode 100644 scripts/MiniMax-M3/docs/proxy_and_disagg.md create mode 100644 scripts/MiniMax-M3/lib/cfg.py create mode 100644 scripts/MiniMax-M3/lib/check_accuracy.py create mode 100644 scripts/MiniMax-M3/lib/clean_node.sh create mode 100644 scripts/MiniMax-M3/lib/lib_inferencex.sh create mode 100644 scripts/MiniMax-M3/lib/prompts.json create mode 100644 scripts/MiniMax-M3/lib/run_disagg.sh create mode 100644 scripts/MiniMax-M3/lib/topo.py create mode 100644 scripts/MiniMax-M3/model.yaml create mode 100644 scripts/MiniMax-M3/run_atom_disagg.sh create mode 100644 scripts/MiniMax-M3/utils/bench_serving/KNOWN_LIMITATION.md create mode 100644 scripts/MiniMax-M3/utils/bench_serving/backend_request_func.py create mode 100644 scripts/MiniMax-M3/utils/bench_serving/benchmark_serving.py create mode 100644 scripts/MiniMax-M3/utils/bench_serving/benchmark_utils.py create mode 100644 scripts/MiniMax-M3/utils/bench_serving/encoding_dsv4.py diff --git a/README.md b/README.md index da9004f1..047ca027 100644 --- a/README.md +++ b/README.md @@ -26,6 +26,7 @@ Below are blueprints of supported models along with their documentation. | [**SGLang disaggregated P/D inference with WideEP/LargeEP**](scripts/sglang_disagg/README.MD) | Distributed Inference P/D disaggregation with SGLang with WideEP/LargeEP | DeepSeek-V3, DeepSeek-R1 | | [**DeepSeek-V4-Pro disaggregated P/D inference**](scripts/DeepSeekV4/README.md) | Distributed Inference P/D disaggregation for DeepSeek-V4-Pro (ATOM+mooncake / SGLang+MoRI) | DeepSeek-V4-Pro | | [**Qwen3.5-35B-A3B aggregated inference**](scripts/Qwen3.5/README.md) | Single-node aggregated serving/benchmark (SGLang/vLLM); recommended default SGLang+MXFP4+TP1 (8 instances/node) | Qwen3.5-35B-A3B (MXFP4, FP8) | +| [**MiniMax-M3 disaggregated P/D inference**](scripts/MiniMax-M3/README.md) | Distributed Inference P/D disaggregation for MiniMax-M3 on MI355X (ATOM+mooncake, MXFP4); top-throughput InferenceX config | MiniMax-M3 (MXFP4) | | [**KVCache Transfer Bench**](scripts/kvcache_transfer_bench/README.md) | Inter-node Transfer Benchmark | no specific models | ## Table of Contents diff --git a/scripts/MiniMax-M3/README.md b/scripts/MiniMax-M3/README.md new file mode 100644 index 00000000..870e21f5 --- /dev/null +++ b/scripts/MiniMax-M3/README.md @@ -0,0 +1,72 @@ +# MiniMax-M3 — disaggregated (P/D) serving/benchmark harness + +Multi-node **prefill/decode disaggregated** serving for MiniMax-M3 (MXFP4) on AMD Instinct +MI355X (gfx950), ported from InferenceX (`minimaxm3-fp4-mi355x-atom-disagg`). + +## Why disagg +MiniMax-M3 is served **disaggregated**: prefill workers (kv_producer) and decode workers +(kv_consumer) on separate nodes, KV streamed over RDMA, fronted by ONE router. Bench drives the +router → a TRUE aggregate number (no ×N extrapolation). + +**All workers are TP4, no EP.** The single-node M3 recipe found plain TP4 beats both TP8 and +TP4/EP4 on tok/s/GPU on gfx950, so prefill and decode both use TP4 and we tune the +prefill:decode worker ratio (xP:yD) instead of TP. STP only (no MTP): `DECODE_MTP_SIZE=0`. + +## Engine & router +| Engine | Router | KV transfer | Status | +|--------|--------|-------------|--------| +| **atom-disagg** | **atomesh** | **mooncake** (RDMA) | PRIMARY — InferenceX-validated M3 MXFP4 config | + +Precision: **MXFP4** (`amd/MiniMax-M3-MXFP4`). Image: `rocm/atom-dev:nightly_202607011530`. + +## Files +| File | Role | +|------|------| +| `cluster.yaml` | nodes, gpus/node, reservation, router/prefill/decode/handshake ports, paths | +| `model.yaml` | M3 def + shared `defaults{}` + atom-disagg image/router/kv + `topologies{}` | +| `run_atom_disagg.sh` | thin wrapper → `lib/run_disagg.sh` | +| `lib/run_disagg.sh` | orchestrator: clean nodes → launch prefill/decode servers → start router → bench → teardown | +| `lib/topo.py` | resolve a topology (1p1d / 2p1d_dpa) → node-role placement (prefill/decode, ports) | +| `lib/cfg.py` | yaml reader | +| `lib/clean_node.sh` | pre-flight GPU clean | +| `atom_disagg/` | vendored InferenceX ATOM server path (`server_atom.sh`, `env_atom.sh`, `bench.sh`, `setup_deps.sh`) + `launch.sh` SLURM-equivalent | +| `utils/bench_serving/` | vendored InferenceX bench client | +| `docs/proxy_and_disagg.md` | per-engine proxy/router/KV mechanics + topologies | + +## Topologies (from minimaxm3-fp4-mi355x-atom-disagg) +| name | nodes | layout | shapes | conc | +|------|-------|--------|--------|------| +| `1p1d` | 2 | 1 prefill TP4 + 1 decode TP4 (STP) | 8192/1024, 1024/1024 | 1–256 | +| `2p1d_dpa` | 3 | 2 prefill + 1 decode, TP4 + DP-attn | 8192/1024 | 256–1024 | + +All layouts keep `xP + yD <= 3` for the 3-node MI355X disagg pool (1 node/worker at TP4). + +## Usage +```bash +# ATOM disagg, 1-prefill/1-decode (2 nodes), full bench sweep through the router: +MODEL=minimaxm3-fp4 TOPO=1p1d ./run_atom_disagg.sh +MODEL=minimaxm3-fp4 TOPO=2p1d_dpa ./run_atom_disagg.sh # 3 nodes, high-conc 8k1k tail + +# Just serve (hold the endpoint up; drive it yourself): +MODEL=minimaxm3-fp4 TOPO=1p1d ACTION=serve ./run_atom_disagg.sh +``` +Bench drives `http://:8000` (OpenAI-compatible). Results land under `results/`. + +## Reference performance (InferenceX, MI355X) +From the InferenceX `minimaxm3-fp4-mi355x-atom-disagg` sweep (PR #2000), this is the +**top-throughput** MiniMax-M3 disagg config on MI355X: + +| shape | layout | peak tok/s/GPU | conc | +|-------|--------|----------------|------| +| 8k/1k | 2P1D + DPA | **~10,829** | 1024 | +| 8k/1k | 1P1D | ~7,594 | 256 | +| 1k/1k | 1P1D | ~2,703 | 256 | + +Dashboard: + +## Status / TODO +- **SCAFFOLDED** from the InferenceX config; **needs a live gfx950 dry-run** (RDMA/mooncake + handshake + atomesh router) once `amd/MiniMax-M3-MXFP4` weights are staged. +- Fill in `cluster.yaml` (reservation, `models_root`, node list) before running. +- MXFP4 quick-reduce knob `AITER_QUICK_REDUCE_QUANTIZATION=INT4` is applied by the launcher; + STP-only is enforced via `DECODE_MTP_SIZE=0`. diff --git a/scripts/MiniMax-M3/atom_disagg/accuracy_ctx.py b/scripts/MiniMax-M3/atom_disagg/accuracy_ctx.py new file mode 100644 index 00000000..0a5f41f9 --- /dev/null +++ b/scripts/MiniMax-M3/atom_disagg/accuracy_ctx.py @@ -0,0 +1,73 @@ +#!/usr/bin/env python3 +# accuracy_ctx.py — context-varied accuracy probe for the MiniMax-M3 ATOM disagg pipeline. +# +# Why: a PD-disaggregated pipeline can be correct at short context but corrupt the KV that +# mooncake transfers prefill->decode at long context. A 5-short-prompt gate misses that. +# This embeds a known "needle" fact inside filler of increasing size and checks the router +# (which spans prefill+decode+mooncake) returns it verbatim — exercising the FULL pipeline +# at each context length. +# +# Usage: ROUTER_URL, MODEL, OUT(json), [CTX_SIZES="512 2048 8192 16384"] via env. +# Exits 0 if all context sizes pass, 1 otherwise. +import json, os, sys, time, urllib.request + +ROUTER = os.environ.get("ROUTER_URL", "http://0.0.0.0:8000") +MODEL = os.environ["MODEL"] +OUT = os.environ.get("OUT", "/out/accuracy_ctx.json") +# approx tokens of filler context to test at (the needle adds a few more) +CTX_SIZES = [int(x) for x in os.environ.get("CTX_SIZES", "512 2048 8192 16384").split()] + +# A deterministic needle: a unique code the model must echo back. Greedy (temp 0). +SECRET = "PLATYPUS-7731-AZURE" +FILLER_SENTENCE = ("The quarterly logistics report notes that warehouse throughput " + "remained within nominal parameters across all regional hubs. ") + +def make_prompt(approx_tokens): + # ~1.3 tokens/word for English; build filler then plant the needle in the MIDDLE + words_needed = int(approx_tokens / 1.3) + filler_words = FILLER_SENTENCE.split() + reps = max(1, words_needed // len(filler_words)) + block = (" ".join(filler_words) + " ") * reps + half = len(block) // 2 + needle = f" IMPORTANT: the access code is {SECRET}. Remember it exactly. " + ctx = block[:half] + needle + block[half:] + q = ("\n\nQuestion: What is the exact access code mentioned above? " + "Reply with ONLY the code, nothing else.") + return ctx + q + +def ask(prompt): + body = json.dumps({ + "model": MODEL, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": 32, "temperature": 0.0, + }).encode() + req = urllib.request.Request(f"{ROUTER}/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=300) as r: + d = json.load(r) + return d["choices"][0]["message"]["content"], time.time() - t0 + +results = []; allpass = True +for ctx in CTX_SIZES: + try: + prompt = make_prompt(ctx) + text, dt = ask(prompt) + found = SECRET in text.upper() + ok = found + results.append({"ctx_tokens": ctx, "ok": ok, "found_needle": found, + "reply": text.strip()[:80], "latency_s": round(dt, 2), + "prompt_chars": len(prompt)}) + print(f" ctx~{ctx:>6}tok: {'PASS' if ok else 'FAIL'} " + f"(needle={'yes' if found else 'NO'}, {dt:.1f}s) reply='{text.strip()[:40]}'") + allpass = allpass and ok + except Exception as e: + results.append({"ctx_tokens": ctx, "ok": False, "error": str(e)[:120]}) + print(f" ctx~{ctx:>6}tok: ERROR {str(e)[:80]}") + allpass = False + +os.makedirs(os.path.dirname(OUT), exist_ok=True) +json.dump({"model": MODEL, "router": ROUTER, "all_pass": allpass, "results": results}, + open(OUT, "w"), indent=2) +print(f"[accuracy_ctx] {'ALL PASS' if allpass else 'FAILURES'} -> {OUT}") +sys.exit(0 if allpass else 1) diff --git a/scripts/MiniMax-M3/atom_disagg/bench.sh b/scripts/MiniMax-M3/atom_disagg/bench.sh new file mode 100644 index 00000000..53442f10 --- /dev/null +++ b/scripts/MiniMax-M3/atom_disagg/bench.sh @@ -0,0 +1,47 @@ +#!/bin/bash +# bench.sh — driven by server_atom.sh on node 0 (router node), AFTER atomesh is /v1/models-ready. +# Reproduces the InferenceX call signature exactly: +# bash bench.sh \ +# "" +# +# Drives the OpenAI-compatible ATOM router at localhost:$ROUTER_PORT with the vendored +# InferenceX benchmark_serving client (same one used for the Qwen sweep) — one result JSON +# per concurrency. This is a TRUE aggregate (one endpoint, real PD + mooncake KV), no x N. +set -uo pipefail +xP="$1"; yD="$2"; PF_GPUS="$3"; DEC_GPUS="$4" +MODEL_DIR="$5"; MODEL_NAME="$6"; LOG_DIR="$7" +ISL="$8"; OSL="$9"; CONC_LIST="${10}"; REQ_RATE="${11}"; RANGE_RATIO="${12}"; PROMPTS_MULT="${13}" +# server_atom.sh passes the conc list 'x'-separated (e.g. 4x8x16). Normalise to spaces. +CONC_LIST=$(echo "$CONC_LIST" | tr 'x' ' ') + +ROUTER_PORT="${ROUTER_PORT:-8000}" +BENCH_CLIENT="${BENCH_CLIENT:-/scripts/utils/bench_serving/benchmark_serving.py}" +OUT="${LOG_DIR}"; mkdir -p "$OUT" 2>/dev/null || true + +echo "[bench] router=localhost:${ROUTER_PORT} model=${MODEL_NAME} isl=${ISL} osl=${OSL} conc='${CONC_LIST}' (xP=${xP} yD=${yD})" + +# --- ACCURACY GATE (context-varied) before perf ------------------------------------------- +# Validate the FULL PD pipeline (prefill+mooncake KV+decode) returns correct output at varied +# context sizes. Runs ONCE per launch (guarded by a marker), not per-shape. Non-fatal: logs +# result but does not abort the perf sweep (so we still get numbers even if a long-ctx probe fails). +ACC_OUT="${OUT}/accuracy_ctx.json" +if [[ ! -f "$ACC_OUT" ]]; then + echo "[accuracy] context-varied probe via router (sizes: ${CTX_SIZES:-512 2048 8192 16384})" + ROUTER_URL="http://localhost:${ROUTER_PORT}" MODEL="${MODEL_DIR}/${MODEL_NAME}" \ + OUT="$ACC_OUT" CTX_SIZES="${CTX_SIZES:-512 2048 8192 16384}" \ + python3 /ws/accuracy_ctx.py || echo "[accuracy] WARN: some context sizes failed (see $ACC_OUT)" +fi + +for c in $CONC_LIST; do + echo "[bench] --- concurrency $c ---" + python3 "$BENCH_CLIENT" \ + --model "${MODEL_DIR}/${MODEL_NAME}" --backend vllm \ + --base-url "http://localhost:${ROUTER_PORT}" --dataset-name random \ + --random-input-len "$ISL" --random-output-len "$OSL" --random-range-ratio "$RANGE_RATIO" \ + --num-prompts "$(( c * PROMPTS_MULT ))" --num-warmups "$(( c * 2 ))" \ + --max-concurrency "$c" --request-rate "$REQ_RATE" --ignore-eos --save-result \ + --percentile-metrics ttft,tpot,itl,e2el --result-dir "$OUT" \ + --result-filename "disagg_${xP}p${yD}d_isl${ISL}_osl${OSL}_c${c}.json" \ + || echo "[bench] FAILED isl=$ISL osl=$OSL conc=$c" +done +echo "[bench] done -> $OUT" diff --git a/scripts/MiniMax-M3/atom_disagg/env_atom.sh b/scripts/MiniMax-M3/atom_disagg/env_atom.sh new file mode 100644 index 00000000..e0135e87 --- /dev/null +++ b/scripts/MiniMax-M3/atom_disagg/env_atom.sh @@ -0,0 +1,75 @@ +#!/bin/bash +# ATOM/mooncake-specific environment setup for multi-node disaggregated serving. +# +# Sourced by server_atom.sh in place of env.sh (which is SGLang/MoRI-specific). +# +# REQUIRED ENVIRONMENT VARIABLES: +# IBDEVICES - RDMA/InfiniBand device names (e.g., ionic_0,ionic_1,...) +# Set by runner or auto-detected from hostname. + +set -x + +export PYTHONUNBUFFERED=1 +export PYTHONDONTWRITEBYTECODE=1 + +# ============================================================================= +# IBDEVICES detection (same as env.sh) +# ============================================================================= + +if [[ -z "$IBDEVICES" ]]; then + DETECTED=$(ibv_devinfo 2>/dev/null | grep "hca_id:" | awk '{print $2}' | paste -sd',') + if [[ -n "$DETECTED" ]]; then + export IBDEVICES="$DETECTED" + echo "[INFO] Auto-detected IBDEVICES=$IBDEVICES via ibv_devinfo on $(hostname -s)" + else + # ATOM uses mooncake proxy_ip/handshake_port for KV transfer — IBDEVICES is + # not passed as a server argument (unlike SGLang --disaggregation-ib-device). + # Log a warning but do not fail; mooncake will use its own RDMA device selection. + echo "[WARN] Unable to detect RDMA devices via ibv_devinfo; IBDEVICES unset (non-fatal for ATOM/mooncake)" >&2 + fi +else + echo "[INFO] Using IBDEVICES=$IBDEVICES (set by runner or environment)" +fi +export IBDEVICES + +export SAFETENSORS_FAST_GPU=1 +export VLLM_LOG_LEVEL=WARNING +export ATOM_LOG_LEVEL=WARNING +export AITER_LOG_LEVEL=WARNING +export LOG_LEVEL=WARNING +export LOGLEVEL=WARNING + +# ============================================================================= +# ATOM/mooncake-specific environment +# ============================================================================= + +# mooncake RDMA KV transfer library path +# The ATOM disagg image (rocm/atom-dev:nightly_202607011530) ships Python 3.12, NOT 3.10. +# The original InferenceX path (.../python3.10/...) does not exist here -> mooncake native libs +# fail to load. Resolve the actual site-packages dir dynamically so this is image-version agnostic. +_MOONCAKE_DIR=$(python3 -c "import os,mooncake;print(os.path.dirname(mooncake.__file__))" 2>/dev/null) +export LD_LIBRARY_PATH=${_MOONCAKE_DIR:-/opt/venv/lib/python3.12/site-packages/mooncake}:/opt/rocm/lib:${LD_LIBRARY_PATH:-} + + +# ATOM_HOST_IP is set per-node in server_atom.sh (= host_ip, used as handshake IP) + +# aiter logging (WARNING to reduce noise; use DEBUG for troubleshooting) +export AITER_LOG_LEVEL=WARNING + +# MiniMax-M3 MXFP4 needs no model-specific env here: the MXFP4 quick-reduce knob +# (AITER_QUICK_REDUCE_QUANTIZATION=INT4) is exported by the launcher via model.yaml `env`, +# and server_atom.sh applies it for all non-DSv4 models. STP-only (DECODE_MTP_SIZE=0) is +# passed as a server arg, not an env var. This block is kept for parity with the DSv4 recipe. +if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then + # ATOM MoE gather/scatter interleave optimization + export ATOM_MOE_GU_ITLV=1 + # Disable bf16->fp8 MoE bound (only for DeepSeek-V4-Pro) + export AITER_BF16_FP8_MOE_BOUND=0 +fi + +# Clear stale ATOM cache on startup (server_atom.sh handles this via rm -rf) +# No env var needed; documented here for reference. + +set +x + +echo "[INFO] ATOM env: IBDEVICES=$IBDEVICES LD_LIBRARY_PATH includes mooncake" diff --git a/scripts/MiniMax-M3/atom_disagg/launch.sh b/scripts/MiniMax-M3/atom_disagg/launch.sh new file mode 100644 index 00000000..90fb181b --- /dev/null +++ b/scripts/MiniMax-M3/atom_disagg/launch.sh @@ -0,0 +1,132 @@ +#!/bin/bash +# launch.sh — multi-node ATOM PD-disagg launcher for MiniMax-M3 MXFP4 (the `job.slurm`-equiv). +# Drives the VENDORED, verbatim InferenceX server_atom.sh across N nodes via docker+srun. +# +# Faithful to the InferenceX chain: +# minimaxm3_..._atom-disagg.sh -> submit.sh -> job.slurm -> server_atom.sh +# We replace submit.sh+job.slurm (the slurm plumbing) with this; server_atom.sh runs unmodified. +# +# Env (with defaults): +# TOPO=1p1d|2p1d_dpa topology (sets xP/yD/dp-attn) [default 1p1d] +# ISL, OSL bench shape [default 8192 1024] +# CONC_LIST space-separated concurrencies [topology default] +# ACTION=bench|serve|dry bench (default) | hold endpoint | print cmds only +# +# Topologies (from minimaxm3-fp4-mi355x-atom-disagg). All workers TP4 (no EP); STP only. +# 1p1d : xP=1 yD=1, 2 nodes, TP4, dp-attn OFF (conc 1-256; shapes 8192/1024 & 1024/1024) +# 2p1d_dpa : xP=2 yD=1, 3 nodes, TP4, dp-attn ON (conc 256-1024; shape 8192/1024) +set -uo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # scripts/MiniMax-M3/atom_disagg +ROOT="$(cd "$HERE/.." && pwd)" # scripts/MiniMax-M3 +cfg(){ python3 "$ROOT/lib/cfg.py" "$@"; } +CLUSTER="$ROOT/cluster.yaml" + +RESV=$(cfg "$CLUSTER" cluster reservation) +MROOT=$(cfg "$CLUSTER" cluster models_root) +RROOT=$(cfg "$CLUSTER" cluster results_root) +RPORT=$(cfg "$CLUSTER" cluster router_port) +PPORT=$(cfg "$CLUSTER" cluster prefill_port) +DPORT=$(cfg "$CLUSTER" cluster decode_port) +HPORT=$(cfg "$CLUSTER" cluster handshake_port) +IMG=$(cfg "$ROOT/model.yaml" engine atom-disagg image) +MODEL_NAME="MiniMax-M3-MXFP4" + +TOPO="${TOPO:-1p1d}"; ACTION="${ACTION:-bench}" +ISL="${ISL:-8192}"; OSL="${OSL:-1024}" +case "$TOPO" in + 1p1d) xP=1; yD=1; DPATTN=0; DEF_CONC="1 2 4 8 16 32 64 128 256" ;; + 2p1d_dpa) xP=2; yD=1; DPATTN=1; DEF_CONC="256 512 768 1024" ;; + *) echo "ERROR: TOPO must be 1p1d|2p1d_dpa (got $TOPO)"; exit 1 ;; +esac +CONC_LIST="${CONC_LIST:-$DEF_CONC}" +# server_atom.sh expects BENCH_MAX_CONCURRENCY 'x'-separated (it does `tr 'x' '\n'|sort -n|tail -1` +# to derive the decode server's --max-num-seqs). Our bench.sh wants space-separated. Keep both. +CONC_X=$(echo "$CONC_LIST" | tr ' ' 'x') +TP=4 # TP4 throughout (M3 on gfx950: plain TP4 beats TP8 / TP4-EP) +NNODES=$(( xP + yD )) # 1 node/worker at TP4 on 8-GPU nodes + +# --- node selection: explicit NODES= override (REQUIRED for parallel launches to avoid +# two launches racing to grab the same 'first free' nodes), else auto-pick free ones. ---- +if [[ -n "${NODES:-}" ]]; then + IFS=',' read -ra NODES <<< "$NODES" + [[ ${#NODES[@]} -lt $NNODES ]] && { echo "ERROR: NODES= has ${#NODES[@]}, need $NNODES"; exit 1; } + NODES=("${NODES[@]:0:$NNODES}") +else + mapfile -t ALL < <(scontrol show hostnames "$(scontrol show res "$RESV" 2>/dev/null | grep -oE 'Nodes=[^ ]+' | head -1 | cut -d= -f2)") + BUSY=$(squeue -u "$USER" -h -o "%N" 2>/dev/null | scontrol show hostnames 2>/dev/null | sort -u) + NODES=(); for n in "${ALL[@]}"; do grep -qx "$n" <<<"$BUSY" || NODES+=("$n"); [[ ${#NODES[@]} -ge $NNODES ]] && break; done + [[ ${#NODES[@]} -lt $NNODES ]] && { echo "ERROR: need $NNODES free nodes, have ${#NODES[@]}"; exit 1; } +fi + +STAMP=$(python3 -c "import time;print(time.strftime('%Y%m%d_%H%M%S'))") +RUN="$RROOT/minimaxm3_atom_${TOPO}_isl${ISL}_${STAMP}"; mkdir -p "$RUN"; chmod 777 "$RUN" +echo "[launch] TOPO=$TOPO xP=$xP yD=$yD nodes=${NODES[*]} shape=${ISL}/${OSL} conc='$CONC_LIST'" +echo "[launch] image=$IMG run=$RUN" + +# --- resolve node IPs (IPADDRS = prefill nodes first, then decode; server_atom.sh slices it) - +declare -a IPS=() +for n in "${NODES[@]}"; do + ip=$(srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=1G \ + bash -c "ip route get 1.1.1.1 2>/dev/null | awk '/src/{print \$7}'" 2>/dev/null \ + | grep -oE '([0-9]+\.){3}[0-9]+' | head -1) + IPS+=("$ip"); echo " $n -> $ip" +done +IPADDRS=$(IFS=,; echo "${IPS[*]}") +NODE0="${NODES[0]}"; NODE0_IP="${IPS[0]}" + +# --- pre-flight clean every node ------------------------------------------------------------ +for n in "${NODES[@]}"; do + srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=0 \ + bash "$ROOT/lib/clean_node.sh" > "$RUN/clean_${n}.log" 2>&1 && echo " [clean] $n OK" || echo " [clean] $n WARN" +done + +DRY=0; [[ "$ACTION" == "dry" ]] && DRY=1 + +# --- RDMA device flags (THE fix for mooncake 30s timeouts) ---------------------------------- +# The ATOM disagg image (rocm/atom-dev:nightly_202607011530) carries its own bnxt_re RDMA stack. +# The fix for mooncake handshake timeouts is exposing --device=/dev/infiniband + --ulimit +# memlock=-1, NOT bind-mounting host libs. DO NOT bind-mount host libibverbs over the +# container's: host and container can have incompatible IBVERBS_PRIVATE ABIs, so mounting host +# libs BREAKS the container's RDMA tools. Expose devices + memlock only. +RDMA_DEVICES="--device=/dev/infiniband --ulimit memlock=-1 --ulimit stack=67108864 --init" +RDMA_MOUNT_SNIPPET='RM=""' # no host-lib mounts (container stack is self-sufficient) + +# --- launch server_atom.sh on each node, NODE_RANK = its index ------------------------------ +# server_atom.sh logic: rank 0 = prefill0 + atomesh router (+ bench); rank=NODE_OFFSET = decode. We pass the SAME env to every node; the script self-selects by rank. +rank=0 +for n in "${NODES[@]}"; do + rdir="$RUN/node${rank}_${n}"; mkdir -p "$rdir"; chmod 777 "$rdir" + srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --exclusive bash -c " + docker rm -f mm3_atom_n${rank} >/dev/null 2>&1 || true + $RDMA_MOUNT_SNIPPET # builds \$RM (host RDMA lib bind-mounts) on THIS node + echo \"[rdma] mounts: \$RM\" + docker run --rm --name mm3_atom_n${rank} --entrypoint bash \ + --device=/dev/kfd --device=/dev/dri $RDMA_DEVICES \$RM -v /sys:/sys \ + --group-add video --ipc host --shm-size 128G \ + --network host --privileged --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e ATOM_WS_PATH=/ws -e NODE_RANK=${rank} -e NODE0_ADDR=${NODE0_IP} -e IPADDRS='${IPADDRS}' \ + -e MODEL_DIR=${MROOT} -e MODEL_NAME=${MODEL_NAME} \ + -e xP=${xP} -e yD=${yD} -e GPUS_PER_NODE=8 \ + -e PREFILL_TP_SIZE=${TP} -e DECODE_TP_SIZE=${TP} \ + -e PREFILL_ENABLE_DP=$([[ $DPATTN == 1 ]] && echo true || echo false) -e PREFILL_ENABLE_EP=1 \ + -e DECODE_ENABLE_DP=$([[ $DPATTN == 1 ]] && echo true || echo false) -e DECODE_ENABLE_EP=1 \ + -e PREFILL_PORT=${PPORT} -e DECODE_PORT=${DPORT} -e ROUTER_PORT=${RPORT} -e HANDSHAKE_PORT=${HPORT} \ + -e MEM_FRAC_STATIC=0.8 -e KV_CACHE_DTYPE=auto -e BLOCK_SIZE=128 -e MAX_NUM_SEQS=256 \ + -e MAX_MODEL_LEN=32768 -e MAX_NUM_BATCHED_TOKENS=32768 \ + -e SPEC_DECODING= -e DECODE_MTP_SIZE=0 \ + -e AITER_QUICK_REDUCE_QUANTIZATION=INT4 \ + -e BENCH_INPUT_LEN=${ISL} -e BENCH_OUTPUT_LEN=${OSL} -e BENCH_MAX_CONCURRENCY='${CONC_X}' \ + -e BENCH_REQUEST_RATE=inf -e BENCH_RANDOM_RANGE_RATIO=1 -e BENCH_NUM_PROMPTS_MULTIPLIER=10 \ + -e DRY_RUN=${DRY} -e SLURM_JOB_ID=${STAMP} \ + -v ${MROOT}:${MROOT} -v ${HERE}:/ws -v ${ROOT}:/scripts -v ${rdir}:/run_logs/slurm_job-${STAMP} \ + ${IMG} /ws/server_atom.sh + " > "$rdir/srun.log" 2>&1 & + echo " [launch] node${rank} ($n) server_atom.sh dispatched" + rank=$(( rank + 1 )) +done + +echo "[launch] all $NNODES nodes dispatched. node0=$NODE0 hosts atomesh router :$RPORT + drives bench." +echo "[launch] tail node0: $RUN/node0_${NODE0}/srun.log" +wait +echo "[launch] done -> $RUN (results: disagg_*.json under node0_*/)" diff --git a/scripts/MiniMax-M3/atom_disagg/server_atom.sh b/scripts/MiniMax-M3/atom_disagg/server_atom.sh new file mode 100644 index 00000000..ccc86403 --- /dev/null +++ b/scripts/MiniMax-M3/atom_disagg/server_atom.sh @@ -0,0 +1,585 @@ +#!/bin/bash +# ATOM Disaggregated Server Launcher +# ============================================================================= +# Uses atom.entrypoints.openai_server with mooncake RDMA KV transfer. +# Mirrors server_sglang.sh topology (dynamic xP/yD) but adapts to ATOM's +# explicit kv-transfer-config and atomesh router. +# +# Key differences from server_sglang.sh: +# - Engine: atom.entrypoints.openai_server (not sglang.launch_server) +# - KV transfer: mooncake (--kv-transfer-config JSON) +# - Router: atomesh (not sglang_router) +# - Prefill port: $PREFILL_PORT (default 8010) / Decode port: $DECODE_PORT (default 8020) +# - Router port: $ROUTER_PORT (default 8000) +# ============================================================================= + +# ============================================================================= +# Environment Configuration +# ============================================================================= + +NODE0_ADDR="${NODE0_ADDR:-localhost}" +NODE_RANK="${NODE_RANK:-0}" +MODEL_DIR="${MODEL_DIR:-}" +MODEL_NAME="${MODEL_NAME:-}" + +xP="${xP:-1}" +yD="${yD:-1}" + +IPADDRS="${IPADDRS:-localhost}" + +# Parallelism +PREFILL_TP_SIZE="${PREFILL_TP_SIZE:-8}" +PREFILL_ENABLE_EP="${PREFILL_ENABLE_EP}" +PREFILL_ENABLE_DP="${PREFILL_ENABLE_DP}" +DECODE_TP_SIZE="${DECODE_TP_SIZE:-8}" +DECODE_ENABLE_EP="${DECODE_ENABLE_EP}" +DECODE_ENABLE_DP="${DECODE_ENABLE_DP}" + +# MTP +SPEC_DECODING="${SPEC_DECODING:-}" +DECODE_MTP_SIZE="${DECODE_MTP_SIZE:-1}" + +# ATOM server ports (different from SGLang which uses 8000 for all) +PREFILL_PORT="${PREFILL_PORT:-8010}" +DECODE_PORT="${DECODE_PORT:-8020}" +ROUTER_PORT="${ROUTER_PORT:-8000}" +HANDSHAKE_PORT="${HANDSHAKE_PORT:-6301}" + +# ATOM server tuning (from reference script defaults) +MEM_FRAC_STATIC="${MEM_FRAC_STATIC:-0.85}" +KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" +BLOCK_SIZE="${BLOCK_SIZE:-16}" +MAX_NUM_SEQS="${MAX_NUM_SEQS:-256}" +MAX_MODEL_LEN="${MAX_MODEL_LEN:-}" +MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-}" +EXTRA_SERVER_ARGS="${EXTRA_SERVER_ARGS:-}" + +# Benchmark Configuration +BENCH_INPUT_LEN="${BENCH_INPUT_LEN:-1024}" +BENCH_OUTPUT_LEN="${BENCH_OUTPUT_LEN:-1024}" +BENCH_RANDOM_RANGE_RATIO="${BENCH_RANDOM_RANGE_RATIO:-1}" +BENCH_REQUEST_RATE="${BENCH_REQUEST_RATE:-inf}" +BENCH_NUM_PROMPTS_MULTIPLIER="${BENCH_NUM_PROMPTS_MULTIPLIER:-10}" +BENCH_MAX_CONCURRENCY="${BENCH_MAX_CONCURRENCY:-512}" + +DRY_RUN="${DRY_RUN:-0}" +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" + +# ============================================================================= +# Dependencies and Environment Setup +# ============================================================================= + +source $ATOM_WS_PATH/setup_deps.sh +source $ATOM_WS_PATH/env_atom.sh + +host_ip=$(ip route get 1.1.1.1 2>/dev/null | awk '/src/ {print $7}') +if [[ -z "$host_ip" ]]; then + host_ip=$(hostname -I 2>/dev/null | awk '{print $1}') +fi +host_name=$(hostname) + +# ============================================================================= +# Cluster Topology Configuration +# ============================================================================= + +IFS=',' read -ra IP_ARRAY <<< "$IPADDRS" + +PREFILL_NODES_PER_WORKER=$(((PREFILL_TP_SIZE + GPUS_PER_NODE - 1) / GPUS_PER_NODE)) +DECODE_NODES_PER_WORKER=$(((DECODE_TP_SIZE + GPUS_PER_NODE - 1) / GPUS_PER_NODE)) +NODE_OFFSET=$((PREFILL_NODES_PER_WORKER * xP)) + +# Build prefill IP list and atomesh --prefill args +PREFILL_ARGS="" +PREFILL_IPS=() +for i in $(seq 0 $((xP - 1))); do + idx=$((i * PREFILL_NODES_PER_WORKER)) + PREFILL_IPS[$i]="${IP_ARRAY[$idx]}" + PREFILL_ARGS="$PREFILL_ARGS --prefill http://${IP_ARRAY[$idx]}:${PREFILL_PORT}" +done + +# Build decode IP list and atomesh --decode args +DECODE_ARGS="" +DECODE_IPS=() +for i in $(seq 0 $((yD - 1))); do + idx=$((i * DECODE_NODES_PER_WORKER + NODE_OFFSET)) + DECODE_IPS[$i]="${IP_ARRAY[$idx]}" + DECODE_ARGS="$DECODE_ARGS --decode http://${IP_ARRAY[$idx]}:${DECODE_PORT}" +done + +PREFILL_ENABLE_EP="${PREFILL_ENABLE_EP}" +PREFILL_ENABLE_DP="${PREFILL_ENABLE_DP}" +DECODE_ENABLE_EP="${DECODE_ENABLE_EP}" +DECODE_ENABLE_DP="${DECODE_ENABLE_DP}" + +# Parallel args +PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE") #TP +if [ "$PREFILL_ENABLE_DP" = "true" ]; then + if [ "$PREFILL_ENABLE_EP" -gt 1 ]; then #DPA+EP + PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" --enable-expert-parallel --enable-dp-attention ) + else #TP+DPA+TBO + if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then + PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" --enable-dp-attention --enable-tbo ) + export GPU_MAX_HW_QUEUES=5 + export ATOM_CPU_AFFINITY=1 + else #TP+DPA + PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" --enable-dp-attention ) + fi + fi +fi + +# (srok), split DPA & TBO cases +DECODE_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE") #TP +if [ "$DECODE_ENABLE_DP" = "true" ]; then + if [ "$DECODE_ENABLE_EP" -gt 1 ]; then #DPA+EP + DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" --enable-expert-parallel --enable-dp-attention ) + else #TP+DPA+TBO + if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then + DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" --enable-dp-attention --enable-tbo ) + export GPU_MAX_HW_QUEUES=5 + export ATOM_CPU_AFFINITY=1 + else #TP+DPA + DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" --enable-dp-attention ) + fi + fi +fi + +# MTP args +SPEC_ARGS=() #TP +if [ "$SPEC_DECODING" = "mtp" ]; then + SPEC_ARGS=(--method mtp --num-speculative-tokens "$DECODE_MTP_SIZE") +fi + +# HF overrides (single-quoted JSON preserved through eval) +HF_OVERRIDES_ARG="" +if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then + HF_OVERRIDES_ARG="--hf-overrides '{\"use_index_cache\":true,\"index_topk_freq\":4}'" +fi + +# KV cache dtype (skip if unset or 'auto') +KV_CACHE_ARG="" +if [[ -n "$KV_CACHE_DTYPE" && "$KV_CACHE_DTYPE" != "auto" ]]; then + KV_CACHE_ARG="--kv_cache_dtype ${KV_CACHE_DTYPE}" +fi + +# Optional model length / batched-token cap +MODEL_LEN_ARGS="" +if [[ -n "$MAX_MODEL_LEN" ]]; then + MODEL_LEN_ARGS="${MODEL_LEN_ARGS} --max-model-len ${MAX_MODEL_LEN}" +fi +if [[ -n "$MAX_NUM_BATCHED_TOKENS" ]]; then + MODEL_LEN_ARGS="${MODEL_LEN_ARGS} --max-num-batched-tokens ${MAX_NUM_BATCHED_TOKENS}" +fi + +if [[ "$MODEL_NAME" != "DeepSeek-V4-Pro" ]]; then + export AITER_QUICK_REDUCE_QUANTIZATION=INT4 +fi + +cat < prefill node 0 + router +# rank 1 .. (NODE_OFFSET-1) -> remaining prefill nodes +# rank NODE_OFFSET .. -> decode nodes +# ============================================================================= +if [ "$NODE_RANK" -eq 0 ]; then + # ────────────────────────────────────────────────────────────────────────── + # Node 0: prefill server (producer) + atomesh router + # ────────────────────────────────────────────────────────────────────────── + echo "NODE INFO =======================================" + echo "${host_name}:${host_ip} is Prefill Node 0 + Router" + echo "Prefill TP=${PREFILL_TP_SIZE}, Decode TP=${DECODE_TP_SIZE}" + echo "Prefill servers: ${PREFILL_ARGS}" + echo "Decode servers: ${DECODE_ARGS}" + echo "================================================" + + PREFILL_CMD="python3 -m atom.entrypoints.openai_server \ + --model ${MODEL_DIR}/${MODEL_NAME} \ + --host 0.0.0.0 --server-port ${PREFILL_PORT} \ + --trust-remote-code \ + ${PREFILL_PARALLEL_ARGS[*]} \ + ${SPEC_ARGS[*]} \ + ${KV_CACHE_ARG} \ + --block-size ${BLOCK_SIZE} \ + --gpu-memory-utilization ${MEM_FRAC_STATIC} \ + --max-num-seqs ${MAX_NUM_SEQS} \ + ${MODEL_LEN_ARGS} \ + --no-enable_prefix_caching \ + ${HF_OVERRIDES_ARG} \ + --kv-transfer-config '{\"kv_role\":\"kv_producer\",\"kv_connector\":\"mooncake\",\"proxy_ip\":\"${host_ip}\",\"handshake_port\":${HANDSHAKE_PORT}}' \ + ${EXTRA_SERVER_ARGS}" + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $PREFILL_CMD" + else + set -x + eval "$PREFILL_CMD" \ + 2>&1 | tee /run_logs/slurm_job-${SLURM_JOB_ID}/prefill0_${host_name}.log & + set +x + prefill0_pid=$! + fi + + # Wait for all prefill and decode servers to be ready + WAIT_SERVER_TIMEOUT="${WAIT_SERVER_TIMEOUT:-2500}" + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting for all servers to be up (timeout=${WAIT_SERVER_TIMEOUT}s)..." + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait for prefill/decode /health endpoints" + else + _deadline=$(( $(date +%s) + WAIT_SERVER_TIMEOUT )) + for _ip in "${PREFILL_IPS[@]}"; do + echo "[wait] prefill http://${_ip}:${PREFILL_PORT}/health" + while ! curl -sf --max-time 10 "http://${_ip}:${PREFILL_PORT}/health" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_deadline ]]; then + echo "[wait][FAIL] prefill ${_ip}:${PREFILL_PORT} not ready after ${WAIT_SERVER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] prefill ${_ip}:${PREFILL_PORT} ready" + done + for _ip in "${DECODE_IPS[@]}"; do + echo "[wait] decode http://${_ip}:${DECODE_PORT}/health" + while ! curl -sf --max-time 10 "http://${_ip}:${DECODE_PORT}/health" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_deadline ]]; then + echo "[wait][FAIL] decode ${_ip}:${DECODE_PORT} not ready after ${WAIT_SERVER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] decode ${_ip}:${DECODE_PORT} ready" + done + fi + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "All servers up. Starting atomesh router..." + + ROUTER_CMD="/usr/local/bin/atomesh launch \ + --host 0.0.0.0 --port ${ROUTER_PORT} \ + --pd-disaggregation \ + ${PREFILL_ARGS} \ + ${DECODE_ARGS} \ + --policy random \ + --backend atom \ + --log-level info \ + --disable-health-check \ + --disable-circuit-breaker \ + --prometheus-port 29100" + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $ROUTER_CMD" + else + ROUTER_LOG_FILE="/tmp/slurm_job-${SLURM_JOB_ID}_router_${host_name}.log" + set -x + eval "$ROUTER_CMD" 2>&1 | tee "$ROUTER_LOG_FILE" & + set +x + proxy_pid=$! + + # Wait for router to accept connections + WAIT_ROUTER_TIMEOUT="${WAIT_ROUTER_TIMEOUT:-300}" + echo "[wait] router http://0.0.0.0:${ROUTER_PORT}/v1/models (timeout=${WAIT_ROUTER_TIMEOUT}s)" + _router_deadline=$(( $(date +%s) + WAIT_ROUTER_TIMEOUT )) + while ! curl -sf --max-time 10 "http://0.0.0.0:${ROUTER_PORT}/v1/models" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_router_deadline ]]; then + echo "[wait][FAIL] router ${ROUTER_PORT}/v1/models not ready after ${WAIT_ROUTER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] router /v1/models ready" + + echo "Router is ready for benchmarking" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Ready for benchmarking on ${host_name}:${host_ip}" + + cd $ATOM_WS_PATH + + export IS_MTP="false" + if [ "$SPEC_DECODING" = "mtp" ]; then + export IS_MTP="true" + fi + + BENCH_CMD="bash $ATOM_WS_PATH/bench.sh ${xP} ${yD} $((PREFILL_TP_SIZE*xP)) $((DECODE_TP_SIZE*yD)) \ + $MODEL_DIR $MODEL_NAME /run_logs/slurm_job-${SLURM_JOB_ID} ${BENCH_INPUT_LEN} \ + ${BENCH_OUTPUT_LEN} \"${BENCH_MAX_CONCURRENCY}\" ${BENCH_REQUEST_RATE} \ + ${BENCH_RANDOM_RANGE_RATIO} ${BENCH_NUM_PROMPTS_MULTIPLIER}" + + if [[ "${EVAL_ONLY:-false}" == "true" ]]; then + echo "EVAL_ONLY mode: skipping throughput benchmark" + elif [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $BENCH_CMD" + else + set -x + eval "$BENCH_CMD" + set +x + fi + + # Run evaluation if requested (before killing router) + if [[ "${RUN_EVAL:-false}" == "true" ]]; then + echo "Running lm-eval evaluation on Node 0..." + + # Health check: verify the router is still serving before running eval. + EVAL_HEALTH_OK=false + for _attempt in 1 2 3; do + if curl -sf --max-time 10 "http://0.0.0.0:${ROUTER_PORT}/health" >/dev/null 2>&1; then + EVAL_HEALTH_OK=true + break + fi + echo "Eval health check attempt $_attempt failed, retrying in 10s..." + sleep 10 + done + + if [[ "$EVAL_HEALTH_OK" != "true" ]]; then + echo "WARNING: Router health check failed after 3 attempts. Skipping eval." + else + pushd /workspace + + source /workspace/benchmarks/benchmark_lib.sh + + if [[ -n "${EVAL_CONC:-}" ]]; then + export EVAL_CONCURRENT_REQUESTS="${EVAL_CONC}" + else + export EVAL_CONCURRENT_REQUESTS=$(echo "$BENCH_MAX_CONCURRENCY" | tr 'x' '\n' | sort -n | tail -1) + fi + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: run_eval --framework lm-eval --port ${ROUTER_PORT} (conc=${EVAL_CONCURRENT_REQUESTS})" + else + MODEL_NAME="${MODEL_DIR}/${MODEL_NAME}" run_eval --framework lm-eval --port ${ROUTER_PORT} + eval_rc=$? + + if [[ $eval_rc -ne 0 ]]; then + echo "ERROR: run_eval exited rc=$eval_rc; skipping metadata write and eval artifact staging" >&2 + EVAL_FAILED=1 + else + export TP="${PREFILL_TP_SIZE}" + export CONC="${EVAL_CONCURRENT_REQUESTS}" + export PREFILL_TP="${PREFILL_TP_SIZE}" + export PREFILL_EP=1 + export PREFILL_NUM_WORKERS="${xP}" + export DECODE_TP="${DECODE_TP_SIZE}" + export DECODE_EP=1 + export DECODE_NUM_WORKERS="${yD}" + export ISL="${BENCH_INPUT_LEN}" + export OSL="${BENCH_OUTPUT_LEN}" + + MODEL_NAME="${MODEL_DIR}/${MODEL_NAME}" append_lm_eval_summary + + EVAL_COPY_DIR="/run_logs/slurm_job-${SLURM_JOB_ID}/eval_results" + mkdir -p "$EVAL_COPY_DIR" + for f in meta_env.json; do + [ -e "/workspace/$f" ] && cp -f "/workspace/$f" "$EVAL_COPY_DIR/" + done + find /workspace -maxdepth 1 -name 'results*.json' -exec cp -f {} "$EVAL_COPY_DIR/" \; + find /workspace -maxdepth 1 -name 'sample*.jsonl' -exec cp -f {} "$EVAL_COPY_DIR/" \; + + echo "Eval completed. Artifacts staged in $EVAL_COPY_DIR" + fi + fi + + popd + fi + fi + + # Copy results + LOGS_OUTPUT="${BENCHMARK_LOGS_DIR:-/run_logs}/logs" + mkdir -p "$LOGS_OUTPUT" + if [[ "$DRY_RUN" -eq 0 ]]; then + cp -r /run_logs/slurm_job-${SLURM_JOB_ID} "$LOGS_OUTPUT/" + echo "Copied results to $LOGS_OUTPUT/slurm_job-${SLURM_JOB_ID}" + fi + + echo "Waiting 60s before killing router and prefill server..." + sleep 60 + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Killing router and prefill server" + if [[ "$DRY_RUN" -eq 0 ]]; then + kill $proxy_pid + kill $prefill0_pid + fi + + if [[ "${EVAL_FAILED:-0}" -eq 1 ]]; then + echo "ERROR: eval failed; exiting node-0 with rc=1" + exit 1 + fi + +elif [ "$NODE_RANK" -gt 0 ] && [ "$NODE_RANK" -lt "$NODE_OFFSET" ]; then + # ────────────────────────────────────────────────────────────────────────── + # Prefill nodes 1..N (kv_producer) + # ────────────────────────────────────────────────────────────────────────── + echo "${host_name}:${host_ip} is Prefill Node (rank ${NODE_RANK})" + + # Determine which prefill worker this node belongs to, and its headnode IP + prefill_worker_idx=$((NODE_RANK / PREFILL_NODES_PER_WORKER)) + PREFILL_HEADNODE_IP="${PREFILL_IPS[$prefill_worker_idx]}" + + PREFILL_CMD="python3 -m atom.entrypoints.openai_server \ + --model ${MODEL_DIR}/${MODEL_NAME} \ + --host 0.0.0.0 --server-port ${PREFILL_PORT} \ + --trust-remote-code \ + ${PREFILL_PARALLEL_ARGS[*]} \ + ${SPEC_ARGS[*]} \ + ${KV_CACHE_ARG} \ + --block-size ${BLOCK_SIZE} \ + --gpu-memory-utilization ${MEM_FRAC_STATIC} \ + --max-num-seqs ${MAX_NUM_SEQS} \ + ${MODEL_LEN_ARGS} \ + --no-enable_prefix_caching \ + ${HF_OVERRIDES_ARG} \ + --kv-transfer-config '{\"kv_role\":\"kv_producer\",\"kv_connector\":\"mooncake\",\"proxy_ip\":\"${host_ip}\",\"handshake_port\":${HANDSHAKE_PORT}}' \ + ${EXTRA_SERVER_ARGS}" + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $PREFILL_CMD" + else + set -x + eval "$PREFILL_CMD" \ + 2>&1 | tee /run_logs/slurm_job-${SLURM_JOB_ID}/prefill_${host_name}.log & + set +x + prefill_pid=$! + trap 'echo "Caught signal, killing prefill (pid=$prefill_pid)"; kill $prefill_pid 2>/dev/null; exit 0' SIGTERM SIGINT + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting for router to be up..." + WAIT_ROUTER_TIMEOUT="${WAIT_ROUTER_TIMEOUT:-2800}" + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait for router ${NODE0_ADDR}:${ROUTER_PORT}/health" + else + _router_deadline=$(( $(date +%s) + WAIT_ROUTER_TIMEOUT )) + while ! curl -sf --max-time 10 "http://${NODE0_ADDR}:${ROUTER_PORT}/health" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_router_deadline ]]; then + echo "[wait][FAIL] router ${NODE0_ADDR}:${ROUTER_PORT} not ready after ${WAIT_ROUTER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] router ${NODE0_ADDR}:${ROUTER_PORT} ready" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting until router closes..." + trap 'echo "Caught signal, killing prefill (pid=$prefill_pid)"; kill $prefill_pid 2>/dev/null; exit 0' SIGTERM SIGINT + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait until router ${NODE0_ADDR}:${ROUTER_PORT} closes" + else + while curl -sf --max-time 10 "http://${NODE0_ADDR}:${ROUTER_PORT}/health" >/dev/null 2>&1; do + sleep 10 & + wait $! + done + echo "[wait] router ${NODE0_ADDR}:${ROUTER_PORT} closed" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Killing prefill server (rank ${NODE_RANK})" + if [[ "$DRY_RUN" -eq 0 ]]; then kill $prefill_pid 2>/dev/null; fi + +else + # ────────────────────────────────────────────────────────────────────────── + # Decode nodes (kv_consumer) + # ────────────────────────────────────────────────────────────────────────── + RANK=$((NODE_RANK - NODE_OFFSET)) + echo "${host_name}:${host_ip} is Decode Node (rank ${RANK})" + + _MAX_CONC=$(echo "$BENCH_MAX_CONCURRENCY" | tr 'x' '\n' | sort -n | tail -1) + if [[ "$_MAX_CONC" -gt 2048 ]]; then + CUDAGRAPH_SIZES='[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124,128,132,136,140,144,148,152,156,160,164,168,172,176,180,184,188,192,196,200,204,208,212,216,220,224,228,232,236,240,244,248,252,256,512,1024,2048,4096]' + elif [[ "$_MAX_CONC" -gt 1024 ]]; then + CUDAGRAPH_SIZES='[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124,128,132,136,140,144,148,152,156,160,164,168,172,176,180,184,188,192,196,200,204,208,212,216,220,224,228,232,236,240,244,248,252,256,512,1024,2048]' + elif [[ "$_MAX_CONC" -gt 512 ]]; then + CUDAGRAPH_SIZES='[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124,128,132,136,140,144,148,152,156,160,164,168,172,176,180,184,188,192,196,200,204,208,212,216,220,224,228,232,236,240,244,248,252,256,512,768,1024]' + else + CUDAGRAPH_SIZES='[1,2,4,8,12,16,20,24,28,32,36,40,44,48,52,56,60,64,68,72,76,80,84,88,92,96,100,104,108,112,116,120,124,128,132,136,140,144,148,152,156,160,164,168,172,176,180,184,188,192,196,200,204,208,212,216,220,224,228,232,236,240,244,248,252,256,512]' + fi + + if [[ "$BENCH_INPUT_LEN" == "1024" && "$BENCH_OUTPUT_LEN" == "1024" ]]; then + DECODE_MAX_NUM_SEQS="${_MAX_CONC}" + else + DECODE_MAX_NUM_SEQS="${MAX_NUM_SEQS}" + fi + + DECODE_CMD="python3 -m atom.entrypoints.openai_server \ + --model ${MODEL_DIR}/${MODEL_NAME} \ + --host 0.0.0.0 --server-port ${DECODE_PORT} \ + --trust-remote-code \ + ${DECODE_PARALLEL_ARGS[*]} \ + ${SPEC_ARGS[*]} \ + ${KV_CACHE_ARG} \ + --block-size ${BLOCK_SIZE} \ + --gpu-memory-utilization ${MEM_FRAC_STATIC} \ + --max-num-seqs ${DECODE_MAX_NUM_SEQS} \ + ${MODEL_LEN_ARGS} \ + --no-enable_prefix_caching \ + ${HF_OVERRIDES_ARG} \ + --kv-transfer-config '{\"kv_role\":\"kv_consumer\",\"kv_connector\":\"mooncake\",\"proxy_ip\":\"${host_ip}\",\"handshake_port\":${HANDSHAKE_PORT}}' \ + --cudagraph-capture-sizes "${CUDAGRAPH_SIZES}" \ + ${EXTRA_SERVER_ARGS}" + + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: $DECODE_CMD" + else + set -x + eval "$DECODE_CMD" \ + 2>&1 | tee /run_logs/slurm_job-${SLURM_JOB_ID}/decode_${host_name}.log & + set +x + decode_pid=$! + trap 'echo "Caught signal, killing decode (pid=$decode_pid)"; kill $decode_pid 2>/dev/null; exit 0' SIGTERM SIGINT + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting for router to be up..." + WAIT_ROUTER_TIMEOUT="${WAIT_ROUTER_TIMEOUT:-2800}" + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait for router ${NODE0_ADDR}:${ROUTER_PORT}/health" + else + _router_deadline=$(( $(date +%s) + WAIT_ROUTER_TIMEOUT )) + while ! curl -sf --max-time 10 "http://${NODE0_ADDR}:${ROUTER_PORT}/health" >/dev/null 2>&1; do + if [[ $(date +%s) -ge $_router_deadline ]]; then + echo "[wait][FAIL] router ${NODE0_ADDR}:${ROUTER_PORT} not ready after ${WAIT_ROUTER_TIMEOUT}s" >&2 + exit 1 + fi + sleep 10 + done + echo "[wait][OK] router ${NODE0_ADDR}:${ROUTER_PORT} ready" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Waiting until router closes..." + trap 'echo "Caught signal, killing decode (pid=$decode_pid)"; kill $decode_pid 2>/dev/null; exit 0' SIGTERM SIGINT + if [[ "$DRY_RUN" -eq 1 ]]; then + echo "DRY RUN: wait until router ${NODE0_ADDR}:${ROUTER_PORT} closes" + else + while curl -sf --max-time 10 "http://${NODE0_ADDR}:${ROUTER_PORT}/health" >/dev/null 2>&1; do + sleep 10 & + wait $! + done + echo "[wait] router ${NODE0_ADDR}:${ROUTER_PORT} closed" + fi + + echo "[-------]" NODE $NODE_RANK "[--------]" + echo "Killing decode server (rank ${RANK})" + if [[ "$DRY_RUN" -eq 0 ]]; then kill $decode_pid 2>/dev/null; fi +fi + +echo "Script completed successfully" +exit 0 diff --git a/scripts/MiniMax-M3/atom_disagg/setup_deps.sh b/scripts/MiniMax-M3/atom_disagg/setup_deps.sh new file mode 100644 index 00000000..76b993c1 --- /dev/null +++ b/scripts/MiniMax-M3/atom_disagg/setup_deps.sh @@ -0,0 +1,11 @@ +#!/bin/bash +# setup_deps.sh — sourced by server_atom.sh at startup. +# +# In the InferenceX CI this installs/patches runtime deps into the container. The M3 ATOM +# disagg image (rocm/atom-dev:nightly_202607011530) is expected to ship everything needed: +# atomesh (/usr/local/bin/atomesh), mooncake (py3.12 site-packages + +# /usr/local/bin/mooncake_{master,client}), atom.entrypoints.openai_server. +# +# So this is intentionally a NO-OP stub. If a future image is missing a dep, add the +# install here. Kept as a separate file so server_atom.sh runs unmodified (verbatim vendor). +: diff --git a/scripts/MiniMax-M3/cluster.yaml b/scripts/MiniMax-M3/cluster.yaml new file mode 100644 index 00000000..71e48e6e --- /dev/null +++ b/scripts/MiniMax-M3/cluster.yaml @@ -0,0 +1,28 @@ +# cluster.yaml — MiniMax-M3 disagg cluster definition. SCALE LEVER for multi-node PD. +# M3 MXFP4 disagg spans MULTIPLE nodes per serving instance: xP prefill nodes + yD decode +# nodes + 1 router (on prefill node 0). All workers TP4 => 1 node/worker on 8-GPU nodes, +# so 1p1d=2 nodes and 2p1d_dpa=3 nodes (the InferenceX mi355x-disagg pool is 3 nodes). + +reservation: +gpus_per_node: 8 # gfx950 (MI355X) +hardware: MI355X + +# Model staging root (same path on every node; local NVMe recommended). +# Stage amd/MiniMax-M3-MXFP4 to /MiniMax-M3-MXFP4 (matches model.yaml dir). +models_root: /path/to/models + +results_root: ./results + +# Disagg ports (match InferenceX server_atom.sh defaults) +router_port: 8000 +prefill_port: 8010 +decode_port: 8020 +handshake_port: 6301 + +# Node pool. The launcher picks the first +# (xP*pf_nodes_per_worker + yD*dec_nodes_per_worker) nodes for a given topology. +# 2p1d_dpa needs 3 nodes; list at least that many. +nodes: + - + - + - diff --git a/scripts/MiniMax-M3/docs/proxy_and_disagg.md b/scripts/MiniMax-M3/docs/proxy_and_disagg.md new file mode 100644 index 00000000..2dc6825b --- /dev/null +++ b/scripts/MiniMax-M3/docs/proxy_and_disagg.md @@ -0,0 +1,52 @@ +# MiniMax-M3 disaggregated serving — proxy/router mechanics (from InferenceX) + +MiniMax-M3 (MXFP4) is served **prefill/decode disaggregated (PD-disagg)** across nodes on AMD +MI355X (gfx950): prefill workers (kv_producer) and decode workers (kv_consumer) on separate +nodes, KV streamed over RDMA, fronted by ONE atomesh router. Bench drives the router → a true +aggregate throughput number. + +**All workers are TP4, no EP.** The single-node M3 recipe found plain TP4 beats both TP8 and +TP4/EP4 on tok/s/GPU on gfx950. We tune the prefill:decode worker ratio (xP:yD) instead of TP. +Serving is **STP only** (no MTP): `DECODE_MTP_SIZE=0`. + +## The proxy/router this recipe uses + +| Engine | Router/proxy | KV-transfer backend | Ports | InferenceX source | +|--------|--------------|---------------------|-------|-------------------| +| **ATOM** | **atomesh** (`/usr/local/bin/atomesh launch --pd-disaggregation`) | **mooncake** (RDMA, `--kv-transfer-config` JSON) | router 8000, prefill 8010, decode 8020, handshake 6301 | `server_atom.sh` (`minimaxm3-fp4-mi355x-atom-disagg`) | + +This folder ships the **ATOM MXFP4** path only (the top-throughput MiniMax-M3 disagg config on +MI355X). The InferenceX config family also includes vLLM disagg (MoRI-IO) and MXFP8 variants; +those are out of scope here. + +## ATOM disagg mechanics + +- **Topology**: `xP` prefill workers + `yD` decode workers, each TP4 (1 worker = 1 node on an + 8-GPU node, since ceil(4/8)=1). + - `1p1d` = 2 nodes; `2p1d_dpa` = 2 prefill + 1 decode with DP-attention = 3 nodes (high conc). +- **Prefill server** (kv_producer): + ``` + python3 -m atom.entrypoints.openai_server --model --host 0.0.0.0 --server-port 8010 \ + -tp 4 [--enable-dp-attention] --block-size 128 \ + --gpu-memory-utilization 0.8 --max-num-seqs N --max-model-len 32768 \ + --max-num-batched-tokens 32768 --no-enable_prefix_caching \ + --kv-transfer-config '{"kv_role":"kv_producer","kv_connector":"mooncake","proxy_ip":"","handshake_port":6301}' + ``` +- **Decode server** (kv_consumer): same, `--server-port 8020`, `kv_role":"kv_consumer"`. +- **Router** (on prefill node 0, after all servers healthy): + ``` + atomesh launch --host 0.0.0.0 --port 8000 --pd-disaggregation \ + --prefill http://:8010 [...] --decode http://:8020 [...] \ + --policy random --backend atom --disable-health-check --disable-circuit-breaker + ``` +- **Bench**: drive `http://:8000` (OpenAI-compatible) at aggregate concurrency. +- **M3-specific env**: `AITER_QUICK_REDUCE_QUANTIZATION=INT4` (MXFP4 quick-reduce), applied for + all non-DSv4 models by `server_atom.sh`. No `--enable-tbo` (that is a DSv4-only path). +- KV transfer uses mooncake's own RDMA device selection (IBDEVICES not passed as a server arg). + +## Topologies to benchmark (from minimaxm3-fp4-mi355x-atom-disagg) +| shape | conc | topo | nodes | +|-------|------|------|-------| +| 8192/1024 | 1–256 | 1P1D TP4 | 2 | +| 1024/1024 | 1–256 | 1P1D TP4 | 2 | +| 8192/1024 | 256–1024 | 2P1D + DP-attn TP4 | 3 | diff --git a/scripts/MiniMax-M3/lib/cfg.py b/scripts/MiniMax-M3/lib/cfg.py new file mode 100644 index 00000000..9ced83da --- /dev/null +++ b/scripts/MiniMax-M3/lib/cfg.py @@ -0,0 +1,34 @@ +#!/usr/bin/env python3 +"""cfg.py — tiny yaml field reader for the shell scripts. +Usage: + cfg.py cluster -> cluster.yaml top-level key + cfg.py model -> models.. + cfg.py engine -> engines.. +Lists print space-separated; scalars print as-is. +""" +import sys, yaml + +def emit(v): + if isinstance(v, list): print(" ".join(str(x) for x in v)) + elif isinstance(v, dict): print(" ".join(f"{k}={v}" for k, v in v.items())) + else: print("" if v is None else v) + +def main(): + f = sys.argv[1]; kind = sys.argv[2] + d = yaml.safe_load(open(f)) + if kind == "cluster": + emit(d.get(sys.argv[3])) + elif kind == "model": + emit(d["models"][sys.argv[3]].get(sys.argv[4])) + elif kind == "engine": + emit(d["engines"][sys.argv[3]].get(sys.argv[4])) + elif kind == "default": + # default : per-model override wins over global defaults{} + model = sys.argv[3]; key = sys.argv[4] + mv = d.get("models", {}).get(model, {}).get(key) + emit(mv if mv is not None else d.get("defaults", {}).get(key)) + else: + sys.exit(f"unknown kind {kind}") + +if __name__ == "__main__": + main() diff --git a/scripts/MiniMax-M3/lib/check_accuracy.py b/scripts/MiniMax-M3/lib/check_accuracy.py new file mode 100644 index 00000000..e89571ce --- /dev/null +++ b/scripts/MiniMax-M3/lib/check_accuracy.py @@ -0,0 +1,59 @@ +#!/usr/bin/env python3 +# Sanity/accuracy check: send 5 real prompts to a running OpenAI-compatible server, +# verify coherent + expected output. Catches corrupt-weight / bad-quant failures that +# throughput benchmarks (which use --ignore-eos on random tokens) miss. +# +# Usage: PORT, MODEL, OUT(json) via env. Exits 0 if all coherent, 1 otherwise. +import json, os, sys, time, urllib.request + +PORT = os.environ.get("PORT", "8000") +MODEL = os.environ["MODEL"] +OUT = os.environ.get("OUT", "/results/sanity.json") +HERE = os.path.dirname(os.path.abspath(__file__)) +prompts = json.load(open(os.path.join(HERE, "prompts.json"))) + +def chat(prompt, max_tokens): + body = json.dumps({ + "model": MODEL, + "messages": [{"role": "user", "content": prompt}], + "max_tokens": max_tokens, + "temperature": 0.0, + }).encode() + req = urllib.request.Request(f"http://0.0.0.0:{PORT}/v1/chat/completions", + data=body, headers={"Content-Type": "application/json"}) + t0 = time.time() + with urllib.request.urlopen(req, timeout=120) as r: + d = json.load(r) + return d["choices"][0]["message"]["content"], time.time() - t0 + +results = [] +allpass = True +for p in prompts: + try: + text, dt = chat(p["prompt"], p["max_tokens"]) + low = text.lower() + hits = [k for k in p["expect_contains"] if k.lower() in low] + # coherent = non-empty, not pure repetition/garbage + coherent = len(text.strip()) > 0 and len(set(text.split())) > 1 or len(text.strip()) <= 8 + ok = len(hits) > 0 and coherent + results.append({"id": p["id"], "ok": ok, "hits": hits, + "expected": p["expect_contains"], "latency_s": round(dt, 2), + "output": text.strip()[:200]}) + if not ok: + allpass = False + except Exception as e: + results.append({"id": p["id"], "ok": False, "error": str(e)[:150]}) + allpass = False + +summary = {"model": MODEL, "all_pass": allpass, + "passed": sum(1 for r in results if r.get("ok")), + "total": len(results), "results": results} +os.makedirs(os.path.dirname(OUT), exist_ok=True) +json.dump(summary, open(OUT, "w"), indent=2) + +print(f"=== SANITY: {summary['passed']}/{summary['total']} prompts passed ===") +for r in results: + mark = "PASS" if r.get("ok") else "FAIL" + out = r.get("output", r.get("error", "")) + print(f" [{mark}] {r['id']:12s} -> {out[:90]!r}") +sys.exit(0 if allpass else 1) diff --git a/scripts/MiniMax-M3/lib/clean_node.sh b/scripts/MiniMax-M3/lib/clean_node.sh new file mode 100644 index 00000000..6d1aea5a --- /dev/null +++ b/scripts/MiniMax-M3/lib/clean_node.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# clean_node.sh — ensure a node's GPUs are in a CLEAN state before deploying an engine. +# Run via srun on the target node (exclusive). Kills leftover serving containers/processes +# that squat VRAM (the #1 cause of OOM-at-init), then verifies VRAM is actually free. +# +# Usage (inside srun on the node): clean_node.sh [max_used_gb_threshold] +# default threshold = 5 GB; exits non-zero if any GPU still above it after cleanup. +set -uo pipefail +THRESH_GB="${1:-5}" + +echo "[clean] $(hostname): removing ALL containers + stray GPU procs (reservation is exclusive to us)" +# The reservation is slurm-exclusive to us, but Docker containers are NOT slurm-tracked, +# so orphaned/stale containers from prior runs can linger and squat VRAM. Remove them ALL. +cids=$(docker ps -aq 2>/dev/null) +[ -n "$cids" ] && { docker kill $cids >/dev/null 2>&1; docker rm -f $cids >/dev/null 2>&1; } || true +# belt-and-suspenders: kill stray engine procs that may hold VRAM +pkill -9 -f "vllm|atom.entrypoints|sglang.launch|EngineCore|sglang.launch_server" 2>/dev/null || true +sleep 5 + +# 3) verify VRAM free on every GPU. +# rocm-smi prints one line per GPU: "GPU[N] : VRAM Total Used Memory (B): ". +# Take the LAST integer on each line (the byte count), not the GPU index. +echo "[clean] verifying VRAM free (threshold ${THRESH_GB}GB)" +busy=0; i=0 +while IFS= read -r line; do + b=$(echo "$line" | grep -oE "[0-9]+" | tail -1) + [ -z "$b" ] && continue + gb=$(( b / 1000000000 )) + if [ "$gb" -gt "$THRESH_GB" ]; then echo " GPU$i: ${gb}GB USED (>threshold)"; busy=1; else echo " GPU$i: ${gb}GB ok"; fi + i=$((i+1)) +done < <(rocm-smi --showmeminfo vram 2>/dev/null | grep -i "Used Memory") +if [ "$busy" -ne 0 ]; then + echo "[clean] WARNING: some GPUs still hold VRAM after cleanup — listing holders:" + rocm-smi --showpids 2>/dev/null | grep -iE "vllm|atom|sglang|EngineCore" | head + exit 1 +fi +echo "[clean] $(hostname): all GPUs clean" +exit 0 diff --git a/scripts/MiniMax-M3/lib/lib_inferencex.sh b/scripts/MiniMax-M3/lib/lib_inferencex.sh new file mode 100644 index 00000000..95e063a3 --- /dev/null +++ b/scripts/MiniMax-M3/lib/lib_inferencex.sh @@ -0,0 +1,1281 @@ +#!/usr/bin/env bash + +# Shared benchmarking utilities for InferenceX + +# Keep Python bytecode out of the mounted workspace. Benchmark jobs often run as +# root inside containers, and root-owned cache directories break future checkout +# cleanup on self-hosted runners. +export PYTHONDONTWRITEBYTECODE=1 +export PYTHONPYCACHEPREFIX="${PYTHONPYCACHEPREFIX:-/tmp/inferencex-pycache}" +mkdir -p "$PYTHONPYCACHEPREFIX" 2>/dev/null || true + +# Inference server port shared by every benchmark recipe. Launchers that need +# a non-default value (e.g. launch_mi355x-amds.sh derives PORT from RUNNER_NAME +# to avoid collisions across concurrent gh-runners on a shared host) set PORT +# themselves before sourcing this file; the `:-` fallback only kicks in when +# nothing upstream set it. +export PORT="${PORT:-8888}" + +# -------------------------------- +# GPU monitoring helpers +# -------------------------------- + +GPU_MONITOR_PID="" +GPU_METRICS_CSV="/workspace/gpu_metrics.csv" +export GPU_METRICS_CSV + +# Start background GPU monitoring that logs metrics every second to CSV. +# Auto-detects NVIDIA (nvidia-smi) or AMD (amd-smi) GPUs. +# Usage: start_gpu_monitor [--output /path/to/output.csv] [--interval 1] +start_gpu_monitor() { + local output="$GPU_METRICS_CSV" + local interval=1 + + while [[ $# -gt 0 ]]; do + case $1 in + --output) output="$2"; shift 2 ;; + --interval) interval="$2"; shift 2 ;; + *) shift ;; + esac + done + + GPU_METRICS_CSV="$output" + export GPU_METRICS_CSV + + if command -v nvidia-smi &>/dev/null; then + nvidia-smi --query-gpu=timestamp,index,power.draw,temperature.gpu,clocks.current.sm,clocks.current.memory,utilization.gpu,utilization.memory \ + --format=csv -l "$interval" > "$output" 2>/dev/null & + GPU_MONITOR_PID=$! + echo "[GPU Monitor] Started NVIDIA (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" + elif command -v amd-smi &>/dev/null; then + # Use amd-smi native watch mode (-w) which includes timestamps automatically. + # Pipe through awk to: skip preamble lines, keep first CSV header, skip repeated headers. + amd-smi metric -p -c -t -u -w "$interval" --csv 2>/dev/null \ + | awk '/^timestamp,/{if(!h){print;h=1};next} h{print}' > "$output" & + GPU_MONITOR_PID=$! + echo "[GPU Monitor] Started AMD (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" + else + echo "[GPU Monitor] No GPU monitoring tool found (nvidia-smi or amd-smi), skipping" + return 0 + fi +} + +# Stop the background GPU monitor and report file size. +stop_gpu_monitor() { + if [[ -n "$GPU_MONITOR_PID" ]] && kill -0 "$GPU_MONITOR_PID" 2>/dev/null; then + kill "$GPU_MONITOR_PID" 2>/dev/null + wait "$GPU_MONITOR_PID" 2>/dev/null || true + echo "[GPU Monitor] Stopped (PID=$GPU_MONITOR_PID)" + if [[ -f "$GPU_METRICS_CSV" ]]; then + local lines + lines=$(wc -l < "$GPU_METRICS_CSV") + echo "[GPU Monitor] Collected $lines rows -> $GPU_METRICS_CSV" + fi + fi + GPU_MONITOR_PID="" +} + +# Check if required environment variables are set +# Usage: check_env_vars VAR1 VAR2 VAR3 ... +# Exits with code 1 if any variable is not set +check_env_vars() { + local missing_vars=() + + for var_name in "$@"; do + if [[ -z "${!var_name:-}" ]]; then + missing_vars+=("$var_name") + fi + done + + if [[ ${#missing_vars[@]} -gt 0 ]]; then + echo "Error: The following required environment variables are not set:" + for var in "${missing_vars[@]}"; do + echo " - $var" + done + exit 1 + fi +} + +# Wait for server to be ready by polling the health endpoint +# All parameters are required +# Parameters: +# --port: Server port +# --server-log: Path to server log file +# --server-pid: Server process ID (required) +# --sleep-interval: Sleep interval between health checks (optional, default: 5) +wait_for_server_ready() { + set +x + local port="" + local server_log="" + local server_pid="" + local sleep_interval=5 + + # Parse arguments + while [[ $# -gt 0 ]]; do + case $1 in + --port) + port="$2" + shift 2 + ;; + --server-log) + server_log="$2" + shift 2 + ;; + --server-pid) + server_pid="$2" + shift 2 + ;; + --sleep-interval) + sleep_interval="$2" + shift 2 + ;; + *) + echo "Unknown parameter: $1" + return 1 + ;; + esac + done + + # Validate required parameters + if [[ -z "$port" ]]; then + echo "Error: --port is required" + return 1 + fi + if [[ -z "$server_log" ]]; then + echo "Error: --server-log is required" + return 1 + fi + if [[ -z "$server_pid" ]]; then + echo "Error: --server-pid is required" + return 1 + fi + + # Wait for server log file to be created (container startup may delay this) + while [ ! -f "$server_log" ]; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "Server died before creating log file. Exiting." + exit 1 + fi + sleep 1 + done + + # Show logs until server is ready + tail -f -n +1 "$server_log" & + local TAIL_PID=$! + until curl --output /dev/null --silent --fail http://0.0.0.0:$port/health; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "Server died before becoming healthy. Exiting." + kill $TAIL_PID + exit 1 + fi + sleep "$sleep_interval" + done + kill $TAIL_PID +} + +# Run benchmark serving with standardized parameters +# All parameters are required except --endpoint, --use-chat-template, --dsv4, and --trust-remote-code +# Parameters: +# --model: Model name +# --port: Server port +# --backend: Backend type - e.g., 'vllm' or 'openai' +# --endpoint: Optional API endpoint override +# --input-len: Random input sequence length +# --output-len: Random output sequence length +# --random-range-ratio: Random range ratio +# --num-prompts: Number of prompts +# --max-concurrency: Max concurrency +# --result-filename: Result filename without extension +# --result-dir: Result directory +# --use-chat-template: Optional flag to enable chat template +# --dsv4: Optional flag to use the DeepSeek-V4 chat template +# (encoding_dsv4.py) instead of the tokenizer's built-in jinja +# template. Implies --use-chat-template. +# --trust-remote-code: Optional flag to trust remote code from HuggingFace +# --server-pid: Optional server process ID to monitor during benchmark +run_benchmark_serving() { + # In eval-only mode, skip the throughput benchmark entirely. + if [ "${EVAL_ONLY}" = "true" ]; then + echo "EVAL_ONLY mode: skipping throughput benchmark" + return 0 + fi + + set +x + local model="" + local port="" + local backend="" + local endpoint="" + local input_len="" + local output_len="" + local random_range_ratio="" + local num_prompts="" + local max_concurrency="" + local result_filename="" + local result_dir="" + local workspace_dir="" + local use_chat_template=false + local dsv4=false + local trust_remote_code=false + local server_pid="" + local tokenizer="" + + while [[ $# -gt 0 ]]; do + case $1 in + --model) + model="$2" + shift 2 + ;; + --port) + port="$2" + shift 2 + ;; + --backend) + backend="$2" + shift 2 + ;; + --endpoint) + endpoint="$2" + shift 2 + ;; + --input-len) + input_len="$2" + shift 2 + ;; + --output-len) + output_len="$2" + shift 2 + ;; + --random-range-ratio) + random_range_ratio="$2" + shift 2 + ;; + --num-prompts) + num_prompts="$2" + shift 2 + ;; + --max-concurrency) + max_concurrency="$2" + shift 2 + ;; + --result-filename) + result_filename="$2" + shift 2 + ;; + --result-dir) + result_dir="$2" + shift 2 + ;; + --bench-serving-dir) + workspace_dir="$2" + shift 2 + ;; + --use-chat-template) + use_chat_template=true + shift + ;; + --dsv4) + dsv4=true + use_chat_template=true + shift + ;; + --trust-remote-code) + trust_remote_code=true + shift + ;; + --server-pid) + server_pid="$2" + shift 2 + ;; + --tokenizer) + tokenizer="$2" + shift 2 + ;; + *) + echo "Unknown parameter: $1" + return 1 + ;; + esac + done + + # Validate all required parameters + if [[ -z "$model" ]]; then + echo "Error: --model is required" + return 1 + fi + if [[ -z "$port" ]]; then + echo "Error: --port is required" + return 1 + fi + if [[ -z "$backend" ]]; then + echo "Error: --backend is required" + return 1 + fi + if [[ -z "$input_len" ]]; then + echo "Error: --input-len is required" + return 1 + fi + if [[ -z "$output_len" ]]; then + echo "Error: --output-len is required" + return 1 + fi + if [[ -z "$random_range_ratio" ]]; then + echo "Error: --random-range-ratio is required" + return 1 + fi + if [[ -z "$num_prompts" ]]; then + echo "Error: --num-prompts is required" + return 1 + fi + if [[ -z "$max_concurrency" ]]; then + echo "Error: --max-concurrency is required" + return 1 + fi + if [[ -z "$result_filename" ]]; then + echo "Error: --result-filename is required" + return 1 + fi + if [[ -z "$result_dir" ]]; then + echo "Error: --result-dir is required" + return 1 + fi + + if [[ -z "$workspace_dir" ]]; then + workspace_dir=$(pwd) + fi + + # Profiling support: when PROFILE=1, ensure profiler dir exists, add --profile flag, + # and cap num_prompts to keep traces small. + local profile_flag=() + if [[ "${PROFILE:-}" == "1" ]]; then + local _prof_dir="${SGLANG_TORCH_PROFILER_DIR:-${VLLM_TORCH_PROFILER_DIR:-}}" + if [[ -n "$_prof_dir" ]]; then + mkdir -p "$_prof_dir" + fi + profile_flag+=(--profile) + num_prompts="$max_concurrency" + fi + + # Build benchmark command + local benchmark_cmd=( + python3 "$workspace_dir/utils/bench_serving/benchmark_serving.py" + --model "$model" + --backend "$backend" + --base-url "http://0.0.0.0:$port" + --dataset-name random + --random-input-len "$input_len" + --random-output-len "$output_len" + --random-range-ratio "$random_range_ratio" + --num-prompts "$num_prompts" + --max-concurrency "$max_concurrency" + --request-rate inf + --ignore-eos + "${profile_flag[@]}" + --save-result + --num-warmups "$((2 * max_concurrency))" \ + --percentile-metrics 'ttft,tpot,itl,e2el' + --result-dir "$result_dir" + --result-filename "$result_filename.json" + ) + + if [[ -n "$endpoint" ]]; then + benchmark_cmd+=(--endpoint "$endpoint") + fi + + # Add --use-chat-template if requested + if [[ "$use_chat_template" == true ]]; then + benchmark_cmd+=(--use-chat-template) + fi + + # Add --dsv4 if requested (requires --use-chat-template, which we + # auto-enable when --dsv4 is passed in). + if [[ "$dsv4" == true ]]; then + benchmark_cmd+=(--dsv4) + fi + + # Add --trust-remote-code if requested + if [[ "$trust_remote_code" == true ]]; then + benchmark_cmd+=(--trust-remote-code) + fi + + if [[ -n "$tokenizer" ]]; then + benchmark_cmd+=(--tokenizer "$tokenizer") + fi + + # Run benchmark with optional server monitoring + set -x + if [[ -n "$server_pid" ]]; then + # Run benchmark in background and monitor server health + "${benchmark_cmd[@]}" & + local benchmark_pid=$! + + # Monitor loop: check both benchmark and server status + while kill -0 "$benchmark_pid" 2>/dev/null; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "ERROR: Server process $server_pid died during benchmark" + kill "$benchmark_pid" 2>/dev/null + wait "$benchmark_pid" 2>/dev/null + set +x + return 1 + fi + sleep 2 + done + + # Benchmark finished, get its exit code + wait "$benchmark_pid" + local benchmark_exit_code=$? + else + # No server monitoring, run benchmark directly + "${benchmark_cmd[@]}" + local benchmark_exit_code=$? + fi + set +x + + # If profiling, move trace to relay-upload location + if [[ "${PROFILE:-}" == "1" ]]; then + move_profile_trace_for_relay + fi + + return $benchmark_exit_code +} + + +# -------------------------------- +# Profiling trace helpers +# -------------------------------- + +_find_latest_profile_trace() { + local latest="" + local dir="" candidate="" base="" + local -a search_roots=() + + for dir in "$@"; do + search_roots=() + if [[ -d "$dir" ]]; then + search_roots+=("$dir") + fi + if [[ -d "$dir/profiles" ]]; then + search_roots+=("$dir/profiles") + fi + if [[ ${#search_roots[@]} -eq 0 ]]; then + continue + fi + + while IFS= read -r -d '' candidate; do + base="$(basename "$candidate")" + if [[ "$base" == profile_*.trace.json.gz ]]; then + continue + fi + if [[ -z "$latest" || "$candidate" -nt "$latest" ]]; then + latest="$candidate" + fi + done < <( + find "${search_roots[@]}" -maxdepth 1 -type f \ + \( -name "*.trace.json" -o -name "*.trace.json.gz" -o -name "*trace*.json" -o -name "*trace*.json.gz" -o -name "*profile*.json" -o -name "*profile*.json.gz" \) \ + -print0 2>/dev/null + ) + done + + printf '%s' "$latest" +} + +# Move profiler trace into a stable workspace path for workflow relay/upload. +move_profile_trace_for_relay() { + if [[ "${PROFILE:-}" != "1" ]]; then + return 0 + fi + + if [[ -z "${RESULT_FILENAME:-}" ]]; then + echo "[PROFILE] RESULT_FILENAME is not set; skipping relay trace staging." >&2 + return 0 + fi + + local sglang_dir="${SGLANG_TORCH_PROFILER_DIR:-/workspace}" + local vllm_dir="${VLLM_TORCH_PROFILER_DIR:-/workspace}" + local -a search_dirs=() + local dir="" existing="" + local seen=0 + + for dir in "$sglang_dir" "$vllm_dir" "/workspace"; do + if [[ -z "$dir" ]]; then + continue + fi + seen=0 + for existing in "${search_dirs[@]}"; do + if [[ "$existing" == "$dir" ]]; then + seen=1 + break + fi + done + if [[ "$seen" -eq 0 ]]; then + search_dirs+=("$dir") + fi + done + + local trace_file="" + local wait_attempts=10 + for (( i=1; i<=wait_attempts; i++ )); do + trace_file="$(_find_latest_profile_trace "${search_dirs[@]}")" + if [[ -n "$trace_file" ]]; then + break + fi + sleep 10 + done + + if [[ -z "$trace_file" ]]; then + echo "[PROFILE] No trace found for relay under: ${search_dirs[*]}" >&2 + return 0 + fi + + local dest_trace="/workspace/profile_${RESULT_FILENAME}.trace.json.gz" + if [[ "$trace_file" == *.gz ]]; then + cp -f "$trace_file" "$dest_trace" + else + gzip -c "$trace_file" > "$dest_trace" + fi + + echo "[PROFILE] Relay trace prepared: $dest_trace (source: $trace_file)" +} + + +# ------------------------------ +# Eval (lm-eval-harness) helpers +# ------------------------------ + +_install_lm_eval_deps() { + # torchvision causes circular imports in ATOM; TRT-LLM/SGLang need it at module level. + if [[ "${IMAGE:-}" == *atom* ]]; then + python3 -m pip uninstall -y torchvision 2>/dev/null || true + fi + python3 -m pip install -q --no-cache-dir --break-system-packages "lm-eval[api]" || true + local lm_eval_ref="b315ef3b05176acc9732bb7fdec116abe1ecc476" + if command -v git >/dev/null 2>&1; then + if ! python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "git+https://github.com/EleutherAI/lm-evaluation-harness.git@${lm_eval_ref}"; then + python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true + fi + else + python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true + fi +} + +# Patch lm-eval filters to be robust to empty strings via sitecustomize +_patch_lm_eval() { + local patch_dir + patch_dir="$(mktemp -d)" + cat > "$patch_dir/sitecustomize.py" <<'PY' +# --- Patch LocalChatCompletion.parse_generations to handle empty content with reasoning_content --- +import re, sys, unicodedata, json +from lm_eval.filters import extraction as ex +from lm_eval.models.openai_completions import LocalChatCompletion as _LCC + +def _le_parse_generations(outputs, **kwargs): + res = [] + if not isinstance(outputs, list): + outputs = [outputs] + for out in (outputs or []): + try: + choices = out.get("choices", []) + tmp = ["" for _ in choices] + for choice in choices: + idx = choice.get("index", 0) + msg = (choice.get("message") or {}) + content = msg.get("content") + if content in (None, "", []): + content = msg.get("reasoning_content") or "" + tmp[idx] = content + except Exception: + tmp = [""] + res.extend(tmp) + return res + +# Keep staticmethod semantics +_LCC.parse_generations = staticmethod(_le_parse_generations) + +# --- Patch TemplateAPI.apply_chat_template to avoid injecting "type": "text" for TRT --- +try: + from lm_eval.models import api_models as _api_models + _TemplateAPI = _api_models.TemplateAPI + _JsonChatStr = _api_models.JsonChatStr +except Exception: + _TemplateAPI = None + _JsonChatStr = None + +if _TemplateAPI is not None and _JsonChatStr is not None: + _orig_apply_chat_template = _TemplateAPI.apply_chat_template + + def _patched_apply_chat_template( + self, + chat_history, + add_generation_prompt: bool = True, + ): + """Applies a chat template to a list of chat history between user and model.""" + if self.tokenizer_backend == "huggingface" and self.tokenized_requests: + return self.tokenizer.apply_chat_template( + chat_history, + tokenize=False, + add_generation_prompt=add_generation_prompt, + continue_final_message=not add_generation_prompt, + ) + elif self.tokenizer_backend == "remote" and self.tokenized_requests: + return chat_history + else: + # NOTE: we no longer inject `"type": "text"` when tokenizer is None / non-HF + return _JsonChatStr( + json.dumps( + [{**item} for item in chat_history], + ensure_ascii=False, + ) + ) + + _TemplateAPI.apply_chat_template = _patched_apply_chat_template +PY + export PYTHONPATH="${patch_dir}:${PYTHONPATH:-}" +} + +get_native_max_context_length() { + local model_path="$1" + # Prefer MODEL_PATH (local model directory) when available, since the + # argument may be a served-model name that is neither a valid HF repo + # ID nor a local path (e.g. "deepseek-r1-fp4" on the B300 cluster). + if [ -n "${MODEL_PATH:-}" ] && [ -d "${MODEL_PATH}" ]; then + model_path="${MODEL_PATH}" + fi + python3 -c " +try: + from transformers import AutoConfig + config = AutoConfig.from_pretrained('${model_path}', trust_remote_code=True) + for attr in ['max_position_embeddings', 'max_sequence_length', 'seq_length', 'n_positions']: + if hasattr(config, attr): + print(getattr(config, attr)) + break + else: + print(0) +except Exception: + print(0) +" +} + +# Compute the context length for eval-only mode. +# Uses the requested benchmark context capped at the model's native max. +# Sets EVAL_MAX_MODEL_LEN (needed by run_lm_eval). +# Echoes the computed value for scripts to capture. +# +# Usage: local ctx=$(compute_eval_context_length "$MODEL" "${current_ctx}") +compute_eval_context_length() { + local model="$1" + local benchmark_ctx="${2:-0}" + local native_max + native_max=$(get_native_max_context_length "$model") + native_max="${native_max:-0}" + + if [ "$benchmark_ctx" -eq 0 ] 2>/dev/null; then + benchmark_ctx="${native_max:-0}" + fi + local eval_ctx=$(( benchmark_ctx * 1 )) + if [ "$native_max" -gt 0 ] 2>/dev/null && [ "$eval_ctx" -gt "$native_max" ]; then + eval_ctx="$native_max" + fi + # If eval_ctx is still 0 (both benchmark_ctx and native_max were 0), fall back + if [ "$eval_ctx" -le 0 ] 2>/dev/null; then + echo "WARN: compute_eval_context_length could not determine context length for $model" >&2 + eval_ctx="${MAX_MODEL_LEN:-16384}" + fi + EVAL_MAX_MODEL_LEN="$eval_ctx" + echo "$eval_ctx" +} + +# Convenience wrapper: compute eval context from ISL/OSL and export EVAL_MAX_MODEL_LEN. +# Call directly (not in a subshell) so the export persists. +# Scripts then wire $EVAL_MAX_MODEL_LEN into whichever server variable they need. +setup_eval_context() { + EVAL_MAX_MODEL_LEN=$(compute_eval_context_length "$MODEL" "$((ISL + OSL + 256))") + export EVAL_MAX_MODEL_LEN +} + +run_lm_eval() { + local port="${PORT:-8888}" + local tasks_dir="${EVAL_TASKS_DIR:-utils/evals/gsm8k.yaml}" + local results_dir="${EVAL_RESULT_DIR:-$(mktemp -d /tmp/eval_out-XXXXXX)}" + local eval_context_len="${EVAL_MAX_MODEL_LEN:-16384}" + local temperature=0 + local top_p=1 + local concurrent_requests="${EVAL_CONCURRENT_REQUESTS:-${CONC:-64}}" + + while [[ $# -gt 0 ]]; do + case $1 in + --port) port="$2"; shift 2 ;; + --task) tasks_dir="$2"; shift 2 ;; + --results-dir) results_dir="$2"; shift 2 ;; + --gen-max-tokens) eval_context_len="$2"; shift 2 ;; + --temperature) temperature="$2"; shift 2 ;; + --top-p) top_p="$2"; shift 2 ;; + *) echo "Unknown parameter: $1"; return 1 ;; + esac + done + + if [ "${INFERENCEX_LM_EVAL_RUNTIME_READY:-false}" != "true" ]; then + _install_lm_eval_deps + _patch_lm_eval + export INFERENCEX_LM_EVAL_RUNTIME_READY=true + fi + + local openai_server_base="http://0.0.0.0:${port}" + local openai_chat_base="${openai_server_base}/v1/chat/completions" + export OPENAI_API_KEY=${OPENAI_API_KEY:-EMPTY} + MODEL_NAME=${MODEL_NAME:-$MODEL} # Prefer MODEL_NAME, else MODEL + + # Cap output tokens: must fit within context window (leave room for input), + # and avoid excessive KV cache reservation per request on TRT. + local max_output_tokens=$(( eval_context_len > 4096 ? eval_context_len - 4096 : eval_context_len / 2 )) + if [ "$max_output_tokens" -gt 16384 ]; then + max_output_tokens=16384 + fi + echo "Eval budget: eval_context_len=${eval_context_len}, max_output_tokens=${max_output_tokens}" + + # Export for append_lm_eval_summary to pick up + export EVAL_RESULT_DIR="$results_dir" + set -x + python3 -m lm_eval --model local-chat-completions --apply_chat_template \ + --tasks "${tasks_dir}" \ + --output_path "${results_dir}" \ + --log_samples \ + --model_args "model=${MODEL_NAME},base_url=${openai_chat_base},api_key=${OPENAI_API_KEY},eos_string=,max_retries=5,num_concurrent=${concurrent_requests},timeout=1800,tokenized_requests=False,max_length=${eval_context_len}" \ + --gen_kwargs "max_tokens=${max_output_tokens},temperature=${temperature},top_p=${top_p}" + local eval_exit=$? + set +x + return $eval_exit +} + +_stage_lm_eval_artifacts() { + local results_dir="$1" + local eval_conc="$2" + local moved=0 + local failed=0 + local jf base stem extension target suffix + + if [ ! -d "${results_dir}" ]; then + echo "WARN: eval result directory '${results_dir}' does not exist" >&2 + return 1 + fi + + while IFS= read -r -d '' jf; do + base=$(basename "$jf") + case "$base" in + meta_env.json) + continue + ;; + *.jsonl) + stem="${base%.jsonl}" + extension=".jsonl" + ;; + *.json) + stem="${base%.json}" + extension=".json" + ;; + *) + continue + ;; + esac + + target="./${stem}_conc${eval_conc}${extension}" + suffix=2 + while [ -e "$target" ]; do + target="./${stem}_conc${eval_conc}_${suffix}${extension}" + suffix=$((suffix + 1)) + done + + if mv -f "$jf" "$target"; then + moved=1 + else + echo "WARN: failed to stage eval artifact ${jf}" >&2 + failed=1 + fi + done < <( + find "${results_dir}" -type f \ + \( -name "*.json" -o -name "*.jsonl" \) -print0 2>/dev/null + ) + + rm -rf --one-file-system "${results_dir}" 2>/dev/null \ + || rm -rf "${results_dir}" \ + || true + + if [ "$moved" -eq 0 ]; then + echo "WARN: no eval artifacts were produced for concurrency ${eval_conc}" >&2 + return 1 + fi + return "$failed" +} + +_eval_concs_to_json() { + local values="$1" + local value + local joined="" + + for value in $values; do + if [[ ! "$value" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: invalid eval concurrency '${value}'" >&2 + return 1 + fi + if [ -n "$joined" ]; then + joined="${joined}, " + fi + joined="${joined}${value}" + done + + printf '[%s]' "$joined" +} + +append_lm_eval_summary() { + local batch_concs="${EVAL_BATCHED_CONCS:-}" + local results_dir="${EVAL_RESULT_DIR:-}" + local out_dir="${results_dir}" + local meta_json + local metadata_conc="${CONC:-1}" + local batch_metadata="" + + if [ -n "$batch_concs" ]; then + meta_json="./meta_env.json" + metadata_conc="${batch_concs%% *}" + + local eval_concs_json completed_concs_json failed_concs_json + eval_concs_json=$(_eval_concs_to_json "$batch_concs") || return 1 + completed_concs_json=$( + _eval_concs_to_json "${EVAL_BATCHED_COMPLETED_CONCS:-}" + ) || return 1 + failed_concs_json=$( + _eval_concs_to_json "${EVAL_BATCHED_FAILED_CONCS:-}" + ) || return 1 + printf -v batch_metadata \ + ' "eval_concs": %s,\n "completed_eval_concs": %s,\n "failed_eval_concs": %s,\n' \ + "$eval_concs_json" \ + "$completed_concs_json" \ + "$failed_concs_json" + else + if [ -z "${results_dir}" ]; then + echo "WARN: EVAL_RESULT_DIR is empty; skipping artifact collection" >&2 + return 1 + fi + if [ ! -d "${out_dir}" ]; then + echo "WARN: EVAL_RESULT_DIR='${out_dir}' does not exist; skipping artifact collection" >&2 + return 1 + fi + meta_json="${out_dir}/meta_env.json" + fi + + # Write minimal meta for collectors that expect it + local model_name="${MODEL_NAME:-$MODEL}" + local is_multinode_json="false" + if [ "${IS_MULTINODE:-false}" = "true" ]; then + is_multinode_json="true" + fi + + local prefill_tp="${PREFILL_TP:-${TP:-1}}" + local prefill_ep="${PREFILL_EP:-${EP_SIZE:-1}}" + local prefill_num_workers="${PREFILL_NUM_WORKERS:-1}" + local decode_tp="${DECODE_TP:-${TP:-1}}" + local decode_ep="${DECODE_EP:-${EP_SIZE:-1}}" + local decode_num_workers="${DECODE_NUM_WORKERS:-1}" + + local dp_json="false" + if [ "${DP_ATTENTION:-false}" = "true" ]; then dp_json="true"; fi + local prefill_dp_json="$dp_json" + if [ "${PREFILL_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then + prefill_dp_json="true" + else + prefill_dp_json="false" + fi + local decode_dp_json="$dp_json" + if [ "${DECODE_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then + decode_dp_json="true" + else + decode_dp_json="false" + fi + + # Derive framework/precision from env, fallback to parsing RESULT_FILENAME + # RESULT_FILENAME format (from workflow): + # ___tp<...>_ep<...>_dpa_<...>_conc<...>_ + local fw="${FRAMEWORK:-}" + local prec="${PRECISION:-}" + if [[ -z "$fw" || -z "$prec" ]]; then + if [[ -n "${RESULT_FILENAME}" ]]; then + # Extract the two fields immediately before "_tp" + # Handles arbitrary underscores in exp_name by matching from the end + local parsed + parsed=$(echo "${RESULT_FILENAME}" | sed -n 's/.*_\([^_][^_]*\)_\([^_][^_]*\)_tp.*/\1 \2/p') + local p1="${parsed%% *}" + local p2="${parsed#* }" + if [[ -z "$prec" && -n "$p1" && "$p1" != "$parsed" ]]; then + prec="$p1" + fi + if [[ -z "$fw" && -n "$p2" && "$p2" != "$parsed" ]]; then + fw="$p2" + fi + fi + fi + cat > "${meta_json}" <&2 + fi + if [ -d "${out_dir}" ]; then + while IFS= read -r -d '' jf; do + base=$(basename "$jf") + if [ "$base" != "meta_env.json" ]; then + mv -f "$jf" ./ || echo "WARN: failed to move ${jf}" >&2 + fi + done < <(find "${out_dir}" -type f -name "*.json*" -print0 2>/dev/null) + fi + + # Best-effort cleanup of the temp directory + if [ -n "${out_dir}" ] && [ -d "${out_dir}" ]; then + rm -rf --one-file-system "${out_dir}" || rm -rf "${out_dir}" || true + fi + + echo "Moved eval artifacts to: $(pwd)" +} + +# ------------------------------ +# Unified eval entrypoint +# ------------------------------ + +run_eval() { + local framework="${EVAL_FRAMEWORK:-lm-eval}" + local forwarded=() + + while [[ $# -gt 0 ]]; do + case "$1" in + --framework) framework="$2"; shift 2 ;; + *) forwarded+=("$1"); shift ;; + esac + done + + # Compute EVAL_MAX_MODEL_LEN if not already set by the calling script + if [ -z "${EVAL_MAX_MODEL_LEN:-}" ]; then + compute_eval_context_length "$MODEL" "${MAX_MODEL_LEN:-0}" > /dev/null + fi + + unset EVAL_BATCHED_CONCS + unset EVAL_BATCHED_COMPLETED_CONCS + unset EVAL_BATCHED_FAILED_CONCS + + local requested_concs="${EVAL_CONCURRENT_REQUESTS:-}" + local eval_concs=() + if [ -n "$requested_concs" ]; then + read -r -a eval_concs <<< "$requested_concs" + fi + + if [ "${#eval_concs[@]}" -gt 1 ]; then + if [[ "$framework" != "lm-eval" && "$framework" != "lm_eval" ]]; then + echo "ERROR: batched eval concurrency is only supported for lm-eval" >&2 + return 1 + fi + + local eval_conc results_dir eval_rc stage_rc + local completed_concs=() + local failed_concs=() + + for eval_conc in "${eval_concs[@]}"; do + if [[ ! "$eval_conc" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: invalid eval concurrency '${eval_conc}'" >&2 + return 1 + fi + + if ! results_dir=$(mktemp -d /tmp/eval_out-conc"${eval_conc}"-XXXXXX); then + echo "ERROR: failed to create eval output directory for concurrency ${eval_conc}" >&2 + failed_concs+=("$eval_conc") + continue + fi + + echo "Running lm-eval at concurrency ${eval_conc} using the existing engine" + export EVAL_CONCURRENT_REQUESTS="$eval_conc" + export CONC="$eval_conc" + eval_rc=0 + stage_rc=0 + run_lm_eval "${forwarded[@]}" --results-dir "$results_dir" \ + || eval_rc=$? + _stage_lm_eval_artifacts "$results_dir" "$eval_conc" \ + || stage_rc=$? + + if [ "$eval_rc" -eq 0 ] && [ "$stage_rc" -eq 0 ]; then + completed_concs+=("$eval_conc") + else + echo "ERROR: lm-eval failed at concurrency ${eval_conc} (eval_rc=${eval_rc}, stage_rc=${stage_rc})" >&2 + failed_concs+=("$eval_conc") + fi + done + + export EVAL_CONCURRENT_REQUESTS="$requested_concs" + export EVAL_RESULT_DIR="" + export EVAL_BATCHED_CONCS="${eval_concs[*]}" + export EVAL_BATCHED_COMPLETED_CONCS="${completed_concs[*]}" + export EVAL_BATCHED_FAILED_CONCS="${failed_concs[*]}" + + if [ "${#failed_concs[@]}" -gt 0 ]; then + echo "ERROR: batched eval failed for concurrency: ${failed_concs[*]}" >&2 + echo "Deferring failure until post-upload score validation preserves all artifacts" >&2 + fi + return 0 + fi + + local eval_rc=0 + case "$framework" in + lm-eval|lm_eval) run_lm_eval "${forwarded[@]}" || eval_rc=$? ;; + *) echo "Unknown framework '${framework}'"; eval_rc=1 ;; + esac + + if [ "$eval_rc" -ne 0 ]; then + echo "ERROR: run_eval failed with exit code $eval_rc" >&2 + if [ "${EVAL_ONLY}" = "true" ]; then + echo "Eval-only mode: failing after artifact collection" >&2 + return "$eval_rc" + fi + fi + return $eval_rc +} + + +# -------------------------------- +# Agentic trace replay helpers (aiperf driver) +# -------------------------------- + +INFMAX_CONTAINER_WORKSPACE="${INFMAX_CONTAINER_WORKSPACE:-/workspace}" +AGENTIC_DIR="${AGENTIC_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/agentic-benchmark}" +AIPERF_DIR="${AIPERF_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/aiperf}" + +agentic_pip_install() { + local pip_install=(python3 -m pip install) + if python3 -m pip install --help 2>/dev/null | grep -q -- "--break-system-packages"; then + pip_install+=(--break-system-packages) + fi + + "${pip_install[@]}" "$@" +} + +ensure_hf_cli() { + if command -v hf >/dev/null 2>&1; then + return 0 + fi + + # Some lean runtime images used by multinode SGLang include Python but not + # the Hugging Face CLI. Install just the hub CLI before prefetching traces. + agentic_pip_install --quiet "huggingface_hub[cli]>=0.25.0" +} + +resolve_trace_source() { + # Per-recipe override: set WEKA_LOADER_OVERRIDE to one of the aiperf + # public-dataset loader names allowed by the inferencex-agentx-mvp + # scenario. Used by recipes whose servers have non-default context + # caps (e.g. minimaxm2.5 at max_model_len ~256k can't replay the + # unfiltered 052726 corpus and switches to the 256k-capped variant). + local loader="${WEKA_LOADER_OVERRIDE:-semianalysis_cc_traces_weka_with_subagents}" + local dataset + case "$loader" in + semianalysis_cc_traces_weka_with_subagents) + dataset="semianalysisai/cc-traces-weka-with-subagents-052726" + ;; + semianalysis_cc_traces_weka_with_subagents_256k) + dataset="semianalysisai/cc-traces-weka-with-subagents-052726-256k" + ;; + *) + echo "Error: unknown WEKA_LOADER_OVERRIDE='$loader'. Allowed: semianalysis_cc_traces_weka_with_subagents, semianalysis_cc_traces_weka_with_subagents_256k" >&2 + exit 1 + ;; + esac + TRACE_SOURCE_FLAG="--public-dataset $loader" + echo "Loading traces via aiperf public-dataset: $loader ($dataset)" + # Pre-download the dataset into the shared HF_HUB_CACHE (same mount used + # for model weights) so subsequent runs read from cache instead of + # re-downloading every job. + ensure_hf_cli + hf download --repo-type dataset "$dataset" +} + +install_agentic_deps() { + # vllm/vllm-openai container ships without git. pip needs git to + # introspect the aiperf source tree on install. Install on demand; + # no-op when git is already present (e.g. AMD images that ship it). + if ! command -v git >/dev/null 2>&1; then + apt-get update && apt-get install -y git + fi + agentic_pip_install --quiet urllib3 requests 2>/dev/null || true + agentic_pip_install -q -r "$AGENTIC_DIR/requirements.txt" + # Editable install of aiperf from the submodule — gives us the + # `aiperf` CLI plus the inferencex-agentx-mvp scenario plugin. + # + # `--ignore-installed` sidesteps the distutils-uninstall error that + # vLLM containers hit on apt-managed system packages (blinker, etc.) + # when pip's resolver tries to upgrade one of aiperf's transitive + # deps. Installing fresh into the user/site location is safe — the + # system package stays in place and pip's import order picks up our + # newer copy first. + agentic_pip_install -q --ignore-installed -e "$AIPERF_DIR" + # Force-upgrade datasets: containers often ship an older version without + # the `Json` feature type used by the HF traces dataset. `Json` was added + # in datasets 4.7.0 (March 2025). Unpinned installs won't upgrade an + # already-present package. + agentic_pip_install --upgrade "datasets>=4.7.0" +} + +build_replay_cmd() { + # aiperf invocation for the inferencex-agentx-mvp scenario. + # + # Pre-canned assistant replay is the default: recorded assistant responses + # are used for future prompt construction, and live server responses are + # discarded. Set AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES=1 explicitly + # to use live-assistant mode, where the loader emits user-only deltas and + # the worker threads the server's live assistant response back into the + # session. + # + # The scenario plugin locks: --cache-bust first_turn_prefix and + # --trace-idle-gap-cap-seconds 60 (per-trace idle-gap compression + # against parent + subagent request-start timestamps; supersedes the + # legacy --use-think-time-only / --inter-turn-delay-cap-seconds path), + # and auto-injects them — so we do not pass them. See + # utils/aiperf/docs/tutorials/agentx-mvp.md. + local result_dir="$1" + local duration="$DURATION" + + export AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES="${AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES:-0}" + # Dataset configuration (load + reconstruct + inputs.json + mmap) + # routinely takes 4-5 min for the Weka corpus on fast /tmp + # (B300) but can stretch to 14 min on slower /tmp + parallel contention + # (observed on H200 where all 14 R3 jobs hit aiperf's 900s Configure + # Profiling timeout simultaneously). Bump to 1800s to absorb 3x + # worst-case slowdown — the post-setup measurement window is unaffected. + export AIPERF_DATASET_CONFIGURATION_TIMEOUT=1800 + # aiperf validates that SERVICE_PROFILE_CONFIGURE_TIMEOUT >= + # DATASET_CONFIGURATION_TIMEOUT at startup. Bump it in lockstep. + export AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT=1800 + REPLAY_CMD="aiperf profile --scenario inferencex-agentx-mvp" + REPLAY_CMD+=" --url http://localhost:$PORT" + REPLAY_CMD+=" --endpoint /v1/chat/completions" + REPLAY_CMD+=" --endpoint-type chat" + REPLAY_CMD+=" --streaming" + REPLAY_CMD+=" --model $MODEL" + REPLAY_CMD+=" --concurrency $CONC" + REPLAY_CMD+=" --benchmark-duration $duration" + REPLAY_CMD+=" --random-seed 42" + # Fail runs once more than 10% of requests error. This keeps known + # transient low-rate failures from killing long sweeps while still + # catching malformed payloads or server crashes before they get aggregated + # as benchmarkable data. + REPLAY_CMD+=" --failed-request-threshold 0.10" + # Sample each trajectory's warmup start position uniformly from + # [25%, 75%] of the trace's turn count (was hardcoded 0%-70% upstream). + # Avoids starting trajectories right at turn 0 where the KV cache is + # cold and skews early steady-state samples. + REPLAY_CMD+=" --trajectory-start-min-ratio 0.25" + REPLAY_CMD+=" --trajectory-start-max-ratio 0.75" + # Use server-reported usage fields (prompt_tokens / completion_tokens) for + # ISL/OSL instead of client-side tokenizer.encode(). Auto-enables + # stream_options.include_usage on the OpenAI chat endpoint. Skips the + # heavy per-record tokenization in the records pipeline that was pinning + # CPU on minimax-m2.5 at high concurrency. Lossless for vLLM (server + # usage is authoritative). + REPLAY_CMD+=" --use-server-token-count" + # aiperf's dataset manager (separate from the inference parser) loads + # the model's tokenizer for trace-prompt tokenization regardless of + # --use-server-token-count. Models like kimi (amd/Kimi-K2.5-MXFP4, + # moonshotai/Kimi-K2.5) ship a custom tokenizer in their HF repo and + # need trust_remote_code=True to load. Benign for models without + # custom tokenizer code, so we set it unconditionally. + REPLAY_CMD+=" --tokenizer-trust-remote-code" + # Keep replay inputs inside the same context window used to launch the + # server. The WEKA corpus contains a few very long parent/subagent traces; + # if we mmap and replay them against a smaller-context server they become + # deterministic 4xxs and can still pressure the engine while queued. + if [ -n "${MAX_MODEL_LEN:-}" ] && [ "$MAX_MODEL_LEN" != "0" ]; then + REPLAY_CMD+=" --max-context-length $MAX_MODEL_LEN" + fi + # Default --num-dataset-entries is 100; the with-subagents Weka corpus + # has 472. Cap at 472 so all unique traces are loaded (the loader treats + # this as a ``min(cap, available)`` ceiling, not a target — see + # semianalysis_cc_traces_weka.py). + REPLAY_CMD+=" --num-dataset-entries 472" + # 1-second timeslices on the server-metrics scrape so the post-run + # plotter has per-window time series (KV usage, cache hit rate, + # throughput, etc.). Matches kv-cache-tester's poll_interval=1.0 + # snapshot cadence so metrics_plots.png is visually comparable. + # Without this, aiperf only emits aggregate stats and the 6x2 panels + # collapse to flat lines. + REPLAY_CMD+=" --slice-duration 1.0" + REPLAY_CMD+=" --output-artifact-dir $result_dir/aiperf_artifacts" + # The inferencex-agentx-mvp scenario enforces a 900s minimum + # benchmark duration. For smoke tests with shorter durations, opt + # into --unsafe-override (the run's submission_valid will be flagged + # false; that's expected for non-canonical runs). + if [ "$duration" -lt 900 ] || [ "${AIPERF_UNSAFE_OVERRIDE:-false}" = "true" ]; then + REPLAY_CMD+=" --unsafe-override" + fi + REPLAY_CMD+=" $TRACE_SOURCE_FLAG" +} + +write_agentic_result_json() { + # Aggregate aiperf's profile_export.{json,jsonl} + server_metrics_export.json + # into $AGENTIC_OUTPUT_DIR/$RESULT_FILENAME.json. The workflow's existing + # retry-based existence check is the single success gate. + local result_dir="$1" + RESULT_DIR="$result_dir" AGENTIC_OUTPUT_DIR="${AGENTIC_OUTPUT_DIR:-$INFMAX_CONTAINER_WORKSPACE}" \ + python3 "$INFMAX_CONTAINER_WORKSPACE/utils/process_agentic_result.py" + + # Generate metrics_plots.png from the same aiperf artifacts. Best-effort: + # don't fail the launcher if plot generation has trouble (e.g. matplotlib + # missing in a stripped-down image). The agg JSON is the success gate. + python3 "$INFMAX_CONTAINER_WORKSPACE/utils/generate_aiperf_plots.py" "$result_dir" 2>&1 || true +} + +run_agentic_replay_and_write_outputs() { + local result_dir="$1" + local replay_rc + + echo "$REPLAY_CMD" > "$result_dir/benchmark_command.txt" + + set +e + set -x + $REPLAY_CMD 2>&1 | tee "$result_dir/benchmark.log" + replay_rc=${PIPESTATUS[0]} + set +x + set -e + + write_agentic_result_json "$result_dir" + + python3 "$AGENTIC_DIR/scripts/analyze_benchmark_distributions.py" \ + "$result_dir/aiperf_artifacts" -o "$result_dir" 2>&1 || true + + if [ "$replay_rc" -ne 0 ]; then + echo "ERROR: agentic trace replay exited with code $replay_rc after writing available results" >&2 + return "$replay_rc" + fi +} diff --git a/scripts/MiniMax-M3/lib/prompts.json b/scripts/MiniMax-M3/lib/prompts.json new file mode 100644 index 00000000..1b0d3755 --- /dev/null +++ b/scripts/MiniMax-M3/lib/prompts.json @@ -0,0 +1,32 @@ +[ + { + "id": "factual", + "prompt": "What is the capital of France? Answer in one word.", + "expect_contains": ["Paris"], + "max_tokens": 16 + }, + { + "id": "arithmetic", + "prompt": "Compute 17 multiplied by 23. Give only the number.", + "expect_contains": ["391"], + "max_tokens": 24 + }, + { + "id": "code", + "prompt": "Write a Python one-liner that returns the sum of a list named nums. Output only the code.", + "expect_contains": ["sum(", "nums"], + "max_tokens": 48 + }, + { + "id": "reasoning", + "prompt": "If a train travels 60 km in 1.5 hours, what is its average speed in km/h? Give only the number.", + "expect_contains": ["40"], + "max_tokens": 32 + }, + { + "id": "instruction", + "prompt": "List exactly three primary colors, comma-separated, lowercase.", + "expect_contains": ["red", "blue"], + "max_tokens": 24 + } +] diff --git a/scripts/MiniMax-M3/lib/run_disagg.sh b/scripts/MiniMax-M3/lib/run_disagg.sh new file mode 100644 index 00000000..a3dae5d7 --- /dev/null +++ b/scripts/MiniMax-M3/lib/run_disagg.sh @@ -0,0 +1,192 @@ +#!/bin/bash +# run_disagg.sh — MiniMax-M3 MXFP4 PD-disaggregated serving + benchmark across multiple nodes. +# Faithful port of InferenceX server_atom.sh (atomesh + mooncake) for minimaxm3-fp4-mi355x-atom-disagg. +# +# Env in: ENGINE=atom-disagg MODEL=minimaxm3-fp4 TOPO=1p1d|2p1d_dpa [ACTION=bench|serve] +# +# Layout per topology (topo.py resolves nodes): +# prefill worker(s) = kv_producer servers (port 8010) [worker 0's node also hosts the router] +# decode worker(s) = kv_consumer servers (port 8020) +# router (atomesh) on prefill-0, port 8000 <- bench drives THIS endpoint +# +# Each server runs in its own docker container via a per-node `srun --exclusive`. KV flows +# node->node over RDMA (mooncake). Bench runs from a clean container against the router +# (aggregate throughput; no x N extrapolation — this is a real disagg instance). +# All workers TP4, STP only (DECODE_MTP_SIZE=0); DPA topology adds --enable-dp-attention. +set -uo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" # scripts/MiniMax-M3 +CLUSTER="$HERE/cluster.yaml"; MODELS="$HERE/model.yaml" +ENGINE="${ENGINE:-atom-disagg}" +MODEL="${MODEL:-minimaxm3-fp4}"; TOPO="${TOPO:?set TOPO=1p1d|2p1d_dpa}"; ACTION="${ACTION:-bench}" +cfg(){ python3 "$HERE/lib/cfg.py" "$@"; } +topo(){ python3 "$HERE/lib/topo.py" "$@"; } + +RESV=$(cfg "$CLUSTER" cluster reservation) +MROOT=$(cfg "$CLUSTER" cluster models_root) +RROOT=$(cfg "$CLUSTER" cluster results_root) +RPORT=$(cfg "$CLUSTER" cluster router_port) +PPORT=$(cfg "$CLUSTER" cluster prefill_port) +DPORT=$(cfg "$CLUSTER" cluster decode_port) +HPORT=$(cfg "$CLUSTER" cluster handshake_port) +MDIR=$(cfg "$MODELS" model "$MODEL" dir) +IMG=$(cfg "$MODELS" engine "$ENGINE" image) +ROUTER=$(cfg "$MODELS" engine "$ENGINE" router) +KVCONN=$(cfg "$MODELS" engine "$ENGINE" kv_connector) +GMU=$(cfg "$MODELS" default "$MODEL" gpu_memory_util) +MNS=$(cfg "$MODELS" default "$MODEL" max_num_seqs) +KVDT=$(cfg "$MODELS" default "$MODEL" kv_cache_dtype) +BS=$(cfg "$MODELS" default "$MODEL" block_size) +MML=$(cfg "$MODELS" default "$MODEL" max_model_len) +MBT=$(cfg "$MODELS" default "$MODEL" max_num_batched_tokens) +TP=$(cfg "$MODELS" default "$MODEL" tp) +SHAPES=$(topo "$MODELS" "$TOPO" field shapes) +CONC=$(topo "$MODELS" "$TOPO" field conc) +DPATTN=$(python3 -c "import yaml;t=yaml.safe_load(open('$MODELS'))['topologies']['$TOPO'];print('1' if t['decode'].get('dp_attn') else '0')") + +ENV_ARGS=""; for kv in $(cfg "$MODELS" engine "$ENGINE" env); do ENV_ARGS+=" -e $kv"; done + +STAMP=$(python3 -c "import time;print(time.strftime('%Y%m%d_%H%M%S'))") +RUN="$RROOT/${MODEL}_${ENGINE}_${TOPO}_${STAMP}"; mkdir -p "$RUN"; chmod 777 "$RUN" +topo "$CLUSTER" "$MODELS" "$TOPO" roles > "$RUN/roles.tsv" +echo "[disagg] ENGINE=$ENGINE MODEL=$MODEL TOPO=$TOPO router=$ROUTER kv=$KVCONN image=$IMG" +echo "[disagg] roles:"; sed 's/^/ /' "$RUN/roles.tsv" + +PREFILL_NODE0=$(awk -F'\t' '$1=="prefill" && $2=="0"{print $3}' "$RUN/roles.tsv") +PREFILL_IPS=""; DECODE_IPS="" + +ip_of(){ srun --reservation="$RESV" --nodelist="$1" --nodes=1 --ntasks=1 --overlap --mem=1G \ + bash -c "ip route get 1.1.1.1 2>/dev/null | awk '/src/{print \$7}'" 2>/dev/null | grep -oE '([0-9]+\.){3}[0-9]+' | head -1; } + +# --- PRE-FLIGHT: clean every node in this topology ------------------------------------------ +NODES=$(cut -f3 "$RUN/roles.tsv" | sort -u) +echo "[disagg] cleaning nodes: $(echo $NODES | tr '\n' ' ')" +for n in $NODES; do + srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=0 \ + bash "$HERE/lib/clean_node.sh" > "$RUN/clean_${n}.log" 2>&1 && echo " [clean] $n OK" || echo " [clean] $n WARN" +done + +# --- resolve router IP (prefill node 0) ----------------------------------------------------- +ROUTER_IP=$(ip_of "$PREFILL_NODE0") +echo "[disagg] router will be on $PREFILL_NODE0 ($ROUTER_IP):$RPORT" + +# DP-attention arg for M3 (non-DSv4: DPA only, no --enable-tbo which is DSv4-specific). +dpa_arg(){ [[ "$1" == "1" ]] && echo "--enable-dp-attention" || true; } +# Optional model-length caps (M3: max_model_len 32768, max_num_batched_tokens 32768). +LEN_ARGS=""; [[ -n "$MML" ]] && LEN_ARGS="$LEN_ARGS --max-model-len $MML" +[[ -n "$MBT" ]] && LEN_ARGS="$LEN_ARGS --max-num-batched-tokens $MBT" +# KV cache dtype: skip the flag when 'auto' (M3 uses the engine default). +KV_ARG=""; [[ -n "$KVDT" && "$KVDT" != "auto" ]] && KV_ARG="--kv-cache-dtype $KVDT" + +# --- launch one server (prefill|decode) on a node in its own srun+docker -------------------- +launch_server() { + local role="$1" node="$2" port="$3" kvrole="$4" rdir="$RUN/${role}_${node}" + mkdir -p "$rdir"; chmod 777 "$rdir" + local extra="" + [[ "$role" == "prefill" ]] && extra="$(cfg "$MODELS" engine "$ENGINE" prefill_flags)" || extra="$(cfg "$MODELS" engine "$ENGINE" decode_flags)" + [[ "$DPATTN" == "1" ]] && extra="$extra $(dpa_arg 1)" + local SRV="python3 -m atom.entrypoints.openai_server --model $MROOT/$MDIR \ + --host 0.0.0.0 --server-port $port -tp $TP \ + $KV_ARG --block-size $BS --gpu-memory-utilization $GMU --max-num-seqs $MNS \ + $LEN_ARGS $extra \ + --kv-transfer-config '{\"kv_role\":\"$kvrole\",\"kv_connector\":\"$KVCONN\",\"proxy_ip\":\"$ROUTER_IP\",\"handshake_port\":$HPORT}'" + srun --reservation="$RESV" --nodelist="$node" --nodes=1 --ntasks=1 --exclusive bash -c " + docker rm -f mm3_${role} >/dev/null 2>&1 || true; sleep 3 + cid=\$(docker run -d --name mm3_${role} --entrypoint bash \ + --device=/dev/kfd --device=/dev/dri --group-add video --ipc host --shm-size 128G --network host \ + --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --privileged \ + -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $ENV_ARGS \ + -v $MROOT:$MROOT -v $HERE:/scripts -v $rdir:/out \ + $IMG -c '$SRV') + echo \"[$role] container \$cid on $node:$port\" + for i in \$(seq 1 300); do + docker logs mm3_${role} > $rdir/server.log 2>&1 || true + curl -sf http://localhost:$port/health >/dev/null 2>&1 && { echo READY > $rdir/READY; break; } + st=\$(docker inspect -f '{{.State.Status}}' mm3_${role} 2>/dev/null || echo missing) + { [ \"\$st\" = exited ] || [ \"\$st\" = missing ]; } && { echo \"[$role] container \$st\"; break; } + sleep 5 + done + docker logs mm3_${role} > $rdir/server.log 2>&1 || true + # hold the container alive for the duration (router + bench run elsewhere); wait on it + [ -f $rdir/READY ] && docker wait mm3_${role} || { echo '[FAIL] tail:'; tail -25 $rdir/server.log; exit 1; } + " > "$rdir/srun.log" 2>&1 & + echo "$!" +} + +# Build router prefill/decode arg lists (need IPs) +declare -a SRV_PIDS=() +while IFS=$'\t' read -r role widx node port; do + ipn=$(ip_of "$node") + if [[ "$role" == "prefill" ]]; then PREFILL_IPS+="$ipn "; else DECODE_IPS+="$ipn "; fi +done < "$RUN/roles.tsv" + +echo "[disagg] launching servers..." +while IFS=$'\t' read -r role widx node port; do + kvrole="kv_producer"; [[ "$role" == "decode" ]] && kvrole="kv_consumer" + pid=$(launch_server "$role" "$node" "$port" "$kvrole"); SRV_PIDS+=("$pid") + echo " [$role $widx] $node:$port ($kvrole) srun-pid $pid" +done < "$RUN/roles.tsv" + +# --- wait for all servers READY ------------------------------------------------------------- +echo "[disagg] waiting for all servers healthy (timeout ~25min)..." +ALLREADY=0 +for i in $(seq 1 300); do + r=$(find "$RUN" -name READY 2>/dev/null | wc -l); n=$(wc -l < "$RUN/roles.tsv") + [[ "$r" -ge "$n" ]] && { ALLREADY=1; echo " all $n servers READY"; break; } + sleep 5 +done +if [[ "$ALLREADY" != "1" ]]; then echo "[disagg] NOT all servers came up — see $RUN/*/server.log"; fi + +# --- start the atomesh router on prefill node 0 --------------------------------------------- +RDIR="$RUN/router"; mkdir -p "$RDIR"; chmod 777 "$RDIR" +PF_ARGS=""; for ip in $PREFILL_IPS; do PF_ARGS+=" --prefill http://$ip:$PPORT"; done +DC_ARGS=""; for ip in $DECODE_IPS; do DC_ARGS+=" --decode http://$ip:$DPORT"; done +RFLAGS=$(cfg "$MODELS" engine "$ENGINE" router_flags) +ROUTER_CMD="/usr/local/bin/atomesh launch --host 0.0.0.0 --port $RPORT --pd-disaggregation $PF_ARGS $DC_ARGS $RFLAGS" +echo "[disagg] router cmd: $ROUTER_CMD" +srun --reservation="$RESV" --nodelist="$PREFILL_NODE0" --nodes=1 --ntasks=1 --overlap --mem=8G bash -c " + docker rm -f mm3_router >/dev/null 2>&1 || true + cid=\$(docker run -d --name mm3_router --entrypoint bash --network host \ + -v $HERE:/scripts $IMG -c '$ROUTER_CMD') + for i in \$(seq 1 60); do + docker logs mm3_router > $RDIR/router.log 2>&1 || true + curl -sf http://localhost:$RPORT/v1/models >/dev/null 2>&1 && { echo READY > $RDIR/READY; break; } + sleep 5 + done + docker logs mm3_router > $RDIR/router.log 2>&1 || true +" > "$RDIR/srun.log" 2>&1 & + +for i in $(seq 1 60); do [[ -f "$RDIR/READY" ]] && { echo "[disagg] router READY on $ROUTER_IP:$RPORT"; break; }; sleep 5; done +[[ -f "$RDIR/READY" ]] || echo "[disagg] router did NOT become ready — see $RDIR/router.log" + +if [[ "$ACTION" == "serve" ]]; then + echo "[disagg] ACTION=serve — endpoint http://$ROUTER_IP:$RPORT (router). Servers held; Ctrl-C / scancel to stop." + wait; exit 0 +fi + +# --- ACTION=bench: drive the router at each (shape,conc) ------------------------------------- +BENCH_IMG=$(cfg "$MODELS" engine "$ENGINE" bench_image) +BENCH_IMG="${BENCH_IMG:-vllm/vllm-openai-rocm:nightly-9037498c22891e55b594f567fb91d9b4efbf3e99}" +echo "[disagg] benchmarking router http://$ROUTER_IP:$RPORT shapes='$SHAPES' conc='$CONC'" +for shape in $SHAPES; do + ISL=${shape%,*}; OSL=${shape#*,} + for c in $CONC; do + echo " [bench] isl=$ISL osl=$OSL conc=$c" + srun --reservation="$RESV" --nodelist="$PREFILL_NODE0" --nodes=1 --ntasks=1 --overlap --mem=8G bash -c " + docker run --rm --entrypoint bash --network host -v $HERE:/scripts -v $RUN:/out \ + $BENCH_IMG -c 'python3 /scripts/utils/bench_serving/benchmark_serving.py \ + --model $MROOT/$MDIR --backend vllm --base-url http://$ROUTER_IP:$RPORT --dataset-name random \ + --random-input-len $ISL --random-output-len $OSL --random-range-ratio 1.0 \ + --num-prompts \$(( $c * 10 )) --num-warmups \$(( $c * 2 )) --max-concurrency $c \ + --request-rate inf --ignore-eos --save-result --percentile-metrics ttft,tpot,itl,e2el \ + --result-dir /out --result-filename disagg_isl${ISL}_osl${OSL}_c${c}.json' + " > "$RUN/bench_isl${ISL}_osl${OSL}_c${c}.log" 2>&1 || echo " [bench] FAILED isl=$ISL conc=$c" + done +done + +# --- teardown ------------------------------------------------------------------------------- +echo "[disagg] teardown: stopping servers + router" +for n in $NODES; do + srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=1G \ + bash -c "docker rm -f mm3_prefill mm3_decode mm3_router >/dev/null 2>&1" 2>/dev/null || true +done +echo "[disagg] done -> $RUN" diff --git a/scripts/MiniMax-M3/lib/topo.py b/scripts/MiniMax-M3/lib/topo.py new file mode 100644 index 00000000..d767eda2 --- /dev/null +++ b/scripts/MiniMax-M3/lib/topo.py @@ -0,0 +1,52 @@ +#!/usr/bin/env python3 +"""topo.py — resolve a MiniMax-M3 disagg topology into a node-role placement. + +Usage: + topo.py roles + -> TSV lines: roleworker_idxnodeport + role in {prefill, decode}; node 0 of prefill also hosts the router. + topo.py field + -> scalar/list field of the topology (xP, yD, shapes, conc, ...) + +Topology drives node count: nodes_needed = xP*pf_nodes_per_worker + yD*dec_nodes_per_worker, +with nodes_per_worker = ceil(TP / gpus_per_node) (TP4 on 8-GPU node => 1 node/worker). +""" +import sys, yaml, math + +def load(f): return yaml.safe_load(open(f)) + +def roles(cluster_f, model_f, topo_name): + c = load(cluster_f); m = load(model_f) + gpn = c["gpus_per_node"]; nodes = c["nodes"] + t = m["topologies"][topo_name] + xP, yD = t["xP"], t["yD"] + pf_tp = t["prefill"]["tp"]; dec_tp = t["decode"]["tp"] + pf_npw = math.ceil(pf_tp / gpn); dec_npw = math.ceil(dec_tp / gpn) + need = xP * pf_npw + yD * dec_npw + if need > len(nodes): + sys.exit(f"ERROR: topo {topo_name} needs {need} nodes, cluster has {len(nodes)}") + pf_port = c["prefill_port"]; dec_port = c["decode_port"] + out = []; i = 0 + for w in range(xP): + out.append(f"prefill\t{w}\t{nodes[i]}\t{pf_port}"); i += pf_npw + for w in range(yD): + out.append(f"decode\t{w}\t{nodes[i]}\t{dec_port}"); i += dec_npw + return out + +def main(): + a = sys.argv + if a[2:3] == ["roles"] or (len(a) >= 4 and a[3] == "roles"): + # form: topo.py cluster model topo roles -> but support both orderings + pass + if len(a) >= 5 and a[4] == "roles": + print("\n".join(roles(a[1], a[2], a[3]))); return + if len(a) >= 5 and a[3] == "field": + m = load(a[1]); t = m["topologies"][a[2]]; v = t.get(a[4]) + if isinstance(v, list): print(" ".join(str(x) for x in v)) + elif isinstance(v, dict): print(" ".join(f"{k}={vv}" for k, vv in v.items())) + else: print("" if v is None else v) + return + sys.exit("usage: topo.py roles | topo.py field ") + +if __name__ == "__main__": + main() diff --git a/scripts/MiniMax-M3/model.yaml b/scripts/MiniMax-M3/model.yaml new file mode 100644 index 00000000..f579595c --- /dev/null +++ b/scripts/MiniMax-M3/model.yaml @@ -0,0 +1,64 @@ +# model.yaml — MiniMax-M3 model def + ATOM disagg image/flags/router. +# +# MiniMax-M3 (MXFP4): served PD-disaggregated on AMD Instinct MI355X (gfx950). +# All workers TP4 (no EP): the single-node M3 recipe found plain TP4 beats both TP8 and +# TP4/EP4 on tok/s/GPU for this model on gfx950, so prefill and decode both use TP4 and we +# tune the prefill:decode worker ratio (xP:yD) instead of TP. STP only (DECODE_MTP_SIZE=0). +# +# Source of truth: InferenceX `minimaxm3-fp4-mi355x-atom-disagg` +# (SemiAnalysisAI/InferenceX configs/amd-master.yaml, PR #2000). + +defaults: + gpu_memory_util: 0.8 # ATOM --gpu-memory-utilization (InferenceX MEM_FRAC_STATIC 0.8) + max_num_seqs: 256 + kv_cache_dtype: auto # M3 MXFP4 uses the engine default KV dtype (unlike DSv4 fp8) + block_size: 128 # InferenceX ATOM M3 BLOCK_SIZE=128 + max_model_len: 32768 # InferenceX ATOM M3 MAX_MODEL_LEN=32768 + max_num_batched_tokens: 32768 + tp: 4 # per prefill/decode worker (TP4 throughout) + +models: + minimaxm3-fp4: + dir: MiniMax-M3-MXFP4 # under cluster.models_root (staged to scratch); from amd/MiniMax-M3-MXFP4 + hf_repo: amd/MiniMax-M3-MXFP4 + arch: minimax_m3 + precision: fp4 + status: SCAFFOLDED # ported from InferenceX; needs a live gfx950 dry-run + +# ---- engines: disagg image + router + KV-transfer backend ---------------------------------- +engines: + atom-disagg: + # InferenceX-VALIDATED M3 MXFP4 disagg path. atomesh router + mooncake RDMA KV. + image: rocm/atom-dev:nightly_202607011530 + # bench client runs from a clean vLLM image (ATOM's own image can crash the bench client) + bench_image: vllm/vllm-openai-rocm:nightly-9037498c22891e55b594f567fb91d9b4efbf3e99 + router: atomesh # /usr/local/bin/atomesh launch --pd-disaggregation + router_version: "04b120d3040e0dd1a6915e88a3a3c1f588e6684a" + kv_connector: mooncake # --kv-transfer-config '{"kv_connector":"mooncake",...}' + entrypoint: "python3 -m atom.entrypoints.openai_server" + env: + SAFETENSORS_FAST_GPU: "1" + AITER_QUICK_REDUCE_QUANTIZATION: "INT4" # applied for non-DSv4 models in InferenceX server_atom.sh + # engine-unique server flags (shared knobs come from defaults{}) + prefill_flags: "--trust-remote-code --no-enable_prefix_caching" + decode_flags: "--trust-remote-code --no-enable_prefix_caching" + router_flags: "--policy random --backend atom --disable-health-check --disable-circuit-breaker" + status: PRIMARY # validated M3 MXFP4 atom-disagg config + +# ---- disagg topologies (from minimaxm3-fp4-mi355x-atom-disagg search-space) ----------------- +# All workers TP4 -> 1 node/worker on 8-GPU nodes. xP + yD <= 3 (3-node mi355x-disagg pool). +topologies: + 1p1d: # 2 nodes: 1 prefill (TP4) + 1 decode (TP4), STP + xP: 1 + yD: 1 + prefill: { tp: 4, ep: 1, dp_attn: false } + decode: { tp: 4, ep: 1, dp_attn: false, mtp_size: 0 } + shapes: ["8192,1024", "1024,1024"] + conc: [1, 2, 4, 8, 16, 32, 64, 128, 256] + 2p1d_dpa: # 3 nodes: 2 prefill (TP4, DP-attn) + 1 decode (TP4, DP-attn), 8k1k tail + xP: 2 + yD: 1 + prefill: { tp: 4, ep: 1, dp_attn: true } + decode: { tp: 4, ep: 1, dp_attn: true, mtp_size: 0 } + shapes: ["8192,1024"] + conc: [256, 512, 768, 1024] diff --git a/scripts/MiniMax-M3/run_atom_disagg.sh b/scripts/MiniMax-M3/run_atom_disagg.sh new file mode 100644 index 00000000..ead44e19 --- /dev/null +++ b/scripts/MiniMax-M3/run_atom_disagg.sh @@ -0,0 +1,5 @@ +#!/bin/bash +# run_atom_disagg.sh — MiniMax-M3 ATOM PD-disagg (atomesh + mooncake). Thin wrapper. +# Usage: MODEL=minimaxm3-fp4 TOPO=1p1d|2p1d_dpa [ACTION=bench|serve] ./run_atom_disagg.sh +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +ENGINE=atom-disagg MODEL="${MODEL:-minimaxm3-fp4}" exec bash "$HERE/lib/run_disagg.sh" diff --git a/scripts/MiniMax-M3/utils/bench_serving/KNOWN_LIMITATION.md b/scripts/MiniMax-M3/utils/bench_serving/KNOWN_LIMITATION.md new file mode 100644 index 00000000..342de691 --- /dev/null +++ b/scripts/MiniMax-M3/utils/bench_serving/KNOWN_LIMITATION.md @@ -0,0 +1,7 @@ +# KNOWN_LIMITATION + +When benchmarking small models (e.g., Gemma 1B or Llama 8B) and at ultra-high querys per second (QPS) and with short input and output lengths, the current InferenceX bench_serving client becomes client-bound. + +InferenceX does not currently benchmark in this regime and has no plans to. Our roadmap skews the opposite direction, larger models, longer ISL/OSL (i.e., agentic workloads), and interactive TTFT & tok/s/user scenarios rather than ultra-high QPS. + +Should the need arise, we shall fix this known limitation & migrate to a multi-process benchmark client. If your needs are benchmarking small models at high QPS at small ISL/OSL lengths, we recommend xyz benchmark instead. diff --git a/scripts/MiniMax-M3/utils/bench_serving/backend_request_func.py b/scripts/MiniMax-M3/utils/bench_serving/backend_request_func.py new file mode 100644 index 00000000..4c8820f8 --- /dev/null +++ b/scripts/MiniMax-M3/utils/bench_serving/backend_request_func.py @@ -0,0 +1,574 @@ +# SPDX-License-Identifier: Apache-2.0 + +import json +import os +import sys +import time +import traceback +from dataclasses import dataclass, field +from typing import List, Optional, Union + +import aiohttp +import huggingface_hub.constants +from tqdm.asyncio import tqdm +from transformers import (AutoTokenizer, PreTrainedTokenizer, + PreTrainedTokenizerFast) + +AIOHTTP_TIMEOUT = aiohttp.ClientTimeout(total=6 * 60 * 60) + + +@dataclass +class RequestFuncInput: + prompt: str + api_url: str + prompt_len: int + output_len: int + model: str + model_name: Optional[str] = None + best_of: int = 1 + logprobs: Optional[int] = None + extra_body: Optional[dict] = None + multi_modal_content: Optional[dict] = None + ignore_eos: bool = False + + +@dataclass +class RequestFuncOutput: + generated_text: str = "" + success: bool = False + latency: float = 0.0 + output_tokens: int = 0 + ttft: float = 0.0 # Time to first token + itl: List[float] = field( + default_factory=list) # List of inter-token latencies + tpot: float = 0.0 # avg next-token latencies + prompt_len: int = 0 + error: str = "" + + +async def async_request_tgi( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith("generate_stream") + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + params = { + "best_of": request_func_input.best_of, + "max_new_tokens": request_func_input.output_len, + "do_sample": True, + "temperature": 0.01, # TGI does not accept 0.0 temperature. + "top_p": 0.99, # TGI does not accept 1.0 top_p. + "truncate": request_func_input.prompt_len, + # TGI does not accept ignore_eos flag. + } + payload = { + "inputs": request_func_input.prompt, + "parameters": params, + } + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + chunk_bytes = chunk_bytes.decode("utf-8") + + # NOTE: Sometimes TGI returns a ping response without + # any data, we should skip it. + if chunk_bytes.startswith(":"): + continue + chunk = chunk_bytes.removeprefix("data:") + + data = json.loads(chunk) + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + most_recent_timestamp = timestamp + + output.latency = most_recent_timestamp - st + output.success = True + output.generated_text = data["generated_text"] + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_trt_llm( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith("generate_stream") + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + assert request_func_input.best_of == 1 + payload = { + "accumulate_tokens": True, + "text_input": request_func_input.prompt, + "temperature": 0.0, + "top_p": 1.0, + "max_tokens": request_func_input.output_len, + "stream": True, + } + if request_func_input.ignore_eos: + payload["min_length"] = request_func_input.output_len + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = chunk_bytes.decode("utf-8").removeprefix( + "data:") + + data = json.loads(chunk) + output.generated_text += data["text_output"] + timestamp = time.perf_counter() + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + most_recent_timestamp = timestamp + + output.latency = most_recent_timestamp - st + output.success = True + + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_deepspeed_mii( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + assert request_func_input.best_of == 1 + + payload = { + "prompt": request_func_input.prompt, + "max_tokens": request_func_input.output_len, + "temperature": 0.01, # deepspeed-mii does not accept 0.0 temp. + "top_p": 1.0, + } + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + # NOTE: DeepSpeed-MII doesn't support streaming as of Jan 28 2024, + # will use 0 as placeholder. + # See https://github.com/microsoft/DeepSpeed-MII/pull/311 + output.ttft = 0 + + st = time.perf_counter() + try: + async with session.post(url=request_func_input.api_url, + json=payload) as response: + if response.status == 200: + parsed_resp = await response.json() + output.latency = time.perf_counter() - st + output.generated_text = parsed_resp["text"][0] + output.success = True + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_openai_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + ("completions", "profile") + ), "OpenAI Completions API URL must end with 'completions' or 'profile'." + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + payload = { + "model": request_func_input.model_name \ + if request_func_input.model_name else request_func_input.model, + "prompt": request_func_input.prompt, + "temperature": 0.0, + "best_of": request_func_input.best_of, + "max_tokens": request_func_input.output_len, + "logprobs": request_func_input.logprobs, + "stream": True, + "stream_options": { + "include_usage": True, + }, + } + if request_func_input.ignore_eos: + payload["ignore_eos"] = request_func_input.ignore_eos + if request_func_input.extra_body: + payload.update(request_func_input.extra_body) + headers = { + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}" + } + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload, + headers=headers) as response: + if response.status == 200: + first_chunk_received = False + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = chunk_bytes.decode("utf-8").removeprefix( + "data: ") + if chunk != "[DONE]": + data = json.loads(chunk) + + # NOTE: Some completion API might have a last + # usage summary response without a token so we + # want to check a token was generated + if choices := data.get("choices"): + # Note that text could be empty here + # e.g. for special tokens + text = choices[0].get("text") + timestamp = time.perf_counter() + # First token + if not first_chunk_received: + first_chunk_received = True + ttft = time.perf_counter() - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + most_recent_timestamp = timestamp + generated_text += text or "" + elif usage := data.get("usage"): + output.output_tokens = usage.get( + "completion_tokens") + if first_chunk_received: + output.success = True + else: + output.success = False + output.error = ( + "Never received a valid chunk to calculate TTFT." + "This response will be marked as failed!") + output.generated_text = generated_text + output.latency = most_recent_timestamp - st + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +async def async_request_openai_chat_completions( + request_func_input: RequestFuncInput, + pbar: Optional[tqdm] = None, +) -> RequestFuncOutput: + api_url = request_func_input.api_url + assert api_url.endswith( + "chat/completions" + ), "OpenAI Chat Completions API URL must end with 'chat/completions'." + + async with aiohttp.ClientSession(trust_env=True, + timeout=AIOHTTP_TIMEOUT) as session: + content = request_func_input.prompt + if request_func_input.multi_modal_content: + content = [{"type": "text", "text": request_func_input.prompt}] + content.append(request_func_input.multi_modal_content) + payload = { + "model": request_func_input.model_name \ + if request_func_input.model_name else request_func_input.model, + "messages": [ + { + "role": "user", + "content": content + }, + ], + "temperature": 0.0, + "max_completion_tokens": request_func_input.output_len, + "stream": True, + "stream_options": { + "include_usage": True, + }, + } + if request_func_input.ignore_eos: + payload["ignore_eos"] = request_func_input.ignore_eos + if request_func_input.extra_body: + payload.update(request_func_input.extra_body) + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {os.environ.get('OPENAI_API_KEY')}", + } + + output = RequestFuncOutput() + output.prompt_len = request_func_input.prompt_len + + generated_text = "" + ttft = 0.0 + st = time.perf_counter() + most_recent_timestamp = st + try: + async with session.post(url=api_url, json=payload, + headers=headers) as response: + if response.status == 200: + async for chunk_bytes in response.content: + chunk_bytes = chunk_bytes.strip() + if not chunk_bytes: + continue + + chunk = chunk_bytes.decode("utf-8").removeprefix( + "data: ") + if chunk != "[DONE]": + timestamp = time.perf_counter() + data = json.loads(chunk) + + if choices := data.get("choices"): + content = choices[0]["delta"].get("content") + # First token + if ttft == 0.0: + ttft = timestamp - st + output.ttft = ttft + + # Decoding phase + else: + output.itl.append(timestamp - + most_recent_timestamp) + + generated_text += content or "" + elif usage := data.get("usage"): + output.output_tokens = usage.get( + "completion_tokens") + + most_recent_timestamp = timestamp + + output.generated_text = generated_text + output.success = True + output.latency = most_recent_timestamp - st + else: + output.error = response.reason or "" + output.success = False + except Exception: + output.success = False + exc_info = sys.exc_info() + output.error = "".join(traceback.format_exception(*exc_info)) + + if pbar: + pbar.update(1) + return output + + +def get_model(pretrained_model_name_or_path: str) -> str: + if os.getenv('VLLM_USE_MODELSCOPE', 'False').lower() == 'true': + from modelscope import snapshot_download + + model_path = snapshot_download( + model_id=pretrained_model_name_or_path, + local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE, + ignore_file_pattern=[".*.pt", ".*.safetensors", ".*.bin"]) + + return model_path + return pretrained_model_name_or_path + + +def _fix_tokenizer_for_sglang(tokenizer, model_path): + """Fix transformers v5 tokenizer to match sglang server-side behavior. + + Root cause: transformers v5 (>= 5.0) changed how tokenizers are loaded. + Specifically, LlamaTokenizerFast.__init__ in v5 rebuilds the pre_tokenizer + and decoder from scratch using class-specific components, discarding the + originals from tokenizer.json. For models like DeepSeek-R1 that declare + LlamaTokenizerFast but actually use a ByteLevel/Sequence tokenizer + architecture, v5 incorrectly replaces the original Sequence pre_tokenizer + with Metaspace, and the original ByteLevel decoder with Sequence. + See: https://github.com/sgl-project/sglang/blob/9238bd08a2895fa3b7ec79ea567e5c27ac951343/python/sglang/srt/utils/hf_transformers_utils.py#L836 + + The sglang server applies fixes for this in hf_transformers_utils.py + (_fix_v5_tokenizer_components and _fix_v5_add_bos_eos_token), but the + benchmark client loads the tokenizer directly via AutoTokenizer without + these fixes. This mismatch causes the client to encode text differently + from the server -- e.g. a 7000-token prompt on the client becomes ~35000 + tokens on the server, leading to ~5x TTFT inflation and false performance + regressions in benchmarks. + + This function replicates the same fixes so the benchmark client tokenizes + identically to the sglang server. It is a no-op on transformers v4. + """ + import json + from pathlib import Path + + def _resolve(filename): + """Return a filesystem path for `filename`, whether `model_path` is a + local directory or an HF Hub repo id. Returns None and logs a warning + on failure so we don't silently fail to apply the v5 fix.""" + local = Path(model_path) / filename + if local.is_file(): + return str(local) + try: + from huggingface_hub import hf_hub_download + return hf_hub_download(repo_id=model_path, filename=filename) + except Exception as e: + print( + f"v5 tokenizer fix: cannot resolve {filename} for {model_path!r} " + f"({type(e).__name__}: {e}); fix will not apply.", + flush=True, + ) + return None + + backend = getattr(tokenizer, "_tokenizer", None) + if backend is not None: + tok_file = _resolve("tokenizer.json") + if tok_file is not None: + from tokenizers import Tokenizer as RawTokenizer + raw = RawTokenizer.from_file(tok_file) + raw_pre = type(raw.pre_tokenizer).__name__ if raw.pre_tokenizer else None + loaded_pre = type(backend.pre_tokenizer).__name__ if backend.pre_tokenizer else None + if raw_pre and loaded_pre and raw_pre != loaded_pre: + print( + f"v5 tokenizer fix: {model_path} pre_tokenizer {loaded_pre} -> {raw_pre}, " + f"decoder {type(backend.decoder).__name__ if backend.decoder else None} -> " + f"{type(raw.decoder).__name__ if raw.decoder else None}", + flush=True, + ) + backend.pre_tokenizer = raw.pre_tokenizer + backend.decoder = raw.decoder + + config_file = _resolve("tokenizer_config.json") + if config_file is not None: + with open(config_file) as f: + config = json.load(f) + tok_class = config.get("tokenizer_class", "") + bos_eos_classes = { + "LlamaTokenizer", "LlamaTokenizerFast", + "CodeLlamaTokenizer", "CodeLlamaTokenizerFast", + "GemmaTokenizer", "GemmaTokenizerFast", "CohereTokenizerFast", + } + if tok_class in bos_eos_classes: + defaults = {"add_bos_token": True, "add_eos_token": False} + changed = False + for attr in ("add_bos_token", "add_eos_token"): + val = config.get(attr) + if val is None: + val = defaults.get(attr, False) + if getattr(tokenizer, attr, None) != val: + setattr(tokenizer, f"_{attr}", val) + changed = True + if changed and hasattr(tokenizer, "update_post_processor"): + tokenizer.update_post_processor() + + return tokenizer + + +def get_tokenizer( + pretrained_model_name_or_path: str, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + **kwargs, +) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: + if pretrained_model_name_or_path is not None and not os.path.exists( + pretrained_model_name_or_path): + pretrained_model_name_or_path = get_model( + pretrained_model_name_or_path) + if tokenizer_mode == "slow": + if kwargs.get("use_fast", False): + raise ValueError( + "Cannot use the fast tokenizer in slow tokenizer mode.") + kwargs["use_fast"] = False + if tokenizer_mode == "mistral": + try: + from vllm.transformers_utils.tokenizer import MistralTokenizer + except ImportError as e: + raise ImportError("MistralTokenizer requires vllm package.\n" + "Please install it with `pip install vllm` " + "to use mistral tokenizer mode.") from e + return MistralTokenizer.from_pretrained( + str(pretrained_model_name_or_path)) + else: + tokenizer = AutoTokenizer.from_pretrained( + pretrained_model_name_or_path, + trust_remote_code=trust_remote_code, + **kwargs, + ) + return _fix_tokenizer_for_sglang(tokenizer, pretrained_model_name_or_path) + + +ASYNC_REQUEST_FUNCS = { + "tgi": async_request_tgi, + "vllm": async_request_openai_completions, + "lmdeploy": async_request_openai_completions, + "deepspeed-mii": async_request_deepspeed_mii, + "openai": async_request_openai_completions, + "openai-chat": async_request_openai_chat_completions, + "tensorrt-llm": async_request_trt_llm, + "scalellm": async_request_openai_completions, + "sglang": async_request_openai_completions, +} diff --git a/scripts/MiniMax-M3/utils/bench_serving/benchmark_serving.py b/scripts/MiniMax-M3/utils/bench_serving/benchmark_serving.py new file mode 100644 index 00000000..7cc806ec --- /dev/null +++ b/scripts/MiniMax-M3/utils/bench_serving/benchmark_serving.py @@ -0,0 +1,1276 @@ +# SPDX-License-Identifier: Apache-2.0 +r"""Benchmark online serving throughput. + +On the server side, run one of the following commands: + vLLM OpenAI API server + vllm serve \ + --swap-space 16 \ + --disable-log-requests + + (TGI backend) + ./launch_tgi_server.sh + +On the client side, run: + python benchmarks/benchmark_serving.py \ + --backend \ + --model \ + --dataset-name sharegpt \ + --dataset-path \ + --request-rate \ # By default is inf + --num-prompts # By default is 1000 + + when using tgi backend, add + --endpoint /generate_stream + to the end of the command above. +""" +import argparse +import asyncio +import base64 +import contextlib +import gc +import io +import json +import os +import random +import time +import warnings +from dataclasses import dataclass +from datetime import datetime +from multiprocessing import Pool, cpu_count +from typing import Any, AsyncGenerator, Collection, Dict, List, Optional, Tuple + +import numpy as np +from backend_request_func import (ASYNC_REQUEST_FUNCS, RequestFuncInput, + RequestFuncOutput) +from tqdm.asyncio import tqdm +from transformers import PreTrainedTokenizerBase + +try: + from backend_request_func import get_tokenizer +except ImportError: + from vllm.transformers_utils.tokenizer import get_tokenizer + +try: + from vllm.utils import FlexibleArgumentParser +except ImportError: + from argparse import ArgumentParser as FlexibleArgumentParser + +from benchmark_utils import convert_to_pytorch_benchmark_format +from encoding_dsv4 import encode_messages as dsv4_encode_messages + +MILLISECONDS_TO_SECONDS_CONVERSION = 1000 + + +@dataclass +class BenchmarkMetrics: + completed: int + total_input: int + total_output: int + request_throughput: float + request_goodput: float + output_throughput: float + total_token_throughput: float + mean_ttft_ms: float + median_ttft_ms: float + std_ttft_ms: float + percentiles_ttft_ms: List[Tuple[float, float]] + mean_tpot_ms: float + median_tpot_ms: float + std_tpot_ms: float + percentiles_tpot_ms: List[Tuple[float, float]] + mean_itl_ms: float + median_itl_ms: float + std_itl_ms: float + percentiles_itl_ms: List[Tuple[float, float]] + # E2EL stands for end-to-end latency per request. + # It is the time taken on the client side from sending + # a request to receiving a complete response. + mean_e2el_ms: float + median_e2el_ms: float + std_e2el_ms: float + percentiles_e2el_ms: List[Tuple[float, float]] + + +# --- Multiprocessing helpers for sample_random_requests --- +_worker_tokenizer = None + + +def _load_tokenizer(tokenizer_id, tokenizer_mode, trust_remote_code): + """Load tokenizer for random-prompt generation. + + vLLM's get_tokenizer can raise AttributeError when transformers removes + LlamaTokenizer.all_special_tokens_extended (e.g. Qwen3.5 with newer + transformers). Prefer backend_request_func.get_tokenizer on fallback so + client tokenization stays aligned with the sglang server (#1381, #1428). + """ + try: + return get_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) + except AttributeError as exc: + if "all_special_tokens_extended" not in str(exc): + raise + try: + from backend_request_func import get_tokenizer as _backend_get_tokenizer + return _backend_get_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) + except ImportError: + from transformers import AutoTokenizer + use_fast = tokenizer_mode != "slow" + return AutoTokenizer.from_pretrained( + tokenizer_id, + trust_remote_code=trust_remote_code, + use_fast=use_fast, + ) + + +def _init_tokenizer_worker(tokenizer_id, tokenizer_mode, trust_remote_code): + """Initialize tokenizer once per worker process.""" + global _worker_tokenizer + _worker_tokenizer = _load_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) + + +def _apply_chat_template(prompt, tokenizer, dsv4): + """Render a single user message into the appropriate chat-template prompt. + + When `dsv4` is True we use the self-contained DeepSeek-V4 encoder + (encoding_dsv4.encode_messages) which emits the + ... framing the model expects. Otherwise we + fall back to the tokenizer's built-in jinja chat template. + """ + if dsv4: + return dsv4_encode_messages( + [{"role": "user", "content": prompt}], + thinking_mode="thinking", + ) + return tokenizer.apply_chat_template( + [{"role": "user", "content": prompt}], + add_generation_prompt=True, + tokenize=False, + ) + + +def _process_prompt_chunk(chunk_args): + """Generate a chunk of random prompts in a worker process.""" + (indices, prefix_token_ids, input_lens, output_lens, offsets, + prefix_len, vocab_size, use_chat_template, dsv4, seed) = chunk_args + + rng = np.random.RandomState(seed) + tokenizer = _worker_tokenizer + + results = [] + for local_idx, global_idx in enumerate(indices): + tgt_prompt_len = prefix_len + input_lens[local_idx] + prompt_token_ids = prefix_token_ids + [ + (offsets[local_idx] + global_idx + j) % vocab_size + for j in range(input_lens[local_idx]) + ] + prompt = tokenizer.decode(prompt_token_ids) + + max_retries = 10 + for _ in range(max_retries): + prompt_token_ids = tokenizer.encode(prompt, add_special_tokens=False) + if len(prompt_token_ids) < tgt_prompt_len: + num_extras = tgt_prompt_len - len(prompt_token_ids) + prompt_token_ids.extend( + rng.randint(0, vocab_size, size=num_extras).tolist()) + elif len(prompt_token_ids) > tgt_prompt_len: + prompt_token_ids = prompt_token_ids[:tgt_prompt_len] + else: + break + prompt = tokenizer.decode(prompt_token_ids) + + if use_chat_template: + prompt = _apply_chat_template(prompt, tokenizer, dsv4) + + prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False)) + mismatch = prompt_len - tgt_prompt_len + results.append((prompt, prompt_len, output_lens[local_idx], None, mismatch)) + + return results + + +def sample_random_requests( + prefix_len: int, + input_len: int, + output_len: int, + num_prompts: int, + range_ratio: float, + tokenizer: PreTrainedTokenizerBase, + use_chat_template: bool = False, + dsv4: bool = False, + tokenizer_id: Optional[str] = None, + tokenizer_mode: str = "auto", + trust_remote_code: bool = False, + num_workers: int = 0, +) -> List[Tuple[str, int, int]]: + vocab_size = tokenizer.vocab_size + prefix_token_ids = np.random.randint(0, vocab_size, size=prefix_len).tolist() + + if dsv4 and not use_chat_template: + raise ValueError("--dsv4 requires --use-chat-template to be set.") + + if use_chat_template: + chat_template_dummy = _apply_chat_template("a", tokenizer, dsv4) + tokenized_chat_template_dummy = tokenizer.encode(chat_template_dummy, add_special_tokens=False) + chat_template_len = len(tokenized_chat_template_dummy) - 1 + input_len = input_len - chat_template_len + + def sample_uniform(seq_len): + lower = int(seq_len * range_ratio) + upper = seq_len + seq_lens = np.random.randint(lower, upper+1, size=num_prompts).tolist() + return seq_lens + + input_lens = sample_uniform(input_len) + output_lens = sample_uniform(output_len) + offsets = np.random.randint(0, vocab_size, size=num_prompts) + + # Create a local RNG for retry-loop padding so that neither serial nor + # parallel path consumes global np.random draws beyond this point. + # This ensures downstream code (e.g. gamma draws for inter-arrival times) + # sees identical global RNG state regardless of num_workers. + local_rng = np.random.RandomState( + np.random.get_state()[1][:4].tolist() # derive seed from current state without advancing it + ) + + # Decide whether to use multiprocessing + if num_workers <= 0: + num_workers = min(cpu_count() or 1, 8) + use_parallel = num_workers > 1 and tokenizer_id is not None + + if use_parallel: + # Split work into chunks, one per worker + chunk_size = (num_prompts + num_workers - 1) // num_workers + chunk_args_list = [] + for w in range(num_workers): + start = w * chunk_size + end = min(start + chunk_size, num_prompts) + if start >= num_prompts: + break + chunk_args_list.append(( + list(range(start, end)), + prefix_token_ids, + input_lens[start:end], + output_lens[start:end], + offsets[start:end].tolist(), + prefix_len, + vocab_size, + use_chat_template, + dsv4, + int(local_rng.randint(0, 2**31)), + )) + + actual_workers = len(chunk_args_list) + print(f"Generating {num_prompts} prompts using {actual_workers} worker processes...") + t0 = time.perf_counter() + with Pool( + processes=actual_workers, + initializer=_init_tokenizer_worker, + initargs=(tokenizer_id, tokenizer_mode, trust_remote_code), + ) as pool: + chunk_results = pool.map(_process_prompt_chunk, chunk_args_list) + + input_requests = [] + mismatches = [] + for chunk in chunk_results: + for prompt, prompt_len, out_len, mm_content, mismatch in chunk: + input_requests.append((prompt, prompt_len, out_len, mm_content)) + mismatches.append(mismatch) + elapsed = time.perf_counter() - t0 + print(f"Prompt generation completed in {elapsed:.1f}s") + else: + # Original serial path — also uses local_rng for retry-loop padding + # to keep global RNG consumption identical to the parallel path. + if tokenizer_id is None and num_workers > 1: + print("Warning: tokenizer_id not provided, falling back to serial prompt generation.") + input_requests = [] + mismatches = [] + for i in range(num_prompts): + tgt_prompt_len = prefix_len + input_lens[i] + prompt_token_ids = prefix_token_ids + [(offsets[i] + i + j) % vocab_size for j in range(input_lens[i])] + prompt = tokenizer.decode(prompt_token_ids) + + max_retries = 10 + for _ in range(max_retries): + prompt_token_ids = tokenizer.encode(prompt, add_special_tokens=False) + if len(prompt_token_ids) < tgt_prompt_len: + num_extras = tgt_prompt_len - len(prompt_token_ids) + prompt_token_ids.extend(local_rng.randint(0, vocab_size, size=num_extras).tolist()) + elif len(prompt_token_ids) > tgt_prompt_len: + prompt_token_ids = prompt_token_ids[:tgt_prompt_len] + else: + break + prompt = tokenizer.decode(prompt_token_ids) + + if use_chat_template: + prompt = _apply_chat_template(prompt, tokenizer, dsv4) + + prompt_len = len(tokenizer.encode(prompt, add_special_tokens=False)) + mismatches.append(prompt_len - tgt_prompt_len) + input_requests.append((prompt, prompt_len, output_lens[i], None)) + + header_str = f'{"-"*19} Input/Output Length Statistics {"-"*19}' + print(header_str) + print( + f' input_lens : ' + f'min={min(r[1] for r in input_requests):<4d} ' + f'max={max(r[1] for r in input_requests):<4d} ' + f'mean={np.mean([r[1] for r in input_requests]):<7.2f} ' + f'avg_token_mismatch={np.mean(mismatches):<5.2f} ' + ) + print( + f' output_lens: ' + f'min={min(r[2] for r in input_requests):<4d} ' + f'max={max(r[2] for r in input_requests):<4d} ' + f'mean={np.mean([r[2] for r in input_requests]):<7.2f} ' + ) + print('-' * len(header_str), '\n') + + return input_requests + + +async def get_request( + input_requests: List[Tuple[str, int, int]], + request_rate: float, + burstiness: float = 1.0, +) -> AsyncGenerator[Tuple[str, int, int], None]: + """ + Asynchronously generates requests at a specified rate + with OPTIONAL burstiness. + + Args: + input_requests: + A list of input requests, each represented as a tuple. + request_rate: + The rate at which requests are generated (requests/s). + burstiness (optional): + The burstiness factor of the request generation. + Only takes effect when request_rate is not inf. + Default value is 1, which follows a Poisson process. + Otherwise, the request intervals follow a gamma distribution. + A lower burstiness value (0 < burstiness < 1) results + in more bursty requests, while a higher burstiness value + (burstiness > 1) results in a more uniform arrival of requests. + """ + input_requests = iter(input_requests) + + # Calculate scale parameter theta to maintain the desired request_rate. + assert burstiness > 0, ( + f"A positive burstiness factor is expected, but given {burstiness}.") + theta = 1.0 / (request_rate * burstiness) + + for request in input_requests: + yield request + + if request_rate == float("inf"): + # If the request rate is infinity, then we don't need to wait. + continue + + # Sample the request interval from the gamma distribution. + # If burstiness is 1, it follows exponential distribution. + interval = np.random.gamma(shape=burstiness, scale=theta) + # The next request will be sent after the interval. + await asyncio.sleep(interval) + + +def calculate_metrics( + input_requests: List[Tuple[str, int, int]], + outputs: List[RequestFuncOutput], + dur_s: float, + tokenizer: PreTrainedTokenizerBase, + selected_percentile_metrics: List[str], + selected_percentiles: List[float], + goodput_config_dict: Dict[str, float], +) -> Tuple[BenchmarkMetrics, List[int]]: + actual_output_lens: List[int] = [] + total_input = 0 + completed = 0 + good_completed = 0 + itls: List[float] = [] + tpots: List[float] = [] + all_tpots: List[float] = [] + ttfts: List[float] = [] + e2els: List[float] = [] + for i in range(len(outputs)): + if outputs[i].success: + output_len = outputs[i].output_tokens + + if output_len is None: + # We use the tokenizer to count the number of output tokens + # for some serving backends instead of looking at + # len(outputs[i].itl) since multiple output tokens may be + # bundled together + # Note : this may inflate the output token count slightly + output_len = len( + tokenizer(outputs[i].generated_text, + add_special_tokens=False).input_ids) + actual_output_lens.append(output_len) + total_input += input_requests[i][1] + tpot = 0 + if output_len > 1: + latency_minus_ttft = outputs[i].latency - outputs[i].ttft + tpot = latency_minus_ttft / (output_len - 1) + tpots.append(tpot) + # Note: if output_len <= 1, we regard tpot as 0 for goodput + all_tpots.append(tpot) + itls += outputs[i].itl + ttfts.append(outputs[i].ttft) + e2els.append(outputs[i].latency) + completed += 1 + else: + actual_output_lens.append(0) + + if goodput_config_dict: + valid_metrics = [] + slo_values = [] + + if "ttft" in goodput_config_dict: + valid_metrics.append(ttfts) + slo_values.append(goodput_config_dict["ttft"] / + MILLISECONDS_TO_SECONDS_CONVERSION) + if "tpot" in goodput_config_dict: + valid_metrics.append(all_tpots) + slo_values.append(goodput_config_dict["tpot"] / + MILLISECONDS_TO_SECONDS_CONVERSION) + if "e2el" in goodput_config_dict: + valid_metrics.append(e2els) + slo_values.append(goodput_config_dict["e2el"] / + MILLISECONDS_TO_SECONDS_CONVERSION) + + for req_metric in zip(*valid_metrics): + is_good_req = all([s >= r for s, r in zip(slo_values, req_metric)]) + if is_good_req: + good_completed += 1 + + if completed == 0: + warnings.warn( + "All requests failed. This is likely due to a misconfiguration " + "on the benchmark arguments.", + stacklevel=2) + metrics = BenchmarkMetrics( + completed=completed, + total_input=total_input, + total_output=sum(actual_output_lens), + request_throughput=completed / dur_s, + request_goodput=good_completed / dur_s, + output_throughput=sum(actual_output_lens) / dur_s, + total_token_throughput=(total_input + sum(actual_output_lens)) / dur_s, + mean_ttft_ms=np.mean(ttfts or 0) * + 1000, # ttfts is empty if streaming is not supported by backend + std_ttft_ms=np.std(ttfts or 0) * 1000, + median_ttft_ms=np.median(ttfts or 0) * 1000, + percentiles_ttft_ms=[(p, np.percentile(ttfts or 0, p) * 1000) + for p in selected_percentiles], + mean_tpot_ms=np.mean(tpots or 0) * 1000, + std_tpot_ms=np.std(tpots or 0) * 1000, + median_tpot_ms=np.median(tpots or 0) * 1000, + percentiles_tpot_ms=[(p, np.percentile(tpots or 0, p) * 1000) + for p in selected_percentiles], + mean_itl_ms=np.mean(itls or 0) * 1000, + std_itl_ms=np.std(itls or 0) * 1000, + median_itl_ms=np.median(itls or 0) * 1000, + percentiles_itl_ms=[(p, np.percentile(itls or 0, p) * 1000) + for p in selected_percentiles], + mean_e2el_ms=np.mean(e2els or 0) * 1000, + std_e2el_ms=np.std(e2els or 0) * 1000, + median_e2el_ms=np.median(e2els or 0) * 1000, + percentiles_e2el_ms=[(p, np.percentile(e2els or 0, p) * 1000) + for p in selected_percentiles], + ) + + return metrics, actual_output_lens + + +async def benchmark( + backend: str, + api_url: str, + base_url: str, + model_id: str, + model_name: str, + tokenizer: PreTrainedTokenizerBase, + input_requests: List[Tuple[str, int, int]], + logprobs: Optional[int], + best_of: int, + request_rate: float, + burstiness: float, + disable_tqdm: bool, + num_warmups: int, + profile: bool, + selected_percentile_metrics: List[str], + selected_percentiles: List[str], + ignore_eos: bool, + goodput_config_dict: Dict[str, float], + max_concurrency: Optional[int], + lora_modules: Optional[List[str]], +): + if backend in ASYNC_REQUEST_FUNCS: + request_func = ASYNC_REQUEST_FUNCS[backend] + else: + raise ValueError(f"Unknown backend: {backend}") + + print("Starting initial single prompt test run...") + test_prompt, test_prompt_len, test_output_len, test_mm_content = ( + input_requests[0]) + if backend != "openai-chat" and test_mm_content is not None: + # multi-modal benchmark is only available on OpenAI Chat backend. + raise ValueError( + "Multi-modal content is only supported on 'openai-chat' backend.") + test_input = RequestFuncInput( + model=model_id, + model_name=model_name, + prompt=test_prompt, + api_url=api_url, + prompt_len=test_prompt_len, + output_len=test_output_len, + logprobs=logprobs, + best_of=best_of, + multi_modal_content=test_mm_content, + ignore_eos=ignore_eos, + ) + + if num_warmups > 0: + print(f"Warming up with {num_warmups} requests...") + warmup_pbar = None if disable_tqdm else tqdm(total=num_warmups) + warmup_semaphore = asyncio.Semaphore(max_concurrency) if max_concurrency else None + + async def warmup_limited_req_fn(): + if warmup_semaphore is None: + return await request_func(request_func_input=test_input, pbar=warmup_pbar) + async with warmup_semaphore: + return await request_func(request_func_input=test_input, pbar=warmup_pbar) + + warmup_tasks = [] + for _ in range(num_warmups): + task = asyncio.create_task(warmup_limited_req_fn()) + warmup_tasks.append(task) + _ = await asyncio.gather(*warmup_tasks) + + if warmup_pbar is not None: + warmup_pbar.close() + print("Warmup completed.") + + if lora_modules: + # For each input request, choose a LoRA module at random. + lora_modules = iter( + [random.choice(lora_modules) for _ in range(len(input_requests))]) + + if profile: + print("Starting profiler...") + profile_input = RequestFuncInput(model=model_id, + model_name=model_name, + prompt=test_prompt, + api_url=base_url + "/start_profile", + prompt_len=test_prompt_len, + output_len=test_output_len, + extra_body={"num_steps": 1, "merge_profiles": True, "profile_by_stage": True}, + logprobs=logprobs, + best_of=best_of, + multi_modal_content=test_mm_content, + ignore_eos=ignore_eos) + profile_output = await request_func(request_func_input=profile_input) + if profile_output.success: + print("Profiler started") + + if burstiness == 1.0: + distribution = "Poisson process" + else: + distribution = "Gamma distribution" + + print(f"Traffic request rate: {request_rate}") + print(f"Burstiness factor: {burstiness} ({distribution})") + print(f"Maximum request concurrency: {max_concurrency}") + + pbar = None if disable_tqdm else tqdm(total=len(input_requests)) + + semaphore = (asyncio.Semaphore(max_concurrency) + if max_concurrency else None) + + async def limited_request_func(request_func_input, pbar): + if semaphore is None: + return await request_func(request_func_input=request_func_input, + pbar=pbar) + async with semaphore: + return await request_func(request_func_input=request_func_input, + pbar=pbar) + + print("Starting main benchmark run...") + + benchmark_start_time = time.perf_counter() + benchmark_start_time_unix = time.time() + tasks: List[asyncio.Task] = [] + async for request in get_request(input_requests, request_rate, burstiness): + prompt, prompt_len, output_len, mm_content = request + req_model_id, req_model_name = model_id, model_name + if lora_modules: + req_lora_module = next(lora_modules) + req_model_id, req_model_name = req_lora_module, req_lora_module + + request_func_input = RequestFuncInput(model=req_model_id, + model_name=req_model_name, + prompt=prompt, + api_url=api_url, + prompt_len=prompt_len, + output_len=output_len, + logprobs=logprobs, + best_of=best_of, + multi_modal_content=mm_content, + ignore_eos=ignore_eos) + tasks.append( + asyncio.create_task( + limited_request_func(request_func_input=request_func_input, + pbar=pbar))) + outputs: List[RequestFuncOutput] = await asyncio.gather(*tasks) + + if profile: + print("Stopping profiler...") + profile_input = RequestFuncInput( + model=model_id, + prompt=test_prompt, + api_url=base_url + "/stop_profile", + prompt_len=test_prompt_len, + output_len=test_output_len, + logprobs=logprobs, + best_of=best_of, + ) + profile_output = await request_func(request_func_input=profile_input) + if profile_output.success: + print("Profiler stopped") + + if pbar is not None: + pbar.close() + + benchmark_duration = time.perf_counter() - benchmark_start_time + benchmark_end_time_unix = time.time() + + metrics, actual_output_lens = calculate_metrics( + input_requests=input_requests, + outputs=outputs, + dur_s=benchmark_duration, + tokenizer=tokenizer, + selected_percentile_metrics=selected_percentile_metrics, + selected_percentiles=selected_percentiles, + goodput_config_dict=goodput_config_dict, + ) + + print("{s:{c}^{n}}".format(s=' Serving Benchmark Result ', n=50, c='=')) + print("{:<40} {:<10}".format("Successful requests:", metrics.completed)) + print("{:<40} {:<10.2f}".format("Benchmark duration (s):", + benchmark_duration)) + print("{:<40} {:<10}".format("Total input tokens:", metrics.total_input)) + print("{:<40} {:<10}".format("Total generated tokens:", + metrics.total_output)) + print("{:<40} {:<10.2f}".format("Request throughput (req/s):", + metrics.request_throughput)) + if goodput_config_dict: + print("{:<40} {:<10.2f}".format("Request goodput (req/s):", + metrics.request_goodput)) + print("{:<40} {:<10.2f}".format("Output token throughput (tok/s):", + metrics.output_throughput)) + print("{:<40} {:<10.2f}".format("Total Token throughput (tok/s):", + metrics.total_token_throughput)) + + result = { + "duration": benchmark_duration, + "benchmark_start_time_unix": benchmark_start_time_unix, + "benchmark_end_time_unix": benchmark_end_time_unix, + "completed": metrics.completed, + "total_input_tokens": metrics.total_input, + "total_output_tokens": metrics.total_output, + "request_throughput": metrics.request_throughput, + "request_goodput:": + metrics.request_goodput if goodput_config_dict else None, + "output_throughput": metrics.output_throughput, + "total_token_throughput": metrics.total_token_throughput, + "input_lens": [output.prompt_len for output in outputs], + "output_lens": actual_output_lens, + "ttfts": [output.ttft for output in outputs], + "itls": [output.itl for output in outputs], + "generated_texts": [output.generated_text for output in outputs], + "errors": [output.error for output in outputs], + } + + def process_one_metric( + # E.g., "ttft" + metric_attribute_name: str, + # E.g., "TTFT" + metric_name: str, + # E.g., "Time to First Token" + metric_header: str, + ): + # This function prints and adds statistics of the specified + # metric. + if metric_attribute_name not in selected_percentile_metrics: + return + print("{s:{c}^{n}}".format(s=metric_header, n=50, c='-')) + print("{:<40} {:<10.2f}".format( + f"Mean {metric_name} (ms):", + getattr(metrics, f"mean_{metric_attribute_name}_ms"))) + print("{:<40} {:<10.2f}".format( + f"Median {metric_name} (ms):", + getattr(metrics, f"median_{metric_attribute_name}_ms"))) + result[f"mean_{metric_attribute_name}_ms"] = getattr( + metrics, f"mean_{metric_attribute_name}_ms") + result[f"median_{metric_attribute_name}_ms"] = getattr( + metrics, f"median_{metric_attribute_name}_ms") + result[f"std_{metric_attribute_name}_ms"] = getattr( + metrics, f"std_{metric_attribute_name}_ms") + for p, value in getattr(metrics, + f"percentiles_{metric_attribute_name}_ms"): + p_word = str(int(p)) if int(p) == p else str(p) + print("{:<40} {:<10.2f}".format(f"P{p_word} {metric_name} (ms):", + value)) + result[f"p{p_word}_{metric_attribute_name}_ms"] = value + + process_one_metric("ttft", "TTFT", "Time to First Token") + process_one_metric("tpot", "TPOT", + "Time per Output Token (excl. 1st token)") + process_one_metric("itl", "ITL", "Inter-token Latency") + process_one_metric("e2el", "E2EL", "End-to-end Latency") + + print("=" * 50) + + return result + + +def check_goodput_args(args): + # Check and parse goodput arguments + goodput_config_dict = {} + VALID_NAMES = ["ttft", "tpot", "e2el"] + if args.goodput: + goodput_config_dict = parse_goodput(args.goodput) + for slo_name, slo_val in goodput_config_dict.items(): + if slo_name not in VALID_NAMES: + raise ValueError( + f"Invalid metric name found, {slo_name}: {slo_val}. " + "The service level objective name should be one of " + f"{str(VALID_NAMES)}. ") + if slo_val < 0: + raise ValueError( + f"Invalid value found, {slo_name}: {slo_val}. " + "The service level objective value should be " + "non-negative.") + return goodput_config_dict + + +def parse_goodput(slo_pairs): + goodput_config_dict = {} + try: + for slo_pair in slo_pairs: + slo_name, slo_val = slo_pair.split(":") + goodput_config_dict[slo_name] = float(slo_val) + except ValueError as err: + raise argparse.ArgumentTypeError( + "Invalid format found for service level objectives. " + "Specify service level objectives for goodput as \"KEY:VALUE\" " + "pairs, where the key is a metric name, and the value is a " + "number in milliseconds.") from err + return goodput_config_dict + + +def save_to_pytorch_benchmark_format(args: argparse.Namespace, + results: Dict[str, Any], + file_name: str) -> None: + metrics = [ + "median_ttft_ms", "mean_ttft_ms", "std_ttft_ms", "p99_ttft_ms", + "mean_tpot_ms", "median_tpot_ms", "std_tpot_ms", "p99_tpot_ms", + "median_itl_ms", "mean_itl_ms", "std_itl_ms", "p99_itl_ms" + ] + # These raw data might be useful, but they are rather big. They can be added + # later if needed + ignored_metrics = ["ttfts", "itls", "generated_texts", "errors"] + pt_records = convert_to_pytorch_benchmark_format( + args=args, + metrics={k: [results[k]] + for k in metrics}, + extra_info={ + k: results[k] + for k in results if k not in metrics and k not in ignored_metrics + }) + if pt_records: + # Don't use json suffix here as we don't want CI to pick it up + pt_file = f"{os.path.splitext(file_name)[0]}.pytorch.json" + with open(pt_file, "w") as f: + json.dump(pt_records, f) + + +def main(args: argparse.Namespace): + print(args) + random.seed(args.seed) + np.random.seed(args.seed) + + backend = args.backend + model_id = args.model + model_name = args.served_model_name + tokenizer_id = args.tokenizer if args.tokenizer is not None else args.model + tokenizer_mode = args.tokenizer_mode + + if args.base_url is not None: + api_url = f"{args.base_url}{args.endpoint}" + base_url = f"{args.base_url}" + else: + api_url = f"http://{args.host}:{args.port}{args.endpoint}" + base_url = f"http://{args.host}:{args.port}" + + tokenizer = _load_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=args.trust_remote_code, + ) + + + if args.dataset_name == "random": + input_requests = sample_random_requests( + prefix_len=args.random_prefix_len, + input_len=args.random_input_len, + output_len=args.random_output_len, + num_prompts=args.num_prompts, + range_ratio=args.random_range_ratio, + tokenizer=tokenizer, + use_chat_template=args.use_chat_template, + dsv4=args.dsv4, + tokenizer_id=tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=args.trust_remote_code, + num_workers=args.random_num_workers, + ) + + else: + raise ValueError(f"Unknown dataset: {args.dataset_name}") + + goodput_config_dict = check_goodput_args(args) + + # Avoid GC processing "static" data - reduce pause times. + gc.collect() + gc.freeze() + + benchmark_result = asyncio.run( + benchmark( + backend=backend, + api_url=api_url, + base_url=base_url, + model_id=model_id, + model_name=model_name, + tokenizer=tokenizer, + input_requests=input_requests, + logprobs=args.logprobs, + best_of=args.best_of, + request_rate=args.request_rate, + burstiness=args.burstiness, + disable_tqdm=args.disable_tqdm, + num_warmups=args.num_warmups, + profile=args.profile, + selected_percentile_metrics=args.percentile_metrics.split(","), + selected_percentiles=[ + float(p) for p in args.metric_percentiles.split(",") + ], + ignore_eos=args.ignore_eos, + goodput_config_dict=goodput_config_dict, + max_concurrency=args.max_concurrency, + lora_modules=args.lora_modules, + )) + + # Save config and results to json + if args.save_result: + result_json: Dict[str, Any] = {} + + # Setup + current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") + result_json["date"] = current_dt + result_json["backend"] = backend + result_json["model_id"] = model_id + result_json["tokenizer_id"] = tokenizer_id + result_json["best_of"] = args.best_of + result_json["num_prompts"] = args.num_prompts + + # Metadata + if args.metadata: + for item in args.metadata: + if "=" in item: + kvstring = item.split("=") + result_json[kvstring[0].strip()] = kvstring[1].strip() + else: + raise ValueError( + "Invalid metadata format. Please use KEY=VALUE format." + ) + + # Traffic + result_json["request_rate"] = (args.request_rate if args.request_rate + < float("inf") else "inf") + result_json["burstiness"] = args.burstiness + result_json["max_concurrency"] = args.max_concurrency + + # Merge with benchmark result + result_json = {**result_json, **benchmark_result} + + if not args.save_detailed: + # Remove fields with too many data points + for field in [ + "ttfts", + "itls", + "generated_texts", + "errors", + ]: + if field in result_json: + del result_json[field] + if field in benchmark_result: + del benchmark_result[field] + + # Save to file + base_model_id = model_id.split("/")[-1] + max_concurrency_str = (f"-concurrency{args.max_concurrency}" + if args.max_concurrency is not None else "") + file_name = f"{backend}-{args.request_rate}qps{max_concurrency_str}-{base_model_id}-{current_dt}.json" #noqa + if args.result_filename: + file_name = args.result_filename + if args.result_dir: + file_name = os.path.join(args.result_dir, file_name) + with open(file_name, "w", encoding='utf-8') as outfile: + json.dump(result_json, outfile) + save_to_pytorch_benchmark_format(args, result_json, file_name) + + max_failure_rate = 0.05 + completed = benchmark_result["completed"] + failure_rate = 1 - completed / args.num_prompts + if failure_rate > max_failure_rate: + raise SystemExit( + f"FAIL: request failure rate {failure_rate:.1%} exceeds " + f"{max_failure_rate:.0%} threshold " + f"({completed}/{args.num_prompts} completed)" + ) + + +if __name__ == "__main__": + parser = FlexibleArgumentParser( + description="Benchmark the online serving throughput.") + parser.add_argument( + "--backend", + type=str, + default="vllm", + choices=list(ASYNC_REQUEST_FUNCS.keys()), + ) + parser.add_argument( + "--base-url", + type=str, + default=None, + help="Server or API base url if not using http host and port.", + ) + # Use 127.0.0.1 here instead of localhost to force the use of ipv4 + parser.add_argument("--host", type=str, default="127.0.0.1") + parser.add_argument("--port", type=int, default=8000) + parser.add_argument( + "--endpoint", + type=str, + default="/v1/completions", + help="API endpoint.", + ) + parser.add_argument( + "--dataset-name", + type=str, + default="sharegpt", + choices=["random"], + help="Name of the dataset to benchmark on.", + ) + parser.add_argument("--dataset-path", + type=str, + default=None, + help="Path to the sharegpt/sonnet dataset. " + "Or the huggingface dataset ID if using HF dataset.") + parser.add_argument( + "--max-concurrency", + type=int, + default=None, + help="Maximum number of concurrent requests. This can be used " + "to help simulate an environment where a higher level component " + "is enforcing a maximum number of concurrent requests. While the " + "--request-rate argument controls the rate at which requests are " + "initiated, this argument will control how many are actually allowed " + "to execute at a time. This means that when used in combination, the " + "actual request rate may be lower than specified with --request-rate, " + "if the server is not processing requests fast enough to keep up.") + + parser.add_argument( + "--model", + type=str, + required=True, + help="Name of the model.", + ) + parser.add_argument( + "--tokenizer", + type=str, + help= + "Name or path of the tokenizer, if not using the default tokenizer.", # noqa: E501 + ) + parser.add_argument( + "--best-of", + type=int, + default=1, + help="Generates `best_of` sequences per prompt and " + "returns the best one.", + ) + parser.add_argument("--use-beam-search", action="store_true") + parser.add_argument( + "--num-prompts", + type=int, + default=1000, + help="Number of prompts to process.", + ) + parser.add_argument( + "--logprobs", + type=int, + default=None, + help=("Number of logprobs-per-token to compute & return as part of " + "the request. If unspecified, then either (1) if beam search " + "is disabled, no logprobs are computed & a single dummy " + "logprob is returned for each token; or (2) if beam search " + "is enabled 1 logprob per token is computed"), + ) + parser.add_argument( + "--request-rate", + type=float, + default=float("inf"), + help="Number of requests per second. If this is inf, " + "then all the requests are sent at time 0. " + "Otherwise, we use Poisson process or gamma distribution " + "to synthesize the request arrival times.", + ) + parser.add_argument( + "--burstiness", + type=float, + default=1.0, + help="Burstiness factor of the request generation. " + "Only take effect when request_rate is not inf. " + "Default value is 1, which follows Poisson process. " + "Otherwise, the request intervals follow a gamma distribution. " + "A lower burstiness value (0 < burstiness < 1) results in more " + "bursty requests. A higher burstiness value (burstiness > 1) " + "results in a more uniform arrival of requests.", + ) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument( + "--trust-remote-code", + action="store_true", + help="Trust remote code from huggingface", + ) + parser.add_argument( + "--disable-tqdm", + action="store_true", + help="Specify to disable tqdm progress bar.", + ) + parser.add_argument( + "--profile", + action="store_true", + help="Use Torch Profiler. The endpoint must be launched with " + "VLLM_TORCH_PROFILER_DIR to enable profiler.", + ) + parser.add_argument( + "--save-result", + action="store_true", + help="Specify to save benchmark results to a json file", + ) + parser.add_argument( + "--save-detailed", + action="store_true", + default=False, + help="When saving results, include detailed per-request data " + "(input_lens, output_lens, ttfts, itls, generated_texts, errors). " + "By default, only aggregated metrics are saved to reduce file size.", + ) + parser.add_argument( + "--metadata", + metavar="KEY=VALUE", + nargs="*", + help="Key-value pairs (e.g, --metadata version=0.3.3 tp=1) " + "for metadata of this run to be saved in the result JSON file " + "for record keeping purposes.", + ) + parser.add_argument( + "--result-dir", + type=str, + default=None, + help="Specify directory to save benchmark json results." + "If not specified, results are saved in the current directory.", + ) + parser.add_argument( + "--result-filename", + type=str, + default=None, + help="Specify the filename to save benchmark json results." + "If not specified, results will be saved in " + "{backend}-{args.request_rate}qps-{base_model_id}-{current_dt}.json" + " format.", + ) + parser.add_argument( + "--ignore-eos", + action="store_true", + help="Set ignore_eos flag when sending the benchmark request." + "Warning: ignore_eos is not supported in deepspeed_mii and tgi.") + parser.add_argument( + "--percentile-metrics", + type=str, + default="ttft,tpot,itl,e2el", + help="Comma-seperated list of selected metrics to report percentils. " + "This argument specifies the metrics to report percentiles. " + "Allowed metric names are \"ttft\", \"tpot\", \"itl\", \"e2el\". " + "Default value is \"ttft,tpot,itl,e2el\".") + parser.add_argument( + "--metric-percentiles", + type=str, + default="90,99,99.9", + help="Comma-seperated list of percentiles for selected metrics. " + "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". " + "Default value is \"90,99,99.9\". " + "Use \"--percentile-metrics\" to select metrics.", + ) + parser.add_argument( + "--goodput", + nargs="+", + required=False, + help="Specify service level objectives for goodput as \"KEY:VALUE\" " + "pairs, where the key is a metric name, and the value is in " + "milliseconds. Multiple \"KEY:VALUE\" pairs can be provided, " + "separated by spaces. Allowed request level metric names are " + "\"ttft\", \"tpot\", \"e2el\". For more context on the definition of " + "goodput, refer to DistServe paper: https://arxiv.org/pdf/2401.09670 " + "and the blog: https://hao-ai-lab.github.io/blogs/distserve") + + # group for dataset specific arguments + sonnet_group = parser.add_argument_group("sonnet dataset options") + sonnet_group.add_argument( + "--sonnet-input-len", + type=int, + default=550, + help= + "Number of input tokens per request, used only for sonnet dataset.", + ) + sonnet_group.add_argument( + "--sonnet-output-len", + type=int, + default=150, + help= + "Number of output tokens per request, used only for sonnet dataset.", + ) + sonnet_group.add_argument( + "--sonnet-prefix-len", + type=int, + default=200, + help= + "Number of prefix tokens per request, used only for sonnet dataset.", + ) + + sharegpt_group = parser.add_argument_group("sharegpt dataset options") + sharegpt_group.add_argument( + "--sharegpt-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output length " + "from the ShareGPT dataset.") + + random_group = parser.add_argument_group("random dataset options") + random_group.add_argument( + "--random-input-len", + type=int, + default=1024, + help= + "Number of input tokens per request, used only for random sampling.", + ) + random_group.add_argument( + "--random-output-len", + type=int, + default=128, + help= + "Number of output tokens per request, used only for random sampling.", + ) + random_group.add_argument( + "--random-range-ratio", + type=float, + default=1.0, + help="Range of sampled ratio of input/output length, " + "used only for random sampling.", + ) + random_group.add_argument( + "--random-prefix-len", + type=int, + default=0, + help="Number of fixed prefix tokens before random " + " context. The length range of context in a random " + " request is [random-prefix-len, " + " random-prefix-len + random-prefix-len * random-range-ratio).") + random_group.add_argument( + "--use-chat-template", + action="store_true", + help="Use chat template to format the prompt.", + ) + random_group.add_argument( + '--random-num-workers', + type=int, + default=0, + help="Number of worker processes for parallel random prompt generation. " + "Only used with --dataset-name random. " + "0 (default) = auto (min(cpu_count, 8)). 1 = serial (no multiprocessing).", + ) + + dsv4_group = parser.add_argument_group("DeepSeek-V4 chat template options") + dsv4_group.add_argument( + "--dsv4", + action="store_true", + help="Use the DeepSeek-V4 chat template (encoding_dsv4.py) instead of " + "the tokenizer's built-in jinja chat template. Requires " + "--use-chat-template to also be set. Applies to the random dataset.", + ) + + hf_group = parser.add_argument_group("hf dataset options") + hf_group.add_argument("--hf-subset", + type=str, + default=None, + help="Subset of the HF dataset.") + hf_group.add_argument("--hf-split", + type=str, + default=None, + help="Split of the HF dataset.") + hf_group.add_argument( + "--hf-output-len", + type=int, + default=None, + help="Output length for each request. Overrides the output lengths " + "from the sampled HF dataset.", + ) + + parser.add_argument( + '--tokenizer-mode', + type=str, + default="auto", + choices=['auto', 'slow', 'mistral', 'custom'], + help='The tokenizer mode.\n\n* "auto" will use the ' + 'fast tokenizer if available.\n* "slow" will ' + 'always use the slow tokenizer. \n* ' + '"mistral" will always use the `mistral_common` tokenizer. \n*' + '"custom" will use --tokenizer to select the preregistered tokenizer.') + + parser.add_argument("--served-model-name", + type=str, + default=None, + help="The model name used in the API. " + "If not specified, the model name will be the " + "same as the ``--model`` argument. ") + + parser.add_argument("--lora-modules", + nargs='+', + default=None, + help="A subset of LoRA module names passed in when " + "launching the server. For each request, the " + "script chooses a LoRA module at random.") + + parser.add_argument('--num-warmups', type=int, default=0) + + args = parser.parse_args() + main(args) diff --git a/scripts/MiniMax-M3/utils/bench_serving/benchmark_utils.py b/scripts/MiniMax-M3/utils/bench_serving/benchmark_utils.py new file mode 100644 index 00000000..dc6d31f6 --- /dev/null +++ b/scripts/MiniMax-M3/utils/bench_serving/benchmark_utils.py @@ -0,0 +1,39 @@ +# SPDX-License-Identifier: Apache-2.0 + +import argparse +import os +from typing import Any, Dict, List + + +def convert_to_pytorch_benchmark_format(args: argparse.Namespace, + metrics: Dict[str, List], + extra_info: Dict[str, Any]) -> List: + """ + Save the benchmark results in the format used by PyTorch OSS benchmark with + on metric per record + https://github.com/pytorch/pytorch/wiki/How-to-integrate-with-PyTorch-OSS-benchmark-database + """ + records = [] + if not os.environ.get("SAVE_TO_PYTORCH_BENCHMARK_FORMAT", False): + return records + + for name, benchmark_values in metrics.items(): + record = { + "benchmark": { + "name": "vLLM benchmark", + "extra_info": { + "args": vars(args), + }, + }, + "model": { + "name": args.model, + }, + "metric": { + "name": name, + "benchmark_values": benchmark_values, + "extra_info": extra_info, + }, + } + records.append(record) + + return records \ No newline at end of file diff --git a/scripts/MiniMax-M3/utils/bench_serving/encoding_dsv4.py b/scripts/MiniMax-M3/utils/bench_serving/encoding_dsv4.py new file mode 100644 index 00000000..cdb68f23 --- /dev/null +++ b/scripts/MiniMax-M3/utils/bench_serving/encoding_dsv4.py @@ -0,0 +1,744 @@ +""" +DeepSeek-V4 Encoding + +A self-contained implementation for encoding/decoding DeepSeek-V4 chat messages +with tool calling, thinking mode, and quick instruction task support. +""" + +from typing import Any, Dict, List, Union, Optional, Tuple +import copy +import json +import re + +# ============================================================ +# Special Tokens +# ============================================================ + +bos_token: str = "<｜begin▁of▁sentence｜>" +eos_token: str = "<｜end▁of▁sentence｜>" +thinking_start_token: str = "" +thinking_end_token: str = "" +dsml_token: str = "｜DSML｜" + +USER_SP_TOKEN = "<｜User｜>" +ASSISTANT_SP_TOKEN = "<｜Assistant｜>" +LATEST_REMINDER_SP_TOKEN = "<｜latest_reminder｜>" + +# Task special tokens for internal classification tasks +DS_TASK_SP_TOKENS = { + "action": "<｜action｜>", + "query": "<｜query｜>", + "authority": "<｜authority｜>", + "domain": "<｜domain｜>", + "title": "<｜title｜>", + "read_url": "<｜read_url｜>", +} +VALID_TASKS = set(DS_TASK_SP_TOKENS.keys()) + +# ============================================================ +# Templates +# ============================================================ + +system_msg_template: str = "{content}" +user_msg_template: str = "{content}" +latest_reminder_msg_template: str = "{content}" +assistant_msg_template: str = "{reasoning}{content}{tool_calls}" + eos_token +assistant_msg_wo_eos_template: str = "{reasoning}{content}{tool_calls}" +thinking_template: str = "{reasoning_content}" + +response_format_template: str = ( + "## Response Format:\n\nYou MUST strictly adhere to the following schema to reply:\n{schema}" +) +tool_call_template: str = ( + "<{dsml_token}invoke name=\"{name}\">\n{arguments}\n" +) +tool_calls_template = ( + "<{dsml_token}{tc_block_name}>\n{tool_calls}\n" +) +tool_calls_block_name: str = "tool_calls" + +tool_output_template: str = ( + "{content}" +) + +REASONING_EFFORT_MAX = ( + "Reasoning Effort: Absolute maximum with no shortcuts permitted.\n" + "You MUST be very thorough in your thinking and comprehensively decompose the problem to resolve the root cause, rigorously stress-testing your logic against all potential paths, edge cases, and adversarial scenarios.\n" + "Explicitly write out your entire deliberation process, documenting every intermediate step, considered alternative, and rejected hypothesis to ensure absolutely no assumption is left unchecked.\n\n" +) + +TOOLS_TEMPLATE = """## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<{dsml_token}tool_calls>" block like the following: + +<{dsml_token}tool_calls> +<{dsml_token}invoke name="$TOOL_NAME"> +<{dsml_token}parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<{dsml_token}invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by {thinking_start_token}), you MUST output your complete reasoning inside {thinking_start_token}...{thinking_end_token} BEFORE any tool calls or final response. + +Otherwise, output directly after {thinking_end_token} with tool calls or final response. + +### Available Tool Schemas + +{tool_schemas} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +""" + +# ============================================================ +# Utility Functions +# ============================================================ + +def to_json(value: Any) -> str: + """Serialize a value to JSON string.""" + try: + return json.dumps(value, ensure_ascii=False) + except: + return json.dumps(value, ensure_ascii=True) + + +def tools_from_openai_format(tools): + """Extract function definitions from OpenAI-format tool list.""" + return [tool["function"] for tool in tools] + + +def tool_calls_from_openai_format(tool_calls): + """Convert OpenAI-format tool calls to internal format.""" + return [ + { + "name": tool_call["function"]["name"], + "arguments": tool_call["function"]["arguments"], + } + for tool_call in tool_calls + ] + + +def tool_calls_to_openai_format(tool_calls): + """Convert internal tool calls to OpenAI format.""" + return [ + { + "type": "function", + "function": { + "name": tool_call["name"], + "arguments": tool_call["arguments"], + } + } + for tool_call in tool_calls + ] + + +def encode_arguments_to_dsml(tool_call: Dict[str, str]) -> str: + """ + Encode tool call arguments into DSML parameter format. + + Args: + tool_call: Dict with "name" and "arguments" (JSON string) keys. + + Returns: + DSML-formatted parameter string. + """ + p_dsml_template = '<{dsml_token}parameter name="{key}" string="{is_str}">{value}' + P_dsml_strs = [] + + try: + arguments = json.loads(tool_call["arguments"]) + except Exception as err: + arguments = {"arguments": tool_call["arguments"]} + + for k, v in arguments.items(): + p_dsml_str = p_dsml_template.format( + dsml_token=dsml_token, + key=k, + is_str="true" if isinstance(v, str) else "false", + value=v if isinstance(v, str) else to_json(v), + ) + P_dsml_strs.append(p_dsml_str) + + return "\n".join(P_dsml_strs) + + +def decode_dsml_to_arguments(tool_name: str, tool_args: Dict[str, Tuple[str, str]]) -> Dict[str, str]: + """ + Decode DSML parameters back to a tool call dict. + + Args: + tool_name: Name of the tool. + tool_args: Dict mapping param_name -> (value, is_string_flag). + + Returns: + Dict with "name" and "arguments" (JSON string) keys. + """ + def _decode_value(key: str, value: str, string: str): + if string == "true": + value = to_json(value) + return f"{to_json(key)}: {value}" + + tool_args_json = "{" + ", ".join([_decode_value(k, v, string=is_str) for k, (v, is_str) in tool_args.items()]) + "}" + return dict(name=tool_name, arguments=tool_args_json) + + +def render_tools(tools: List[Dict[str, Union[str, Dict[str, Any]]]]) -> str: + """ + Render tool schemas into the system prompt format. + + Args: + tools: List of tool schema dicts (each with name, description, parameters). + + Returns: + Formatted tools section string. + """ + tools_json = [to_json(t) for t in tools] + + return TOOLS_TEMPLATE.format( + tool_schemas="\n".join(tools_json), + dsml_token=dsml_token, + thinking_start_token=thinking_start_token, + thinking_end_token=thinking_end_token, + ) + + +def find_last_user_index(messages: List[Dict[str, Any]]) -> int: + """Find the index of the last user/developer message.""" + last_user_index = -1 + for idx in range(len(messages) - 1, -1, -1): + if messages[idx].get("role") in ["user", "developer"]: + last_user_index = idx + break + return last_user_index + + +# ============================================================ +# Message Rendering +# ============================================================ + +def render_message(index: int, messages: List[Dict[str, Any]], thinking_mode: str, drop_thinking: bool = True, reasoning_effort: Optional[str] = None) -> str: + """ + Render a single message at the given index into its encoded string form. + + This is the core function that converts each message in the conversation + into the DeepSeek-V4 format. + + Args: + index: Index of the message to render. + messages: Full list of messages in the conversation. + thinking_mode: Either "chat" or "thinking". + drop_thinking: Whether to drop reasoning content from earlier turns. + reasoning_effort: Optional reasoning effort level ("max", "high", or None). + + Returns: + Encoded string for this message. + """ + assert 0 <= index < len(messages) + assert thinking_mode in ["chat", "thinking"], f"Invalid thinking_mode `{thinking_mode}`" + + prompt = "" + msg = messages[index] + last_user_idx = find_last_user_index(messages) + + role = msg.get("role") + content = msg.get("content") + tools = msg.get("tools") + response_format = msg.get("response_format") + tool_calls = msg.get("tool_calls") + reasoning_content = msg.get("reasoning_content") + wo_eos = msg.get("wo_eos", False) + + if tools: + tools = tools_from_openai_format(tools) + if tool_calls: + tool_calls = tool_calls_from_openai_format(tool_calls) + + # Reasoning effort prefix (only at index 0 in thinking mode with max effort) + assert reasoning_effort in ['max', None, 'high'], f"Invalid reasoning effort: {reasoning_effort}" + if index == 0 and thinking_mode == "thinking" and reasoning_effort == 'max': + prompt += REASONING_EFFORT_MAX + + if role == "system": + prompt += system_msg_template.format(content=content or "") + if tools: + prompt += "\n\n" + render_tools(tools) + if response_format: + prompt += "\n\n" + response_format_template.format(schema=to_json(response_format)) + + elif role == "developer": + assert content, f"Invalid message for role `{role}`: {msg}" + + content_developer = USER_SP_TOKEN + content_developer += content + + if tools: + content_developer += "\n\n" + render_tools(tools) + if response_format: + content_developer += "\n\n" + response_format_template.format(schema=to_json(response_format)) + + prompt += user_msg_template.format(content=content_developer) + + elif role == "user": + prompt += USER_SP_TOKEN + + # Handle content blocks (tool results mixed with text) + content_blocks = msg.get("content_blocks") + if content_blocks: + parts = [] + for block in content_blocks: + block_type = block.get("type") + if block_type == "text": + parts.append(block.get("text", "")) + elif block_type == "tool_result": + tool_content = block.get("content", "") + if isinstance(tool_content, list): + text_parts = [] + for b in tool_content: + if b.get("type") == "text": + text_parts.append(b.get("text", "")) + else: + text_parts.append(f"[Unsupported {b.get('type')}]") + tool_content = "\n\n".join(text_parts) + parts.append(tool_output_template.format(content=tool_content)) + else: + parts.append(f"[Unsupported {block_type}]") + prompt += "\n\n".join(parts) + else: + prompt += content or "" + + elif role == "latest_reminder": + prompt += LATEST_REMINDER_SP_TOKEN + latest_reminder_msg_template.format(content=content) + + elif role == "tool": + raise NotImplementedError("deepseek_v4 merges tool messages into user; please preprocess with merge_tool_messages()") + + elif role == "assistant": + thinking_part = "" + tc_content = "" + + if tool_calls: + tc_list = [ + tool_call_template.format( + dsml_token=dsml_token, + name=tc.get("name"), + arguments=encode_arguments_to_dsml(tc) + ) + for tc in tool_calls + ] + tc_content += '\n\n' + tool_calls_template.format( + dsml_token=dsml_token, + tool_calls="\n".join(tc_list), + tc_block_name=tool_calls_block_name, + ) + + summary_content = content or "" + rc = reasoning_content or "" + + # Check if previous message has a task - if so, this is a task output (no thinking) + prev_has_task = index - 1 >= 0 and messages[index - 1].get("task") is not None + + if thinking_mode == "thinking" and not prev_has_task: + if not drop_thinking or index > last_user_idx: + thinking_part = thinking_template.format(reasoning_content=rc) + thinking_end_token + else: + thinking_part = "" + + if wo_eos: + prompt += assistant_msg_wo_eos_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) + else: + prompt += assistant_msg_template.format( + reasoning=thinking_part, + content=summary_content, + tool_calls=tc_content, + ) + else: + raise NotImplementedError(f"Unknown role: {role}") + + # Append transition tokens based on what follows + if index + 1 < len(messages) and messages[index + 1].get("role") not in ["assistant", "latest_reminder"]: + return prompt + + task = messages[index].get("task") + if task is not None: + # Task special token for internal classification tasks + assert task in VALID_TASKS, f"Invalid task: '{task}'. Valid tasks are: {list(VALID_TASKS)}" + task_sp_token = DS_TASK_SP_TOKENS[task] + + if task != "action": + # Non-action tasks: append task sp token directly after the message + prompt += task_sp_token + else: + # Action task: append Assistant + thinking token + action sp token + prompt += ASSISTANT_SP_TOKEN + prompt += thinking_end_token if thinking_mode != "thinking" else thinking_start_token + prompt += task_sp_token + + elif messages[index].get("role") in ["user", "developer"]: + # Normal generation: append Assistant + thinking token + prompt += ASSISTANT_SP_TOKEN + if not drop_thinking and thinking_mode == "thinking": + prompt += thinking_start_token + elif drop_thinking and thinking_mode == "thinking" and index >= last_user_idx: + prompt += thinking_start_token + else: + prompt += thinking_end_token + + return prompt + + +# ============================================================ +# Preprocessing +# ============================================================ + +def merge_tool_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Merge tool messages into the preceding user message using content_blocks format. + + DeepSeek-V4 does not have a standalone "tool" role; instead, tool results + are encoded as blocks within user messages. + + This function converts a standard OpenAI-format conversation (with separate + "tool" role messages) into V4 format where tool results are merged into + user messages. + + Args: + messages: List of message dicts in OpenAI format. + + Returns: + Processed message list with tool messages merged into user messages. + """ + merged: List[Dict[str, Any]] = [] + + for msg in messages: + msg = copy.deepcopy(msg) + role = msg.get("role") + + if role == "tool": + # Convert tool message to a user message with tool_result block + tool_block = { + "type": "tool_result", + "tool_use_id": msg.get("tool_call_id", ""), + "content": msg.get("content", ""), + } + # Merge into previous message if it's already a user (merged tool) + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1]: + merged[-1]["content_blocks"].append(tool_block) + else: + merged.append({ + "role": "user", + "content_blocks": [tool_block], + }) + elif role == "user": + text_block = {"type": "text", "text": msg.get("content", "")} + if merged and merged[-1].get("role") == "user" and "content_blocks" in merged[-1] and merged[-1].get("task") is None: + merged[-1]["content_blocks"].append(text_block) + else: + new_msg = { + "role": "user", + "content": msg.get("content", ""), + "content_blocks": [text_block], + } + # Preserve extra fields (task, wo_eos, mask, etc.) + for key in ("task", "wo_eos", "mask"): + if key in msg: + new_msg[key] = msg[key] + merged.append(new_msg) + else: + merged.append(msg) + + return merged + + +def sort_tool_results_by_call_order(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Sort tool_result blocks within user messages by the order of tool_calls + in the preceding assistant message. + + Args: + messages: Preprocessed message list (after merge_tool_messages). + + Returns: + Message list with sorted tool result blocks. + """ + last_tool_call_order: Dict[str, int] = {} + + for msg in messages: + role = msg.get("role") + if role == "assistant" and msg.get("tool_calls"): + last_tool_call_order = {} + for idx, tc in enumerate(msg["tool_calls"]): + tc_id = tc.get("id") or tc.get("function", {}).get("id", "") + if tc_id: + last_tool_call_order[tc_id] = idx + + elif role == "user" and msg.get("content_blocks"): + tool_blocks = [b for b in msg["content_blocks"] if b.get("type") == "tool_result"] + if len(tool_blocks) > 1 and last_tool_call_order: + sorted_blocks = sorted( + tool_blocks, + key=lambda b: last_tool_call_order.get(b.get("tool_use_id", ""), 0) + ) + sorted_idx = 0 + new_blocks = [] + for block in msg["content_blocks"]: + if block.get("type") == "tool_result": + new_blocks.append(sorted_blocks[sorted_idx]) + sorted_idx += 1 + else: + new_blocks.append(block) + msg["content_blocks"] = new_blocks + + return messages + + +# ============================================================ +# Main Encoding Function +# ============================================================ + +def encode_messages( + messages: List[Dict[str, Any]], + thinking_mode: str, + context: Optional[List[Dict[str, Any]]] = None, + drop_thinking: bool = True, + add_default_bos_token: bool = True, + reasoning_effort: Optional[str] = None, +) -> str: + """ + Encode a list of messages into the DeepSeek-V4 prompt format. + + This is the main entry point for encoding conversations. It handles: + - BOS token insertion + - Thinking mode with optional reasoning content dropping + - Tool message merging into user messages + - Multi-turn conversation context + + Args: + messages: List of message dicts to encode. + thinking_mode: Either "chat" or "thinking". + context: Optional preceding context messages (already encoded prefix). + drop_thinking: If True, drop reasoning_content from earlier assistant turns + (only keep reasoning for messages after the last user message). + add_default_bos_token: Whether to prepend BOS token at conversation start. + reasoning_effort: Optional reasoning effort level ("max", "high", or None). + + Returns: + The encoded prompt string. + """ + context = context if context else [] + + # Preprocess: merge tool messages and sort tool results + messages = merge_tool_messages(messages) + messages = sort_tool_results_by_call_order(context + messages)[len(context):] + if context: + context = merge_tool_messages(context) + context = sort_tool_results_by_call_order(context) + + full_messages = context + messages + + prompt = bos_token if add_default_bos_token and len(context) == 0 else "" + + # Resolve drop_thinking: if any message has tools defined, don't drop thinking + effective_drop_thinking = drop_thinking + if any(m.get("tools") for m in full_messages): + effective_drop_thinking = False + + if thinking_mode == "thinking" and effective_drop_thinking: + full_messages = _drop_thinking_messages(full_messages) + # After dropping, recalculate how many messages to render + # (context may have shrunk too) + num_to_render = len(full_messages) - len(_drop_thinking_messages(context)) + context_len = len(full_messages) - num_to_render + else: + num_to_render = len(messages) + context_len = len(context) + + for idx in range(num_to_render): + prompt += render_message( + idx + context_len, + full_messages, + thinking_mode=thinking_mode, + drop_thinking=effective_drop_thinking, + reasoning_effort=reasoning_effort, + ) + + return prompt + + +def _drop_thinking_messages(messages: List[Dict[str, Any]]) -> List[Dict[str, Any]]: + """ + Drop reasoning_content and non-essential messages before the last user message. + + Behavior: + - Messages with role in ["user", "system", "tool", "latest_reminder"] are always kept. + - Messages at or after the last user index are always kept. + - Assistant messages before the last user get reasoning_content removed. + - Developer messages before the last user are dropped entirely. + """ + last_user_idx = find_last_user_index(messages) + result = [] + keep_roles = {"user", "system", "tool", "latest_reminder", "direct_search_results"} + + for idx, msg in enumerate(messages): + role = msg.get("role") + if role in keep_roles or idx >= last_user_idx: + result.append(msg) + elif role == "assistant": + msg = copy.copy(msg) + msg.pop("reasoning_content", None) + result.append(msg) + # developer and other roles before last_user_idx are dropped + + return result + + +# ============================================================ +# Parsing (Decoding model output) +# ============================================================ + +def _read_until_stop(index: int, text: str, stop: List[str]) -> Tuple[int, str, Optional[str]]: + """ + Read text from index until one of the stop strings is found. + + Returns: + Tuple of (new_index, content_before_stop, matched_stop_string_or_None). + """ + min_pos = len(text) + matched_stop = None + + for s in stop: + pos = text.find(s, index) + if pos != -1 and pos < min_pos: + min_pos = pos + matched_stop = s + + if matched_stop: + content = text[index:min_pos] + return min_pos + len(matched_stop), content, matched_stop + else: + content = text[index:] + return len(text), content, None + + +def parse_tool_calls(index: int, text: str) -> Tuple[int, Optional[str], List[Dict[str, str]]]: + """ + Parse DSML tool calls from text starting at the given index. + + Args: + index: Starting position in text. + text: The full text to parse. + + Returns: + Tuple of (new_index, last_stop_token, list_of_tool_call_dicts). + Each tool call dict has "name" and "arguments" keys. + """ + tool_calls: List[Dict[str, Any]] = [] + stop_token = None + tool_calls_end_token = f"" + + while index < len(text): + index, _, stop_token = _read_until_stop(index, text, [f"<{dsml_token}invoke", tool_calls_end_token]) + if _ != ">\n": + raise ValueError(f"Tool call format error: expected '>\\n' but got '{_}'") + + if stop_token == tool_calls_end_token: + break + + if stop_token is None: + raise ValueError("Missing special token in tool calls") + + index, tool_name_content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n$', tool_name_content, flags=re.DOTALL) + if len(p_tool_name) != 1: + raise ValueError(f"Tool name format error: '{tool_name_content}'") + tool_name = p_tool_name[0] + + tool_args: Dict[str, Tuple[str, str]] = {} + while stop_token == f"<{dsml_token}parameter": + index, param_content, stop_token = _read_until_stop(index, text, [f"/{dsml_token}parameter"]) + + param_kv = re.findall(r'^ name="(.*?)" string="(true|false)">(.*?)<$', param_content, flags=re.DOTALL) + if len(param_kv) != 1: + raise ValueError(f"Parameter format error: '{param_content}'") + param_name, string, param_value = param_kv[0] + + if param_name in tool_args: + raise ValueError(f"Duplicate parameter name: '{param_name}'") + tool_args[param_name] = (param_value, string) + + index, content, stop_token = _read_until_stop(index, text, [f"<{dsml_token}parameter", f"\n": + raise ValueError(f"Parameter format error: expected '>\\n' but got '{content}'") + + tool_call = decode_dsml_to_arguments(tool_name=tool_name, tool_args=tool_args) + tool_calls.append(tool_call) + + return index, stop_token, tool_calls + + +def parse_message_from_completion_text(text: str, thinking_mode: str) -> Dict[str, Any]: + """ + Parse a model completion text into a structured assistant message. + + This function takes the raw text output from the model (a single assistant turn) + and extracts: + - reasoning_content (thinking block) + - content (summary/response) + - tool_calls (if any) + + NOTE: This function is designed to parse only correctly formatted strings and + will raise ValueError for malformed output. + + Args: + text: The raw completion text (including EOS token). + thinking_mode: Either "chat" or "thinking". + + Returns: + Dict with keys: "role", "content", "reasoning_content", "tool_calls". + tool_calls are in OpenAI format. + """ + summary_content, reasoning_content, tool_calls = "", "", [] + index, stop_token = 0, None + tool_calls_start_token = f"\n\n<{dsml_token}{tool_calls_block_name}" + + is_thinking = thinking_mode == "thinking" + is_tool_calling = False + + if is_thinking: + index, content_delta, stop_token = _read_until_stop(index, text, [thinking_end_token, tool_calls_start_token]) + reasoning_content = content_delta + assert stop_token == thinking_end_token, "Invalid thinking format: missing " + + index, content_delta, stop_token = _read_until_stop(index, text, [eos_token, tool_calls_start_token]) + summary_content = content_delta + if stop_token == tool_calls_start_token: + is_tool_calling = True + else: + assert stop_token == eos_token, "Invalid format: missing EOS token" + + if is_tool_calling: + index, stop_token, tool_calls = parse_tool_calls(index, text) + + index, tool_ends_text, stop_token = _read_until_stop(index, text, [eos_token]) + assert not tool_ends_text, "Unexpected content after tool calls" + + assert len(text) == index and stop_token in [eos_token, None], "Unexpected content at end" + + for sp_token in [bos_token, eos_token, thinking_start_token, thinking_end_token, dsml_token]: + assert sp_token not in summary_content and sp_token not in reasoning_content, \ + f"Unexpected special token '{sp_token}' in content" + + return { + "role": "assistant", + "content": summary_content, + "reasoning_content": reasoning_content, + "tool_calls": tool_calls_to_openai_format(tool_calls) + } From 32856d633a9e66cfb6848471cbdd15af1cfe597f Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Thu, 16 Jul 2026 07:22:21 +0000 Subject: [PATCH 06/12] MiniMax-M3: vendor InferenceX ATOM MXFP4 disagg code verbatim Replace the hand-adapted MiniMax-M3 harness with the actual InferenceX ATOM MXFP4 disaggregated benchmark, copied unmodified from SemiAnalysisAI/InferenceX @ a174f20 (benchmarks/multi_node/amd_utils + utils/bench_serving). Config key minimaxm3-fp4-mi355x-atom-disagg (atomesh + mooncake, all-TP4). Includes models_atom.yaml (authoritative MiniMax-M3-MXFP4 flags), server_atom.sh, env_atom.sh, bench.sh, setup_deps.sh, job.slurm/submit.sh/sync.py, the bench_serving client, and the extracted config block under configs/. Co-Authored-By: Claude --- scripts/MiniMax-M3/README.md | 111 +- scripts/MiniMax-M3/amd_utils/bench.sh | 112 ++ .../{atom_disagg => amd_utils}/env_atom.sh | 38 +- scripts/MiniMax-M3/amd_utils/job.slurm | 603 ++++++++ scripts/MiniMax-M3/amd_utils/models_atom.yaml | 75 + scripts/MiniMax-M3/amd_utils/server.sh | 23 + .../{atom_disagg => amd_utils}/server_atom.sh | 174 ++- scripts/MiniMax-M3/amd_utils/setup_deps.sh | 214 +++ scripts/MiniMax-M3/amd_utils/submit.sh | 246 ++++ scripts/MiniMax-M3/amd_utils/sync.py | 220 +++ .../MiniMax-M3/atom_disagg/accuracy_ctx.py | 73 - scripts/MiniMax-M3/atom_disagg/bench.sh | 47 - scripts/MiniMax-M3/atom_disagg/launch.sh | 132 -- scripts/MiniMax-M3/atom_disagg/setup_deps.sh | 11 - scripts/MiniMax-M3/cluster.yaml | 28 - .../minimaxm3-fp4-mi355x-atom-disagg.yaml | 74 + scripts/MiniMax-M3/docs/proxy_and_disagg.md | 52 - scripts/MiniMax-M3/lib/cfg.py | 34 - scripts/MiniMax-M3/lib/check_accuracy.py | 59 - scripts/MiniMax-M3/lib/clean_node.sh | 38 - scripts/MiniMax-M3/lib/lib_inferencex.sh | 1281 ----------------- scripts/MiniMax-M3/lib/prompts.json | 32 - scripts/MiniMax-M3/lib/run_disagg.sh | 192 --- scripts/MiniMax-M3/lib/topo.py | 52 - scripts/MiniMax-M3/model.yaml | 64 - scripts/MiniMax-M3/run_atom_disagg.sh | 5 - .../utils/bench_serving/benchmark_serving.py | 23 +- 27 files changed, 1739 insertions(+), 2274 deletions(-) create mode 100644 scripts/MiniMax-M3/amd_utils/bench.sh rename scripts/MiniMax-M3/{atom_disagg => amd_utils}/env_atom.sh (61%) create mode 100644 scripts/MiniMax-M3/amd_utils/job.slurm create mode 100644 scripts/MiniMax-M3/amd_utils/models_atom.yaml create mode 100644 scripts/MiniMax-M3/amd_utils/server.sh rename scripts/MiniMax-M3/{atom_disagg => amd_utils}/server_atom.sh (81%) create mode 100644 scripts/MiniMax-M3/amd_utils/setup_deps.sh create mode 100644 scripts/MiniMax-M3/amd_utils/submit.sh create mode 100644 scripts/MiniMax-M3/amd_utils/sync.py delete mode 100644 scripts/MiniMax-M3/atom_disagg/accuracy_ctx.py delete mode 100644 scripts/MiniMax-M3/atom_disagg/bench.sh delete mode 100644 scripts/MiniMax-M3/atom_disagg/launch.sh delete mode 100644 scripts/MiniMax-M3/atom_disagg/setup_deps.sh delete mode 100644 scripts/MiniMax-M3/cluster.yaml create mode 100644 scripts/MiniMax-M3/configs/minimaxm3-fp4-mi355x-atom-disagg.yaml delete mode 100644 scripts/MiniMax-M3/docs/proxy_and_disagg.md delete mode 100644 scripts/MiniMax-M3/lib/cfg.py delete mode 100644 scripts/MiniMax-M3/lib/check_accuracy.py delete mode 100644 scripts/MiniMax-M3/lib/clean_node.sh delete mode 100644 scripts/MiniMax-M3/lib/lib_inferencex.sh delete mode 100644 scripts/MiniMax-M3/lib/prompts.json delete mode 100644 scripts/MiniMax-M3/lib/run_disagg.sh delete mode 100644 scripts/MiniMax-M3/lib/topo.py delete mode 100644 scripts/MiniMax-M3/model.yaml delete mode 100644 scripts/MiniMax-M3/run_atom_disagg.sh diff --git a/scripts/MiniMax-M3/README.md b/scripts/MiniMax-M3/README.md index 870e21f5..cbfe76fa 100644 --- a/scripts/MiniMax-M3/README.md +++ b/scripts/MiniMax-M3/README.md @@ -1,72 +1,39 @@ -# MiniMax-M3 — disaggregated (P/D) serving/benchmark harness - -Multi-node **prefill/decode disaggregated** serving for MiniMax-M3 (MXFP4) on AMD Instinct -MI355X (gfx950), ported from InferenceX (`minimaxm3-fp4-mi355x-atom-disagg`). - -## Why disagg -MiniMax-M3 is served **disaggregated**: prefill workers (kv_producer) and decode workers -(kv_consumer) on separate nodes, KV streamed over RDMA, fronted by ONE router. Bench drives the -router → a TRUE aggregate number (no ×N extrapolation). - -**All workers are TP4, no EP.** The single-node M3 recipe found plain TP4 beats both TP8 and -TP4/EP4 on tok/s/GPU on gfx950, so prefill and decode both use TP4 and we tune the -prefill:decode worker ratio (xP:yD) instead of TP. STP only (no MTP): `DECODE_MTP_SIZE=0`. - -## Engine & router -| Engine | Router | KV transfer | Status | -|--------|--------|-------------|--------| -| **atom-disagg** | **atomesh** | **mooncake** (RDMA) | PRIMARY — InferenceX-validated M3 MXFP4 config | - -Precision: **MXFP4** (`amd/MiniMax-M3-MXFP4`). Image: `rocm/atom-dev:nightly_202607011530`. - -## Files -| File | Role | -|------|------| -| `cluster.yaml` | nodes, gpus/node, reservation, router/prefill/decode/handshake ports, paths | -| `model.yaml` | M3 def + shared `defaults{}` + atom-disagg image/router/kv + `topologies{}` | -| `run_atom_disagg.sh` | thin wrapper → `lib/run_disagg.sh` | -| `lib/run_disagg.sh` | orchestrator: clean nodes → launch prefill/decode servers → start router → bench → teardown | -| `lib/topo.py` | resolve a topology (1p1d / 2p1d_dpa) → node-role placement (prefill/decode, ports) | -| `lib/cfg.py` | yaml reader | -| `lib/clean_node.sh` | pre-flight GPU clean | -| `atom_disagg/` | vendored InferenceX ATOM server path (`server_atom.sh`, `env_atom.sh`, `bench.sh`, `setup_deps.sh`) + `launch.sh` SLURM-equivalent | -| `utils/bench_serving/` | vendored InferenceX bench client | -| `docs/proxy_and_disagg.md` | per-engine proxy/router/KV mechanics + topologies | - -## Topologies (from minimaxm3-fp4-mi355x-atom-disagg) -| name | nodes | layout | shapes | conc | -|------|-------|--------|--------|------| -| `1p1d` | 2 | 1 prefill TP4 + 1 decode TP4 (STP) | 8192/1024, 1024/1024 | 1–256 | -| `2p1d_dpa` | 3 | 2 prefill + 1 decode, TP4 + DP-attn | 8192/1024 | 256–1024 | - -All layouts keep `xP + yD <= 3` for the 3-node MI355X disagg pool (1 node/worker at TP4). - -## Usage -```bash -# ATOM disagg, 1-prefill/1-decode (2 nodes), full bench sweep through the router: -MODEL=minimaxm3-fp4 TOPO=1p1d ./run_atom_disagg.sh -MODEL=minimaxm3-fp4 TOPO=2p1d_dpa ./run_atom_disagg.sh # 3 nodes, high-conc 8k1k tail - -# Just serve (hold the endpoint up; drive it yourself): -MODEL=minimaxm3-fp4 TOPO=1p1d ACTION=serve ./run_atom_disagg.sh -``` -Bench drives `http://:8000` (OpenAI-compatible). Results land under `results/`. - -## Reference performance (InferenceX, MI355X) -From the InferenceX `minimaxm3-fp4-mi355x-atom-disagg` sweep (PR #2000), this is the -**top-throughput** MiniMax-M3 disagg config on MI355X: - -| shape | layout | peak tok/s/GPU | conc | -|-------|--------|----------------|------| -| 8k/1k | 2P1D + DPA | **~10,829** | 1024 | -| 8k/1k | 1P1D | ~7,594 | 256 | -| 1k/1k | 1P1D | ~2,703 | 256 | - -Dashboard: - -## Status / TODO -- **SCAFFOLDED** from the InferenceX config; **needs a live gfx950 dry-run** (RDMA/mooncake - handshake + atomesh router) once `amd/MiniMax-M3-MXFP4` weights are staged. -- Fill in `cluster.yaml` (reservation, `models_root`, node list) before running. -- MXFP4 quick-reduce knob `AITER_QUICK_REDUCE_QUANTIZATION=INT4` is applied by the launcher; - STP-only is enforced via `DECODE_MTP_SIZE=0`. +# MiniMax-M3 — MXFP4 disaggregated (P/D) ATOM benchmark (from InferenceX) + +Vendored **verbatim** from [SemiAnalysisAI/InferenceX](https://github.com/SemiAnalysisAI/InferenceX) +— the ATOM MXFP4 prefill/decode disaggregated benchmark for MiniMax-M3 on AMD MI355X (gfx950), +config key **`minimaxm3-fp4-mi355x-atom-disagg`** (atomesh router + mooncake RDMA KV transfer). + +This is the top-throughput MiniMax-M3 disagg config on MI355X (InferenceX): ~10,829 tok/s/GPU at +8k/1k (2P1D + DP-attn). Dashboard: + +## Contents (copied unchanged from InferenceX @ `a174f20`) +| Path | InferenceX source | +|------|-------------------| +| `amd_utils/server_atom.sh` | `benchmarks/multi_node/amd_utils/server_atom.sh` — ATOM disagg launcher (reads `models_atom.yaml`) | +| `amd_utils/models_atom.yaml` | `benchmarks/multi_node/amd_utils/models_atom.yaml` — per-model ATOM flags (incl. `MiniMax-M3-MXFP4`) | +| `amd_utils/server.sh` | dispatcher → `server_atom.sh` when `ENGINE=atom-disagg` | +| `amd_utils/env_atom.sh` | ATOM/mooncake env + RDMA device detection | +| `amd_utils/setup_deps.sh` | in-container dep setup | +| `amd_utils/bench.sh` | bench driver (calls the bench client against the atomesh router) | +| `amd_utils/job.slurm`, `submit.sh`, `sync.py` | SLURM multi-node plumbing | +| `utils/bench_serving/` | InferenceX benchmark client (`benchmark_serving.py` + backends) | +| `configs/minimaxm3-fp4-mi355x-atom-disagg.yaml` | the config entry extracted from `configs/amd-master.yaml` | + +## The MiniMax-M3-MXFP4 ATOM config (`models_atom.yaml`) +All workers TP4 (no EP; TP4 beats TP8/TP4-EP on gfx950). Key flags: +- `env: AITER_QUICK_REDUCE_QUANTIZATION=INT4 ATOM_FORCE_ATTN_TRITON=1` +- `kv_cache_flags: --kv_cache_dtype fp8` +- `online_quant_config: ptpc_fp8` (exclude lm_head/embed/vision/moe) +- `block_size 128`, `mem_frac_static 0.8`, `max_model_len 32768`, `max_num_batched_tokens 32768` + +## Topologies (`configs/minimaxm3-fp4-mi355x-atom-disagg.yaml`) +| shape | conc | layout | nodes | +|-------|------|--------|-------| +| 8192/1024 | 1–256 | 1P1D TP4 (STP, `DECODE_MTP_SIZE=0`) | 2 | +| 8192/1024 | 256–1024 | 2P1D TP4 + DP-attn | 3 | +| 1024/1024 | 1–256 | 1P1D TP4 (STP) | 2 | + +## Provenance +Source commit: `a174f20d26983d499e0eac315ee39f0c9aa48262` (SemiAnalysisAI/InferenceX). Files are +unmodified. Update by re-copying from the same paths in a newer InferenceX revision. diff --git a/scripts/MiniMax-M3/amd_utils/bench.sh b/scripts/MiniMax-M3/amd_utils/bench.sh new file mode 100644 index 00000000..c0dbb60a --- /dev/null +++ b/scripts/MiniMax-M3/amd_utils/bench.sh @@ -0,0 +1,112 @@ +#!/bin/bash +# Dual-Engine Disaggregated Benchmark Runner +# +# ENGINE=sglang (default): SGLang benchmark +# ENGINE=vllm: vLLM benchmark +# +# Produces JSON result files via benchmark_serving.py so that the CI pipeline +# can collect and process results. +# +# Usage: bash bench.sh \ +# \ +# + +ENGINE="${ENGINE:-sglang-disagg}" + +n_prefill=$1 +n_decode=$2 +prefill_gpus=$3 +decode_gpus=$4 +model_path=$5 +model_name=$6 +MODEL_PATH="${MODEL_PATH:-${model_path}/${model_name}}" +# vllm-disagg uses --served-model-name MODEL_NAME; sglang defaults to MODEL_PATH +if [[ "$ENGINE" == "vllm-disagg" ]]; then + BENCH_MODEL="${MODEL_NAME:-${MODEL_PATH}}" +else + BENCH_MODEL="${MODEL_PATH}" +fi +log_path=$7 + +chosen_isl=${8:-1024} +chosen_osl=${9:-1024} +concurrency_list=${10:-"512x1"} +if [[ "$ENGINE" == "vllm-disagg" ]]; then + chosen_req_rate=${11:-inf} +else + chosen_req_rate=${11:-1} +fi +random_range_ratio=${12:-0.8} +num_prompts_multiplier=${13:-10} + +IFS='x' read -r -a chosen_concurrencies <<< "$concurrency_list" + +ROUTER_PORT="${ROUTER_PORT:-30000}" + +export TRANSFORMERS_VERBOSITY=error +export TOKENIZERS_PARALLELISM=false + +echo "Config ${chosen_isl}; ${chosen_osl}; ${chosen_concurrencies[0]}; ${chosen_req_rate}" + +profile_folder="${log_path}/${ENGINE}_isl_${chosen_isl}_osl_${chosen_osl}" +mkdir -p "$profile_folder" + +source "$(dirname "$0")/../../benchmark_lib.sh" + +REPO_ROOT="$(cd "$(dirname "$0")/../../.." && pwd)" + +for max_concurrency in "${chosen_concurrencies[@]}"; do + + export_file="${profile_folder}/concurrency_${max_concurrency}_req_rate_${chosen_req_rate}_gpus_$((prefill_gpus+decode_gpus))_ctx_${prefill_gpus}_gen_${decode_gpus}" + + num_prompts=$(( max_concurrency * num_prompts_multiplier )) + if [[ "$num_prompts" -lt 16 ]]; then + num_prompts=16 + fi + + echo "profile_folder: $profile_folder" + echo "max_concurrency: $max_concurrency" + echo "chosen_req_rate: $chosen_req_rate" + echo "MODEL_PATH: $MODEL_PATH" + echo "ROUTER_PORT: $ROUTER_PORT" + echo "chosen_isl: $chosen_isl" + echo "chosen_osl: $chosen_osl" + echo "num_prompts: $num_prompts" + echo "export_file: $export_file" + + # Engine-specific extra flags + extra_flags="" + if [[ "$ENGINE" == "vllm-disagg" ]]; then + extra_flags="--trust-remote-code --tokenizer $MODEL_PATH" + else + if [ "$IS_MTP" = "true" ]; then + if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then + extra_flags="--dsv4" + else + extra_flags="--use-chat-template" + fi + fi + fi + + run_benchmark_serving \ + --bench-serving-dir "$REPO_ROOT" \ + --model "$BENCH_MODEL" \ + --port "$ROUTER_PORT" \ + --backend openai \ + --input-len "$chosen_isl" \ + --output-len "$chosen_osl" \ + --random-range-ratio "$random_range_ratio" \ + --num-prompts "$num_prompts" \ + --max-concurrency "$max_concurrency" \ + --result-filename "$export_file" \ + --result-dir /workspace/ \ + $extra_flags + + echo "-----------------------------------------" + + # vLLM: cooldown between rounds for idle KV block reaper + if [[ "$ENGINE" == "vllm-disagg" ]]; then + echo "[BENCH] Cooldown: waiting 10s for idle KV block reaper..." + sleep 10 + fi +done diff --git a/scripts/MiniMax-M3/atom_disagg/env_atom.sh b/scripts/MiniMax-M3/amd_utils/env_atom.sh similarity index 61% rename from scripts/MiniMax-M3/atom_disagg/env_atom.sh rename to scripts/MiniMax-M3/amd_utils/env_atom.sh index e0135e87..089594cb 100644 --- a/scripts/MiniMax-M3/atom_disagg/env_atom.sh +++ b/scripts/MiniMax-M3/amd_utils/env_atom.sh @@ -32,44 +32,24 @@ else fi export IBDEVICES -export SAFETENSORS_FAST_GPU=1 -export VLLM_LOG_LEVEL=WARNING -export ATOM_LOG_LEVEL=WARNING -export AITER_LOG_LEVEL=WARNING -export LOG_LEVEL=WARNING -export LOGLEVEL=WARNING - # ============================================================================= # ATOM/mooncake-specific environment # ============================================================================= # mooncake RDMA KV transfer library path -# The ATOM disagg image (rocm/atom-dev:nightly_202607011530) ships Python 3.12, NOT 3.10. -# The original InferenceX path (.../python3.10/...) does not exist here -> mooncake native libs -# fail to load. Resolve the actual site-packages dir dynamically so this is image-version agnostic. -_MOONCAKE_DIR=$(python3 -c "import os,mooncake;print(os.path.dirname(mooncake.__file__))" 2>/dev/null) -export LD_LIBRARY_PATH=${_MOONCAKE_DIR:-/opt/venv/lib/python3.12/site-packages/mooncake}:/opt/rocm/lib:${LD_LIBRARY_PATH:-} +export LD_LIBRARY_PATH=/opt/venv/lib/python3.10/site-packages/mooncake:/opt/rocm/lib:${LD_LIBRARY_PATH:-} - -# ATOM_HOST_IP is set per-node in server_atom.sh (= host_ip, used as handshake IP) +# faster model loading (safetensors only) +export SAFETENSORS_FAST_GPU=1 # aiter logging (WARNING to reduce noise; use DEBUG for troubleshooting) +export VLLM_LOG_LEVEL=WARNING +export ATOM_LOG_LEVEL=WARNING export AITER_LOG_LEVEL=WARNING - -# MiniMax-M3 MXFP4 needs no model-specific env here: the MXFP4 quick-reduce knob -# (AITER_QUICK_REDUCE_QUANTIZATION=INT4) is exported by the launcher via model.yaml `env`, -# and server_atom.sh applies it for all non-DSv4 models. STP-only (DECODE_MTP_SIZE=0) is -# passed as a server arg, not an env var. This block is kept for parity with the DSv4 recipe. -if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then - # ATOM MoE gather/scatter interleave optimization - export ATOM_MOE_GU_ITLV=1 - # Disable bf16->fp8 MoE bound (only for DeepSeek-V4-Pro) - export AITER_BF16_FP8_MOE_BOUND=0 -fi - -# Clear stale ATOM cache on startup (server_atom.sh handles this via rm -rf) -# No env var needed; documented here for reference. +export LOG_LEVEL=WARNING +export LOGLEVEL=WARNING set +x -echo "[INFO] ATOM env: IBDEVICES=$IBDEVICES LD_LIBRARY_PATH includes mooncake" +# ATOM_HOST_IP is set per-node in server_atom.sh (= host_ip, used as handshake IP) +echo "[INFO] ATOM env: IBDEVICES=$IBDEVICES LD_LIBRARY_PATH includes mooncake" \ No newline at end of file diff --git a/scripts/MiniMax-M3/amd_utils/job.slurm b/scripts/MiniMax-M3/amd_utils/job.slurm new file mode 100644 index 00000000..4638e5ca --- /dev/null +++ b/scripts/MiniMax-M3/amd_utils/job.slurm @@ -0,0 +1,603 @@ +#!/bin/bash +#SBATCH --job-name=disagg-bench +#SBATCH -N 3 # Overridden by submit.sh -N flag +#SBATCH -n 3 # Overridden by submit.sh -n flag +#SBATCH --ntasks-per-node=1 +#SBATCH --spread-job +#SBATCH --gres=gpu:8 +#SBATCH --time=24:00:00 +# --output and --error are set by submit.sh via BENCHMARK_LOGS_DIR + +ENGINE="${ENGINE:-sglang-disagg}" + +echo "=== Job Start Time ===" +echo "UTC Time: $(TZ=UTC date '+%Y-%m-%d %H:%M:%S %Z')" +echo "PST Time: $(TZ=America/Los_Angeles date '+%Y-%m-%d %H:%M:%S %Z')" +echo "ENGINE: $ENGINE" +echo "=======================" +echo "" + +# ============================================================================= +# Model Validation +# ============================================================================= + +# Use $(pwd) not BASH_SOURCE — sbatch copies the script to /var/spool/slurmd/ +# at runtime, but the CWD remains the submit-time directory (amd_utils/). +if [[ "$ENGINE" == "vllm-disagg" ]]; then + MODELS_YAML="$(pwd)/models_vllm.yaml" +elif [[ "$ENGINE" == "atom-disagg" ]]; then + MODELS_YAML="$(pwd)/models_atom.yaml" +else + MODELS_YAML="$(pwd)/models.yaml" +fi + +if [[ ! -f "$MODELS_YAML" ]]; then + echo "Error: models YAML not found at $MODELS_YAML" + exit 1 +fi + +if [[ -z "${DOCKER_IMAGE_NAME:-}" ]]; then + echo "Error: DOCKER_IMAGE_NAME is not set." + exit 1 +fi + +MODEL_NAME="${MODEL_NAME:-None}" +if ! grep -q "^${MODEL_NAME}:" "$MODELS_YAML"; then + echo "Error: Model '$MODEL_NAME' not found in $MODELS_YAML" + echo "Available models:" + grep -E '^[A-Za-z]' "$MODELS_YAML" | sed 's/:.*$//' | sed 's/^/ - /' + exit 1 +fi +echo "Model found: $MODEL_NAME" + +RUN_FILE="server.sh" +echo "Runfile set: $RUN_FILE" + +# DI_REPO_DIR points to the repo root. +# $(pwd) is amd_utils/ (the sbatch submit dir); go up 3 levels to reach the repo root. +export DI_REPO_DIR=$(cd "$(pwd)/../../.." && pwd) + +xP="${xP:-1}" +yD="${yD:-1}" + +# Benchmark configuration +BENCH_INPUT_LEN="${BENCH_INPUT_LEN:-1024}" +BENCH_OUTPUT_LEN="${BENCH_OUTPUT_LEN:-1024}" +BENCH_RANDOM_RANGE_RATIO="${BENCH_RANDOM_RANGE_RATIO:-1}" +BENCH_NUM_PROMPTS_MULTIPLIER="${BENCH_NUM_PROMPTS_MULTIPLIER:-10}" +BENCH_MAX_CONCURRENCY="${BENCH_MAX_CONCURRENCY:-512}" +BENCH_REQUEST_RATE="${BENCH_REQUEST_RATE:-inf}" + +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" + +# Engine-specific defaults +PREFILL_ENABLE_EP="${PREFILL_ENABLE_EP:-false}" +PREFILL_ENABLE_DP="${PREFILL_ENABLE_DP:-false}" +DECODE_ENABLE_EP="${DECODE_ENABLE_EP:-false}" +DECODE_ENABLE_DP="${DECODE_ENABLE_DP:-false}" +PREFILL_TP_SIZE="${PREFILL_TP_SIZE:-8}" +DECODE_TP_SIZE="${DECODE_TP_SIZE:-8}" +DECODE_MTP_SIZE=${DECODE_MTP_SIZE:-0} + +# Router selection: "vllm-router" (external container) or "moriio" (in-container proxy) +ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}" +ROUTER_PORT="${ROUTER_PORT:-30000}" +PROXY_PING_PORT="${PROXY_PING_PORT:-36367}" + +# ============================================================================= +# Model Path Resolution +# ============================================================================= + +# MODEL_DIR detection: prefer env var, fall back to hostname detection +if [[ -z "$MODEL_DIR" ]]; then + NODENAME=$(hostname -s) + if [[ $NODENAME == GPU* ]] || [[ $NODENAME == smci355-ccs-aus* ]]; then + MODEL_DIR="/nfsdata" + elif [[ $NODENAME == mia1* ]]; then + MODEL_DIR="/it-share/data" + else + MODEL_DIR="/nfsdata" + fi + echo "[INFO] Auto-detected MODEL_DIR=$MODEL_DIR from hostname $(hostname -s)" +fi +export MODEL_DIR + +if [[ "$ENGINE" == "vllm-disagg" ]]; then + # vLLM: Extract hf_dir from models.yaml, search multiple paths, resolve HF cache snapshots + DISK_DIR_NAME=$(awk '/^'"$MODEL_NAME"':/{found=1; next} + found && /^[^ ]/{exit} + found && /hf_dir:/{gsub(/[" ]/, "", $2); print $2; exit}' "$MODELS_YAML") + DISK_DIR_NAME="${DISK_DIR_NAME:-$MODEL_NAME}" + echo "Looking for model: $MODEL_NAME (disk dir: $DISK_DIR_NAME)" + + resolve_hf_cache_path() { + local base_path=$1 + if [[ -d "${base_path}/snapshots" ]]; then + local snapshot=$(ls -1 "${base_path}/snapshots" 2>/dev/null | head -1) + if [[ -n "$snapshot" ]]; then + echo "${base_path}/snapshots/${snapshot}" + return 0 + fi + fi + echo "$base_path" + return 1 + } + + MODEL_PATH="" + SEARCH_PATHS=( + "${MODEL_DIR}/${DISK_DIR_NAME}" + "${MODEL_DIR}/${MODEL_NAME}" + "/nfsdata/hf_hub_cache-0/${DISK_DIR_NAME}" + "/nfsdata/hf_hub_cache-0/${MODEL_NAME}" + ) + + for search_path in "${SEARCH_PATHS[@]}"; do + if [[ -d "$search_path" ]]; then + RESOLVED=$(resolve_hf_cache_path "$search_path") + MODEL_PATH="$RESOLVED" + echo "Found MODEL_PATH: $MODEL_PATH" + break + fi + done + + if [[ -z "$MODEL_PATH" ]]; then + echo "FATAL: Model '$MODEL_NAME' not found. Searched:" + for p in "${SEARCH_PATHS[@]}"; do echo " - $p"; done + exit 1 + fi + echo "Final MODEL_PATH: $MODEL_PATH" +else + # SGLang: Validate model path across all allocated nodes + echo "Looking for model: $MODEL_NAME" + echo "Checking model availability across all allocated nodes..." + + ALL_NODES=$(scontrol show hostnames "$SLURM_JOB_NODELIST") + TOTAL_NODES=$(echo "$ALL_NODES" | wc -l) + echo "Total allocated nodes: $TOTAL_NODES" + echo "Nodes: $(echo "$ALL_NODES" | tr '\n' ' ')" + + check_model_path() { + local path=$1 + local check_name=$2 + echo "Checking $check_name: $path" + srun --nodes=$SLURM_NNODES --ntasks=$SLURM_NNODES /bin/bash -c " + if [ -d '$path' ]; then + echo \"\$(hostname): Found $path\" + exit 0 + else + echo \"\$(hostname): Missing $path\" + exit 1 + fi + " + local exit_code=$? + if [ $exit_code -eq 0 ]; then + echo "$check_name available on ALL nodes" + return 0 + else + echo "$check_name NOT available on all nodes" + return 1 + fi + } + + if check_model_path "$MODEL_DIR/$MODEL_NAME" "$MODEL_DIR"; then + MODEL_PATH="$MODEL_DIR/$MODEL_NAME" + echo "Selected MODEL_PATH: $MODEL_PATH (available on all nodes)" + else + echo "FATAL ERROR: Model '$MODEL_NAME' not found on ALL allocated nodes in:" + echo " - $MODEL_DIR/$MODEL_NAME" + exit 1 + fi + echo "Final MODEL_PATH: $MODEL_PATH" +fi + +# ============================================================================= +# Node Selection +# ============================================================================= + +NUM_NODES=$((xP + yD)) +echo "NUM_NODES: $NUM_NODES (xP=$xP + yD=$yD)" + +FULL_NODELIST=$(scontrol show hostnames "$SLURM_JOB_NODELIST") +SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n $NUM_NODES) +SELECTED_NODELIST_STR=$(echo "$SELECTED_NODES" | tr '\n' ',' | sed 's/,$//') + +# Docker privilege detection — evaluated per-node since group membership varies. +# Exported as a snippet so every srun participant resolves it locally. +export DOCKER_CMD_DETECT='if docker ps &>/dev/null 2>&1; then DOCKER_CMD=docker; else DOCKER_CMD="sudo docker"; fi' + +# Update SLURM environment variables +export SLURM_NNODES=$NUM_NODES +export SLURM_NTASKS=$NUM_NODES +export SLURM_JOB_NUM_NODES=$NUM_NODES +export SLURM_NPROCS=$NUM_NODES +export SLURM_JOB_NODELIST="$SELECTED_NODELIST_STR" +export SLURM_NODELIST="$SELECTED_NODELIST_STR" +export SLURM_TASKS_PER_NODE="1(x$NUM_NODES)" +export SLURM_NTASKS_PER_NODE=1 + +echo "" +echo "Selected nodes: $SELECTED_NODELIST_STR" + +# ============================================================================= +# IP Resolution +# ============================================================================= + +USER_NAME=$(whoami) +MASTER_NODE=$(echo "$SELECTED_NODES" | head -n 1) +NODE0_ADDR=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$MASTER_NODE" bash -c 'ip route get 1.1.1.1') +NODE0_ADDR=$(echo "$NODE0_ADDR" | awk '/src/ {print $7}') + +IPS=() +for NODE in $SELECTED_NODES; do + IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'ip route get 1.1.1.1') + IP=$(echo "$IP" | awk '/src/ {print $7}') + IPS+=("$IP") +done + +echo "Node IPs: ${IPS[*]}" + +DOCKER_MOUNT_PATH="/workspace" +WS_PATH="${DOCKER_MOUNT_PATH}/benchmarks/multi_node/amd_utils" + +NNODES=$NUM_NODES + +echo "MASTER_NODE: ${MASTER_NODE}" +echo "NODE0_ADDR: ${NODE0_ADDR}" +echo "NNODES: ${NNODES}" +echo "REPO DIR: ${DI_REPO_DIR}" +echo "USER: ${USER_NAME}" + +# Reduce log spam +export TQDM_MININTERVAL=20 + +# Translate the host-resolved MODEL_PATH to the Docker mount namespace +DOCKER_MODEL_PATH="${MODEL_PATH/#$MODEL_DIR//models}" + +export DI_REPO_DIR=$DI_REPO_DIR +export WS_PATH=$WS_PATH +export NNODES=$NNODES +export NODE0_ADDR=$NODE0_ADDR +export MODEL_PATH=$MODEL_PATH +export MODEL_DIR=$MODEL_DIR +export xP=$xP +export yD=$yD +export MODEL_NAME=$MODEL_NAME +export USER_NAME=$USER_NAME +export IPADDRS="$(echo "${IPS[*]}" | sed 's/ /,/g')" +export GPUS_PER_NODE=$GPUS_PER_NODE +export BENCH_INPUT_LEN=$BENCH_INPUT_LEN +export BENCH_OUTPUT_LEN=$BENCH_OUTPUT_LEN +export BENCH_RANDOM_RANGE_RATIO=$BENCH_RANDOM_RANGE_RATIO +export BENCH_NUM_PROMPTS_MULTIPLIER=$BENCH_NUM_PROMPTS_MULTIPLIER +export BENCH_MAX_CONCURRENCY=$BENCH_MAX_CONCURRENCY +export BENCH_REQUEST_RATE=$BENCH_REQUEST_RATE +export DRY_RUN="${DRY_RUN:-0}" +export BENCHMARK_LOGS_DIR="${BENCHMARK_LOGS_DIR:-$(pwd)/benchmark_logs}" +export KEEP_CONTAINERS="${KEEP_CONTAINERS:-0}" +export ENGINE=$ENGINE + +# Eval-related env vars (threaded from submit.sh) +export RUN_EVAL="${RUN_EVAL:-false}" +export EVAL_ONLY="${EVAL_ONLY:-false}" +export EVAL_CONC="${EVAL_CONC:-}" +export FRAMEWORK="${FRAMEWORK:-}" +export PRECISION="${PRECISION:-}" +export MODEL_PREFIX="${MODEL_PREFIX:-}" +export RUNNER_TYPE="${RUNNER_TYPE:-}" +export RESULT_FILENAME="${RESULT_FILENAME:-}" +export SPEC_DECODING="${SPEC_DECODING:-}" +export IS_MULTINODE="${IS_MULTINODE:-false}" + +SANITIZED_USER=$(echo "$USER_NAME" | tr -c 'a-zA-Z0-9_.-' '_') +export DOCKER_CONT_NAME="container_${ENGINE}_${SANITIZED_USER}_${MODEL_NAME}_${SLURM_JOB_ID}" + +# vLLM external router container. +# NOTE: vllm/vllm-router only retains ~16 recent nightlies on Docker Hub; older +# dated tags are garbage-collected (manifest unknown) +VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly-20260629-e667ebb}" +ROUTER_CONT_NAME="router_vllm_${SANITIZED_USER}_${SLURM_JOB_ID}" +export RUN_FILE_FULL="$WS_PATH/${RUN_FILE}" + +SELECTED_NODELIST_SRUN=$(echo "$SELECTED_NODES" | paste -sd,) + +cleanup() { + echo "[${SLURM_JOB_ID}] termination received on $(hostname); cleaning up container + stale logs..." + # Backstop: on scancel/timeout/step-hang the foreground `exec docker run` + # client is killed before --rm can fire, so the container (and its GPU/ports/ + # :30000) leaks and the step never returns. Force-remove THIS job's container + # on every allocated node. Scoped to $DOCKER_CONT_NAME so it never touches + # other users' containers. (Ported from InferenceY 51ebfa88.) + srun --nodelist="$SELECTED_NODELIST_SRUN" \ + bash -c 'eval "$DOCKER_CMD_DETECT"; $DOCKER_CMD rm -f '"$DOCKER_CONT_NAME"' 2>/dev/null || true' 2>/dev/null || true + rm -rf ${SLURM_SUBMIT_DIR}/logs 2>/dev/null || true + echo "[${SLURM_JOB_ID}] cleanup done." +} + +trap cleanup INT TERM HUP + +# Force NFS cache refresh on all nodes +echo "Refreshing NFS caches on all nodes..." +srun --nodelist="$SELECTED_NODELIST_SRUN" bash -c ' + sync + ls -la '"$DI_REPO_DIR"'/benchmarks/multi_node/amd_utils > /dev/null 2>&1 + stat '"$DI_REPO_DIR"'/benchmarks/multi_node/amd_utils/server.sh > /dev/null 2>&1 + cat '"$DI_REPO_DIR"'/benchmarks/multi_node/amd_utils/server.sh > /dev/null 2>&1 + echo 3 | sudo tee /proc/sys/vm/drop_caches > /dev/null 2>&1 || true + echo "NFS cache refreshed on $(hostname)" +' + +# ============================================================================= +# Build engine-specific Docker environment variables +# ============================================================================= + +# Common env vars (always passed) +DOCKER_ENV_COMMON=( + -e SLURM_JOB_ID=\$SLURM_JOB_ID + -e SLURM_JOB_NODELIST=\$SLURM_JOB_NODELIST + -e NNODES=\$NNODES + -e NODE_RANK=\$SLURM_PROCID + -e NODE0_ADDR=\$NODE0_ADDR + -e MODEL_DIR=/models + -e MODEL_NAME=\$MODEL_NAME + -e GPUS_PER_NODE=\$GPUS_PER_NODE + -e xP=\$xP + -e yD=\$yD + -e IPADDRS=\$IPADDRS + -e BENCH_INPUT_LEN=\$BENCH_INPUT_LEN + -e BENCH_OUTPUT_LEN=\$BENCH_OUTPUT_LEN + -e BENCH_RANDOM_RANGE_RATIO=\$BENCH_RANDOM_RANGE_RATIO + -e BENCH_NUM_PROMPTS_MULTIPLIER=\$BENCH_NUM_PROMPTS_MULTIPLIER + -e BENCH_MAX_CONCURRENCY=\$BENCH_MAX_CONCURRENCY + -e BENCH_REQUEST_RATE=\$BENCH_REQUEST_RATE + -e TQDM_MININTERVAL=\$TQDM_MININTERVAL + -e DRY_RUN=\$DRY_RUN + -e BENCHMARK_LOGS_DIR=/benchmark_logs + -e ENGINE=\$ENGINE + -e WS_PATH=${WS_PATH} + -e RUN_EVAL=\$RUN_EVAL + -e EVAL_ONLY=\$EVAL_ONLY + -e \"EVAL_CONC=\$EVAL_CONC\" + -e FRAMEWORK=\$FRAMEWORK + -e PRECISION=\$PRECISION + -e MODEL_PREFIX=\$MODEL_PREFIX + -e RUNNER_TYPE=\$RUNNER_TYPE + -e RESULT_FILENAME=\$RESULT_FILENAME + -e SPEC_DECODING=\$SPEC_DECODING + -e PREFILL_TP_SIZE=\$PREFILL_TP_SIZE + -e PREFILL_ENABLE_EP=\$PREFILL_ENABLE_EP + -e PREFILL_ENABLE_DP=\$PREFILL_ENABLE_DP + -e DECODE_TP_SIZE=\$DECODE_TP_SIZE + -e DECODE_ENABLE_EP=\$DECODE_ENABLE_EP + -e DECODE_ENABLE_DP=\$DECODE_ENABLE_DP + -e DECODE_MTP_SIZE=\$DECODE_MTP_SIZE + -e IS_MULTINODE=\$IS_MULTINODE +) + +# Engine-specific env vars +if [[ "$ENGINE" == "vllm-disagg" ]]; then + DOCKER_ENV_ENGINE=( + -e VLLM_WS_PATH=${WS_PATH} + -e MODEL_PATH=$DOCKER_MODEL_PATH + -e UCX_TLS=tcp,self,shm,rocm_ipc,rocm_copy,cma + -e UCX_SOCKADDR_TLS_PRIORITY=tcp + -e UCX_MEMTYPE_CACHE=y + -e UCX_RNDV_SCHEME=get_zcopy + -e UCX_RNDV_THRESH=4k + -e UCX_ROCM_IPC_MIN_ZCOPY=0 + -e UCX_LOG_LEVEL=warn + -e HSA_ENABLE_SDMA=1 + -e PROXY_STREAM_IDLE_TIMEOUT=\${PROXY_STREAM_IDLE_TIMEOUT:-300} + -e PYTHONPYCACHEPREFIX=/tmp/pycache + ) +elif [[ "$ENGINE" == "atom-disagg" ]]; then + DOCKER_ENV_ENGINE=( + -e ATOM_WS_PATH=${WS_PATH} + -e PREFILL_PORT=${PREFILL_PORT:-8010} + -e DECODE_PORT=${DECODE_PORT:-8020} + -e ROUTER_PORT=${ROUTER_PORT:-30000} + -e HANDSHAKE_PORT=${HANDSHAKE_PORT:-6301} + -e MEM_FRAC_STATIC=${MEM_FRAC_STATIC:-0.85} + -e KV_CACHE_DTYPE=${KV_CACHE_DTYPE:-fp8} + -e BLOCK_SIZE=${BLOCK_SIZE:-16} + -e MAX_NUM_SEQS=${MAX_NUM_SEQS:-256} + -e MAX_MODEL_LEN=${MAX_MODEL_LEN:-} + -e MAX_NUM_BATCHED_TOKENS=${MAX_NUM_BATCHED_TOKENS:-} + -e EXTRA_SERVER_ARGS=\${EXTRA_SERVER_ARGS:-} + -e IBDEVICES=${IBDEVICES:-} + ) +else + DOCKER_ENV_ENGINE=( + -e SGLANG_WS_PATH=${WS_PATH} + ) +fi + +# Engine-specific container filter for pre-clean +CONT_FILTER="name=^container_${ENGINE}_" + +srun \ + --nodelist="$SELECTED_NODELIST_SRUN" \ + --kill-on-bad-exit=1 \ + --signal=TERM@30 \ + --unbuffered \ + bash -lc " +set -euo pipefail + +echo \"Rank \$SLURM_PROCID on \$(hostname)\" + +# Per-node docker privilege detection +eval \"\$DOCKER_CMD_DETECT\" +echo \"[docker-detect] rank \$SLURM_PROCID: DOCKER_CMD=\$DOCKER_CMD\" + +# Enable out-of-tree RDMA library mounts for atom-disagg (mooncake requires host RDMA stack) +RDMA_MOUNTS=() +if [[ "$ENGINE" == "atom-disagg" ]]; then + +# When the container base OS differs from the host (e.g. Ubuntu 24.04 image +# on a 22.04 host), the container's bundled libibverbs/libionic may be +# ABI-incompatible with the host kernel drivers. Detect the NIC type and +# bind-mount the host's out-of-tree RDMA userspace libraries into the +# container so the RDMA stack always matches the running kernel. +_detect_nic_type() { + if [[ -n \"\${MORI_NIC_TYPE:-}\" ]]; then echo \"\$MORI_NIC_TYPE\"; return; fi + local bnxt=0 mlx5=0 ionic=0 + if [[ -d /sys/class/infiniband ]]; then + for dev in /sys/class/infiniband/*; do + local name; name=\$(basename \"\$dev\") + case \"\$name\" in + bnxt_re*) ((bnxt++)) ;; mlx5*) ((mlx5++)) ;; ionic*) ((ionic++)) ;; + *) + local drv; drv=\$(basename \"\$(readlink -f \"\$dev/device/driver\" 2>/dev/null)\" 2>/dev/null || true) + case \"\$drv\" in bnxt*) ((bnxt++)) ;; mlx5*) ((mlx5++)) ;; ionic*) ((ionic++)) ;; esac ;; + esac + done + fi + if (( bnxt >= mlx5 && bnxt >= ionic && bnxt > 0 )); then echo bnxt + elif (( ionic >= mlx5 && ionic > 0 )); then echo ionic + else echo mlx5; fi +} + +_find_host_ibverbs() { + for c in /usr/lib64/libibverbs.so.1 /lib/x86_64-linux-gnu/libibverbs.so.1 /usr/lib/x86_64-linux-gnu/libibverbs.so.1.14.39.0 /usr/lib/x86_64-linux-gnu/libibverbs.so.1; do + local r; r=\$(readlink -f \"\$c\" 2>/dev/null || true) + [[ \"\$r\" == *libibverbs.so.1.14.57.0 ]] && continue + if [[ -f \"\$r\" ]]; then echo \"\$r\"; return; fi + done +} + +_NIC_TYPE=\$(_detect_nic_type) +echo \"[rdma] NIC type: \${_NIC_TYPE} on \$(hostname)\" + +if [[ \"\$_NIC_TYPE\" == \"ionic\" || \"\$_NIC_TYPE\" == \"bnxt\" ]]; then + _host_ibv=\$(_find_host_ibverbs) + if [[ -n \"\$_host_ibv\" ]]; then + RDMA_MOUNTS+=(-v \"\$_host_ibv:/lib/x86_64-linux-gnu/libibverbs.so.1\") + fi +fi + +if [[ \"\$_NIC_TYPE\" == \"ionic\" ]]; then + for _dir in /usr/local/lib /usr/lib/x86_64-linux-gnu; do + for _lib in \"\$_dir\"/libionic*.so; do + [[ -f \"\$_lib\" ]] || continue + _real=\$(readlink -f \"\$_lib\") + [[ -f \"\$_real\" ]] && RDMA_MOUNTS+=(-v \"\$_real:\$_real\") + RDMA_MOUNTS+=(-v \"\$_lib:/usr/lib/x86_64-linux-gnu/\$(basename \"\$_lib\")\") + done + done + if [[ -d /usr/lib/x86_64-linux-gnu/libibverbs ]]; then + for _lib in /usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav*.so; do + [[ -f \"\$_lib\" ]] && RDMA_MOUNTS+=(-v \"\$_lib:\$_lib\") + done + fi + [[ -d /etc/libibverbs.d ]] && RDMA_MOUNTS+=(-v /etc/libibverbs.d:/etc/libibverbs.d:ro) +elif [[ \"\$_NIC_TYPE\" == \"bnxt\" ]]; then + for _lib in /usr/local/lib/libbnxt_re-rdmav*.so; do + [[ -f \"\$_lib\" ]] && RDMA_MOUNTS+=(-v \"\$_lib:/usr/lib/x86_64-linux-gnu/libibverbs/\$(basename \"\$_lib\")\") + done + for _lib in /usr/local/lib/libbnxt_re.so; do + [[ -f \"\$_lib\" ]] && RDMA_MOUNTS+=(-v \"\$_lib:/usr/lib/x86_64-linux-gnu/\$(basename \"\$_lib\")\") + done + [[ -d /etc/libibverbs.d ]] && RDMA_MOUNTS+=(-v /etc/libibverbs.d:/etc/libibverbs.d:ro) +fi + +if [[ \${#RDMA_MOUNTS[@]} -gt 0 ]]; then + echo \"[rdma] bind-mounts: \${RDMA_MOUNTS[*]}\" +else + echo \"[rdma] no out-of-tree RDMA mounts needed\" +fi +fi # end: if ENGINE == atom-disagg + +# Pre-clean (idempotent) +\$DOCKER_CMD ps -aq --filter \"$CONT_FILTER\" | xargs -r \$DOCKER_CMD rm -f || true +\$DOCKER_CMD ps -aq | xargs -r \$DOCKER_CMD stop || true + +# Start vLLM external router container on node 0 +if [[ \"$ENGINE\" == \"vllm-disagg\" && \"$ROUTER_TYPE\" == \"vllm-router\" && \"\$SLURM_PROCID\" == \"0\" ]]; then + \$DOCKER_CMD rm -f \"$ROUTER_CONT_NAME\" 2>/dev/null || true + \$DOCKER_CMD run -d \ + --name \"$ROUTER_CONT_NAME\" \ + --network host \ + --ulimit nofile=1048576:1048576 \ + -v /tmp:/run_logs \ + \"$VLLM_ROUTER_IMAGE\" \ + bash -lc \"mkdir -p /run_logs/slurm_job-${SLURM_JOB_ID} && exec vllm-router \ + --vllm-pd-disaggregation \ + --kv-connector moriio \ + --vllm-discovery-address 0.0.0.0:${PROXY_PING_PORT} \ + --port ${ROUTER_PORT} \ + --host 0.0.0.0 \ + --policy consistent_hash \ + --prefill-policy consistent_hash \ + --decode-policy consistent_hash \ + --log-level info 2>&1 | tee /run_logs/slurm_job-${SLURM_JOB_ID}/vllm_router_\$(hostname).log \" +fi + +# Skip exec on vllm-disagg rank 0 so we can stop the router after the main +# container exits. Without this, decode nodes block forever waiting for the +# router port to close (the router is a separate container). +MAYBE_EXEC=exec +if [[ \"$ENGINE\" == \"vllm-disagg\" && \"$ROUTER_TYPE\" == \"vllm-router\" && \"\$SLURM_PROCID\" == \"0\" ]]; then + MAYBE_EXEC= + set +e +fi + +\$MAYBE_EXEC \$DOCKER_CMD run \ + --init \ + --stop-timeout 10 \ + --device /dev/dri \ + --device /dev/kfd \ + --device /dev/infiniband \ + --device=/dev/infiniband/rdma_cm \ + --device=/dev/infiniband/uverbs0 \ + --device=/dev/infiniband/uverbs1 \ + --device=/dev/infiniband/uverbs2 \ + --device=/dev/infiniband/uverbs3 \ + --device=/dev/infiniband/uverbs4 \ + --device=/dev/infiniband/uverbs5 \ + --device=/dev/infiniband/uverbs6 \ + --device=/dev/infiniband/uverbs7 \ + --ulimit memlock=-1 \ + --ulimit stack=67108864 \ + --ulimit nofile=1048576:1048576 \ + --network host \ + --ipc host \ + --group-add video \ + --cap-add SYS_PTRACE \ + --security-opt seccomp=unconfined \ + --privileged \ + -v /sys:/sys \ + $(command -v nicctl >/dev/null 2>&1 && echo "-v $(which nicctl):/usr/sbin/nicctl") \ + -v ${MODEL_DIR}:/models \ + -v \$HOME/.ssh:/root/.ssh \ + --shm-size 128G \ + -v /tmp:/run_logs \ + -v ${BENCHMARK_LOGS_DIR}:/benchmark_logs \ + -v ${DI_REPO_DIR}:${DOCKER_MOUNT_PATH} \ + ${EXTRA_DOCKER_MOUNTS:-} \ + \${RDMA_MOUNTS[@]+"\${RDMA_MOUNTS[@]}"} \ + ${DOCKER_ENV_COMMON[*]} \ + ${DOCKER_ENV_ENGINE[*]} \ + --name \"$DOCKER_CONT_NAME\" \ + --entrypoint \"\" \ + \"$DOCKER_IMAGE_NAME\" bash -lc ' + set -o pipefail + mkdir -p /run_logs/slurm_job-'\"\$SLURM_JOB_ID\"' + '"$RUN_FILE_FULL"' 2>&1 | tee /run_logs/slurm_job-'\"\$SLURM_JOB_ID\"'/server_\$(hostname).log + ' + +# Only reached when exec was skipped (vllm-disagg rank 0) +DOCKER_EXIT_CODE=\$? +echo \"[rank 0] Main container exited (rc=\$DOCKER_EXIT_CODE). Stopping vllm-router...\" +\$DOCKER_CMD rm -f \"$ROUTER_CONT_NAME\" 2>/dev/null || true +exit \$DOCKER_EXIT_CODE +" + +if [[ "${KEEP_CONTAINERS}" != "1" ]]; then + srun --nodelist="$SELECTED_NODELIST_SRUN" bash -c 'eval "$DOCKER_CMD_DETECT"; $DOCKER_CMD rm -f '"$DOCKER_CONT_NAME"' 2>/dev/null || true' + + # Clean up vLLM external router container on node 0 + if [[ "$ENGINE" == "vllm-disagg" && "$ROUTER_TYPE" == "vllm-router" ]]; then + srun --nodes=1 --ntasks=1 --nodelist="$MASTER_NODE" bash -c ' + eval "$DOCKER_CMD_DETECT"; $DOCKER_CMD rm -f '"$ROUTER_CONT_NAME"' 2>/dev/null || true + ' + fi +fi diff --git a/scripts/MiniMax-M3/amd_utils/models_atom.yaml b/scripts/MiniMax-M3/amd_utils/models_atom.yaml new file mode 100644 index 00000000..b2490a42 --- /dev/null +++ b/scripts/MiniMax-M3/amd_utils/models_atom.yaml @@ -0,0 +1,75 @@ +# Model-specific ATOM server configurations for disaggregated inference. +# +# Each top-level key is a MODEL_NAME value (must match the directory name under MODEL_DIR). +# +# To add a new model: add a new top-level entry following the same schema. +# No script changes are required. +# +# Schema: +# : +# env: str # Space-separated KEY=VALUE pairs exported unconditionally +# tp_dp_flags: str # Shared TP+DPA flags (fallback when prefill/decode-specific keys are absent) +# prefill_tp_dp_flags: str # TP+DPA flags for prefill only (overrides tp_dp_flags) +# decode_tp_dp_flags: str # TP+DPA flags for decode only (overrides tp_dp_flags) +# tp_dp_env: str # Space-separated KEY=VALUE pairs exported only in TP+DPA mode +# ep_dp_flags: str # Shared EP+DPA flags (fallback when prefill/decode-specific keys are absent) +# prefill_ep_dp_flags: str # EP+DPA flags for prefill only (overrides ep_dp_flags) +# decode_ep_dp_flags: str # EP+DPA flags for decode only (overrides ep_dp_flags) +# ep_dp_env: str # Space-separated KEY=VALUE pairs exported only in EP+DPA mode +# mtp_flags: str # Flags passed to SPEC_ARGS before $DECODE_MTP_SIZE (e.g. "--method mtp --num-speculative-tokens") +# kv_cache_flags: str # Full --kv_cache_dtype flag string (e.g. "--kv_cache_dtype fp8", or "" for none) +# online_quant_config: str # JSON string passed to --online_quant_config (used when DPA is disabled) +# online_quant_dpa_config: str # JSON string passed to --online_quant_config when DPA is enabled (falls back to online_quant_config) +# block_size: str # --block-size value (overrides server_atom.sh default of 16) +# mem_frac_static: str # --gpu-memory-utilization value (overrides default of 0.85) +# max_model_len: str # --max-model-len value (overrides default of unset) +# max_num_seqs: str # --max-num-seqs value (overrides default of 256) +# max_num_batched_tokens: str # --max-num-batched-tokens value (overrides default of unset) + +DeepSeek-V4-Pro: + env: "ATOM_MOE_GU_ITLV=1 AITER_BF16_FP8_MOE_BOUND=0" + kv_cache_flags: "--kv_cache_dtype fp8" + tp_dp_flags: "--enable-dp-attention --enable-tbo" + prefill_tp_dp_flags: "--enable-dp-attention --enable-tbo" + decode_tp_dp_flags: "--enable-dp-attention --enable-tbo" + tp_dp_env: "GPU_MAX_HW_QUEUES=5 ATOM_CPU_AFFINITY=1" + ep_dp_flags: "--enable-expert-parallel --enable-dp-attention" + prefill_ep_dp_flags: "--enable-expert-parallel --enable-dp-attention" + decode_ep_dp_flags: "--enable-expert-parallel --enable-dp-attention" + mtp_flags: "--method mtp --num-speculative-tokens" + +MiniMax-M3-MXFP4: + env: "AITER_QUICK_REDUCE_QUANTIZATION=INT4 ATOM_FORCE_ATTN_TRITON=1" + kv_cache_flags: "--kv_cache_dtype fp8" + tp_dp_flags: "--enable-dp-attention" + prefill_tp_dp_flags: "--enable-dp-attention --enable-tbo prefill" + decode_tp_dp_flags: "--enable-dp-attention" + ep_dp_flags: "--enable-expert-parallel --enable-dp-attention" + prefill_ep_dp_flags: "--enable-expert-parallel --enable-dp-attention" + decode_ep_dp_flags: "--enable-expert-parallel --enable-dp-attention" + mtp_flags: "--method eagle3 --draft-model Inferact/MiniMax-M3-EAGLE3 --num-speculative-tokens" + online_quant_config: '{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","vision_tower","multi_modal_projector","patch_merge_mlp","*block_sparse_moe"]}' + online_quant_dpa_config: '{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","vision_tower","multi_modal_projector","patch_merge_mlp","*block_sparse_moe"]}' + block_size: "128" + mem_frac_static: "0.8" + max_model_len: "32768" + max_num_seqs: "256" + max_num_batched_tokens: "32768" + +MiniMax-M3-MXFP8: + env: "AITER_QUICK_REDUCE_QUANTIZATION=INT4 ATOM_FORCE_ATTN_TRITON=1" + kv_cache_flags: "--kv_cache_dtype fp8" + tp_dp_flags: "--enable-dp-attention" + prefill_tp_dp_flags: "--enable-dp-attention --enable-tbo prefill" + decode_tp_dp_flags: "--enable-dp-attention" + ep_dp_flags: "--enable-expert-parallel --enable-dp-attention" + prefill_ep_dp_flags: "--enable-expert-parallel --enable-dp-attention" + decode_ep_dp_flags: "--enable-expert-parallel --enable-dp-attention" + mtp_flags: "--method eagle3 --draft-model Inferact/MiniMax-M3-EAGLE3 --num-speculative-tokens" + online_quant_config: '{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","vision_tower","multi_modal_projector","patch_merge_mlp","*block_sparse_moe"]}' + online_quant_dpa_config: '{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","vision_tower","multi_modal_projector","patch_merge_mlp","*.gate.*","*.block_sparse_moe.experts*"]}' + block_size: "128" + mem_frac_static: "0.8" + max_model_len: "32768" + max_num_seqs: "256" + max_num_batched_tokens: "32768" diff --git a/scripts/MiniMax-M3/amd_utils/server.sh b/scripts/MiniMax-M3/amd_utils/server.sh new file mode 100644 index 00000000..b62ca581 --- /dev/null +++ b/scripts/MiniMax-M3/amd_utils/server.sh @@ -0,0 +1,23 @@ +#!/bin/bash +# Multi-Engine Disaggregated Server Dispatcher +# ============================================================================= +# Dispatches to the engine-specific server launcher based on ENGINE env var. +# ENGINE=sglang-disagg (default) -> server_sglang.sh (SGLang + MoRI) +# ENGINE=vllm-disagg -> server_vllm.sh (vLLM + Nixl/MoRI-IO) +# ENGINE=atom-disagg -> server_atom.sh (ATOM + mooncake) +# ============================================================================= + +ENGINE="${ENGINE:-sglang-disagg}" +WS_PATH="${WS_PATH:-${SGLANG_WS_PATH:-${VLLM_WS_PATH:-${ATOM_WS_PATH:-$(dirname "${BASH_SOURCE[0]}")}}}}" +export WS_PATH ENGINE + +echo "[DISPATCHER] ENGINE=$ENGINE WS_PATH=$WS_PATH" + +if [[ "$ENGINE" == "vllm-disagg" ]]; then + source "$WS_PATH/server_vllm.sh" +elif [[ "$ENGINE" == "atom-disagg" ]]; then + export ATOM_WS_PATH="$WS_PATH" + source "$WS_PATH/server_atom.sh" +else + source "$WS_PATH/server_sglang.sh" +fi diff --git a/scripts/MiniMax-M3/atom_disagg/server_atom.sh b/scripts/MiniMax-M3/amd_utils/server_atom.sh similarity index 81% rename from scripts/MiniMax-M3/atom_disagg/server_atom.sh rename to scripts/MiniMax-M3/amd_utils/server_atom.sh index ccc86403..23c9dde9 100644 --- a/scripts/MiniMax-M3/atom_disagg/server_atom.sh +++ b/scripts/MiniMax-M3/amd_utils/server_atom.sh @@ -36,8 +36,7 @@ DECODE_ENABLE_EP="${DECODE_ENABLE_EP}" DECODE_ENABLE_DP="${DECODE_ENABLE_DP}" # MTP -SPEC_DECODING="${SPEC_DECODING:-}" -DECODE_MTP_SIZE="${DECODE_MTP_SIZE:-1}" +DECODE_MTP_SIZE="${DECODE_MTP_SIZE:-0}" # ATOM server ports (different from SGLang which uses 8000 for all) PREFILL_PORT="${PREFILL_PORT:-8010}" @@ -45,13 +44,7 @@ DECODE_PORT="${DECODE_PORT:-8020}" ROUTER_PORT="${ROUTER_PORT:-8000}" HANDSHAKE_PORT="${HANDSHAKE_PORT:-6301}" -# ATOM server tuning (from reference script defaults) -MEM_FRAC_STATIC="${MEM_FRAC_STATIC:-0.85}" -KV_CACHE_DTYPE="${KV_CACHE_DTYPE:-fp8}" -BLOCK_SIZE="${BLOCK_SIZE:-16}" -MAX_NUM_SEQS="${MAX_NUM_SEQS:-256}" -MAX_MODEL_LEN="${MAX_MODEL_LEN:-}" -MAX_NUM_BATCHED_TOKENS="${MAX_NUM_BATCHED_TOKENS:-}" +# ATOM server tuning — defaults applied after YAML load (env var > YAML > shell default) EXTRA_SERVER_ARGS="${EXTRA_SERVER_ARGS:-}" # Benchmark Configuration @@ -72,12 +65,61 @@ GPUS_PER_NODE="${GPUS_PER_NODE:-8}" source $ATOM_WS_PATH/setup_deps.sh source $ATOM_WS_PATH/env_atom.sh +# Raise FD limit — lm-eval with high num_concurrent can exhaust the default 1024 +ulimit -n 65536 2>/dev/null || ulimit -n 8192 2>/dev/null || true +echo "ulimit -n (open files): $(ulimit -n)" + host_ip=$(ip route get 1.1.1.1 2>/dev/null | awk '/src/ {print $7}') if [[ -z "$host_ip" ]]; then host_ip=$(hostname -I 2>/dev/null | awk '{print $1}') fi host_name=$(hostname) +# ============================================================================= +# Model-Specific Configuration from YAML +# ============================================================================= +# Load model-specific config from YAML (single parse for all fields) +set -x +_yaml_tmp=$(mktemp) +python3 << PYEOF > "$_yaml_tmp" +import yaml +with open('${ATOM_WS_PATH}/models_atom.yaml') as f: + m = yaml.safe_load(f).get('${MODEL_NAME}', {}) +def sh(v): return v.replace("'", "'\\''") +print(f"MODEL_ENVS='{sh(m.get('env', ''))}'") +_tp_dp = m.get('tp_dp_flags', '') +print(f"PREFILL_MODEL_TP_DP_FLAGS='{sh(m.get('prefill_tp_dp_flags', _tp_dp))}'") +print(f"DECODE_MODEL_TP_DP_FLAGS='{sh(m.get('decode_tp_dp_flags', _tp_dp))}'") +_ep_dp = m.get('ep_dp_flags', '') +print(f"PREFILL_MODEL_EP_DP_FLAGS='{sh(m.get('prefill_ep_dp_flags', _ep_dp))}'") +print(f"DECODE_MODEL_EP_DP_FLAGS='{sh(m.get('decode_ep_dp_flags', _ep_dp))}'") +print(f"MODEL_TP_DP_ENV='{sh(m.get('tp_dp_env', ''))}'") +print(f"MODEL_EP_DP_ENV='{sh(m.get('ep_dp_env', ''))}'") +print(f"MODEL_MTP_FLAGS='{sh(m.get('mtp_flags', ''))}'") +print(f"MODEL_KV_ARG='{sh(m.get('kv_cache_flags', ''))}'") +print(f"_ONLINE_QUANT_CONFIG='{sh(m.get('online_quant_config', ''))}'") +print(f"_ONLINE_QUANT_DPA_CONFIG='{sh(m.get('online_quant_dpa_config', m.get('online_quant_config', '')))}'") +print(f"_YAML_BLOCK_SIZE='{sh(m.get('block_size', ''))}'") +print(f"_YAML_MEM_FRAC_STATIC='{sh(m.get('mem_frac_static', ''))}'") +print(f"_YAML_MAX_MODEL_LEN='{sh(m.get('max_model_len', ''))}'") +print(f"_YAML_MAX_NUM_SEQS='{sh(m.get('max_num_seqs', ''))}'") +print(f"_YAML_MAX_NUM_BATCHED_TOKENS='{sh(m.get('max_num_batched_tokens', ''))}'") +PYEOF +# shellcheck source=/dev/null +source "$_yaml_tmp" +rm -f "$_yaml_tmp" +unset _yaml_tmp + +# Apply server-tuning: YAML > env var > shell default +# (job.slurm injects BLOCK_SIZE/MEM_FRAC_STATIC/MAX_NUM_SEQS with hardcoded +# defaults into the Docker env, so env-first would always shadow the YAML.) +BLOCK_SIZE="${_YAML_BLOCK_SIZE:-${BLOCK_SIZE:-16}}" +MEM_FRAC_STATIC="${_YAML_MEM_FRAC_STATIC:-${MEM_FRAC_STATIC:-0.85}}" +MAX_MODEL_LEN="${_YAML_MAX_MODEL_LEN:-${MAX_MODEL_LEN:-}}" +MAX_NUM_SEQS="${_YAML_MAX_NUM_SEQS:-${MAX_NUM_SEQS:-256}}" +MAX_NUM_BATCHED_TOKENS="${_YAML_MAX_NUM_BATCHED_TOKENS:-${MAX_NUM_BATCHED_TOKENS:-}}" +unset _YAML_BLOCK_SIZE _YAML_MEM_FRAC_STATIC _YAML_MAX_MODEL_LEN _YAML_MAX_NUM_SEQS _YAML_MAX_NUM_BATCHED_TOKENS + # ============================================================================= # Cluster Topology Configuration # ============================================================================= @@ -111,56 +153,56 @@ PREFILL_ENABLE_DP="${PREFILL_ENABLE_DP}" DECODE_ENABLE_EP="${DECODE_ENABLE_EP}" DECODE_ENABLE_DP="${DECODE_ENABLE_DP}" + + + # Parallel args PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE") #TP +ONLINE_QUANT_ARG="" if [ "$PREFILL_ENABLE_DP" = "true" ]; then - if [ "$PREFILL_ENABLE_EP" -gt 1 ]; then #DPA+EP - PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" --enable-expert-parallel --enable-dp-attention ) - else #TP+DPA+TBO - if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then - PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" --enable-dp-attention --enable-tbo ) - export GPU_MAX_HW_QUEUES=5 - export ATOM_CPU_AFFINITY=1 - else #TP+DPA - PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" --enable-dp-attention ) - fi + if [ "$PREFILL_ENABLE_EP" = "true" ]; then #EP+DPA + PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" ${PREFILL_MODEL_EP_DP_FLAGS}) + for _dp_env_pair in ${MODEL_EP_DP_ENV}; do export "$_dp_env_pair"; done + else #TP+DPA + PREFILL_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE" ${PREFILL_MODEL_TP_DP_FLAGS}) + for _dp_env_pair in ${MODEL_TP_DP_ENV}; do export "$_dp_env_pair"; done + fi + if [[ -n "$_ONLINE_QUANT_DPA_CONFIG" ]]; then + ONLINE_QUANT_ARG="--online_quant_config '${_ONLINE_QUANT_DPA_CONFIG}'" + fi +else + if [[ -n "$_ONLINE_QUANT_CONFIG" ]]; then + ONLINE_QUANT_ARG="--online_quant_config '${_ONLINE_QUANT_CONFIG}'" fi -fi +fi -# (srok), split DPA & TBO cases -DECODE_PARALLEL_ARGS=(-tp "$PREFILL_TP_SIZE") #TP +DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE") #TP if [ "$DECODE_ENABLE_DP" = "true" ]; then - if [ "$DECODE_ENABLE_EP" -gt 1 ]; then #DPA+EP - DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" --enable-expert-parallel --enable-dp-attention ) - else #TP+DPA+TBO - if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then - DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" --enable-dp-attention --enable-tbo ) - export GPU_MAX_HW_QUEUES=5 - export ATOM_CPU_AFFINITY=1 - else #TP+DPA - DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" --enable-dp-attention ) - fi + if [ "$DECODE_ENABLE_EP" = "true" ]; then #EP+DPA + DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" ${DECODE_MODEL_EP_DP_FLAGS}) + for _dp_env_pair in ${MODEL_EP_DP_ENV}; do export "$_dp_env_pair"; done + else #TP+DPA + DECODE_PARALLEL_ARGS=(-tp "$DECODE_TP_SIZE" ${DECODE_MODEL_TP_DP_FLAGS}) + for _dp_env_pair in ${MODEL_TP_DP_ENV}; do export "$_dp_env_pair"; done fi -fi - -# MTP args -SPEC_ARGS=() #TP -if [ "$SPEC_DECODING" = "mtp" ]; then - SPEC_ARGS=(--method mtp --num-speculative-tokens "$DECODE_MTP_SIZE") fi +unset _dp_env_pair +unset _ONLINE_QUANT_CONFIG _ONLINE_QUANT_DPA_CONFIG -# HF overrides (single-quoted JSON preserved through eval) -HF_OVERRIDES_ARG="" -if [[ "$MODEL_NAME" == "DeepSeek-V4-Pro" ]]; then - HF_OVERRIDES_ARG="--hf-overrides '{\"use_index_cache\":true,\"index_topk_freq\":4}'" -fi +for _env_pair in ${MODEL_ENVS}; do + export "$_env_pair" +done +unset _env_pair -# KV cache dtype (skip if unset or 'auto') -KV_CACHE_ARG="" -if [[ -n "$KV_CACHE_DTYPE" && "$KV_CACHE_DTYPE" != "auto" ]]; then - KV_CACHE_ARG="--kv_cache_dtype ${KV_CACHE_DTYPE}" +# MTP args +SPEC_ARGS=() +if [[ -n "$MODEL_MTP_FLAGS" && "${DECODE_MTP_SIZE:-0}" -gt 0 ]]; then + SPEC_ARGS=(${MODEL_MTP_FLAGS} "$DECODE_MTP_SIZE") fi +# KV cache arg - full flag string from YAML +KV_CACHE_ARG="${MODEL_KV_ARG}" + # Optional model length / batched-token cap MODEL_LEN_ARGS="" if [[ -n "$MAX_MODEL_LEN" ]]; then @@ -170,9 +212,6 @@ if [[ -n "$MAX_NUM_BATCHED_TOKENS" ]]; then MODEL_LEN_ARGS="${MODEL_LEN_ARGS} --max-num-batched-tokens ${MAX_NUM_BATCHED_TOKENS}" fi -if [[ "$MODEL_NAME" != "DeepSeek-V4-Pro" ]]; then - export AITER_QUICK_REDUCE_QUANTIZATION=INT4 -fi cat < recipe deps + amd-quark + UCX/RIXL path exports +# (base image: vllm/vllm-openai-rocm:nightly) +# sglang-disagg -> SGLang aiter gluon patch + per-model installs +# (base image: lmsysorg/sglang-rocm:v0.5.12-rocm720-mi35x-*) +# +# Sourced by server_vllm.sh and server_sglang.sh so PATH / LD_LIBRARY_PATH +# exports persist. Each patch is idempotent: skipped if already applied. +# +# Build steps run in subshells to avoid CWD pollution between installers. +# ============================================================================= + +ROCM_PATH="${ROCM_PATH:-/opt/rocm}" +UCX_HOME="${UCX_HOME:-/usr/local/ucx}" +RIXL_HOME="${RIXL_HOME:-/usr/local/rixl}" + +_SETUP_START=$(date +%s) +_SETUP_INSTALLED=() + +git_clone_retry() { + local url="$1" dest="$2" max_tries=3 try=1 + while (( try <= max_tries )); do + if git clone --quiet "$url" "$dest" 2>/dev/null; then return 0; fi + echo "[SETUP] git clone attempt $try/$max_tries failed for $url, retrying in 10s..." + rm -rf "$dest" + sleep 10 + (( try++ )) + done + echo "[SETUP] git clone failed after $max_tries attempts: $url" + return 1 +} + + +# --------------------------------------------------------------------------- +# 5. Container RDMA/net tools +# - ibv_devinfo comes from ibverbs-utils +# - iproute2 provides the `ip` command +# Used for in-container NIC/RDMA validation and routing checks. +# --------------------------------------------------------------------------- +install_recipe_deps() { + if command -v ibv_devinfo >/dev/null 2>&1 && command -v ip >/dev/null 2>&1; then + echo "[SETUP] Container RDMA/net tools already present" + return 0 + fi + + echo "[SETUP] Installing ibv_devinfo + iproute2 in container..." + apt-get update -q -y && apt-get install -q -y \ + ibverbs-utils iproute2 \ + && rm -rf /var/lib/apt/lists/* + + if ! command -v ibv_devinfo >/dev/null 2>&1 || ! command -v ip >/dev/null 2>&1; then + echo "[SETUP] ERROR: Failed to install ibv_devinfo/iproute2"; exit 1 + fi + _SETUP_INSTALLED+=("ibverbs-utils+iproute2") +} + +# --------------------------------------------------------------------------- +# 6b. amd-quark (MXFP4 quantization support for Kimi-K2.5-MXFP4 and similar) +# Required due to ROCm vLLM missing the quark dependency: +# https://github.com/vllm-project/vllm/issues/35633 +# --------------------------------------------------------------------------- +install_amd_quark() { + if python3 -c "import quark" 2>/dev/null; then + echo "[SETUP] amd-quark already present" + return 0 + fi + + echo "[SETUP] Installing amd-quark for MXFP4 quantization support..." + pip install --quiet amd-quark + + if ! python3 -c "import quark" 2>/dev/null; then + echo "[SETUP] WARN: amd-quark install failed (non-fatal for non-MXFP4 models)" + return 0 + fi + _SETUP_INSTALLED+=("amd-quark") +} + +# --------------------------------------------------------------------------- +# SGLang: Patch aiter gluon pa_mqa_logits — fix 2D → 3D instr_shape for +# Triton ≥ 3.5. +# +# Bug: _gluon_deepgemm_fp8_paged_mqa_logits (the non-preshuffle variant) +# hardcodes AMDMFMALayout(instr_shape=[16, 16]) which fails on Triton +# builds where AMDMFMALayout requires 3D (M, N, K) format. +# +# The two preshuffle variants already conditionally select 2D vs 3D via +# the module-level _Use_2d_instr_shape_mfma_layout flag, but the base +# variant was missed. This patch brings it in line. +# +# Affects: GLM-5 (NSA attention) and any future model that uses +# deepgemm_fp8_paged_mqa_logits with Preshuffle=False. +# --------------------------------------------------------------------------- +patch_gluon_pa_mqa_logits_instr_shape() { + python3 -c ' +import os, sys + +target = "/sgl-workspace/aiter/aiter/ops/triton/gluon/pa_mqa_logits.py" +if not os.path.isfile(target): + print("[SETUP] gluon pa_mqa_logits.py not found, skipping") + sys.exit(0) + +src = open(target).read() + +if "[PATCHED] 3D instr_shape for base gluon variant" in src: + print("[SETUP] gluon pa_mqa_logits 3D instr_shape patch already applied") + sys.exit(0) + +# The buggy code: the base _gluon_deepgemm_fp8_paged_mqa_logits uses 2D +# instr_shape unconditionally. We replace it with a conditional that +# mirrors the preshuffle variants. +old = """\ + mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=CDNA_VERSION, + instr_shape=[16, 16], + transposed=False, + warps_per_cta=[1, NumWarps], + ) + mfma_layout_a: gl.constexpr = gl.DotOperandLayout( + operand_index=0, parent=mfma_layout, k_width=16 + ) + mfma_layout_b: gl.constexpr = gl.DotOperandLayout( + operand_index=1, parent=mfma_layout, k_width=16 + )""" + +new = """\ + # [PATCHED] 3D instr_shape for base gluon variant + if _Use_2d_instr_shape_mfma_layout: + mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=CDNA_VERSION, + instr_shape=[16, 16], + transposed=False, + warps_per_cta=[1, NumWarps], + ) + else: + mfma_layout: gl.constexpr = gl.amd.AMDMFMALayout( + version=CDNA_VERSION, + instr_shape=[16, 16, 32], + transposed=False, + warps_per_cta=[1, NumWarps], + ) + mfma_layout_a: gl.constexpr = gl.DotOperandLayout( + operand_index=0, parent=mfma_layout, k_width=16 + ) + mfma_layout_b: gl.constexpr = gl.DotOperandLayout( + operand_index=1, parent=mfma_layout, k_width=16 + )""" + +if old not in src: + print("[SETUP] WARN: gluon pa_mqa_logits pattern not found — aiter version may have changed") + sys.exit(0) + +# Only replace the FIRST occurrence (the base variant, not preshuffle ones) +new_src = src.replace(old, new, 1) + +open(target, "w").write(new_src) +print("[SETUP] Patched: gluon pa_mqa_logits 3D instr_shape for base variant") +' + _SETUP_INSTALLED+=("gluon-instr-shape-fix") +} + +# --------------------------------------------------------------------------- +# SGLang: Install latest transformers for GLM-5 model type support. +# +# GLM-5 (zai-org/GLM-5-FP8) requires a transformers build that includes +# the glm_moe_dsa model type. The mori images do not ship it. +# Only install if GLM-5 is the active model (avoid overhead otherwise). +# --------------------------------------------------------------------------- +install_transformers_glm5() { + if [[ "$MODEL_NAME" != "GLM-5-FP8" ]]; then + return 0 + fi + + if python3 -c "from transformers import AutoConfig; AutoConfig.from_pretrained('zai-org/GLM-5-FP8', trust_remote_code=True)" 2>/dev/null; then + echo "[SETUP] transformers already supports GLM-5 model type" + return 0 + fi + + echo "[SETUP] Installing transformers with GLM-5 (glm_moe_dsa) support..." + pip install --quiet -U --no-cache-dir \ + "git+https://github.com/huggingface/transformers.git@6ed9ee36f608fd145168377345bfc4a5de12e1e2" + _SETUP_INSTALLED+=("transformers-glm5") +} + +# ============================================================================= +# Run installers (engine-gated) +# ============================================================================= + +if [[ "$ENGINE" == "vllm-disagg" ]]; then + install_recipe_deps + install_amd_quark + + # ========================================================================= + # vLLM: Export UCX/RIXL paths (persists since this file is sourced) + # ========================================================================= + export ROCM_PATH="${ROCM_PATH}" + export UCX_HOME="${UCX_HOME}" + export RIXL_HOME="${RIXL_HOME}" + export PATH="${UCX_HOME}/bin:/usr/local/bin/etcd:/root/.cargo/bin:${PATH}" + export LD_LIBRARY_PATH="${UCX_HOME}/lib:${RIXL_HOME}/lib:${RIXL_HOME}/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}" +else + patch_gluon_pa_mqa_logits_instr_shape + install_transformers_glm5 +fi + +_SETUP_END=$(date +%s) +if [[ ${#_SETUP_INSTALLED[@]} -eq 0 ]]; then + echo "[SETUP] All dependencies already present ($(( _SETUP_END - _SETUP_START ))s wallclock)" +else + echo "[SETUP] Installed: ${_SETUP_INSTALLED[*]} in $(( _SETUP_END - _SETUP_START ))s" +fi diff --git a/scripts/MiniMax-M3/amd_utils/submit.sh b/scripts/MiniMax-M3/amd_utils/submit.sh new file mode 100644 index 00000000..262fd25a --- /dev/null +++ b/scripts/MiniMax-M3/amd_utils/submit.sh @@ -0,0 +1,246 @@ +#!/bin/bash +# +# Cluster Configuration Template for Multi-Node Disaggregated Serving +# +# This script submits a multi-node disaggregated benchmark job to SLURM. +# It must be configured for your specific cluster before use. +# +# ENGINE=sglang (default): SGLang disaggregated serving +# ENGINE=vllm: vLLM disaggregated serving +# +# Router is co-located with the first prefill node (same for both engines), +# so NUM_NODES = PREFILL_NODES + DECODE_NODES. + +usage() { + cat << 'USAGE' +Usage: + bash submit.sh \ + \ + \ + \ + \ + [NODE_LIST] + +Arguments: + PREFILL_NODES Number of prefill nodes + PREFILL_WORKERS Number of prefill workers (usually 1) + DECODE_NODES Number of decode nodes + DECODE_WORKERS Number of decode workers (usually 1) + ISL Input sequence length + OSL Output sequence length + CONCURRENCIES Concurrency levels, delimited by 'x' (e.g., "8x16x32") + REQUEST_RATE Request rate ("inf" for max throughput) + PREFILL_ENABLE_EP true/false or 1/0 (expert parallelism on prefill) + PREFILL_ENABLE_DP true/false or 1/0 (data-parallel attention on prefill) + DECODE_ENABLE_EP true/false or 1/0 (expert parallelism on decode) + DECODE_ENABLE_DP true/false or 1/0 (data-parallel attention on decode) + PREFILL_TP Tensor parallel size per prefill node + DECODE_TP Tensor parallel size per decode node + RANDOM_RANGE_RATIO Random range ratio for benchmark client + NODE_LIST Optional: comma-separated hostnames (must match NUM_NODES) + +Required environment variables: + SLURM_ACCOUNT SLURM account name + SLURM_PARTITION SLURM partition + TIME_LIMIT Job time limit (e.g., "08:00:00") + MODEL_PATH Path to model directory (e.g., /nfsdata) + MODEL_NAME Model name directory + CONTAINER_IMAGE Docker image name (e.g., vllm_disagg_pd:latest) + RUNNER_NAME Runner identifier (for job name) + +Optional environment variables: + DRY_RUN 1 = echo composed server/router launch commands instead of + running them (preview a recipe against a real allocation). +USAGE +} + +check_env() { + local name="$1" + if [[ -z "${!name:-}" ]]; then + echo "Error: ${name} not specified" >&2 + usage >&2 + exit 1 + fi +} + +check_env SLURM_ACCOUNT +check_env SLURM_PARTITION +check_env TIME_LIMIT + +check_env MODEL_PATH +check_env MODEL_NAME +check_env CONTAINER_IMAGE +check_env RUNNER_NAME +check_env FRAMEWORK + +# GPUS_PER_NODE defaults to 8 (MI355X). Set to 4 for MI325X if needed. +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" + +# COMMAND_LINE ARGS +PREFILL_NODES=$1 +PREFILL_WORKERS=${2:-1} +DECODE_NODES=$3 +DECODE_WORKERS=${4:-1} +ISL=$5 +OSL=$6 +CONCURRENCIES=$7 +REQUEST_RATE=$8 +PREFILL_ENABLE_EP=${9:-true} +PREFILL_ENABLE_DP=${10:-true} +DECODE_ENABLE_EP=${11:-true} +DECODE_ENABLE_DP=${12:-true} +PREFILL_TP=${13:-8} +DECODE_TP=${14:-8} +RANDOM_RANGE_RATIO=${15:-0.8} +NODE_LIST=${16} + +NUM_NODES=$((PREFILL_NODES + DECODE_NODES)) +profiler_args="${ISL} ${OSL} ${CONCURRENCIES} ${REQUEST_RATE}" + +# Export variables for the SLURM job +export ENGINE="${FRAMEWORK:-sglang}" +export MODEL_DIR=$MODEL_PATH +export DOCKER_IMAGE_NAME=$CONTAINER_IMAGE +export PROFILER_ARGS=$profiler_args + +# Engine-specific xP/yD semantics and TP exports +if [[ "$ENGINE" == "vllm-disagg" ]]; then + export PROXY_STREAM_IDLE_TIMEOUT=${PROXY_STREAM_IDLE_TIMEOUT:-300} +fi +# xP = prefill workers, yD = decode workers (may span multiple nodes) +export xP=$PREFILL_WORKERS +export yD=$DECODE_WORKERS +export PREFILL_TP_SIZE=$(( $PREFILL_NODES * $PREFILL_TP / $PREFILL_WORKERS )) +export PREFILL_ENABLE_EP=${PREFILL_ENABLE_EP} +export PREFILL_ENABLE_DP=${PREFILL_ENABLE_DP} +export DECODE_TP_SIZE=$(( $DECODE_NODES * $DECODE_TP / $DECODE_WORKERS )) +export DECODE_ENABLE_EP=${DECODE_ENABLE_EP} +export DECODE_ENABLE_DP=${DECODE_ENABLE_DP} +export DECODE_MTP_SIZE=${DECODE_MTP_SIZE} + +export NUM_NODES=$NUM_NODES +export GPUS_PER_NODE=$GPUS_PER_NODE +export MODEL_NAME=$MODEL_NAME +export BENCH_INPUT_LEN=${ISL} +export BENCH_OUTPUT_LEN=${OSL} +export BENCH_NUM_PROMPTS_MULTIPLIER=${BENCH_NUM_PROMPTS_MULTIPLIER:-10} +export BENCH_MAX_CONCURRENCY=${CONCURRENCIES} +export BENCH_REQUEST_RATE=${REQUEST_RATE} +export BENCH_RANDOM_RANGE_RATIO=${RANDOM_RANGE_RATIO:-0.8} + +# DRY_RUN=1 makes server_sglang.sh echo the composed prefill/decode/router launch +# commands instead of executing them (useful for previewing a recipe against a real +# allocation). Threaded here → job.slurm → Docker (-e DRY_RUN) → server_sglang.sh. +# sbatch defaults to --export=ALL, so exporting it is what carries it into the job. +export DRY_RUN="${DRY_RUN:-0}" + +# Eval-related env vars (threaded from workflow → runner → here → job.slurm → Docker) +export RUN_EVAL="${RUN_EVAL:-false}" +export EVAL_ONLY="${EVAL_ONLY:-false}" +export EVAL_CONC="${EVAL_CONC:-}" +export FRAMEWORK="${FRAMEWORK:-}" +export PRECISION="${PRECISION:-}" +export MODEL_PREFIX="${MODEL_PREFIX:-}" +export RUNNER_TYPE="${RUNNER_TYPE:-}" +export RESULT_FILENAME="${RESULT_FILENAME:-}" +export SPEC_DECODING="${SPEC_DECODING:-}" +export IS_MULTINODE="${IS_MULTINODE:-false}" + +# Log directory: must be on NFS (shared filesystem) so the submit host can read SLURM output. +export BENCHMARK_LOGS_DIR="${BENCHMARK_LOGS_DIR:-$(pwd)/benchmark_logs}" +mkdir -p "$BENCHMARK_LOGS_DIR" + +# Optional: pass an explicit node list to sbatch. +NODELIST_OPT=() +if [[ -n "${NODE_LIST//[[:space:]]/}" ]]; then + IFS=',' read -r -a NODE_ARR <<< "$NODE_LIST" + if [[ "${#NODE_ARR[@]}" -ne "$NUM_NODES" ]]; then + echo "Error: NODE_LIST has ${#NODE_ARR[@]} nodes but NUM_NODES=${NUM_NODES}" >&2 + echo "Error: NODE_LIST='${NODE_LIST}'" >&2 + exit 1 + fi + NODELIST_CSV="$(IFS=,; echo "${NODE_ARR[*]}")" + NODELIST_OPT=(--nodelist "$NODELIST_CSV") +fi + +# Optional: exclude specific nodes (e.g. nodes with broken Docker sockets). +# Set SLURM_EXCLUDE_NODES env var to a comma-separated list of hostnames. +EXCLUDE_OPT=() +SLURM_EXCLUDE_NODES="${SLURM_EXCLUDE_NODES:-mia1-p01-g11,mia1-p01-g12,mia1-p01-g15}" +if [[ -n "${SLURM_EXCLUDE_NODES:-}" ]]; then + EXCLUDE_OPT=(--exclude "$SLURM_EXCLUDE_NODES") +fi + +# ============================================================================= +# Reuse existing allocation (skip sbatch) +# ============================================================================= +# When SLURM_REUSE_JOBID is set, run job.slurm directly in the current shell, +# attaching to the existing allocation. Inner `srun` calls pick up the +# allocation via SLURM_JOB_ID; SLURM_OVERLAP=1 lets them share task slots with +# the interactive shell already holding the allocation. +if [[ -n "${SLURM_REUSE_JOBID:-}" ]]; then + REUSE_JID="$SLURM_REUSE_JOBID" + echo "Reusing existing Slurm allocation ${REUSE_JID} (skipping sbatch)" >&2 + + # Resolve allocation's nodelist if not already provided. + ALLOC_NODELIST="${SLURM_JOB_NODELIST:-$(squeue -h -j "$REUSE_JID" -o '%N' 2>/dev/null)}" + if [[ -z "$ALLOC_NODELIST" ]]; then + echo "Error: could not resolve nodelist for job ${REUSE_JID}" >&2 + exit 1 + fi + ALLOC_NNODES=$(scontrol show hostnames "$ALLOC_NODELIST" | wc -l) + if [[ "$ALLOC_NNODES" -lt "$NUM_NODES" ]]; then + echo "Error: allocation ${REUSE_JID} has ${ALLOC_NNODES} nodes, need ${NUM_NODES}" >&2 + exit 1 + fi + + export SLURM_JOB_ID="$REUSE_JID" + export SLURM_JOBID="$REUSE_JID" + export SLURM_JOB_NODELIST="$ALLOC_NODELIST" + export SLURM_NODELIST="$ALLOC_NODELIST" + export SLURM_NNODES="$ALLOC_NNODES" + export SLURM_JOB_NUM_NODES="$ALLOC_NNODES" + export SLURM_NTASKS="$ALLOC_NNODES" + export SLURM_NPROCS="$ALLOC_NNODES" + export SLURM_NTASKS_PER_NODE=1 + export SLURM_TASKS_PER_NODE="1(x${ALLOC_NNODES})" + export SLURM_OVERLAP=1 + export SLURM_SUBMIT_DIR="$(pwd)" + + STDOUT_LOG="${BENCHMARK_LOGS_DIR}/slurm_job-${REUSE_JID}.out" + STDERR_LOG="${BENCHMARK_LOGS_DIR}/slurm_job-${REUSE_JID}.err" + rm -f "$STDOUT_LOG" "$STDERR_LOG" + + nohup bash "$(dirname "$0")/job.slurm" >"$STDOUT_LOG" 2>"$STDERR_LOG" & + INLINE_PID=$! + echo "$INLINE_PID" > "${BENCHMARK_LOGS_DIR}/slurm_job-${REUSE_JID}.pid" + echo "Started job.slurm (pid=${INLINE_PID}); logs: ${STDOUT_LOG}" >&2 + + echo "$REUSE_JID" + exit 0 +fi + +# Construct the sbatch command +sbatch_cmd=( + sbatch + --parsable + --exclusive + -N "$NUM_NODES" + -n "$NUM_NODES" + "${NODELIST_OPT[@]}" + "${EXCLUDE_OPT[@]}" + --time "$TIME_LIMIT" + --partition "$SLURM_PARTITION" + --account "$SLURM_ACCOUNT" + --job-name "$RUNNER_NAME" + --output "${BENCHMARK_LOGS_DIR}/slurm_job-%j.out" + --error "${BENCHMARK_LOGS_DIR}/slurm_job-%j.err" + "$(dirname "$0")/job.slurm" +) + +JOB_ID=$("${sbatch_cmd[@]}") +if [[ $? -ne 0 ]]; then + echo "Error: Failed to submit job with sbatch" >&2 + exit 1 +fi +echo "$JOB_ID" diff --git a/scripts/MiniMax-M3/amd_utils/sync.py b/scripts/MiniMax-M3/amd_utils/sync.py new file mode 100644 index 00000000..96e94c1b --- /dev/null +++ b/scripts/MiniMax-M3/amd_utils/sync.py @@ -0,0 +1,220 @@ +#!/usr/bin/env python3 +""" +Multi-node synchronization utilities for disaggregated inference. + +Subcommands: + barrier - Wait until all specified nodes have opened their ports (TCP barrier) + Optionally wait for HTTP health endpoints to return 200 + wait - Block until a remote port closes (shutdown coordination) +""" + +import socket +import time +import threading +import argparse +import sys +import urllib.request +import urllib.error + + +def is_port_open(ip, port, timeout=2): + """Check if a given IP and port are accessible.""" + with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: + s.settimeout(timeout) + return s.connect_ex((ip, port)) == 0 + + +def check_health(ip, port, path="/health", timeout=2): + """Return True if http://ip:port/path returns HTTP 200.""" + try: + url = f"http://{ip}:{port}{path}" + req = urllib.request.Request(url) + with urllib.request.urlopen(req, timeout=timeout) as resp: + return getattr(resp, "status", 200) == 200 + except (urllib.error.URLError, urllib.error.HTTPError, OSError): + return False + + +# ============================================================================= +# barrier subcommand +# ============================================================================= + +def cmd_barrier(args): + """Wait until all nodes have opened the specified ports.""" + NODE_IPS = [ip.strip() for ip in args.node_ips.split(",") if ip.strip()] + NODE_PORTS = [int(p.strip()) for p in args.node_ports.split(",") if p.strip()] + + if not NODE_IPS: + print("Error: NODE_IPS argument is empty or not set.") + sys.exit(1) + + if len(NODE_PORTS) == 1: + NODE_PORTS *= len(NODE_IPS) + elif len(NODE_PORTS) != len(NODE_IPS): + print("Error: Number of ports must match number of node IPs or only one port should be given for all.") + sys.exit(1) + + server_socket = None + + def open_port(): + nonlocal server_socket + server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1) + server_socket.bind((args.local_ip, args.local_port)) + server_socket.listen(5) + print(f"Port {args.local_port} is now open on {args.local_ip}.") + while True: + conn, addr = server_socket.accept() + conn.close() + + def close_port(): + nonlocal server_socket + if server_socket: + server_socket.close() + print(f"Port {args.local_port} has been closed on {args.local_ip}.") + + if args.enable_port: + threading.Thread(target=open_port, daemon=True).start() + + # Wait for all ports (TCP check) + if args.wait_for_all_ports: + start_time = time.time() + timeout = args.timeout + seen_open = set() + + while True: + status = {(ip, port): is_port_open(ip, port) + for ip, port in zip(NODE_IPS, NODE_PORTS)} + + # Fail fast: a port that was open and is now closed means that + # server died/was killed; don't wait out the full timeout. + regressed = [t for t in seen_open if not status[t]] + if regressed: + print("ERROR: the following ports were open then went down (server died):", flush=True) + for ip, port in regressed: + print(f" - {ip}:{port}", flush=True) + sys.exit(1) + seen_open.update(t for t, ok in status.items() if ok) + + if all(status.values()): + break + + if timeout > 0: + elapsed = time.time() - start_time + if elapsed >= timeout: + not_open = [t for t, ok in status.items() if not ok] + print(f"ERROR: Timeout after {timeout} seconds waiting for ports to open.", flush=True) + print("The following nodes/ports are still not responding:", flush=True) + for ip, port in not_open: + print(f" - {ip}:{port}", flush=True) + sys.exit(1) + remaining = timeout - (time.time() - start_time) + print(f"Waiting for nodes.{NODE_PORTS},{NODE_IPS} . . ({remaining:.0f}s remaining)", flush=True) + else: + print(f"Waiting for nodes.{NODE_PORTS},{NODE_IPS} . .", flush=True) + time.sleep(5) + + # Wait for all health endpoints (HTTP check) + if args.wait_for_all_health: + health_path = args.health_endpoint + start_time = time.time() + timeout = args.timeout + seen_ready = set() + + while True: + status = { + (ip, port): check_health(ip, port, health_path) + for ip, port in zip(NODE_IPS, NODE_PORTS) + } + + # Fail fast: an endpoint that was healthy and is now down means the + # server died/was killed; don't wait out the full timeout. + regressed = [t for t in seen_ready if not status[t]] + if regressed: + print(f"ERROR: the following ({health_path}) were healthy then went down (server died):", flush=True) + for ip, port in regressed: + print(f" - http://{ip}:{port}{health_path}", flush=True) + sys.exit(1) + seen_ready.update(t for t, ok in status.items() if ok) + + if all(status.values()): + break + + if timeout > 0: + elapsed = time.time() - start_time + if elapsed >= timeout: + not_ready = [t for t, ok in status.items() if not ok] + print(f"ERROR: Timeout after {timeout} seconds waiting for health endpoints.", flush=True) + print(f"The following (http://ip:port{health_path}) are still not responding:", flush=True) + for ip, port in not_ready: + print(f" - http://{ip}:{port}{health_path}", flush=True) + sys.exit(1) + remaining = timeout - (time.time() - start_time) + print( + f"Waiting for health on {list(zip(NODE_IPS, NODE_PORTS))} ({health_path}) .. ({remaining:.0f}s remaining)", + flush=True, + ) + else: + print(f"Waiting for health on {list(zip(NODE_IPS, NODE_PORTS))} ({health_path}) ..", flush=True) + time.sleep(30) + + if args.enable_port: + # Keep the port open long enough for slow nodes to pass their barrier. + # The previous 30s was too short when setup times vary by minutes. + grace = max(60, args.timeout // 2) if args.timeout > 0 else 300 + time.sleep(grace) + close_port() + + +# ============================================================================= +# wait subcommand +# ============================================================================= + +def cmd_wait(args): + """Wait while a remote port remains open, exit when it closes.""" + print( + f"Waiting while port {args.remote_port} on {args.remote_ip} is open...", + flush=True, + ) + while is_port_open(args.remote_ip, args.remote_port): + time.sleep(5) + print(f"Port {args.remote_port} on {args.remote_ip} is now closed.", flush=True) + + +# ============================================================================= +# CLI +# ============================================================================= + +def main(): + parser = argparse.ArgumentParser(description="Multi-node synchronization utilities.") + subparsers = parser.add_subparsers(dest="command", required=True) + + # barrier subcommand + bp = subparsers.add_parser("barrier", help="Wait for all nodes to open specified ports.") + bp.add_argument("--local-ip", required=False, help="Local IP address to bind the server.") + bp.add_argument("--local-port", type=int, required=False, help="Port number to bind the server.") + bp.add_argument("--enable-port", action="store_true", help="Enable opening and closing of local port.") + bp.add_argument("--node-ips", required=True, help="Comma-separated list of node IPs.") + bp.add_argument("--node-ports", required=True, help="Comma-separated list of ports to check.") + bp.add_argument("--timeout", type=int, default=600, + help="Timeout in seconds (default: 600). Set to 0 for no timeout.") + bp.add_argument("--wait-for-all-ports", action="store_true", + help="Wait until all node ports are open (TCP).") + bp.add_argument("--wait-for-all-health", action="store_true", + help="Wait until http://ip:port/health returns 200 for all nodes.") + bp.add_argument("--health-endpoint", default="/health", + help="Path for health check (default: /health).") + bp.set_defaults(func=cmd_barrier) + + # wait subcommand + wp = subparsers.add_parser("wait", help="Wait while a remote port remains open.") + wp.add_argument("--remote-ip", required=True, help="Remote server IP address.") + wp.add_argument("--remote-port", type=int, required=True, help="Remote port number.") + wp.set_defaults(func=cmd_wait) + + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/scripts/MiniMax-M3/atom_disagg/accuracy_ctx.py b/scripts/MiniMax-M3/atom_disagg/accuracy_ctx.py deleted file mode 100644 index 0a5f41f9..00000000 --- a/scripts/MiniMax-M3/atom_disagg/accuracy_ctx.py +++ /dev/null @@ -1,73 +0,0 @@ -#!/usr/bin/env python3 -# accuracy_ctx.py — context-varied accuracy probe for the MiniMax-M3 ATOM disagg pipeline. -# -# Why: a PD-disaggregated pipeline can be correct at short context but corrupt the KV that -# mooncake transfers prefill->decode at long context. A 5-short-prompt gate misses that. -# This embeds a known "needle" fact inside filler of increasing size and checks the router -# (which spans prefill+decode+mooncake) returns it verbatim — exercising the FULL pipeline -# at each context length. -# -# Usage: ROUTER_URL, MODEL, OUT(json), [CTX_SIZES="512 2048 8192 16384"] via env. -# Exits 0 if all context sizes pass, 1 otherwise. -import json, os, sys, time, urllib.request - -ROUTER = os.environ.get("ROUTER_URL", "http://0.0.0.0:8000") -MODEL = os.environ["MODEL"] -OUT = os.environ.get("OUT", "/out/accuracy_ctx.json") -# approx tokens of filler context to test at (the needle adds a few more) -CTX_SIZES = [int(x) for x in os.environ.get("CTX_SIZES", "512 2048 8192 16384").split()] - -# A deterministic needle: a unique code the model must echo back. Greedy (temp 0). -SECRET = "PLATYPUS-7731-AZURE" -FILLER_SENTENCE = ("The quarterly logistics report notes that warehouse throughput " - "remained within nominal parameters across all regional hubs. ") - -def make_prompt(approx_tokens): - # ~1.3 tokens/word for English; build filler then plant the needle in the MIDDLE - words_needed = int(approx_tokens / 1.3) - filler_words = FILLER_SENTENCE.split() - reps = max(1, words_needed // len(filler_words)) - block = (" ".join(filler_words) + " ") * reps - half = len(block) // 2 - needle = f" IMPORTANT: the access code is {SECRET}. Remember it exactly. " - ctx = block[:half] + needle + block[half:] - q = ("\n\nQuestion: What is the exact access code mentioned above? " - "Reply with ONLY the code, nothing else.") - return ctx + q - -def ask(prompt): - body = json.dumps({ - "model": MODEL, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": 32, "temperature": 0.0, - }).encode() - req = urllib.request.Request(f"{ROUTER}/v1/chat/completions", - data=body, headers={"Content-Type": "application/json"}) - t0 = time.time() - with urllib.request.urlopen(req, timeout=300) as r: - d = json.load(r) - return d["choices"][0]["message"]["content"], time.time() - t0 - -results = []; allpass = True -for ctx in CTX_SIZES: - try: - prompt = make_prompt(ctx) - text, dt = ask(prompt) - found = SECRET in text.upper() - ok = found - results.append({"ctx_tokens": ctx, "ok": ok, "found_needle": found, - "reply": text.strip()[:80], "latency_s": round(dt, 2), - "prompt_chars": len(prompt)}) - print(f" ctx~{ctx:>6}tok: {'PASS' if ok else 'FAIL'} " - f"(needle={'yes' if found else 'NO'}, {dt:.1f}s) reply='{text.strip()[:40]}'") - allpass = allpass and ok - except Exception as e: - results.append({"ctx_tokens": ctx, "ok": False, "error": str(e)[:120]}) - print(f" ctx~{ctx:>6}tok: ERROR {str(e)[:80]}") - allpass = False - -os.makedirs(os.path.dirname(OUT), exist_ok=True) -json.dump({"model": MODEL, "router": ROUTER, "all_pass": allpass, "results": results}, - open(OUT, "w"), indent=2) -print(f"[accuracy_ctx] {'ALL PASS' if allpass else 'FAILURES'} -> {OUT}") -sys.exit(0 if allpass else 1) diff --git a/scripts/MiniMax-M3/atom_disagg/bench.sh b/scripts/MiniMax-M3/atom_disagg/bench.sh deleted file mode 100644 index 53442f10..00000000 --- a/scripts/MiniMax-M3/atom_disagg/bench.sh +++ /dev/null @@ -1,47 +0,0 @@ -#!/bin/bash -# bench.sh — driven by server_atom.sh on node 0 (router node), AFTER atomesh is /v1/models-ready. -# Reproduces the InferenceX call signature exactly: -# bash bench.sh \ -# "" -# -# Drives the OpenAI-compatible ATOM router at localhost:$ROUTER_PORT with the vendored -# InferenceX benchmark_serving client (same one used for the Qwen sweep) — one result JSON -# per concurrency. This is a TRUE aggregate (one endpoint, real PD + mooncake KV), no x N. -set -uo pipefail -xP="$1"; yD="$2"; PF_GPUS="$3"; DEC_GPUS="$4" -MODEL_DIR="$5"; MODEL_NAME="$6"; LOG_DIR="$7" -ISL="$8"; OSL="$9"; CONC_LIST="${10}"; REQ_RATE="${11}"; RANGE_RATIO="${12}"; PROMPTS_MULT="${13}" -# server_atom.sh passes the conc list 'x'-separated (e.g. 4x8x16). Normalise to spaces. -CONC_LIST=$(echo "$CONC_LIST" | tr 'x' ' ') - -ROUTER_PORT="${ROUTER_PORT:-8000}" -BENCH_CLIENT="${BENCH_CLIENT:-/scripts/utils/bench_serving/benchmark_serving.py}" -OUT="${LOG_DIR}"; mkdir -p "$OUT" 2>/dev/null || true - -echo "[bench] router=localhost:${ROUTER_PORT} model=${MODEL_NAME} isl=${ISL} osl=${OSL} conc='${CONC_LIST}' (xP=${xP} yD=${yD})" - -# --- ACCURACY GATE (context-varied) before perf ------------------------------------------- -# Validate the FULL PD pipeline (prefill+mooncake KV+decode) returns correct output at varied -# context sizes. Runs ONCE per launch (guarded by a marker), not per-shape. Non-fatal: logs -# result but does not abort the perf sweep (so we still get numbers even if a long-ctx probe fails). -ACC_OUT="${OUT}/accuracy_ctx.json" -if [[ ! -f "$ACC_OUT" ]]; then - echo "[accuracy] context-varied probe via router (sizes: ${CTX_SIZES:-512 2048 8192 16384})" - ROUTER_URL="http://localhost:${ROUTER_PORT}" MODEL="${MODEL_DIR}/${MODEL_NAME}" \ - OUT="$ACC_OUT" CTX_SIZES="${CTX_SIZES:-512 2048 8192 16384}" \ - python3 /ws/accuracy_ctx.py || echo "[accuracy] WARN: some context sizes failed (see $ACC_OUT)" -fi - -for c in $CONC_LIST; do - echo "[bench] --- concurrency $c ---" - python3 "$BENCH_CLIENT" \ - --model "${MODEL_DIR}/${MODEL_NAME}" --backend vllm \ - --base-url "http://localhost:${ROUTER_PORT}" --dataset-name random \ - --random-input-len "$ISL" --random-output-len "$OSL" --random-range-ratio "$RANGE_RATIO" \ - --num-prompts "$(( c * PROMPTS_MULT ))" --num-warmups "$(( c * 2 ))" \ - --max-concurrency "$c" --request-rate "$REQ_RATE" --ignore-eos --save-result \ - --percentile-metrics ttft,tpot,itl,e2el --result-dir "$OUT" \ - --result-filename "disagg_${xP}p${yD}d_isl${ISL}_osl${OSL}_c${c}.json" \ - || echo "[bench] FAILED isl=$ISL osl=$OSL conc=$c" -done -echo "[bench] done -> $OUT" diff --git a/scripts/MiniMax-M3/atom_disagg/launch.sh b/scripts/MiniMax-M3/atom_disagg/launch.sh deleted file mode 100644 index 90fb181b..00000000 --- a/scripts/MiniMax-M3/atom_disagg/launch.sh +++ /dev/null @@ -1,132 +0,0 @@ -#!/bin/bash -# launch.sh — multi-node ATOM PD-disagg launcher for MiniMax-M3 MXFP4 (the `job.slurm`-equiv). -# Drives the VENDORED, verbatim InferenceX server_atom.sh across N nodes via docker+srun. -# -# Faithful to the InferenceX chain: -# minimaxm3_..._atom-disagg.sh -> submit.sh -> job.slurm -> server_atom.sh -# We replace submit.sh+job.slurm (the slurm plumbing) with this; server_atom.sh runs unmodified. -# -# Env (with defaults): -# TOPO=1p1d|2p1d_dpa topology (sets xP/yD/dp-attn) [default 1p1d] -# ISL, OSL bench shape [default 8192 1024] -# CONC_LIST space-separated concurrencies [topology default] -# ACTION=bench|serve|dry bench (default) | hold endpoint | print cmds only -# -# Topologies (from minimaxm3-fp4-mi355x-atom-disagg). All workers TP4 (no EP); STP only. -# 1p1d : xP=1 yD=1, 2 nodes, TP4, dp-attn OFF (conc 1-256; shapes 8192/1024 & 1024/1024) -# 2p1d_dpa : xP=2 yD=1, 3 nodes, TP4, dp-attn ON (conc 256-1024; shape 8192/1024) -set -uo pipefail -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # scripts/MiniMax-M3/atom_disagg -ROOT="$(cd "$HERE/.." && pwd)" # scripts/MiniMax-M3 -cfg(){ python3 "$ROOT/lib/cfg.py" "$@"; } -CLUSTER="$ROOT/cluster.yaml" - -RESV=$(cfg "$CLUSTER" cluster reservation) -MROOT=$(cfg "$CLUSTER" cluster models_root) -RROOT=$(cfg "$CLUSTER" cluster results_root) -RPORT=$(cfg "$CLUSTER" cluster router_port) -PPORT=$(cfg "$CLUSTER" cluster prefill_port) -DPORT=$(cfg "$CLUSTER" cluster decode_port) -HPORT=$(cfg "$CLUSTER" cluster handshake_port) -IMG=$(cfg "$ROOT/model.yaml" engine atom-disagg image) -MODEL_NAME="MiniMax-M3-MXFP4" - -TOPO="${TOPO:-1p1d}"; ACTION="${ACTION:-bench}" -ISL="${ISL:-8192}"; OSL="${OSL:-1024}" -case "$TOPO" in - 1p1d) xP=1; yD=1; DPATTN=0; DEF_CONC="1 2 4 8 16 32 64 128 256" ;; - 2p1d_dpa) xP=2; yD=1; DPATTN=1; DEF_CONC="256 512 768 1024" ;; - *) echo "ERROR: TOPO must be 1p1d|2p1d_dpa (got $TOPO)"; exit 1 ;; -esac -CONC_LIST="${CONC_LIST:-$DEF_CONC}" -# server_atom.sh expects BENCH_MAX_CONCURRENCY 'x'-separated (it does `tr 'x' '\n'|sort -n|tail -1` -# to derive the decode server's --max-num-seqs). Our bench.sh wants space-separated. Keep both. -CONC_X=$(echo "$CONC_LIST" | tr ' ' 'x') -TP=4 # TP4 throughout (M3 on gfx950: plain TP4 beats TP8 / TP4-EP) -NNODES=$(( xP + yD )) # 1 node/worker at TP4 on 8-GPU nodes - -# --- node selection: explicit NODES= override (REQUIRED for parallel launches to avoid -# two launches racing to grab the same 'first free' nodes), else auto-pick free ones. ---- -if [[ -n "${NODES:-}" ]]; then - IFS=',' read -ra NODES <<< "$NODES" - [[ ${#NODES[@]} -lt $NNODES ]] && { echo "ERROR: NODES= has ${#NODES[@]}, need $NNODES"; exit 1; } - NODES=("${NODES[@]:0:$NNODES}") -else - mapfile -t ALL < <(scontrol show hostnames "$(scontrol show res "$RESV" 2>/dev/null | grep -oE 'Nodes=[^ ]+' | head -1 | cut -d= -f2)") - BUSY=$(squeue -u "$USER" -h -o "%N" 2>/dev/null | scontrol show hostnames 2>/dev/null | sort -u) - NODES=(); for n in "${ALL[@]}"; do grep -qx "$n" <<<"$BUSY" || NODES+=("$n"); [[ ${#NODES[@]} -ge $NNODES ]] && break; done - [[ ${#NODES[@]} -lt $NNODES ]] && { echo "ERROR: need $NNODES free nodes, have ${#NODES[@]}"; exit 1; } -fi - -STAMP=$(python3 -c "import time;print(time.strftime('%Y%m%d_%H%M%S'))") -RUN="$RROOT/minimaxm3_atom_${TOPO}_isl${ISL}_${STAMP}"; mkdir -p "$RUN"; chmod 777 "$RUN" -echo "[launch] TOPO=$TOPO xP=$xP yD=$yD nodes=${NODES[*]} shape=${ISL}/${OSL} conc='$CONC_LIST'" -echo "[launch] image=$IMG run=$RUN" - -# --- resolve node IPs (IPADDRS = prefill nodes first, then decode; server_atom.sh slices it) - -declare -a IPS=() -for n in "${NODES[@]}"; do - ip=$(srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=1G \ - bash -c "ip route get 1.1.1.1 2>/dev/null | awk '/src/{print \$7}'" 2>/dev/null \ - | grep -oE '([0-9]+\.){3}[0-9]+' | head -1) - IPS+=("$ip"); echo " $n -> $ip" -done -IPADDRS=$(IFS=,; echo "${IPS[*]}") -NODE0="${NODES[0]}"; NODE0_IP="${IPS[0]}" - -# --- pre-flight clean every node ------------------------------------------------------------ -for n in "${NODES[@]}"; do - srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=0 \ - bash "$ROOT/lib/clean_node.sh" > "$RUN/clean_${n}.log" 2>&1 && echo " [clean] $n OK" || echo " [clean] $n WARN" -done - -DRY=0; [[ "$ACTION" == "dry" ]] && DRY=1 - -# --- RDMA device flags (THE fix for mooncake 30s timeouts) ---------------------------------- -# The ATOM disagg image (rocm/atom-dev:nightly_202607011530) carries its own bnxt_re RDMA stack. -# The fix for mooncake handshake timeouts is exposing --device=/dev/infiniband + --ulimit -# memlock=-1, NOT bind-mounting host libs. DO NOT bind-mount host libibverbs over the -# container's: host and container can have incompatible IBVERBS_PRIVATE ABIs, so mounting host -# libs BREAKS the container's RDMA tools. Expose devices + memlock only. -RDMA_DEVICES="--device=/dev/infiniband --ulimit memlock=-1 --ulimit stack=67108864 --init" -RDMA_MOUNT_SNIPPET='RM=""' # no host-lib mounts (container stack is self-sufficient) - -# --- launch server_atom.sh on each node, NODE_RANK = its index ------------------------------ -# server_atom.sh logic: rank 0 = prefill0 + atomesh router (+ bench); rank=NODE_OFFSET = decode. We pass the SAME env to every node; the script self-selects by rank. -rank=0 -for n in "${NODES[@]}"; do - rdir="$RUN/node${rank}_${n}"; mkdir -p "$rdir"; chmod 777 "$rdir" - srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --exclusive bash -c " - docker rm -f mm3_atom_n${rank} >/dev/null 2>&1 || true - $RDMA_MOUNT_SNIPPET # builds \$RM (host RDMA lib bind-mounts) on THIS node - echo \"[rdma] mounts: \$RM\" - docker run --rm --name mm3_atom_n${rank} --entrypoint bash \ - --device=/dev/kfd --device=/dev/dri $RDMA_DEVICES \$RM -v /sys:/sys \ - --group-add video --ipc host --shm-size 128G \ - --network host --privileged --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ - -e ATOM_WS_PATH=/ws -e NODE_RANK=${rank} -e NODE0_ADDR=${NODE0_IP} -e IPADDRS='${IPADDRS}' \ - -e MODEL_DIR=${MROOT} -e MODEL_NAME=${MODEL_NAME} \ - -e xP=${xP} -e yD=${yD} -e GPUS_PER_NODE=8 \ - -e PREFILL_TP_SIZE=${TP} -e DECODE_TP_SIZE=${TP} \ - -e PREFILL_ENABLE_DP=$([[ $DPATTN == 1 ]] && echo true || echo false) -e PREFILL_ENABLE_EP=1 \ - -e DECODE_ENABLE_DP=$([[ $DPATTN == 1 ]] && echo true || echo false) -e DECODE_ENABLE_EP=1 \ - -e PREFILL_PORT=${PPORT} -e DECODE_PORT=${DPORT} -e ROUTER_PORT=${RPORT} -e HANDSHAKE_PORT=${HPORT} \ - -e MEM_FRAC_STATIC=0.8 -e KV_CACHE_DTYPE=auto -e BLOCK_SIZE=128 -e MAX_NUM_SEQS=256 \ - -e MAX_MODEL_LEN=32768 -e MAX_NUM_BATCHED_TOKENS=32768 \ - -e SPEC_DECODING= -e DECODE_MTP_SIZE=0 \ - -e AITER_QUICK_REDUCE_QUANTIZATION=INT4 \ - -e BENCH_INPUT_LEN=${ISL} -e BENCH_OUTPUT_LEN=${OSL} -e BENCH_MAX_CONCURRENCY='${CONC_X}' \ - -e BENCH_REQUEST_RATE=inf -e BENCH_RANDOM_RANGE_RATIO=1 -e BENCH_NUM_PROMPTS_MULTIPLIER=10 \ - -e DRY_RUN=${DRY} -e SLURM_JOB_ID=${STAMP} \ - -v ${MROOT}:${MROOT} -v ${HERE}:/ws -v ${ROOT}:/scripts -v ${rdir}:/run_logs/slurm_job-${STAMP} \ - ${IMG} /ws/server_atom.sh - " > "$rdir/srun.log" 2>&1 & - echo " [launch] node${rank} ($n) server_atom.sh dispatched" - rank=$(( rank + 1 )) -done - -echo "[launch] all $NNODES nodes dispatched. node0=$NODE0 hosts atomesh router :$RPORT + drives bench." -echo "[launch] tail node0: $RUN/node0_${NODE0}/srun.log" -wait -echo "[launch] done -> $RUN (results: disagg_*.json under node0_*/)" diff --git a/scripts/MiniMax-M3/atom_disagg/setup_deps.sh b/scripts/MiniMax-M3/atom_disagg/setup_deps.sh deleted file mode 100644 index 76b993c1..00000000 --- a/scripts/MiniMax-M3/atom_disagg/setup_deps.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/bash -# setup_deps.sh — sourced by server_atom.sh at startup. -# -# In the InferenceX CI this installs/patches runtime deps into the container. The M3 ATOM -# disagg image (rocm/atom-dev:nightly_202607011530) is expected to ship everything needed: -# atomesh (/usr/local/bin/atomesh), mooncake (py3.12 site-packages + -# /usr/local/bin/mooncake_{master,client}), atom.entrypoints.openai_server. -# -# So this is intentionally a NO-OP stub. If a future image is missing a dep, add the -# install here. Kept as a separate file so server_atom.sh runs unmodified (verbatim vendor). -: diff --git a/scripts/MiniMax-M3/cluster.yaml b/scripts/MiniMax-M3/cluster.yaml deleted file mode 100644 index 71e48e6e..00000000 --- a/scripts/MiniMax-M3/cluster.yaml +++ /dev/null @@ -1,28 +0,0 @@ -# cluster.yaml — MiniMax-M3 disagg cluster definition. SCALE LEVER for multi-node PD. -# M3 MXFP4 disagg spans MULTIPLE nodes per serving instance: xP prefill nodes + yD decode -# nodes + 1 router (on prefill node 0). All workers TP4 => 1 node/worker on 8-GPU nodes, -# so 1p1d=2 nodes and 2p1d_dpa=3 nodes (the InferenceX mi355x-disagg pool is 3 nodes). - -reservation: -gpus_per_node: 8 # gfx950 (MI355X) -hardware: MI355X - -# Model staging root (same path on every node; local NVMe recommended). -# Stage amd/MiniMax-M3-MXFP4 to /MiniMax-M3-MXFP4 (matches model.yaml dir). -models_root: /path/to/models - -results_root: ./results - -# Disagg ports (match InferenceX server_atom.sh defaults) -router_port: 8000 -prefill_port: 8010 -decode_port: 8020 -handshake_port: 6301 - -# Node pool. The launcher picks the first -# (xP*pf_nodes_per_worker + yD*dec_nodes_per_worker) nodes for a given topology. -# 2p1d_dpa needs 3 nodes; list at least that many. -nodes: - - - - - - diff --git a/scripts/MiniMax-M3/configs/minimaxm3-fp4-mi355x-atom-disagg.yaml b/scripts/MiniMax-M3/configs/minimaxm3-fp4-mi355x-atom-disagg.yaml new file mode 100644 index 00000000..033ca902 --- /dev/null +++ b/scripts/MiniMax-M3/configs/minimaxm3-fp4-mi355x-atom-disagg.yaml @@ -0,0 +1,74 @@ +# Extracted verbatim from SemiAnalysisAI/InferenceX configs/amd-master.yaml +# Config key: minimaxm3-fp4-mi355x-atom-disagg +minimaxm3-fp4-mi355x-atom-disagg: + image: rocm/atom-dev:nightly_202607011530 + model: amd/MiniMax-M3-MXFP4 + model-prefix: minimaxm3 + runner: mi355x-disagg + precision: fp4 + framework: atom-disagg + router: { name: atomesh, version: "04b120d3040e0dd1a6915e88a3a3c1f588e6684a" } + kv-p2p-transfer: mooncake + multinode: true + disagg: true + scenarios: + fixed-seq-len: + - isl: 8192 + osl: 1024 + search-space: + # 1P1D TP4 + - conc-list: [ 1, 2, 4, 8, 16, 32, 64, 128, 256 ] + prefill: + num-worker: 1 + tp: 4 + ep: 1 + dp-attn: false + additional-settings: + - "PREFILL_NODES=1" + decode: + num-worker: 1 + tp: 4 + ep: 1 + dp-attn: false + additional-settings: + - "DECODE_NODES=1" + - "DECODE_MTP_SIZE=0" + # 2P1D, DPA TP4 + - conc-list: [ 256, 512, 768, 1024 ] + prefill: + num-worker: 2 + tp: 4 + ep: 1 + dp-attn: true + additional-settings: + - "PREFILL_NODES=2" + decode: + num-worker: 1 + tp: 4 + ep: 1 + dp-attn: true + additional-settings: + - "DECODE_NODES=1" + - "DECODE_MTP_SIZE=0" + # 1P1D TP4 + - isl: 1024 + osl: 1024 + search-space: + # 1P1D TP4 + - conc-list: [ 1, 2, 4, 8, 16, 32, 64, 128, 256 ] + prefill: + num-worker: 1 + tp: 4 + ep: 1 + dp-attn: false + additional-settings: + - "PREFILL_NODES=1" + decode: + num-worker: 1 + tp: 4 + ep: 1 + dp-attn: false + additional-settings: + - "DECODE_NODES=1" + - "DECODE_MTP_SIZE=0" + diff --git a/scripts/MiniMax-M3/docs/proxy_and_disagg.md b/scripts/MiniMax-M3/docs/proxy_and_disagg.md deleted file mode 100644 index 2dc6825b..00000000 --- a/scripts/MiniMax-M3/docs/proxy_and_disagg.md +++ /dev/null @@ -1,52 +0,0 @@ -# MiniMax-M3 disaggregated serving — proxy/router mechanics (from InferenceX) - -MiniMax-M3 (MXFP4) is served **prefill/decode disaggregated (PD-disagg)** across nodes on AMD -MI355X (gfx950): prefill workers (kv_producer) and decode workers (kv_consumer) on separate -nodes, KV streamed over RDMA, fronted by ONE atomesh router. Bench drives the router → a true -aggregate throughput number. - -**All workers are TP4, no EP.** The single-node M3 recipe found plain TP4 beats both TP8 and -TP4/EP4 on tok/s/GPU on gfx950. We tune the prefill:decode worker ratio (xP:yD) instead of TP. -Serving is **STP only** (no MTP): `DECODE_MTP_SIZE=0`. - -## The proxy/router this recipe uses - -| Engine | Router/proxy | KV-transfer backend | Ports | InferenceX source | -|--------|--------------|---------------------|-------|-------------------| -| **ATOM** | **atomesh** (`/usr/local/bin/atomesh launch --pd-disaggregation`) | **mooncake** (RDMA, `--kv-transfer-config` JSON) | router 8000, prefill 8010, decode 8020, handshake 6301 | `server_atom.sh` (`minimaxm3-fp4-mi355x-atom-disagg`) | - -This folder ships the **ATOM MXFP4** path only (the top-throughput MiniMax-M3 disagg config on -MI355X). The InferenceX config family also includes vLLM disagg (MoRI-IO) and MXFP8 variants; -those are out of scope here. - -## ATOM disagg mechanics - -- **Topology**: `xP` prefill workers + `yD` decode workers, each TP4 (1 worker = 1 node on an - 8-GPU node, since ceil(4/8)=1). - - `1p1d` = 2 nodes; `2p1d_dpa` = 2 prefill + 1 decode with DP-attention = 3 nodes (high conc). -- **Prefill server** (kv_producer): - ``` - python3 -m atom.entrypoints.openai_server --model --host 0.0.0.0 --server-port 8010 \ - -tp 4 [--enable-dp-attention] --block-size 128 \ - --gpu-memory-utilization 0.8 --max-num-seqs N --max-model-len 32768 \ - --max-num-batched-tokens 32768 --no-enable_prefix_caching \ - --kv-transfer-config '{"kv_role":"kv_producer","kv_connector":"mooncake","proxy_ip":"","handshake_port":6301}' - ``` -- **Decode server** (kv_consumer): same, `--server-port 8020`, `kv_role":"kv_consumer"`. -- **Router** (on prefill node 0, after all servers healthy): - ``` - atomesh launch --host 0.0.0.0 --port 8000 --pd-disaggregation \ - --prefill http://:8010 [...] --decode http://:8020 [...] \ - --policy random --backend atom --disable-health-check --disable-circuit-breaker - ``` -- **Bench**: drive `http://:8000` (OpenAI-compatible) at aggregate concurrency. -- **M3-specific env**: `AITER_QUICK_REDUCE_QUANTIZATION=INT4` (MXFP4 quick-reduce), applied for - all non-DSv4 models by `server_atom.sh`. No `--enable-tbo` (that is a DSv4-only path). -- KV transfer uses mooncake's own RDMA device selection (IBDEVICES not passed as a server arg). - -## Topologies to benchmark (from minimaxm3-fp4-mi355x-atom-disagg) -| shape | conc | topo | nodes | -|-------|------|------|-------| -| 8192/1024 | 1–256 | 1P1D TP4 | 2 | -| 1024/1024 | 1–256 | 1P1D TP4 | 2 | -| 8192/1024 | 256–1024 | 2P1D + DP-attn TP4 | 3 | diff --git a/scripts/MiniMax-M3/lib/cfg.py b/scripts/MiniMax-M3/lib/cfg.py deleted file mode 100644 index 9ced83da..00000000 --- a/scripts/MiniMax-M3/lib/cfg.py +++ /dev/null @@ -1,34 +0,0 @@ -#!/usr/bin/env python3 -"""cfg.py — tiny yaml field reader for the shell scripts. -Usage: - cfg.py cluster -> cluster.yaml top-level key - cfg.py model -> models.. - cfg.py engine -> engines.. -Lists print space-separated; scalars print as-is. -""" -import sys, yaml - -def emit(v): - if isinstance(v, list): print(" ".join(str(x) for x in v)) - elif isinstance(v, dict): print(" ".join(f"{k}={v}" for k, v in v.items())) - else: print("" if v is None else v) - -def main(): - f = sys.argv[1]; kind = sys.argv[2] - d = yaml.safe_load(open(f)) - if kind == "cluster": - emit(d.get(sys.argv[3])) - elif kind == "model": - emit(d["models"][sys.argv[3]].get(sys.argv[4])) - elif kind == "engine": - emit(d["engines"][sys.argv[3]].get(sys.argv[4])) - elif kind == "default": - # default : per-model override wins over global defaults{} - model = sys.argv[3]; key = sys.argv[4] - mv = d.get("models", {}).get(model, {}).get(key) - emit(mv if mv is not None else d.get("defaults", {}).get(key)) - else: - sys.exit(f"unknown kind {kind}") - -if __name__ == "__main__": - main() diff --git a/scripts/MiniMax-M3/lib/check_accuracy.py b/scripts/MiniMax-M3/lib/check_accuracy.py deleted file mode 100644 index e89571ce..00000000 --- a/scripts/MiniMax-M3/lib/check_accuracy.py +++ /dev/null @@ -1,59 +0,0 @@ -#!/usr/bin/env python3 -# Sanity/accuracy check: send 5 real prompts to a running OpenAI-compatible server, -# verify coherent + expected output. Catches corrupt-weight / bad-quant failures that -# throughput benchmarks (which use --ignore-eos on random tokens) miss. -# -# Usage: PORT, MODEL, OUT(json) via env. Exits 0 if all coherent, 1 otherwise. -import json, os, sys, time, urllib.request - -PORT = os.environ.get("PORT", "8000") -MODEL = os.environ["MODEL"] -OUT = os.environ.get("OUT", "/results/sanity.json") -HERE = os.path.dirname(os.path.abspath(__file__)) -prompts = json.load(open(os.path.join(HERE, "prompts.json"))) - -def chat(prompt, max_tokens): - body = json.dumps({ - "model": MODEL, - "messages": [{"role": "user", "content": prompt}], - "max_tokens": max_tokens, - "temperature": 0.0, - }).encode() - req = urllib.request.Request(f"http://0.0.0.0:{PORT}/v1/chat/completions", - data=body, headers={"Content-Type": "application/json"}) - t0 = time.time() - with urllib.request.urlopen(req, timeout=120) as r: - d = json.load(r) - return d["choices"][0]["message"]["content"], time.time() - t0 - -results = [] -allpass = True -for p in prompts: - try: - text, dt = chat(p["prompt"], p["max_tokens"]) - low = text.lower() - hits = [k for k in p["expect_contains"] if k.lower() in low] - # coherent = non-empty, not pure repetition/garbage - coherent = len(text.strip()) > 0 and len(set(text.split())) > 1 or len(text.strip()) <= 8 - ok = len(hits) > 0 and coherent - results.append({"id": p["id"], "ok": ok, "hits": hits, - "expected": p["expect_contains"], "latency_s": round(dt, 2), - "output": text.strip()[:200]}) - if not ok: - allpass = False - except Exception as e: - results.append({"id": p["id"], "ok": False, "error": str(e)[:150]}) - allpass = False - -summary = {"model": MODEL, "all_pass": allpass, - "passed": sum(1 for r in results if r.get("ok")), - "total": len(results), "results": results} -os.makedirs(os.path.dirname(OUT), exist_ok=True) -json.dump(summary, open(OUT, "w"), indent=2) - -print(f"=== SANITY: {summary['passed']}/{summary['total']} prompts passed ===") -for r in results: - mark = "PASS" if r.get("ok") else "FAIL" - out = r.get("output", r.get("error", "")) - print(f" [{mark}] {r['id']:12s} -> {out[:90]!r}") -sys.exit(0 if allpass else 1) diff --git a/scripts/MiniMax-M3/lib/clean_node.sh b/scripts/MiniMax-M3/lib/clean_node.sh deleted file mode 100644 index 6d1aea5a..00000000 --- a/scripts/MiniMax-M3/lib/clean_node.sh +++ /dev/null @@ -1,38 +0,0 @@ -#!/bin/bash -# clean_node.sh — ensure a node's GPUs are in a CLEAN state before deploying an engine. -# Run via srun on the target node (exclusive). Kills leftover serving containers/processes -# that squat VRAM (the #1 cause of OOM-at-init), then verifies VRAM is actually free. -# -# Usage (inside srun on the node): clean_node.sh [max_used_gb_threshold] -# default threshold = 5 GB; exits non-zero if any GPU still above it after cleanup. -set -uo pipefail -THRESH_GB="${1:-5}" - -echo "[clean] $(hostname): removing ALL containers + stray GPU procs (reservation is exclusive to us)" -# The reservation is slurm-exclusive to us, but Docker containers are NOT slurm-tracked, -# so orphaned/stale containers from prior runs can linger and squat VRAM. Remove them ALL. -cids=$(docker ps -aq 2>/dev/null) -[ -n "$cids" ] && { docker kill $cids >/dev/null 2>&1; docker rm -f $cids >/dev/null 2>&1; } || true -# belt-and-suspenders: kill stray engine procs that may hold VRAM -pkill -9 -f "vllm|atom.entrypoints|sglang.launch|EngineCore|sglang.launch_server" 2>/dev/null || true -sleep 5 - -# 3) verify VRAM free on every GPU. -# rocm-smi prints one line per GPU: "GPU[N] : VRAM Total Used Memory (B): ". -# Take the LAST integer on each line (the byte count), not the GPU index. -echo "[clean] verifying VRAM free (threshold ${THRESH_GB}GB)" -busy=0; i=0 -while IFS= read -r line; do - b=$(echo "$line" | grep -oE "[0-9]+" | tail -1) - [ -z "$b" ] && continue - gb=$(( b / 1000000000 )) - if [ "$gb" -gt "$THRESH_GB" ]; then echo " GPU$i: ${gb}GB USED (>threshold)"; busy=1; else echo " GPU$i: ${gb}GB ok"; fi - i=$((i+1)) -done < <(rocm-smi --showmeminfo vram 2>/dev/null | grep -i "Used Memory") -if [ "$busy" -ne 0 ]; then - echo "[clean] WARNING: some GPUs still hold VRAM after cleanup — listing holders:" - rocm-smi --showpids 2>/dev/null | grep -iE "vllm|atom|sglang|EngineCore" | head - exit 1 -fi -echo "[clean] $(hostname): all GPUs clean" -exit 0 diff --git a/scripts/MiniMax-M3/lib/lib_inferencex.sh b/scripts/MiniMax-M3/lib/lib_inferencex.sh deleted file mode 100644 index 95e063a3..00000000 --- a/scripts/MiniMax-M3/lib/lib_inferencex.sh +++ /dev/null @@ -1,1281 +0,0 @@ -#!/usr/bin/env bash - -# Shared benchmarking utilities for InferenceX - -# Keep Python bytecode out of the mounted workspace. Benchmark jobs often run as -# root inside containers, and root-owned cache directories break future checkout -# cleanup on self-hosted runners. -export PYTHONDONTWRITEBYTECODE=1 -export PYTHONPYCACHEPREFIX="${PYTHONPYCACHEPREFIX:-/tmp/inferencex-pycache}" -mkdir -p "$PYTHONPYCACHEPREFIX" 2>/dev/null || true - -# Inference server port shared by every benchmark recipe. Launchers that need -# a non-default value (e.g. launch_mi355x-amds.sh derives PORT from RUNNER_NAME -# to avoid collisions across concurrent gh-runners on a shared host) set PORT -# themselves before sourcing this file; the `:-` fallback only kicks in when -# nothing upstream set it. -export PORT="${PORT:-8888}" - -# -------------------------------- -# GPU monitoring helpers -# -------------------------------- - -GPU_MONITOR_PID="" -GPU_METRICS_CSV="/workspace/gpu_metrics.csv" -export GPU_METRICS_CSV - -# Start background GPU monitoring that logs metrics every second to CSV. -# Auto-detects NVIDIA (nvidia-smi) or AMD (amd-smi) GPUs. -# Usage: start_gpu_monitor [--output /path/to/output.csv] [--interval 1] -start_gpu_monitor() { - local output="$GPU_METRICS_CSV" - local interval=1 - - while [[ $# -gt 0 ]]; do - case $1 in - --output) output="$2"; shift 2 ;; - --interval) interval="$2"; shift 2 ;; - *) shift ;; - esac - done - - GPU_METRICS_CSV="$output" - export GPU_METRICS_CSV - - if command -v nvidia-smi &>/dev/null; then - nvidia-smi --query-gpu=timestamp,index,power.draw,temperature.gpu,clocks.current.sm,clocks.current.memory,utilization.gpu,utilization.memory \ - --format=csv -l "$interval" > "$output" 2>/dev/null & - GPU_MONITOR_PID=$! - echo "[GPU Monitor] Started NVIDIA (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" - elif command -v amd-smi &>/dev/null; then - # Use amd-smi native watch mode (-w) which includes timestamps automatically. - # Pipe through awk to: skip preamble lines, keep first CSV header, skip repeated headers. - amd-smi metric -p -c -t -u -w "$interval" --csv 2>/dev/null \ - | awk '/^timestamp,/{if(!h){print;h=1};next} h{print}' > "$output" & - GPU_MONITOR_PID=$! - echo "[GPU Monitor] Started AMD (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" - else - echo "[GPU Monitor] No GPU monitoring tool found (nvidia-smi or amd-smi), skipping" - return 0 - fi -} - -# Stop the background GPU monitor and report file size. -stop_gpu_monitor() { - if [[ -n "$GPU_MONITOR_PID" ]] && kill -0 "$GPU_MONITOR_PID" 2>/dev/null; then - kill "$GPU_MONITOR_PID" 2>/dev/null - wait "$GPU_MONITOR_PID" 2>/dev/null || true - echo "[GPU Monitor] Stopped (PID=$GPU_MONITOR_PID)" - if [[ -f "$GPU_METRICS_CSV" ]]; then - local lines - lines=$(wc -l < "$GPU_METRICS_CSV") - echo "[GPU Monitor] Collected $lines rows -> $GPU_METRICS_CSV" - fi - fi - GPU_MONITOR_PID="" -} - -# Check if required environment variables are set -# Usage: check_env_vars VAR1 VAR2 VAR3 ... -# Exits with code 1 if any variable is not set -check_env_vars() { - local missing_vars=() - - for var_name in "$@"; do - if [[ -z "${!var_name:-}" ]]; then - missing_vars+=("$var_name") - fi - done - - if [[ ${#missing_vars[@]} -gt 0 ]]; then - echo "Error: The following required environment variables are not set:" - for var in "${missing_vars[@]}"; do - echo " - $var" - done - exit 1 - fi -} - -# Wait for server to be ready by polling the health endpoint -# All parameters are required -# Parameters: -# --port: Server port -# --server-log: Path to server log file -# --server-pid: Server process ID (required) -# --sleep-interval: Sleep interval between health checks (optional, default: 5) -wait_for_server_ready() { - set +x - local port="" - local server_log="" - local server_pid="" - local sleep_interval=5 - - # Parse arguments - while [[ $# -gt 0 ]]; do - case $1 in - --port) - port="$2" - shift 2 - ;; - --server-log) - server_log="$2" - shift 2 - ;; - --server-pid) - server_pid="$2" - shift 2 - ;; - --sleep-interval) - sleep_interval="$2" - shift 2 - ;; - *) - echo "Unknown parameter: $1" - return 1 - ;; - esac - done - - # Validate required parameters - if [[ -z "$port" ]]; then - echo "Error: --port is required" - return 1 - fi - if [[ -z "$server_log" ]]; then - echo "Error: --server-log is required" - return 1 - fi - if [[ -z "$server_pid" ]]; then - echo "Error: --server-pid is required" - return 1 - fi - - # Wait for server log file to be created (container startup may delay this) - while [ ! -f "$server_log" ]; do - if ! kill -0 "$server_pid" 2>/dev/null; then - echo "Server died before creating log file. Exiting." - exit 1 - fi - sleep 1 - done - - # Show logs until server is ready - tail -f -n +1 "$server_log" & - local TAIL_PID=$! - until curl --output /dev/null --silent --fail http://0.0.0.0:$port/health; do - if ! kill -0 "$server_pid" 2>/dev/null; then - echo "Server died before becoming healthy. Exiting." - kill $TAIL_PID - exit 1 - fi - sleep "$sleep_interval" - done - kill $TAIL_PID -} - -# Run benchmark serving with standardized parameters -# All parameters are required except --endpoint, --use-chat-template, --dsv4, and --trust-remote-code -# Parameters: -# --model: Model name -# --port: Server port -# --backend: Backend type - e.g., 'vllm' or 'openai' -# --endpoint: Optional API endpoint override -# --input-len: Random input sequence length -# --output-len: Random output sequence length -# --random-range-ratio: Random range ratio -# --num-prompts: Number of prompts -# --max-concurrency: Max concurrency -# --result-filename: Result filename without extension -# --result-dir: Result directory -# --use-chat-template: Optional flag to enable chat template -# --dsv4: Optional flag to use the DeepSeek-V4 chat template -# (encoding_dsv4.py) instead of the tokenizer's built-in jinja -# template. Implies --use-chat-template. -# --trust-remote-code: Optional flag to trust remote code from HuggingFace -# --server-pid: Optional server process ID to monitor during benchmark -run_benchmark_serving() { - # In eval-only mode, skip the throughput benchmark entirely. - if [ "${EVAL_ONLY}" = "true" ]; then - echo "EVAL_ONLY mode: skipping throughput benchmark" - return 0 - fi - - set +x - local model="" - local port="" - local backend="" - local endpoint="" - local input_len="" - local output_len="" - local random_range_ratio="" - local num_prompts="" - local max_concurrency="" - local result_filename="" - local result_dir="" - local workspace_dir="" - local use_chat_template=false - local dsv4=false - local trust_remote_code=false - local server_pid="" - local tokenizer="" - - while [[ $# -gt 0 ]]; do - case $1 in - --model) - model="$2" - shift 2 - ;; - --port) - port="$2" - shift 2 - ;; - --backend) - backend="$2" - shift 2 - ;; - --endpoint) - endpoint="$2" - shift 2 - ;; - --input-len) - input_len="$2" - shift 2 - ;; - --output-len) - output_len="$2" - shift 2 - ;; - --random-range-ratio) - random_range_ratio="$2" - shift 2 - ;; - --num-prompts) - num_prompts="$2" - shift 2 - ;; - --max-concurrency) - max_concurrency="$2" - shift 2 - ;; - --result-filename) - result_filename="$2" - shift 2 - ;; - --result-dir) - result_dir="$2" - shift 2 - ;; - --bench-serving-dir) - workspace_dir="$2" - shift 2 - ;; - --use-chat-template) - use_chat_template=true - shift - ;; - --dsv4) - dsv4=true - use_chat_template=true - shift - ;; - --trust-remote-code) - trust_remote_code=true - shift - ;; - --server-pid) - server_pid="$2" - shift 2 - ;; - --tokenizer) - tokenizer="$2" - shift 2 - ;; - *) - echo "Unknown parameter: $1" - return 1 - ;; - esac - done - - # Validate all required parameters - if [[ -z "$model" ]]; then - echo "Error: --model is required" - return 1 - fi - if [[ -z "$port" ]]; then - echo "Error: --port is required" - return 1 - fi - if [[ -z "$backend" ]]; then - echo "Error: --backend is required" - return 1 - fi - if [[ -z "$input_len" ]]; then - echo "Error: --input-len is required" - return 1 - fi - if [[ -z "$output_len" ]]; then - echo "Error: --output-len is required" - return 1 - fi - if [[ -z "$random_range_ratio" ]]; then - echo "Error: --random-range-ratio is required" - return 1 - fi - if [[ -z "$num_prompts" ]]; then - echo "Error: --num-prompts is required" - return 1 - fi - if [[ -z "$max_concurrency" ]]; then - echo "Error: --max-concurrency is required" - return 1 - fi - if [[ -z "$result_filename" ]]; then - echo "Error: --result-filename is required" - return 1 - fi - if [[ -z "$result_dir" ]]; then - echo "Error: --result-dir is required" - return 1 - fi - - if [[ -z "$workspace_dir" ]]; then - workspace_dir=$(pwd) - fi - - # Profiling support: when PROFILE=1, ensure profiler dir exists, add --profile flag, - # and cap num_prompts to keep traces small. - local profile_flag=() - if [[ "${PROFILE:-}" == "1" ]]; then - local _prof_dir="${SGLANG_TORCH_PROFILER_DIR:-${VLLM_TORCH_PROFILER_DIR:-}}" - if [[ -n "$_prof_dir" ]]; then - mkdir -p "$_prof_dir" - fi - profile_flag+=(--profile) - num_prompts="$max_concurrency" - fi - - # Build benchmark command - local benchmark_cmd=( - python3 "$workspace_dir/utils/bench_serving/benchmark_serving.py" - --model "$model" - --backend "$backend" - --base-url "http://0.0.0.0:$port" - --dataset-name random - --random-input-len "$input_len" - --random-output-len "$output_len" - --random-range-ratio "$random_range_ratio" - --num-prompts "$num_prompts" - --max-concurrency "$max_concurrency" - --request-rate inf - --ignore-eos - "${profile_flag[@]}" - --save-result - --num-warmups "$((2 * max_concurrency))" \ - --percentile-metrics 'ttft,tpot,itl,e2el' - --result-dir "$result_dir" - --result-filename "$result_filename.json" - ) - - if [[ -n "$endpoint" ]]; then - benchmark_cmd+=(--endpoint "$endpoint") - fi - - # Add --use-chat-template if requested - if [[ "$use_chat_template" == true ]]; then - benchmark_cmd+=(--use-chat-template) - fi - - # Add --dsv4 if requested (requires --use-chat-template, which we - # auto-enable when --dsv4 is passed in). - if [[ "$dsv4" == true ]]; then - benchmark_cmd+=(--dsv4) - fi - - # Add --trust-remote-code if requested - if [[ "$trust_remote_code" == true ]]; then - benchmark_cmd+=(--trust-remote-code) - fi - - if [[ -n "$tokenizer" ]]; then - benchmark_cmd+=(--tokenizer "$tokenizer") - fi - - # Run benchmark with optional server monitoring - set -x - if [[ -n "$server_pid" ]]; then - # Run benchmark in background and monitor server health - "${benchmark_cmd[@]}" & - local benchmark_pid=$! - - # Monitor loop: check both benchmark and server status - while kill -0 "$benchmark_pid" 2>/dev/null; do - if ! kill -0 "$server_pid" 2>/dev/null; then - echo "ERROR: Server process $server_pid died during benchmark" - kill "$benchmark_pid" 2>/dev/null - wait "$benchmark_pid" 2>/dev/null - set +x - return 1 - fi - sleep 2 - done - - # Benchmark finished, get its exit code - wait "$benchmark_pid" - local benchmark_exit_code=$? - else - # No server monitoring, run benchmark directly - "${benchmark_cmd[@]}" - local benchmark_exit_code=$? - fi - set +x - - # If profiling, move trace to relay-upload location - if [[ "${PROFILE:-}" == "1" ]]; then - move_profile_trace_for_relay - fi - - return $benchmark_exit_code -} - - -# -------------------------------- -# Profiling trace helpers -# -------------------------------- - -_find_latest_profile_trace() { - local latest="" - local dir="" candidate="" base="" - local -a search_roots=() - - for dir in "$@"; do - search_roots=() - if [[ -d "$dir" ]]; then - search_roots+=("$dir") - fi - if [[ -d "$dir/profiles" ]]; then - search_roots+=("$dir/profiles") - fi - if [[ ${#search_roots[@]} -eq 0 ]]; then - continue - fi - - while IFS= read -r -d '' candidate; do - base="$(basename "$candidate")" - if [[ "$base" == profile_*.trace.json.gz ]]; then - continue - fi - if [[ -z "$latest" || "$candidate" -nt "$latest" ]]; then - latest="$candidate" - fi - done < <( - find "${search_roots[@]}" -maxdepth 1 -type f \ - \( -name "*.trace.json" -o -name "*.trace.json.gz" -o -name "*trace*.json" -o -name "*trace*.json.gz" -o -name "*profile*.json" -o -name "*profile*.json.gz" \) \ - -print0 2>/dev/null - ) - done - - printf '%s' "$latest" -} - -# Move profiler trace into a stable workspace path for workflow relay/upload. -move_profile_trace_for_relay() { - if [[ "${PROFILE:-}" != "1" ]]; then - return 0 - fi - - if [[ -z "${RESULT_FILENAME:-}" ]]; then - echo "[PROFILE] RESULT_FILENAME is not set; skipping relay trace staging." >&2 - return 0 - fi - - local sglang_dir="${SGLANG_TORCH_PROFILER_DIR:-/workspace}" - local vllm_dir="${VLLM_TORCH_PROFILER_DIR:-/workspace}" - local -a search_dirs=() - local dir="" existing="" - local seen=0 - - for dir in "$sglang_dir" "$vllm_dir" "/workspace"; do - if [[ -z "$dir" ]]; then - continue - fi - seen=0 - for existing in "${search_dirs[@]}"; do - if [[ "$existing" == "$dir" ]]; then - seen=1 - break - fi - done - if [[ "$seen" -eq 0 ]]; then - search_dirs+=("$dir") - fi - done - - local trace_file="" - local wait_attempts=10 - for (( i=1; i<=wait_attempts; i++ )); do - trace_file="$(_find_latest_profile_trace "${search_dirs[@]}")" - if [[ -n "$trace_file" ]]; then - break - fi - sleep 10 - done - - if [[ -z "$trace_file" ]]; then - echo "[PROFILE] No trace found for relay under: ${search_dirs[*]}" >&2 - return 0 - fi - - local dest_trace="/workspace/profile_${RESULT_FILENAME}.trace.json.gz" - if [[ "$trace_file" == *.gz ]]; then - cp -f "$trace_file" "$dest_trace" - else - gzip -c "$trace_file" > "$dest_trace" - fi - - echo "[PROFILE] Relay trace prepared: $dest_trace (source: $trace_file)" -} - - -# ------------------------------ -# Eval (lm-eval-harness) helpers -# ------------------------------ - -_install_lm_eval_deps() { - # torchvision causes circular imports in ATOM; TRT-LLM/SGLang need it at module level. - if [[ "${IMAGE:-}" == *atom* ]]; then - python3 -m pip uninstall -y torchvision 2>/dev/null || true - fi - python3 -m pip install -q --no-cache-dir --break-system-packages "lm-eval[api]" || true - local lm_eval_ref="b315ef3b05176acc9732bb7fdec116abe1ecc476" - if command -v git >/dev/null 2>&1; then - if ! python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ - "git+https://github.com/EleutherAI/lm-evaluation-harness.git@${lm_eval_ref}"; then - python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ - "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true - fi - else - python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ - "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true - fi -} - -# Patch lm-eval filters to be robust to empty strings via sitecustomize -_patch_lm_eval() { - local patch_dir - patch_dir="$(mktemp -d)" - cat > "$patch_dir/sitecustomize.py" <<'PY' -# --- Patch LocalChatCompletion.parse_generations to handle empty content with reasoning_content --- -import re, sys, unicodedata, json -from lm_eval.filters import extraction as ex -from lm_eval.models.openai_completions import LocalChatCompletion as _LCC - -def _le_parse_generations(outputs, **kwargs): - res = [] - if not isinstance(outputs, list): - outputs = [outputs] - for out in (outputs or []): - try: - choices = out.get("choices", []) - tmp = ["" for _ in choices] - for choice in choices: - idx = choice.get("index", 0) - msg = (choice.get("message") or {}) - content = msg.get("content") - if content in (None, "", []): - content = msg.get("reasoning_content") or "" - tmp[idx] = content - except Exception: - tmp = [""] - res.extend(tmp) - return res - -# Keep staticmethod semantics -_LCC.parse_generations = staticmethod(_le_parse_generations) - -# --- Patch TemplateAPI.apply_chat_template to avoid injecting "type": "text" for TRT --- -try: - from lm_eval.models import api_models as _api_models - _TemplateAPI = _api_models.TemplateAPI - _JsonChatStr = _api_models.JsonChatStr -except Exception: - _TemplateAPI = None - _JsonChatStr = None - -if _TemplateAPI is not None and _JsonChatStr is not None: - _orig_apply_chat_template = _TemplateAPI.apply_chat_template - - def _patched_apply_chat_template( - self, - chat_history, - add_generation_prompt: bool = True, - ): - """Applies a chat template to a list of chat history between user and model.""" - if self.tokenizer_backend == "huggingface" and self.tokenized_requests: - return self.tokenizer.apply_chat_template( - chat_history, - tokenize=False, - add_generation_prompt=add_generation_prompt, - continue_final_message=not add_generation_prompt, - ) - elif self.tokenizer_backend == "remote" and self.tokenized_requests: - return chat_history - else: - # NOTE: we no longer inject `"type": "text"` when tokenizer is None / non-HF - return _JsonChatStr( - json.dumps( - [{**item} for item in chat_history], - ensure_ascii=False, - ) - ) - - _TemplateAPI.apply_chat_template = _patched_apply_chat_template -PY - export PYTHONPATH="${patch_dir}:${PYTHONPATH:-}" -} - -get_native_max_context_length() { - local model_path="$1" - # Prefer MODEL_PATH (local model directory) when available, since the - # argument may be a served-model name that is neither a valid HF repo - # ID nor a local path (e.g. "deepseek-r1-fp4" on the B300 cluster). - if [ -n "${MODEL_PATH:-}" ] && [ -d "${MODEL_PATH}" ]; then - model_path="${MODEL_PATH}" - fi - python3 -c " -try: - from transformers import AutoConfig - config = AutoConfig.from_pretrained('${model_path}', trust_remote_code=True) - for attr in ['max_position_embeddings', 'max_sequence_length', 'seq_length', 'n_positions']: - if hasattr(config, attr): - print(getattr(config, attr)) - break - else: - print(0) -except Exception: - print(0) -" -} - -# Compute the context length for eval-only mode. -# Uses the requested benchmark context capped at the model's native max. -# Sets EVAL_MAX_MODEL_LEN (needed by run_lm_eval). -# Echoes the computed value for scripts to capture. -# -# Usage: local ctx=$(compute_eval_context_length "$MODEL" "${current_ctx}") -compute_eval_context_length() { - local model="$1" - local benchmark_ctx="${2:-0}" - local native_max - native_max=$(get_native_max_context_length "$model") - native_max="${native_max:-0}" - - if [ "$benchmark_ctx" -eq 0 ] 2>/dev/null; then - benchmark_ctx="${native_max:-0}" - fi - local eval_ctx=$(( benchmark_ctx * 1 )) - if [ "$native_max" -gt 0 ] 2>/dev/null && [ "$eval_ctx" -gt "$native_max" ]; then - eval_ctx="$native_max" - fi - # If eval_ctx is still 0 (both benchmark_ctx and native_max were 0), fall back - if [ "$eval_ctx" -le 0 ] 2>/dev/null; then - echo "WARN: compute_eval_context_length could not determine context length for $model" >&2 - eval_ctx="${MAX_MODEL_LEN:-16384}" - fi - EVAL_MAX_MODEL_LEN="$eval_ctx" - echo "$eval_ctx" -} - -# Convenience wrapper: compute eval context from ISL/OSL and export EVAL_MAX_MODEL_LEN. -# Call directly (not in a subshell) so the export persists. -# Scripts then wire $EVAL_MAX_MODEL_LEN into whichever server variable they need. -setup_eval_context() { - EVAL_MAX_MODEL_LEN=$(compute_eval_context_length "$MODEL" "$((ISL + OSL + 256))") - export EVAL_MAX_MODEL_LEN -} - -run_lm_eval() { - local port="${PORT:-8888}" - local tasks_dir="${EVAL_TASKS_DIR:-utils/evals/gsm8k.yaml}" - local results_dir="${EVAL_RESULT_DIR:-$(mktemp -d /tmp/eval_out-XXXXXX)}" - local eval_context_len="${EVAL_MAX_MODEL_LEN:-16384}" - local temperature=0 - local top_p=1 - local concurrent_requests="${EVAL_CONCURRENT_REQUESTS:-${CONC:-64}}" - - while [[ $# -gt 0 ]]; do - case $1 in - --port) port="$2"; shift 2 ;; - --task) tasks_dir="$2"; shift 2 ;; - --results-dir) results_dir="$2"; shift 2 ;; - --gen-max-tokens) eval_context_len="$2"; shift 2 ;; - --temperature) temperature="$2"; shift 2 ;; - --top-p) top_p="$2"; shift 2 ;; - *) echo "Unknown parameter: $1"; return 1 ;; - esac - done - - if [ "${INFERENCEX_LM_EVAL_RUNTIME_READY:-false}" != "true" ]; then - _install_lm_eval_deps - _patch_lm_eval - export INFERENCEX_LM_EVAL_RUNTIME_READY=true - fi - - local openai_server_base="http://0.0.0.0:${port}" - local openai_chat_base="${openai_server_base}/v1/chat/completions" - export OPENAI_API_KEY=${OPENAI_API_KEY:-EMPTY} - MODEL_NAME=${MODEL_NAME:-$MODEL} # Prefer MODEL_NAME, else MODEL - - # Cap output tokens: must fit within context window (leave room for input), - # and avoid excessive KV cache reservation per request on TRT. - local max_output_tokens=$(( eval_context_len > 4096 ? eval_context_len - 4096 : eval_context_len / 2 )) - if [ "$max_output_tokens" -gt 16384 ]; then - max_output_tokens=16384 - fi - echo "Eval budget: eval_context_len=${eval_context_len}, max_output_tokens=${max_output_tokens}" - - # Export for append_lm_eval_summary to pick up - export EVAL_RESULT_DIR="$results_dir" - set -x - python3 -m lm_eval --model local-chat-completions --apply_chat_template \ - --tasks "${tasks_dir}" \ - --output_path "${results_dir}" \ - --log_samples \ - --model_args "model=${MODEL_NAME},base_url=${openai_chat_base},api_key=${OPENAI_API_KEY},eos_string=,max_retries=5,num_concurrent=${concurrent_requests},timeout=1800,tokenized_requests=False,max_length=${eval_context_len}" \ - --gen_kwargs "max_tokens=${max_output_tokens},temperature=${temperature},top_p=${top_p}" - local eval_exit=$? - set +x - return $eval_exit -} - -_stage_lm_eval_artifacts() { - local results_dir="$1" - local eval_conc="$2" - local moved=0 - local failed=0 - local jf base stem extension target suffix - - if [ ! -d "${results_dir}" ]; then - echo "WARN: eval result directory '${results_dir}' does not exist" >&2 - return 1 - fi - - while IFS= read -r -d '' jf; do - base=$(basename "$jf") - case "$base" in - meta_env.json) - continue - ;; - *.jsonl) - stem="${base%.jsonl}" - extension=".jsonl" - ;; - *.json) - stem="${base%.json}" - extension=".json" - ;; - *) - continue - ;; - esac - - target="./${stem}_conc${eval_conc}${extension}" - suffix=2 - while [ -e "$target" ]; do - target="./${stem}_conc${eval_conc}_${suffix}${extension}" - suffix=$((suffix + 1)) - done - - if mv -f "$jf" "$target"; then - moved=1 - else - echo "WARN: failed to stage eval artifact ${jf}" >&2 - failed=1 - fi - done < <( - find "${results_dir}" -type f \ - \( -name "*.json" -o -name "*.jsonl" \) -print0 2>/dev/null - ) - - rm -rf --one-file-system "${results_dir}" 2>/dev/null \ - || rm -rf "${results_dir}" \ - || true - - if [ "$moved" -eq 0 ]; then - echo "WARN: no eval artifacts were produced for concurrency ${eval_conc}" >&2 - return 1 - fi - return "$failed" -} - -_eval_concs_to_json() { - local values="$1" - local value - local joined="" - - for value in $values; do - if [[ ! "$value" =~ ^[1-9][0-9]*$ ]]; then - echo "ERROR: invalid eval concurrency '${value}'" >&2 - return 1 - fi - if [ -n "$joined" ]; then - joined="${joined}, " - fi - joined="${joined}${value}" - done - - printf '[%s]' "$joined" -} - -append_lm_eval_summary() { - local batch_concs="${EVAL_BATCHED_CONCS:-}" - local results_dir="${EVAL_RESULT_DIR:-}" - local out_dir="${results_dir}" - local meta_json - local metadata_conc="${CONC:-1}" - local batch_metadata="" - - if [ -n "$batch_concs" ]; then - meta_json="./meta_env.json" - metadata_conc="${batch_concs%% *}" - - local eval_concs_json completed_concs_json failed_concs_json - eval_concs_json=$(_eval_concs_to_json "$batch_concs") || return 1 - completed_concs_json=$( - _eval_concs_to_json "${EVAL_BATCHED_COMPLETED_CONCS:-}" - ) || return 1 - failed_concs_json=$( - _eval_concs_to_json "${EVAL_BATCHED_FAILED_CONCS:-}" - ) || return 1 - printf -v batch_metadata \ - ' "eval_concs": %s,\n "completed_eval_concs": %s,\n "failed_eval_concs": %s,\n' \ - "$eval_concs_json" \ - "$completed_concs_json" \ - "$failed_concs_json" - else - if [ -z "${results_dir}" ]; then - echo "WARN: EVAL_RESULT_DIR is empty; skipping artifact collection" >&2 - return 1 - fi - if [ ! -d "${out_dir}" ]; then - echo "WARN: EVAL_RESULT_DIR='${out_dir}' does not exist; skipping artifact collection" >&2 - return 1 - fi - meta_json="${out_dir}/meta_env.json" - fi - - # Write minimal meta for collectors that expect it - local model_name="${MODEL_NAME:-$MODEL}" - local is_multinode_json="false" - if [ "${IS_MULTINODE:-false}" = "true" ]; then - is_multinode_json="true" - fi - - local prefill_tp="${PREFILL_TP:-${TP:-1}}" - local prefill_ep="${PREFILL_EP:-${EP_SIZE:-1}}" - local prefill_num_workers="${PREFILL_NUM_WORKERS:-1}" - local decode_tp="${DECODE_TP:-${TP:-1}}" - local decode_ep="${DECODE_EP:-${EP_SIZE:-1}}" - local decode_num_workers="${DECODE_NUM_WORKERS:-1}" - - local dp_json="false" - if [ "${DP_ATTENTION:-false}" = "true" ]; then dp_json="true"; fi - local prefill_dp_json="$dp_json" - if [ "${PREFILL_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then - prefill_dp_json="true" - else - prefill_dp_json="false" - fi - local decode_dp_json="$dp_json" - if [ "${DECODE_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then - decode_dp_json="true" - else - decode_dp_json="false" - fi - - # Derive framework/precision from env, fallback to parsing RESULT_FILENAME - # RESULT_FILENAME format (from workflow): - # ___tp<...>_ep<...>_dpa_<...>_conc<...>_ - local fw="${FRAMEWORK:-}" - local prec="${PRECISION:-}" - if [[ -z "$fw" || -z "$prec" ]]; then - if [[ -n "${RESULT_FILENAME}" ]]; then - # Extract the two fields immediately before "_tp" - # Handles arbitrary underscores in exp_name by matching from the end - local parsed - parsed=$(echo "${RESULT_FILENAME}" | sed -n 's/.*_\([^_][^_]*\)_\([^_][^_]*\)_tp.*/\1 \2/p') - local p1="${parsed%% *}" - local p2="${parsed#* }" - if [[ -z "$prec" && -n "$p1" && "$p1" != "$parsed" ]]; then - prec="$p1" - fi - if [[ -z "$fw" && -n "$p2" && "$p2" != "$parsed" ]]; then - fw="$p2" - fi - fi - fi - cat > "${meta_json}" <&2 - fi - if [ -d "${out_dir}" ]; then - while IFS= read -r -d '' jf; do - base=$(basename "$jf") - if [ "$base" != "meta_env.json" ]; then - mv -f "$jf" ./ || echo "WARN: failed to move ${jf}" >&2 - fi - done < <(find "${out_dir}" -type f -name "*.json*" -print0 2>/dev/null) - fi - - # Best-effort cleanup of the temp directory - if [ -n "${out_dir}" ] && [ -d "${out_dir}" ]; then - rm -rf --one-file-system "${out_dir}" || rm -rf "${out_dir}" || true - fi - - echo "Moved eval artifacts to: $(pwd)" -} - -# ------------------------------ -# Unified eval entrypoint -# ------------------------------ - -run_eval() { - local framework="${EVAL_FRAMEWORK:-lm-eval}" - local forwarded=() - - while [[ $# -gt 0 ]]; do - case "$1" in - --framework) framework="$2"; shift 2 ;; - *) forwarded+=("$1"); shift ;; - esac - done - - # Compute EVAL_MAX_MODEL_LEN if not already set by the calling script - if [ -z "${EVAL_MAX_MODEL_LEN:-}" ]; then - compute_eval_context_length "$MODEL" "${MAX_MODEL_LEN:-0}" > /dev/null - fi - - unset EVAL_BATCHED_CONCS - unset EVAL_BATCHED_COMPLETED_CONCS - unset EVAL_BATCHED_FAILED_CONCS - - local requested_concs="${EVAL_CONCURRENT_REQUESTS:-}" - local eval_concs=() - if [ -n "$requested_concs" ]; then - read -r -a eval_concs <<< "$requested_concs" - fi - - if [ "${#eval_concs[@]}" -gt 1 ]; then - if [[ "$framework" != "lm-eval" && "$framework" != "lm_eval" ]]; then - echo "ERROR: batched eval concurrency is only supported for lm-eval" >&2 - return 1 - fi - - local eval_conc results_dir eval_rc stage_rc - local completed_concs=() - local failed_concs=() - - for eval_conc in "${eval_concs[@]}"; do - if [[ ! "$eval_conc" =~ ^[1-9][0-9]*$ ]]; then - echo "ERROR: invalid eval concurrency '${eval_conc}'" >&2 - return 1 - fi - - if ! results_dir=$(mktemp -d /tmp/eval_out-conc"${eval_conc}"-XXXXXX); then - echo "ERROR: failed to create eval output directory for concurrency ${eval_conc}" >&2 - failed_concs+=("$eval_conc") - continue - fi - - echo "Running lm-eval at concurrency ${eval_conc} using the existing engine" - export EVAL_CONCURRENT_REQUESTS="$eval_conc" - export CONC="$eval_conc" - eval_rc=0 - stage_rc=0 - run_lm_eval "${forwarded[@]}" --results-dir "$results_dir" \ - || eval_rc=$? - _stage_lm_eval_artifacts "$results_dir" "$eval_conc" \ - || stage_rc=$? - - if [ "$eval_rc" -eq 0 ] && [ "$stage_rc" -eq 0 ]; then - completed_concs+=("$eval_conc") - else - echo "ERROR: lm-eval failed at concurrency ${eval_conc} (eval_rc=${eval_rc}, stage_rc=${stage_rc})" >&2 - failed_concs+=("$eval_conc") - fi - done - - export EVAL_CONCURRENT_REQUESTS="$requested_concs" - export EVAL_RESULT_DIR="" - export EVAL_BATCHED_CONCS="${eval_concs[*]}" - export EVAL_BATCHED_COMPLETED_CONCS="${completed_concs[*]}" - export EVAL_BATCHED_FAILED_CONCS="${failed_concs[*]}" - - if [ "${#failed_concs[@]}" -gt 0 ]; then - echo "ERROR: batched eval failed for concurrency: ${failed_concs[*]}" >&2 - echo "Deferring failure until post-upload score validation preserves all artifacts" >&2 - fi - return 0 - fi - - local eval_rc=0 - case "$framework" in - lm-eval|lm_eval) run_lm_eval "${forwarded[@]}" || eval_rc=$? ;; - *) echo "Unknown framework '${framework}'"; eval_rc=1 ;; - esac - - if [ "$eval_rc" -ne 0 ]; then - echo "ERROR: run_eval failed with exit code $eval_rc" >&2 - if [ "${EVAL_ONLY}" = "true" ]; then - echo "Eval-only mode: failing after artifact collection" >&2 - return "$eval_rc" - fi - fi - return $eval_rc -} - - -# -------------------------------- -# Agentic trace replay helpers (aiperf driver) -# -------------------------------- - -INFMAX_CONTAINER_WORKSPACE="${INFMAX_CONTAINER_WORKSPACE:-/workspace}" -AGENTIC_DIR="${AGENTIC_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/agentic-benchmark}" -AIPERF_DIR="${AIPERF_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/aiperf}" - -agentic_pip_install() { - local pip_install=(python3 -m pip install) - if python3 -m pip install --help 2>/dev/null | grep -q -- "--break-system-packages"; then - pip_install+=(--break-system-packages) - fi - - "${pip_install[@]}" "$@" -} - -ensure_hf_cli() { - if command -v hf >/dev/null 2>&1; then - return 0 - fi - - # Some lean runtime images used by multinode SGLang include Python but not - # the Hugging Face CLI. Install just the hub CLI before prefetching traces. - agentic_pip_install --quiet "huggingface_hub[cli]>=0.25.0" -} - -resolve_trace_source() { - # Per-recipe override: set WEKA_LOADER_OVERRIDE to one of the aiperf - # public-dataset loader names allowed by the inferencex-agentx-mvp - # scenario. Used by recipes whose servers have non-default context - # caps (e.g. minimaxm2.5 at max_model_len ~256k can't replay the - # unfiltered 052726 corpus and switches to the 256k-capped variant). - local loader="${WEKA_LOADER_OVERRIDE:-semianalysis_cc_traces_weka_with_subagents}" - local dataset - case "$loader" in - semianalysis_cc_traces_weka_with_subagents) - dataset="semianalysisai/cc-traces-weka-with-subagents-052726" - ;; - semianalysis_cc_traces_weka_with_subagents_256k) - dataset="semianalysisai/cc-traces-weka-with-subagents-052726-256k" - ;; - *) - echo "Error: unknown WEKA_LOADER_OVERRIDE='$loader'. Allowed: semianalysis_cc_traces_weka_with_subagents, semianalysis_cc_traces_weka_with_subagents_256k" >&2 - exit 1 - ;; - esac - TRACE_SOURCE_FLAG="--public-dataset $loader" - echo "Loading traces via aiperf public-dataset: $loader ($dataset)" - # Pre-download the dataset into the shared HF_HUB_CACHE (same mount used - # for model weights) so subsequent runs read from cache instead of - # re-downloading every job. - ensure_hf_cli - hf download --repo-type dataset "$dataset" -} - -install_agentic_deps() { - # vllm/vllm-openai container ships without git. pip needs git to - # introspect the aiperf source tree on install. Install on demand; - # no-op when git is already present (e.g. AMD images that ship it). - if ! command -v git >/dev/null 2>&1; then - apt-get update && apt-get install -y git - fi - agentic_pip_install --quiet urllib3 requests 2>/dev/null || true - agentic_pip_install -q -r "$AGENTIC_DIR/requirements.txt" - # Editable install of aiperf from the submodule — gives us the - # `aiperf` CLI plus the inferencex-agentx-mvp scenario plugin. - # - # `--ignore-installed` sidesteps the distutils-uninstall error that - # vLLM containers hit on apt-managed system packages (blinker, etc.) - # when pip's resolver tries to upgrade one of aiperf's transitive - # deps. Installing fresh into the user/site location is safe — the - # system package stays in place and pip's import order picks up our - # newer copy first. - agentic_pip_install -q --ignore-installed -e "$AIPERF_DIR" - # Force-upgrade datasets: containers often ship an older version without - # the `Json` feature type used by the HF traces dataset. `Json` was added - # in datasets 4.7.0 (March 2025). Unpinned installs won't upgrade an - # already-present package. - agentic_pip_install --upgrade "datasets>=4.7.0" -} - -build_replay_cmd() { - # aiperf invocation for the inferencex-agentx-mvp scenario. - # - # Pre-canned assistant replay is the default: recorded assistant responses - # are used for future prompt construction, and live server responses are - # discarded. Set AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES=1 explicitly - # to use live-assistant mode, where the loader emits user-only deltas and - # the worker threads the server's live assistant response back into the - # session. - # - # The scenario plugin locks: --cache-bust first_turn_prefix and - # --trace-idle-gap-cap-seconds 60 (per-trace idle-gap compression - # against parent + subagent request-start timestamps; supersedes the - # legacy --use-think-time-only / --inter-turn-delay-cap-seconds path), - # and auto-injects them — so we do not pass them. See - # utils/aiperf/docs/tutorials/agentx-mvp.md. - local result_dir="$1" - local duration="$DURATION" - - export AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES="${AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES:-0}" - # Dataset configuration (load + reconstruct + inputs.json + mmap) - # routinely takes 4-5 min for the Weka corpus on fast /tmp - # (B300) but can stretch to 14 min on slower /tmp + parallel contention - # (observed on H200 where all 14 R3 jobs hit aiperf's 900s Configure - # Profiling timeout simultaneously). Bump to 1800s to absorb 3x - # worst-case slowdown — the post-setup measurement window is unaffected. - export AIPERF_DATASET_CONFIGURATION_TIMEOUT=1800 - # aiperf validates that SERVICE_PROFILE_CONFIGURE_TIMEOUT >= - # DATASET_CONFIGURATION_TIMEOUT at startup. Bump it in lockstep. - export AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT=1800 - REPLAY_CMD="aiperf profile --scenario inferencex-agentx-mvp" - REPLAY_CMD+=" --url http://localhost:$PORT" - REPLAY_CMD+=" --endpoint /v1/chat/completions" - REPLAY_CMD+=" --endpoint-type chat" - REPLAY_CMD+=" --streaming" - REPLAY_CMD+=" --model $MODEL" - REPLAY_CMD+=" --concurrency $CONC" - REPLAY_CMD+=" --benchmark-duration $duration" - REPLAY_CMD+=" --random-seed 42" - # Fail runs once more than 10% of requests error. This keeps known - # transient low-rate failures from killing long sweeps while still - # catching malformed payloads or server crashes before they get aggregated - # as benchmarkable data. - REPLAY_CMD+=" --failed-request-threshold 0.10" - # Sample each trajectory's warmup start position uniformly from - # [25%, 75%] of the trace's turn count (was hardcoded 0%-70% upstream). - # Avoids starting trajectories right at turn 0 where the KV cache is - # cold and skews early steady-state samples. - REPLAY_CMD+=" --trajectory-start-min-ratio 0.25" - REPLAY_CMD+=" --trajectory-start-max-ratio 0.75" - # Use server-reported usage fields (prompt_tokens / completion_tokens) for - # ISL/OSL instead of client-side tokenizer.encode(). Auto-enables - # stream_options.include_usage on the OpenAI chat endpoint. Skips the - # heavy per-record tokenization in the records pipeline that was pinning - # CPU on minimax-m2.5 at high concurrency. Lossless for vLLM (server - # usage is authoritative). - REPLAY_CMD+=" --use-server-token-count" - # aiperf's dataset manager (separate from the inference parser) loads - # the model's tokenizer for trace-prompt tokenization regardless of - # --use-server-token-count. Models like kimi (amd/Kimi-K2.5-MXFP4, - # moonshotai/Kimi-K2.5) ship a custom tokenizer in their HF repo and - # need trust_remote_code=True to load. Benign for models without - # custom tokenizer code, so we set it unconditionally. - REPLAY_CMD+=" --tokenizer-trust-remote-code" - # Keep replay inputs inside the same context window used to launch the - # server. The WEKA corpus contains a few very long parent/subagent traces; - # if we mmap and replay them against a smaller-context server they become - # deterministic 4xxs and can still pressure the engine while queued. - if [ -n "${MAX_MODEL_LEN:-}" ] && [ "$MAX_MODEL_LEN" != "0" ]; then - REPLAY_CMD+=" --max-context-length $MAX_MODEL_LEN" - fi - # Default --num-dataset-entries is 100; the with-subagents Weka corpus - # has 472. Cap at 472 so all unique traces are loaded (the loader treats - # this as a ``min(cap, available)`` ceiling, not a target — see - # semianalysis_cc_traces_weka.py). - REPLAY_CMD+=" --num-dataset-entries 472" - # 1-second timeslices on the server-metrics scrape so the post-run - # plotter has per-window time series (KV usage, cache hit rate, - # throughput, etc.). Matches kv-cache-tester's poll_interval=1.0 - # snapshot cadence so metrics_plots.png is visually comparable. - # Without this, aiperf only emits aggregate stats and the 6x2 panels - # collapse to flat lines. - REPLAY_CMD+=" --slice-duration 1.0" - REPLAY_CMD+=" --output-artifact-dir $result_dir/aiperf_artifacts" - # The inferencex-agentx-mvp scenario enforces a 900s minimum - # benchmark duration. For smoke tests with shorter durations, opt - # into --unsafe-override (the run's submission_valid will be flagged - # false; that's expected for non-canonical runs). - if [ "$duration" -lt 900 ] || [ "${AIPERF_UNSAFE_OVERRIDE:-false}" = "true" ]; then - REPLAY_CMD+=" --unsafe-override" - fi - REPLAY_CMD+=" $TRACE_SOURCE_FLAG" -} - -write_agentic_result_json() { - # Aggregate aiperf's profile_export.{json,jsonl} + server_metrics_export.json - # into $AGENTIC_OUTPUT_DIR/$RESULT_FILENAME.json. The workflow's existing - # retry-based existence check is the single success gate. - local result_dir="$1" - RESULT_DIR="$result_dir" AGENTIC_OUTPUT_DIR="${AGENTIC_OUTPUT_DIR:-$INFMAX_CONTAINER_WORKSPACE}" \ - python3 "$INFMAX_CONTAINER_WORKSPACE/utils/process_agentic_result.py" - - # Generate metrics_plots.png from the same aiperf artifacts. Best-effort: - # don't fail the launcher if plot generation has trouble (e.g. matplotlib - # missing in a stripped-down image). The agg JSON is the success gate. - python3 "$INFMAX_CONTAINER_WORKSPACE/utils/generate_aiperf_plots.py" "$result_dir" 2>&1 || true -} - -run_agentic_replay_and_write_outputs() { - local result_dir="$1" - local replay_rc - - echo "$REPLAY_CMD" > "$result_dir/benchmark_command.txt" - - set +e - set -x - $REPLAY_CMD 2>&1 | tee "$result_dir/benchmark.log" - replay_rc=${PIPESTATUS[0]} - set +x - set -e - - write_agentic_result_json "$result_dir" - - python3 "$AGENTIC_DIR/scripts/analyze_benchmark_distributions.py" \ - "$result_dir/aiperf_artifacts" -o "$result_dir" 2>&1 || true - - if [ "$replay_rc" -ne 0 ]; then - echo "ERROR: agentic trace replay exited with code $replay_rc after writing available results" >&2 - return "$replay_rc" - fi -} diff --git a/scripts/MiniMax-M3/lib/prompts.json b/scripts/MiniMax-M3/lib/prompts.json deleted file mode 100644 index 1b0d3755..00000000 --- a/scripts/MiniMax-M3/lib/prompts.json +++ /dev/null @@ -1,32 +0,0 @@ -[ - { - "id": "factual", - "prompt": "What is the capital of France? Answer in one word.", - "expect_contains": ["Paris"], - "max_tokens": 16 - }, - { - "id": "arithmetic", - "prompt": "Compute 17 multiplied by 23. Give only the number.", - "expect_contains": ["391"], - "max_tokens": 24 - }, - { - "id": "code", - "prompt": "Write a Python one-liner that returns the sum of a list named nums. Output only the code.", - "expect_contains": ["sum(", "nums"], - "max_tokens": 48 - }, - { - "id": "reasoning", - "prompt": "If a train travels 60 km in 1.5 hours, what is its average speed in km/h? Give only the number.", - "expect_contains": ["40"], - "max_tokens": 32 - }, - { - "id": "instruction", - "prompt": "List exactly three primary colors, comma-separated, lowercase.", - "expect_contains": ["red", "blue"], - "max_tokens": 24 - } -] diff --git a/scripts/MiniMax-M3/lib/run_disagg.sh b/scripts/MiniMax-M3/lib/run_disagg.sh deleted file mode 100644 index a3dae5d7..00000000 --- a/scripts/MiniMax-M3/lib/run_disagg.sh +++ /dev/null @@ -1,192 +0,0 @@ -#!/bin/bash -# run_disagg.sh — MiniMax-M3 MXFP4 PD-disaggregated serving + benchmark across multiple nodes. -# Faithful port of InferenceX server_atom.sh (atomesh + mooncake) for minimaxm3-fp4-mi355x-atom-disagg. -# -# Env in: ENGINE=atom-disagg MODEL=minimaxm3-fp4 TOPO=1p1d|2p1d_dpa [ACTION=bench|serve] -# -# Layout per topology (topo.py resolves nodes): -# prefill worker(s) = kv_producer servers (port 8010) [worker 0's node also hosts the router] -# decode worker(s) = kv_consumer servers (port 8020) -# router (atomesh) on prefill-0, port 8000 <- bench drives THIS endpoint -# -# Each server runs in its own docker container via a per-node `srun --exclusive`. KV flows -# node->node over RDMA (mooncake). Bench runs from a clean container against the router -# (aggregate throughput; no x N extrapolation — this is a real disagg instance). -# All workers TP4, STP only (DECODE_MTP_SIZE=0); DPA topology adds --enable-dp-attention. -set -uo pipefail -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" # scripts/MiniMax-M3 -CLUSTER="$HERE/cluster.yaml"; MODELS="$HERE/model.yaml" -ENGINE="${ENGINE:-atom-disagg}" -MODEL="${MODEL:-minimaxm3-fp4}"; TOPO="${TOPO:?set TOPO=1p1d|2p1d_dpa}"; ACTION="${ACTION:-bench}" -cfg(){ python3 "$HERE/lib/cfg.py" "$@"; } -topo(){ python3 "$HERE/lib/topo.py" "$@"; } - -RESV=$(cfg "$CLUSTER" cluster reservation) -MROOT=$(cfg "$CLUSTER" cluster models_root) -RROOT=$(cfg "$CLUSTER" cluster results_root) -RPORT=$(cfg "$CLUSTER" cluster router_port) -PPORT=$(cfg "$CLUSTER" cluster prefill_port) -DPORT=$(cfg "$CLUSTER" cluster decode_port) -HPORT=$(cfg "$CLUSTER" cluster handshake_port) -MDIR=$(cfg "$MODELS" model "$MODEL" dir) -IMG=$(cfg "$MODELS" engine "$ENGINE" image) -ROUTER=$(cfg "$MODELS" engine "$ENGINE" router) -KVCONN=$(cfg "$MODELS" engine "$ENGINE" kv_connector) -GMU=$(cfg "$MODELS" default "$MODEL" gpu_memory_util) -MNS=$(cfg "$MODELS" default "$MODEL" max_num_seqs) -KVDT=$(cfg "$MODELS" default "$MODEL" kv_cache_dtype) -BS=$(cfg "$MODELS" default "$MODEL" block_size) -MML=$(cfg "$MODELS" default "$MODEL" max_model_len) -MBT=$(cfg "$MODELS" default "$MODEL" max_num_batched_tokens) -TP=$(cfg "$MODELS" default "$MODEL" tp) -SHAPES=$(topo "$MODELS" "$TOPO" field shapes) -CONC=$(topo "$MODELS" "$TOPO" field conc) -DPATTN=$(python3 -c "import yaml;t=yaml.safe_load(open('$MODELS'))['topologies']['$TOPO'];print('1' if t['decode'].get('dp_attn') else '0')") - -ENV_ARGS=""; for kv in $(cfg "$MODELS" engine "$ENGINE" env); do ENV_ARGS+=" -e $kv"; done - -STAMP=$(python3 -c "import time;print(time.strftime('%Y%m%d_%H%M%S'))") -RUN="$RROOT/${MODEL}_${ENGINE}_${TOPO}_${STAMP}"; mkdir -p "$RUN"; chmod 777 "$RUN" -topo "$CLUSTER" "$MODELS" "$TOPO" roles > "$RUN/roles.tsv" -echo "[disagg] ENGINE=$ENGINE MODEL=$MODEL TOPO=$TOPO router=$ROUTER kv=$KVCONN image=$IMG" -echo "[disagg] roles:"; sed 's/^/ /' "$RUN/roles.tsv" - -PREFILL_NODE0=$(awk -F'\t' '$1=="prefill" && $2=="0"{print $3}' "$RUN/roles.tsv") -PREFILL_IPS=""; DECODE_IPS="" - -ip_of(){ srun --reservation="$RESV" --nodelist="$1" --nodes=1 --ntasks=1 --overlap --mem=1G \ - bash -c "ip route get 1.1.1.1 2>/dev/null | awk '/src/{print \$7}'" 2>/dev/null | grep -oE '([0-9]+\.){3}[0-9]+' | head -1; } - -# --- PRE-FLIGHT: clean every node in this topology ------------------------------------------ -NODES=$(cut -f3 "$RUN/roles.tsv" | sort -u) -echo "[disagg] cleaning nodes: $(echo $NODES | tr '\n' ' ')" -for n in $NODES; do - srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=0 \ - bash "$HERE/lib/clean_node.sh" > "$RUN/clean_${n}.log" 2>&1 && echo " [clean] $n OK" || echo " [clean] $n WARN" -done - -# --- resolve router IP (prefill node 0) ----------------------------------------------------- -ROUTER_IP=$(ip_of "$PREFILL_NODE0") -echo "[disagg] router will be on $PREFILL_NODE0 ($ROUTER_IP):$RPORT" - -# DP-attention arg for M3 (non-DSv4: DPA only, no --enable-tbo which is DSv4-specific). -dpa_arg(){ [[ "$1" == "1" ]] && echo "--enable-dp-attention" || true; } -# Optional model-length caps (M3: max_model_len 32768, max_num_batched_tokens 32768). -LEN_ARGS=""; [[ -n "$MML" ]] && LEN_ARGS="$LEN_ARGS --max-model-len $MML" -[[ -n "$MBT" ]] && LEN_ARGS="$LEN_ARGS --max-num-batched-tokens $MBT" -# KV cache dtype: skip the flag when 'auto' (M3 uses the engine default). -KV_ARG=""; [[ -n "$KVDT" && "$KVDT" != "auto" ]] && KV_ARG="--kv-cache-dtype $KVDT" - -# --- launch one server (prefill|decode) on a node in its own srun+docker -------------------- -launch_server() { - local role="$1" node="$2" port="$3" kvrole="$4" rdir="$RUN/${role}_${node}" - mkdir -p "$rdir"; chmod 777 "$rdir" - local extra="" - [[ "$role" == "prefill" ]] && extra="$(cfg "$MODELS" engine "$ENGINE" prefill_flags)" || extra="$(cfg "$MODELS" engine "$ENGINE" decode_flags)" - [[ "$DPATTN" == "1" ]] && extra="$extra $(dpa_arg 1)" - local SRV="python3 -m atom.entrypoints.openai_server --model $MROOT/$MDIR \ - --host 0.0.0.0 --server-port $port -tp $TP \ - $KV_ARG --block-size $BS --gpu-memory-utilization $GMU --max-num-seqs $MNS \ - $LEN_ARGS $extra \ - --kv-transfer-config '{\"kv_role\":\"$kvrole\",\"kv_connector\":\"$KVCONN\",\"proxy_ip\":\"$ROUTER_IP\",\"handshake_port\":$HPORT}'" - srun --reservation="$RESV" --nodelist="$node" --nodes=1 --ntasks=1 --exclusive bash -c " - docker rm -f mm3_${role} >/dev/null 2>&1 || true; sleep 3 - cid=\$(docker run -d --name mm3_${role} --entrypoint bash \ - --device=/dev/kfd --device=/dev/dri --group-add video --ipc host --shm-size 128G --network host \ - --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --privileged \ - -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 $ENV_ARGS \ - -v $MROOT:$MROOT -v $HERE:/scripts -v $rdir:/out \ - $IMG -c '$SRV') - echo \"[$role] container \$cid on $node:$port\" - for i in \$(seq 1 300); do - docker logs mm3_${role} > $rdir/server.log 2>&1 || true - curl -sf http://localhost:$port/health >/dev/null 2>&1 && { echo READY > $rdir/READY; break; } - st=\$(docker inspect -f '{{.State.Status}}' mm3_${role} 2>/dev/null || echo missing) - { [ \"\$st\" = exited ] || [ \"\$st\" = missing ]; } && { echo \"[$role] container \$st\"; break; } - sleep 5 - done - docker logs mm3_${role} > $rdir/server.log 2>&1 || true - # hold the container alive for the duration (router + bench run elsewhere); wait on it - [ -f $rdir/READY ] && docker wait mm3_${role} || { echo '[FAIL] tail:'; tail -25 $rdir/server.log; exit 1; } - " > "$rdir/srun.log" 2>&1 & - echo "$!" -} - -# Build router prefill/decode arg lists (need IPs) -declare -a SRV_PIDS=() -while IFS=$'\t' read -r role widx node port; do - ipn=$(ip_of "$node") - if [[ "$role" == "prefill" ]]; then PREFILL_IPS+="$ipn "; else DECODE_IPS+="$ipn "; fi -done < "$RUN/roles.tsv" - -echo "[disagg] launching servers..." -while IFS=$'\t' read -r role widx node port; do - kvrole="kv_producer"; [[ "$role" == "decode" ]] && kvrole="kv_consumer" - pid=$(launch_server "$role" "$node" "$port" "$kvrole"); SRV_PIDS+=("$pid") - echo " [$role $widx] $node:$port ($kvrole) srun-pid $pid" -done < "$RUN/roles.tsv" - -# --- wait for all servers READY ------------------------------------------------------------- -echo "[disagg] waiting for all servers healthy (timeout ~25min)..." -ALLREADY=0 -for i in $(seq 1 300); do - r=$(find "$RUN" -name READY 2>/dev/null | wc -l); n=$(wc -l < "$RUN/roles.tsv") - [[ "$r" -ge "$n" ]] && { ALLREADY=1; echo " all $n servers READY"; break; } - sleep 5 -done -if [[ "$ALLREADY" != "1" ]]; then echo "[disagg] NOT all servers came up — see $RUN/*/server.log"; fi - -# --- start the atomesh router on prefill node 0 --------------------------------------------- -RDIR="$RUN/router"; mkdir -p "$RDIR"; chmod 777 "$RDIR" -PF_ARGS=""; for ip in $PREFILL_IPS; do PF_ARGS+=" --prefill http://$ip:$PPORT"; done -DC_ARGS=""; for ip in $DECODE_IPS; do DC_ARGS+=" --decode http://$ip:$DPORT"; done -RFLAGS=$(cfg "$MODELS" engine "$ENGINE" router_flags) -ROUTER_CMD="/usr/local/bin/atomesh launch --host 0.0.0.0 --port $RPORT --pd-disaggregation $PF_ARGS $DC_ARGS $RFLAGS" -echo "[disagg] router cmd: $ROUTER_CMD" -srun --reservation="$RESV" --nodelist="$PREFILL_NODE0" --nodes=1 --ntasks=1 --overlap --mem=8G bash -c " - docker rm -f mm3_router >/dev/null 2>&1 || true - cid=\$(docker run -d --name mm3_router --entrypoint bash --network host \ - -v $HERE:/scripts $IMG -c '$ROUTER_CMD') - for i in \$(seq 1 60); do - docker logs mm3_router > $RDIR/router.log 2>&1 || true - curl -sf http://localhost:$RPORT/v1/models >/dev/null 2>&1 && { echo READY > $RDIR/READY; break; } - sleep 5 - done - docker logs mm3_router > $RDIR/router.log 2>&1 || true -" > "$RDIR/srun.log" 2>&1 & - -for i in $(seq 1 60); do [[ -f "$RDIR/READY" ]] && { echo "[disagg] router READY on $ROUTER_IP:$RPORT"; break; }; sleep 5; done -[[ -f "$RDIR/READY" ]] || echo "[disagg] router did NOT become ready — see $RDIR/router.log" - -if [[ "$ACTION" == "serve" ]]; then - echo "[disagg] ACTION=serve — endpoint http://$ROUTER_IP:$RPORT (router). Servers held; Ctrl-C / scancel to stop." - wait; exit 0 -fi - -# --- ACTION=bench: drive the router at each (shape,conc) ------------------------------------- -BENCH_IMG=$(cfg "$MODELS" engine "$ENGINE" bench_image) -BENCH_IMG="${BENCH_IMG:-vllm/vllm-openai-rocm:nightly-9037498c22891e55b594f567fb91d9b4efbf3e99}" -echo "[disagg] benchmarking router http://$ROUTER_IP:$RPORT shapes='$SHAPES' conc='$CONC'" -for shape in $SHAPES; do - ISL=${shape%,*}; OSL=${shape#*,} - for c in $CONC; do - echo " [bench] isl=$ISL osl=$OSL conc=$c" - srun --reservation="$RESV" --nodelist="$PREFILL_NODE0" --nodes=1 --ntasks=1 --overlap --mem=8G bash -c " - docker run --rm --entrypoint bash --network host -v $HERE:/scripts -v $RUN:/out \ - $BENCH_IMG -c 'python3 /scripts/utils/bench_serving/benchmark_serving.py \ - --model $MROOT/$MDIR --backend vllm --base-url http://$ROUTER_IP:$RPORT --dataset-name random \ - --random-input-len $ISL --random-output-len $OSL --random-range-ratio 1.0 \ - --num-prompts \$(( $c * 10 )) --num-warmups \$(( $c * 2 )) --max-concurrency $c \ - --request-rate inf --ignore-eos --save-result --percentile-metrics ttft,tpot,itl,e2el \ - --result-dir /out --result-filename disagg_isl${ISL}_osl${OSL}_c${c}.json' - " > "$RUN/bench_isl${ISL}_osl${OSL}_c${c}.log" 2>&1 || echo " [bench] FAILED isl=$ISL conc=$c" - done -done - -# --- teardown ------------------------------------------------------------------------------- -echo "[disagg] teardown: stopping servers + router" -for n in $NODES; do - srun --reservation="$RESV" --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=1G \ - bash -c "docker rm -f mm3_prefill mm3_decode mm3_router >/dev/null 2>&1" 2>/dev/null || true -done -echo "[disagg] done -> $RUN" diff --git a/scripts/MiniMax-M3/lib/topo.py b/scripts/MiniMax-M3/lib/topo.py deleted file mode 100644 index d767eda2..00000000 --- a/scripts/MiniMax-M3/lib/topo.py +++ /dev/null @@ -1,52 +0,0 @@ -#!/usr/bin/env python3 -"""topo.py — resolve a MiniMax-M3 disagg topology into a node-role placement. - -Usage: - topo.py roles - -> TSV lines: roleworker_idxnodeport - role in {prefill, decode}; node 0 of prefill also hosts the router. - topo.py field - -> scalar/list field of the topology (xP, yD, shapes, conc, ...) - -Topology drives node count: nodes_needed = xP*pf_nodes_per_worker + yD*dec_nodes_per_worker, -with nodes_per_worker = ceil(TP / gpus_per_node) (TP4 on 8-GPU node => 1 node/worker). -""" -import sys, yaml, math - -def load(f): return yaml.safe_load(open(f)) - -def roles(cluster_f, model_f, topo_name): - c = load(cluster_f); m = load(model_f) - gpn = c["gpus_per_node"]; nodes = c["nodes"] - t = m["topologies"][topo_name] - xP, yD = t["xP"], t["yD"] - pf_tp = t["prefill"]["tp"]; dec_tp = t["decode"]["tp"] - pf_npw = math.ceil(pf_tp / gpn); dec_npw = math.ceil(dec_tp / gpn) - need = xP * pf_npw + yD * dec_npw - if need > len(nodes): - sys.exit(f"ERROR: topo {topo_name} needs {need} nodes, cluster has {len(nodes)}") - pf_port = c["prefill_port"]; dec_port = c["decode_port"] - out = []; i = 0 - for w in range(xP): - out.append(f"prefill\t{w}\t{nodes[i]}\t{pf_port}"); i += pf_npw - for w in range(yD): - out.append(f"decode\t{w}\t{nodes[i]}\t{dec_port}"); i += dec_npw - return out - -def main(): - a = sys.argv - if a[2:3] == ["roles"] or (len(a) >= 4 and a[3] == "roles"): - # form: topo.py cluster model topo roles -> but support both orderings - pass - if len(a) >= 5 and a[4] == "roles": - print("\n".join(roles(a[1], a[2], a[3]))); return - if len(a) >= 5 and a[3] == "field": - m = load(a[1]); t = m["topologies"][a[2]]; v = t.get(a[4]) - if isinstance(v, list): print(" ".join(str(x) for x in v)) - elif isinstance(v, dict): print(" ".join(f"{k}={vv}" for k, vv in v.items())) - else: print("" if v is None else v) - return - sys.exit("usage: topo.py roles | topo.py field ") - -if __name__ == "__main__": - main() diff --git a/scripts/MiniMax-M3/model.yaml b/scripts/MiniMax-M3/model.yaml deleted file mode 100644 index f579595c..00000000 --- a/scripts/MiniMax-M3/model.yaml +++ /dev/null @@ -1,64 +0,0 @@ -# model.yaml — MiniMax-M3 model def + ATOM disagg image/flags/router. -# -# MiniMax-M3 (MXFP4): served PD-disaggregated on AMD Instinct MI355X (gfx950). -# All workers TP4 (no EP): the single-node M3 recipe found plain TP4 beats both TP8 and -# TP4/EP4 on tok/s/GPU for this model on gfx950, so prefill and decode both use TP4 and we -# tune the prefill:decode worker ratio (xP:yD) instead of TP. STP only (DECODE_MTP_SIZE=0). -# -# Source of truth: InferenceX `minimaxm3-fp4-mi355x-atom-disagg` -# (SemiAnalysisAI/InferenceX configs/amd-master.yaml, PR #2000). - -defaults: - gpu_memory_util: 0.8 # ATOM --gpu-memory-utilization (InferenceX MEM_FRAC_STATIC 0.8) - max_num_seqs: 256 - kv_cache_dtype: auto # M3 MXFP4 uses the engine default KV dtype (unlike DSv4 fp8) - block_size: 128 # InferenceX ATOM M3 BLOCK_SIZE=128 - max_model_len: 32768 # InferenceX ATOM M3 MAX_MODEL_LEN=32768 - max_num_batched_tokens: 32768 - tp: 4 # per prefill/decode worker (TP4 throughout) - -models: - minimaxm3-fp4: - dir: MiniMax-M3-MXFP4 # under cluster.models_root (staged to scratch); from amd/MiniMax-M3-MXFP4 - hf_repo: amd/MiniMax-M3-MXFP4 - arch: minimax_m3 - precision: fp4 - status: SCAFFOLDED # ported from InferenceX; needs a live gfx950 dry-run - -# ---- engines: disagg image + router + KV-transfer backend ---------------------------------- -engines: - atom-disagg: - # InferenceX-VALIDATED M3 MXFP4 disagg path. atomesh router + mooncake RDMA KV. - image: rocm/atom-dev:nightly_202607011530 - # bench client runs from a clean vLLM image (ATOM's own image can crash the bench client) - bench_image: vllm/vllm-openai-rocm:nightly-9037498c22891e55b594f567fb91d9b4efbf3e99 - router: atomesh # /usr/local/bin/atomesh launch --pd-disaggregation - router_version: "04b120d3040e0dd1a6915e88a3a3c1f588e6684a" - kv_connector: mooncake # --kv-transfer-config '{"kv_connector":"mooncake",...}' - entrypoint: "python3 -m atom.entrypoints.openai_server" - env: - SAFETENSORS_FAST_GPU: "1" - AITER_QUICK_REDUCE_QUANTIZATION: "INT4" # applied for non-DSv4 models in InferenceX server_atom.sh - # engine-unique server flags (shared knobs come from defaults{}) - prefill_flags: "--trust-remote-code --no-enable_prefix_caching" - decode_flags: "--trust-remote-code --no-enable_prefix_caching" - router_flags: "--policy random --backend atom --disable-health-check --disable-circuit-breaker" - status: PRIMARY # validated M3 MXFP4 atom-disagg config - -# ---- disagg topologies (from minimaxm3-fp4-mi355x-atom-disagg search-space) ----------------- -# All workers TP4 -> 1 node/worker on 8-GPU nodes. xP + yD <= 3 (3-node mi355x-disagg pool). -topologies: - 1p1d: # 2 nodes: 1 prefill (TP4) + 1 decode (TP4), STP - xP: 1 - yD: 1 - prefill: { tp: 4, ep: 1, dp_attn: false } - decode: { tp: 4, ep: 1, dp_attn: false, mtp_size: 0 } - shapes: ["8192,1024", "1024,1024"] - conc: [1, 2, 4, 8, 16, 32, 64, 128, 256] - 2p1d_dpa: # 3 nodes: 2 prefill (TP4, DP-attn) + 1 decode (TP4, DP-attn), 8k1k tail - xP: 2 - yD: 1 - prefill: { tp: 4, ep: 1, dp_attn: true } - decode: { tp: 4, ep: 1, dp_attn: true, mtp_size: 0 } - shapes: ["8192,1024"] - conc: [256, 512, 768, 1024] diff --git a/scripts/MiniMax-M3/run_atom_disagg.sh b/scripts/MiniMax-M3/run_atom_disagg.sh deleted file mode 100644 index ead44e19..00000000 --- a/scripts/MiniMax-M3/run_atom_disagg.sh +++ /dev/null @@ -1,5 +0,0 @@ -#!/bin/bash -# run_atom_disagg.sh — MiniMax-M3 ATOM PD-disagg (atomesh + mooncake). Thin wrapper. -# Usage: MODEL=minimaxm3-fp4 TOPO=1p1d|2p1d_dpa [ACTION=bench|serve] ./run_atom_disagg.sh -HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" -ENGINE=atom-disagg MODEL="${MODEL:-minimaxm3-fp4}" exec bash "$HERE/lib/run_disagg.sh" diff --git a/scripts/MiniMax-M3/utils/bench_serving/benchmark_serving.py b/scripts/MiniMax-M3/utils/bench_serving/benchmark_serving.py index 7cc806ec..df6bddb1 100644 --- a/scripts/MiniMax-M3/utils/bench_serving/benchmark_serving.py +++ b/scripts/MiniMax-M3/utils/bench_serving/benchmark_serving.py @@ -48,7 +48,13 @@ try: from backend_request_func import get_tokenizer except ImportError: - from vllm.transformers_utils.tokenizer import get_tokenizer + # vLLM moved get_tokenizer from vllm.transformers_utils.tokenizer to + # vllm.tokenizers; the old path was a deprecation alias in v0.20 and + # was removed by v0.23. Try the new location first, fall back to old. + try: + from vllm.tokenizers import get_tokenizer + except ImportError: + from vllm.transformers_utils.tokenizer import get_tokenizer try: from vllm.utils import FlexibleArgumentParser @@ -103,6 +109,19 @@ def _load_tokenizer(tokenizer_id, tokenizer_mode, trust_remote_code): transformers). Prefer backend_request_func.get_tokenizer on fallback so client tokenization stays aligned with the sglang server (#1381, #1428). """ + if tokenizer_mode == "deepseek_v4": + # HF AutoTokenizer may not recognize deepseek_v4; use vLLM's loader. + try: + from vllm.tokenizers import get_tokenizer as _vllm_get_tokenizer + except ImportError: + from vllm.transformers_utils.tokenizer import ( + get_tokenizer as _vllm_get_tokenizer, + ) + return _vllm_get_tokenizer( + tokenizer_id, + tokenizer_mode=tokenizer_mode, + trust_remote_code=trust_remote_code, + ) try: return get_tokenizer( tokenizer_id, @@ -1249,7 +1268,7 @@ def main(args: argparse.Namespace): '--tokenizer-mode', type=str, default="auto", - choices=['auto', 'slow', 'mistral', 'custom'], + choices=['auto', 'slow', 'mistral', 'custom', 'deepseek_v4'], help='The tokenizer mode.\n\n* "auto" will use the ' 'fast tokenizer if available.\n* "slow" will ' 'always use the slow tokenizer. \n* ' From 5b40488a82e52eca899e39e0c16088d596cb2aae Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Thu, 16 Jul 2026 13:49:45 +0000 Subject: [PATCH 07/12] MiniMax-M3: make ATOM MXFP4 disagg recipe runnable standalone in MAD MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Keep the InferenceX engine code verbatim (benchmarks/multi_node/amd_utils/, benchmarks/benchmark_lib.sh, utils/bench_serving/ @ a174f20 — byte-identical, models_atom.yaml drives all model flags) but preserve InferenceX's relative path layout so bench.sh's ../../benchmark_lib.sh and REPO_ROOT/utils/bench_serving resolve. Add a thin MAD wrapper (run.sh + cluster.yaml) that supplies the per-node env contract + container mounts InferenceX's job.slurm/submit.sh normally provide, mounting the recipe root as /workspace and ATOM_WS_PATH at amd_utils/. README documents prerequisites (amdgpu_peermem for GPUDirect RDMA, weight staging, RDMA fabric), usage (TOPO=1p1d|2p1d_dpa, ACTION=dry), and topologies. Wrapper is syntax-checked; not yet dry-run on live MI355X. Co-Authored-By: Claude --- scripts/MiniMax-M3/README.md | 101 +- .../MiniMax-M3/benchmarks/benchmark_lib.sh | 1624 +++++++++++++++++ .../multi_node}/amd_utils/bench.sh | 0 .../multi_node}/amd_utils/env_atom.sh | 0 .../multi_node}/amd_utils/job.slurm | 0 .../multi_node}/amd_utils/models_atom.yaml | 0 .../multi_node}/amd_utils/server.sh | 0 .../multi_node}/amd_utils/server_atom.sh | 0 .../multi_node}/amd_utils/setup_deps.sh | 0 .../multi_node}/amd_utils/submit.sh | 0 .../multi_node}/amd_utils/sync.py | 0 scripts/MiniMax-M3/cluster.yaml | 25 + scripts/MiniMax-M3/run.sh | 100 + 13 files changed, 1813 insertions(+), 37 deletions(-) create mode 100644 scripts/MiniMax-M3/benchmarks/benchmark_lib.sh rename scripts/MiniMax-M3/{ => benchmarks/multi_node}/amd_utils/bench.sh (100%) mode change 100644 => 100755 rename scripts/MiniMax-M3/{ => benchmarks/multi_node}/amd_utils/env_atom.sh (100%) mode change 100644 => 100755 rename scripts/MiniMax-M3/{ => benchmarks/multi_node}/amd_utils/job.slurm (100%) rename scripts/MiniMax-M3/{ => benchmarks/multi_node}/amd_utils/models_atom.yaml (100%) rename scripts/MiniMax-M3/{ => benchmarks/multi_node}/amd_utils/server.sh (100%) mode change 100644 => 100755 rename scripts/MiniMax-M3/{ => benchmarks/multi_node}/amd_utils/server_atom.sh (100%) mode change 100644 => 100755 rename scripts/MiniMax-M3/{ => benchmarks/multi_node}/amd_utils/setup_deps.sh (100%) mode change 100644 => 100755 rename scripts/MiniMax-M3/{ => benchmarks/multi_node}/amd_utils/submit.sh (100%) mode change 100644 => 100755 rename scripts/MiniMax-M3/{ => benchmarks/multi_node}/amd_utils/sync.py (100%) create mode 100644 scripts/MiniMax-M3/cluster.yaml create mode 100755 scripts/MiniMax-M3/run.sh diff --git a/scripts/MiniMax-M3/README.md b/scripts/MiniMax-M3/README.md index cbfe76fa..c33294b5 100644 --- a/scripts/MiniMax-M3/README.md +++ b/scripts/MiniMax-M3/README.md @@ -1,39 +1,66 @@ -# MiniMax-M3 — MXFP4 disaggregated (P/D) ATOM benchmark (from InferenceX) - -Vendored **verbatim** from [SemiAnalysisAI/InferenceX](https://github.com/SemiAnalysisAI/InferenceX) -— the ATOM MXFP4 prefill/decode disaggregated benchmark for MiniMax-M3 on AMD MI355X (gfx950), -config key **`minimaxm3-fp4-mi355x-atom-disagg`** (atomesh router + mooncake RDMA KV transfer). - -This is the top-throughput MiniMax-M3 disagg config on MI355X (InferenceX): ~10,829 tok/s/GPU at -8k/1k (2P1D + DP-attn). Dashboard: - -## Contents (copied unchanged from InferenceX @ `a174f20`) -| Path | InferenceX source | -|------|-------------------| -| `amd_utils/server_atom.sh` | `benchmarks/multi_node/amd_utils/server_atom.sh` — ATOM disagg launcher (reads `models_atom.yaml`) | -| `amd_utils/models_atom.yaml` | `benchmarks/multi_node/amd_utils/models_atom.yaml` — per-model ATOM flags (incl. `MiniMax-M3-MXFP4`) | -| `amd_utils/server.sh` | dispatcher → `server_atom.sh` when `ENGINE=atom-disagg` | -| `amd_utils/env_atom.sh` | ATOM/mooncake env + RDMA device detection | -| `amd_utils/setup_deps.sh` | in-container dep setup | -| `amd_utils/bench.sh` | bench driver (calls the bench client against the atomesh router) | -| `amd_utils/job.slurm`, `submit.sh`, `sync.py` | SLURM multi-node plumbing | -| `utils/bench_serving/` | InferenceX benchmark client (`benchmark_serving.py` + backends) | -| `configs/minimaxm3-fp4-mi355x-atom-disagg.yaml` | the config entry extracted from `configs/amd-master.yaml` | - -## The MiniMax-M3-MXFP4 ATOM config (`models_atom.yaml`) -All workers TP4 (no EP; TP4 beats TP8/TP4-EP on gfx950). Key flags: -- `env: AITER_QUICK_REDUCE_QUANTIZATION=INT4 ATOM_FORCE_ATTN_TRITON=1` -- `kv_cache_flags: --kv_cache_dtype fp8` -- `online_quant_config: ptpc_fp8` (exclude lm_head/embed/vision/moe) -- `block_size 128`, `mem_frac_static 0.8`, `max_model_len 32768`, `max_num_batched_tokens 32768` +# MiniMax-M3 — MXFP4 disaggregated (P/D) ATOM benchmark + +The ATOM MXFP4 prefill/decode disaggregated benchmark for MiniMax-M3 on AMD MI355X (gfx950), +config **`minimaxm3-fp4-mi355x-atom-disagg`** (atomesh router + mooncake RDMA KV) — the +top-throughput MiniMax-M3 disagg config on InferenceX (~10,829 tok/s/GPU at 8k/1k, 2P1D+DPA). +Dashboard: + +## Layout +The engine code is copied **verbatim** from [SemiAnalysisAI/InferenceX](https://github.com/SemiAnalysisAI/InferenceX) +@ `a174f20` (unmodified — update by re-copying from the same paths). A thin MAD wrapper +(`run.sh` + `cluster.yaml`) supplies the per-node env + mounts that InferenceX's +`job.slurm`/`submit.sh` normally provide, so it runs standalone from MAD. + +| Path | Origin | +|------|--------| +| `benchmarks/multi_node/amd_utils/` | **verbatim** InferenceX (`server_atom.sh`, `models_atom.yaml`, `env_atom.sh`, `bench.sh`, `setup_deps.sh`, `server.sh`, `job.slurm`, `submit.sh`, `sync.py`) | +| `benchmarks/benchmark_lib.sh` | **verbatim** InferenceX (bench.sh sources `../../benchmark_lib.sh`) | +| `utils/bench_serving/` | **verbatim** InferenceX benchmark client | +| `configs/minimaxm3-fp4-mi355x-atom-disagg.yaml` | config entry from `amd-master.yaml` | +| `run.sh`, `cluster.yaml` | **MAD wrapper** (the only non-verbatim files) | + +Model flags (all-TP4, `--kv_cache_dtype fp8`, `online_quant_config` ptpc_fp8, `ATOM_FORCE_ATTN_TRITON=1`, +`AITER_QUICK_REDUCE_QUANTIZATION=INT4`, block 128, mem-frac 0.8, max-model-len 32768) live in +`models_atom.yaml` and are read by `server_atom.sh` — the wrapper sets none of them, so runtime +behavior matches InferenceX exactly. + +## Prerequisites +1. **`amdgpu_peermem` kernel module loaded on every host** — required for GPUDirect RDMA so + mooncake can stream the KV cache GPU→NIC→GPU across nodes. Check: `lsmod | grep amdgpu_peermem`; + load: `sudo modprobe amdgpu_peermem`. Without it, mooncake KV transfer stalls / times out. +2. **RDMA fabric** reachable between nodes (the launcher passes `--device=/dev/infiniband --ulimit memlock=-1`). +3. **Weights staged** at `/MiniMax-M3-MXFP4` (from `amd/MiniMax-M3-MXFP4`), same path on every node. +4. **SLURM** with the nodes/reservation from `cluster.yaml`; Docker + ROCm on each node. +5. Image `rocm/atom-dev:nightly_202607011530` pullable on each node. + +## Usage +Edit `cluster.yaml` (reservation, `models_root`, `nodes`), then: +```bash +# 8k/1k balanced, 2 nodes (1 prefill + 1 decode), conc 1-256: +TOPO=1p1d ISL=8192 OSL=1024 ./run.sh + +# 8k/1k high-throughput, 3 nodes (2 prefill + 1 decode + DP-attn), conc 256-1024: +TOPO=2p1d_dpa ISL=8192 OSL=1024 ./run.sh + +# 1k/1k, 2 nodes: +TOPO=1p1d ISL=1024 OSL=1024 ./run.sh + +# pin nodes / dry-run (print commands, launch nothing): +NODES=n1,n2,n3 TOPO=2p1d_dpa ./run.sh +ACTION=dry TOPO=1p1d ./run.sh +``` +Bench drives the atomesh router (`:8000`, OpenAI-compatible) → true aggregate throughput (no ×N). +Result JSONs land under `results/minimaxm3-fp4_atom___/`. ## Topologies (`configs/minimaxm3-fp4-mi355x-atom-disagg.yaml`) -| shape | conc | layout | nodes | -|-------|------|--------|-------| -| 8192/1024 | 1–256 | 1P1D TP4 (STP, `DECODE_MTP_SIZE=0`) | 2 | -| 8192/1024 | 256–1024 | 2P1D TP4 + DP-attn | 3 | -| 1024/1024 | 1–256 | 1P1D TP4 (STP) | 2 | - -## Provenance -Source commit: `a174f20d26983d499e0eac315ee39f0c9aa48262` (SemiAnalysisAI/InferenceX). Files are -unmodified. Update by re-copying from the same paths in a newer InferenceX revision. +| TOPO | shape | conc | layout | nodes | +|------|-------|------|--------|-------| +| `1p1d` | 8192/1024, 1024/1024 | 1–256 | 1 prefill TP4 + 1 decode TP4 (STP) | 2 | +| `2p1d_dpa` | 8192/1024 | 256–1024 | 2 prefill + 1 decode, TP4 + DP-attn | 3 | + +All-TP4 ⇒ 1 node/worker on 8-GPU nodes; STP only (`DECODE_MTP_SIZE=0`). + +## Status +Engine code is verbatim InferenceX (byte-identical). The MAD wrapper is syntax-checked and its +topology/mount contract mirrors InferenceX `job.slurm`, but has **not** been dry-run on live +MI355X hardware — validate with `ACTION=dry` first, then a small `TOPO=1p1d` run. diff --git a/scripts/MiniMax-M3/benchmarks/benchmark_lib.sh b/scripts/MiniMax-M3/benchmarks/benchmark_lib.sh new file mode 100644 index 00000000..6cba7338 --- /dev/null +++ b/scripts/MiniMax-M3/benchmarks/benchmark_lib.sh @@ -0,0 +1,1624 @@ +#!/usr/bin/env bash + +# Shared benchmarking utilities for InferenceX + +# Keep Python bytecode out of the mounted workspace. Benchmark jobs often run as +# root inside containers, and root-owned cache directories break future checkout +# cleanup on self-hosted runners. +export PYTHONDONTWRITEBYTECODE=1 +export PYTHONPYCACHEPREFIX="${PYTHONPYCACHEPREFIX:-/tmp/inferencex-pycache}" +mkdir -p "$PYTHONPYCACHEPREFIX" 2>/dev/null || true + +# Inference server port shared by every benchmark recipe. Launchers that need +# a non-default value (e.g. launch_mi355x-amds.sh derives PORT from RUNNER_NAME +# to avoid collisions across concurrent gh-runners on a shared host) set PORT +# themselves before sourcing this file; the `:-` fallback only kicks in when +# nothing upstream set it. +export PORT="${PORT:-8888}" + +agentic_kv_offload_enabled() { + if [[ -z "${KV_OFFLOADING+x}" || -z "$KV_OFFLOADING" ]]; then + echo "Error: KV_OFFLOADING must be set for agentic benchmarks" >&2 + exit 1 + fi + [[ "$KV_OFFLOADING" != "none" ]] +} + +require_agentic_kv_offload_none() { + if agentic_kv_offload_enabled; then + echo "Error: expected KV_OFFLOADING=none, got '$KV_OFFLOADING'" >&2 + exit 1 + fi + if [[ -n "${KV_OFFLOAD_BACKEND:-}" ]]; then + echo "Error: KV_OFFLOAD_BACKEND must be empty when KV_OFFLOADING=none" >&2 + exit 1 + fi +} + +require_agentic_kv_offload_backend() { + local expected_backend="$1" + if [[ -z "${KV_OFFLOADING+x}" || -z "$KV_OFFLOADING" ]]; then + echo "Error: KV_OFFLOADING must be set for agentic benchmarks" >&2 + exit 1 + fi + case "$KV_OFFLOADING" in + none) + if [[ -n "${KV_OFFLOAD_BACKEND:-}" ]]; then + echo "Error: KV_OFFLOAD_BACKEND must be empty when KV_OFFLOADING=none" >&2 + exit 1 + fi + return 1 + ;; + dram) + if [[ "${KV_OFFLOAD_BACKEND:-}" != "$expected_backend" ]]; then + echo "Error: expected KV_OFFLOAD_BACKEND=$expected_backend when KV_OFFLOADING=dram, got '${KV_OFFLOAD_BACKEND:-}'" >&2 + exit 1 + fi + if [[ ! "${TOTAL_CPU_DRAM_GB:-}" =~ ^[1-9][0-9]*$ ]]; then + echo "Error: DRAM KV offloading requires a positive TOTAL_CPU_DRAM_GB capacity" >&2 + exit 1 + fi + return 0 + ;; + *) + echo "Error: unsupported KV_OFFLOADING value '$KV_OFFLOADING' (expected one of: none, dram)" >&2 + exit 1 + ;; + esac +} + +# Agentic replays must use the model's native context limit. Ignore inherited +# workflow or shell overrides so neither the server nor AIPerf applies a cap. +_benchmark_caller="${BASH_SOURCE[1]:-}" +if [[ "$_benchmark_caller" == */agentic/* || + "$_benchmark_caller" == */agentic_*.sh || + "${IS_AGENTIC:-0}" == "1" || + "${SCENARIO_TYPE:-}" == "agentic-coding" ]]; then + unset MAX_MODEL_LEN + if [[ -z "${KV_OFFLOADING+x}" || -z "$KV_OFFLOADING" ]]; then + echo "Error: KV_OFFLOADING must be set for agentic benchmarks" >&2 + exit 1 + fi + case "$KV_OFFLOADING" in + none) + if [[ -n "${KV_OFFLOAD_BACKEND:-}" ]]; then + echo "Error: KV_OFFLOAD_BACKEND must be empty when KV_OFFLOADING=none" >&2 + exit 1 + fi + ;; + dram) + if [[ -z "${KV_OFFLOAD_BACKEND:-}" || "${KV_OFFLOAD_BACKEND:-}" == "none" ]]; then + echo "Error: KV_OFFLOAD_BACKEND is required when KV_OFFLOADING=dram" >&2 + exit 1 + fi + if [[ ! "${TOTAL_CPU_DRAM_GB:-}" =~ ^[1-9][0-9]*$ ]]; then + echo "Error: DRAM KV offloading requires a positive configured TOTAL_CPU_DRAM_GB capacity" >&2 + exit 1 + fi + ;; + *) + echo "Error: unsupported KV_OFFLOADING value '$KV_OFFLOADING' (expected one of: none, dram)" >&2 + exit 1 + ;; + esac +fi +unset _benchmark_caller + +# -------------------------------- +# GPU monitoring helpers +# -------------------------------- + +GPU_MONITOR_PID="" +GPU_METRICS_CSV="/workspace/gpu_metrics.csv" +export GPU_METRICS_CSV + +# Start background GPU monitoring that logs metrics every second to CSV. +# Auto-detects NVIDIA (nvidia-smi) or AMD (amd-smi) GPUs. +# Usage: start_gpu_monitor [--output /path/to/output.csv] [--interval 1] +start_gpu_monitor() { + local output="$GPU_METRICS_CSV" + local interval=1 + + while [[ $# -gt 0 ]]; do + case $1 in + --output) output="$2"; shift 2 ;; + --interval) interval="$2"; shift 2 ;; + *) shift ;; + esac + done + + GPU_METRICS_CSV="$output" + export GPU_METRICS_CSV + + if command -v nvidia-smi &>/dev/null; then + nvidia-smi --query-gpu=timestamp,index,power.draw,temperature.gpu,clocks.current.sm,clocks.current.memory,utilization.gpu,utilization.memory \ + --format=csv -l "$interval" > "$output" 2>/dev/null & + GPU_MONITOR_PID=$! + echo "[GPU Monitor] Started NVIDIA (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" + elif command -v amd-smi &>/dev/null; then + # Use amd-smi native watch mode (-w) which includes timestamps automatically. + # Pipe through awk to: skip preamble lines, keep first CSV header, skip repeated headers. + amd-smi metric -p -c -t -u -w "$interval" --csv 2>/dev/null \ + | awk '/^timestamp,/{if(!h){print;h=1};next} h{print}' > "$output" & + GPU_MONITOR_PID=$! + echo "[GPU Monitor] Started AMD (PID=$GPU_MONITOR_PID, interval=${interval}s, output=$output)" + else + echo "[GPU Monitor] No GPU monitoring tool found (nvidia-smi or amd-smi), skipping" + return 0 + fi +} + +# Stop the background GPU monitor and report file size. +stop_gpu_monitor() { + if [[ -n "$GPU_MONITOR_PID" ]] && kill -0 "$GPU_MONITOR_PID" 2>/dev/null; then + kill "$GPU_MONITOR_PID" 2>/dev/null + wait "$GPU_MONITOR_PID" 2>/dev/null || true + echo "[GPU Monitor] Stopped (PID=$GPU_MONITOR_PID)" + if [[ -f "$GPU_METRICS_CSV" ]]; then + local lines + lines=$(wc -l < "$GPU_METRICS_CSV") + echo "[GPU Monitor] Collected $lines rows -> $GPU_METRICS_CSV" + fi + fi + GPU_MONITOR_PID="" +} + +# Return success only while a PID exists and is not a zombie waiting to be +# reaped. `kill -0` alone treats zombies as live processes. +_background_process_is_running() { + local pid="$1" + local state + kill -0 "$pid" 2>/dev/null || return 1 + state=$(ps -o stat= -p "$pid" 2>/dev/null) || return 1 + [[ -n "$state" && "${state:0:1}" != "Z" ]] +} + +_background_process_descendants() { + local parent_pid="$1" + local child_pid + while read -r child_pid; do + [[ -n "$child_pid" ]] || continue + echo "$child_pid" + _background_process_descendants "$child_pid" + done < <(pgrep -P "$parent_pid" 2>/dev/null || true) +} + +# Stop a background service and every process that descended from it. Capture +# descendants before terminating the root because orphaned workers are +# reparented and can otherwise keep a Slurm step alive after the benchmark +# script exits. +stop_background_process_tree() { + local root_pid="${1:-}" + local label="${2:-background process}" + local grace_seconds="${3:-30}" + + if [[ ! "$root_pid" =~ ^[1-9][0-9]*$ ]] || ! _background_process_is_running "$root_pid"; then + return 0 + fi + + local descendants + local child_pid + descendants=$(_background_process_descendants "$root_pid") + + echo "Stopping $label (PID=$root_pid)..." + kill -TERM "$root_pid" 2>/dev/null || true + + local deadline=$((SECONDS + grace_seconds)) + while _background_process_is_running "$root_pid" && [[ $SECONDS -lt $deadline ]]; do + sleep 1 + done + + local forced_stop=false + while read -r child_pid; do + [[ -n "$child_pid" ]] || continue + if _background_process_is_running "$child_pid"; then + if [[ "$forced_stop" == "false" ]]; then + echo "Force-stopping remaining $label processes." + forced_stop=true + fi + echo " PID=$child_pid" + kill -KILL "$child_pid" 2>/dev/null || true + fi + done </dev/null || true + echo "Stopped $label." +} + +# Check if required environment variables are set +# Usage: check_env_vars VAR1 VAR2 VAR3 ... +# Exits with code 1 if any variable is not set +check_env_vars() { + local missing_vars=() + + for var_name in "$@"; do + if [[ -z "${!var_name:-}" ]]; then + missing_vars+=("$var_name") + fi + done + + if [[ ${#missing_vars[@]} -gt 0 ]]; then + echo "Error: The following required environment variables are not set:" + for var in "${missing_vars[@]}"; do + echo " - $var" + done + exit 1 + fi +} + +# Wait for server to be ready by polling the health endpoint +# All parameters are required +# Parameters: +# --port: Server port +# --server-log: Path to server log file +# --server-pid: Server process ID (required) +# --sleep-interval: Sleep interval between health checks (optional, default: 5) +wait_for_server_ready() { + set +x + local port="" + local server_log="" + local server_pid="" + local sleep_interval=5 + + # Parse arguments + while [[ $# -gt 0 ]]; do + case $1 in + --port) + port="$2" + shift 2 + ;; + --server-log) + server_log="$2" + shift 2 + ;; + --server-pid) + server_pid="$2" + shift 2 + ;; + --sleep-interval) + sleep_interval="$2" + shift 2 + ;; + *) + echo "Unknown parameter: $1" + return 1 + ;; + esac + done + + # Validate required parameters + if [[ -z "$port" ]]; then + echo "Error: --port is required" + return 1 + fi + if [[ -z "$server_log" ]]; then + echo "Error: --server-log is required" + return 1 + fi + if [[ -z "$server_pid" ]]; then + echo "Error: --server-pid is required" + return 1 + fi + + # Wait for server log file to be created (container startup may delay this) + while [ ! -f "$server_log" ]; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "Server died before creating log file. Exiting." + exit 1 + fi + sleep 1 + done + + # Show logs until server is ready + tail -f -n +1 "$server_log" & + local TAIL_PID=$! + until curl --output /dev/null --silent --fail http://0.0.0.0:$port/health; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "Server died before becoming healthy. Exiting." + kill $TAIL_PID + exit 1 + fi + sleep "$sleep_interval" + done + kill $TAIL_PID +} + +# Run benchmark serving with standardized parameters +# All parameters are required except --endpoint, --use-chat-template, --dsv4, and --trust-remote-code +# Parameters: +# --model: Model name +# --port: Server port +# --backend: Backend type - e.g., 'vllm' or 'openai' +# --endpoint: Optional API endpoint override +# --input-len: Random input sequence length +# --output-len: Random output sequence length +# --random-range-ratio: Random range ratio +# --num-prompts: Number of prompts +# --max-concurrency: Max concurrency +# --result-filename: Result filename without extension +# --result-dir: Result directory +# --use-chat-template: Optional flag to enable chat template +# --dsv4: Optional flag to use the DeepSeek-V4 chat template +# (encoding_dsv4.py) instead of the tokenizer's built-in jinja +# template. Implies --use-chat-template. +# --trust-remote-code: Optional flag to trust remote code from HuggingFace +# --server-pid: Optional server process ID to monitor during benchmark +run_benchmark_serving() { + # In eval-only mode, skip the throughput benchmark entirely. + if [ "${EVAL_ONLY}" = "true" ]; then + echo "EVAL_ONLY mode: skipping throughput benchmark" + return 0 + fi + + set +x + local model="" + local port="" + local backend="" + local endpoint="" + local input_len="" + local output_len="" + local random_range_ratio="" + local num_prompts="" + local max_concurrency="" + local result_filename="" + local result_dir="" + local workspace_dir="" + local use_chat_template=false + local dsv4=false + local trust_remote_code=false + local server_pid="" + local tokenizer="" + local tokenizer_mode="" + + while [[ $# -gt 0 ]]; do + case $1 in + --model) + model="$2" + shift 2 + ;; + --port) + port="$2" + shift 2 + ;; + --backend) + backend="$2" + shift 2 + ;; + --endpoint) + endpoint="$2" + shift 2 + ;; + --input-len) + input_len="$2" + shift 2 + ;; + --output-len) + output_len="$2" + shift 2 + ;; + --random-range-ratio) + random_range_ratio="$2" + shift 2 + ;; + --num-prompts) + num_prompts="$2" + shift 2 + ;; + --max-concurrency) + max_concurrency="$2" + shift 2 + ;; + --result-filename) + result_filename="$2" + shift 2 + ;; + --result-dir) + result_dir="$2" + shift 2 + ;; + --bench-serving-dir) + workspace_dir="$2" + shift 2 + ;; + --use-chat-template) + use_chat_template=true + shift + ;; + --dsv4) + dsv4=true + use_chat_template=true + shift + ;; + --trust-remote-code) + trust_remote_code=true + shift + ;; + --server-pid) + server_pid="$2" + shift 2 + ;; + --tokenizer) + tokenizer="$2" + shift 2 + ;; + --tokenizer-mode) + tokenizer_mode="$2" + shift 2 + ;; + *) + echo "Unknown parameter: $1" + return 1 + ;; + esac + done + + # Validate all required parameters + if [[ -z "$model" ]]; then + echo "Error: --model is required" + return 1 + fi + if [[ -z "$port" ]]; then + echo "Error: --port is required" + return 1 + fi + if [[ -z "$backend" ]]; then + echo "Error: --backend is required" + return 1 + fi + if [[ -z "$input_len" ]]; then + echo "Error: --input-len is required" + return 1 + fi + if [[ -z "$output_len" ]]; then + echo "Error: --output-len is required" + return 1 + fi + if [[ -z "$random_range_ratio" ]]; then + echo "Error: --random-range-ratio is required" + return 1 + fi + if [[ -z "$num_prompts" ]]; then + echo "Error: --num-prompts is required" + return 1 + fi + if [[ -z "$max_concurrency" ]]; then + echo "Error: --max-concurrency is required" + return 1 + fi + if [[ -z "$result_filename" ]]; then + echo "Error: --result-filename is required" + return 1 + fi + if [[ -z "$result_dir" ]]; then + echo "Error: --result-dir is required" + return 1 + fi + + if [[ -z "$workspace_dir" ]]; then + workspace_dir=$(pwd) + fi + + # Profiling support: when PROFILE=1, ensure profiler dir exists, add --profile flag, + # and cap num_prompts to keep traces small. + local profile_flag=() + if [[ "${PROFILE:-}" == "1" ]]; then + local _prof_dir="${SGLANG_TORCH_PROFILER_DIR:-${VLLM_TORCH_PROFILER_DIR:-}}" + if [[ -n "$_prof_dir" ]]; then + mkdir -p "$_prof_dir" + fi + profile_flag+=(--profile) + num_prompts="$max_concurrency" + fi + + # Build benchmark command + local benchmark_cmd=( + python3 "$workspace_dir/utils/bench_serving/benchmark_serving.py" + --model "$model" + --backend "$backend" + --base-url "http://0.0.0.0:$port" + --dataset-name random + --random-input-len "$input_len" + --random-output-len "$output_len" + --random-range-ratio "$random_range_ratio" + --num-prompts "$num_prompts" + --max-concurrency "$max_concurrency" + --request-rate inf + --ignore-eos + "${profile_flag[@]}" + --save-result + --num-warmups "$((2 * max_concurrency))" \ + --percentile-metrics 'ttft,tpot,itl,e2el' + --result-dir "$result_dir" + --result-filename "$result_filename.json" + ) + + if [[ -n "$endpoint" ]]; then + benchmark_cmd+=(--endpoint "$endpoint") + fi + + # Add --use-chat-template if requested + if [[ "$use_chat_template" == true ]]; then + benchmark_cmd+=(--use-chat-template) + fi + + # Add --dsv4 if requested (requires --use-chat-template, which we + # auto-enable when --dsv4 is passed in). + if [[ "$dsv4" == true ]]; then + benchmark_cmd+=(--dsv4) + fi + + # Add --trust-remote-code if requested + if [[ "$trust_remote_code" == true ]]; then + benchmark_cmd+=(--trust-remote-code) + fi + + if [[ -n "$tokenizer" ]]; then + benchmark_cmd+=(--tokenizer "$tokenizer") + fi + + if [[ -n "$tokenizer_mode" ]]; then + benchmark_cmd+=(--tokenizer-mode "$tokenizer_mode") + fi + + # Run benchmark with optional server monitoring + set -x + if [[ -n "$server_pid" ]]; then + # Run benchmark in background and monitor server health + "${benchmark_cmd[@]}" & + local benchmark_pid=$! + + # Monitor loop: check both benchmark and server status + while kill -0 "$benchmark_pid" 2>/dev/null; do + if ! kill -0 "$server_pid" 2>/dev/null; then + echo "ERROR: Server process $server_pid died during benchmark" + kill "$benchmark_pid" 2>/dev/null + wait "$benchmark_pid" 2>/dev/null + set +x + return 1 + fi + sleep 2 + done + + # Benchmark finished, get its exit code + wait "$benchmark_pid" + local benchmark_exit_code=$? + else + # No server monitoring, run benchmark directly + "${benchmark_cmd[@]}" + local benchmark_exit_code=$? + fi + set +x + + # If profiling, move trace to relay-upload location + if [[ "${PROFILE:-}" == "1" ]]; then + move_profile_trace_for_relay + fi + + return $benchmark_exit_code +} + + +# -------------------------------- +# Profiling trace helpers +# -------------------------------- + +_find_latest_profile_trace() { + local latest="" + local dir="" candidate="" base="" + local -a search_roots=() + + for dir in "$@"; do + search_roots=() + if [[ -d "$dir" ]]; then + search_roots+=("$dir") + fi + if [[ -d "$dir/profiles" ]]; then + search_roots+=("$dir/profiles") + fi + if [[ ${#search_roots[@]} -eq 0 ]]; then + continue + fi + + while IFS= read -r -d '' candidate; do + base="$(basename "$candidate")" + if [[ "$base" == profile_*.trace.json.gz ]]; then + continue + fi + if [[ -z "$latest" || "$candidate" -nt "$latest" ]]; then + latest="$candidate" + fi + done < <( + find "${search_roots[@]}" -maxdepth 1 -type f \ + \( -name "*.trace.json" -o -name "*.trace.json.gz" -o -name "*trace*.json" -o -name "*trace*.json.gz" -o -name "*profile*.json" -o -name "*profile*.json.gz" \) \ + -print0 2>/dev/null + ) + done + + printf '%s' "$latest" +} + +# Move profiler trace into a stable workspace path for workflow relay/upload. +move_profile_trace_for_relay() { + if [[ "${PROFILE:-}" != "1" ]]; then + return 0 + fi + + if [[ -z "${RESULT_FILENAME:-}" ]]; then + echo "[PROFILE] RESULT_FILENAME is not set; skipping relay trace staging." >&2 + return 0 + fi + + local sglang_dir="${SGLANG_TORCH_PROFILER_DIR:-/workspace}" + local vllm_dir="${VLLM_TORCH_PROFILER_DIR:-/workspace}" + local -a search_dirs=() + local dir="" existing="" + local seen=0 + + for dir in "$sglang_dir" "$vllm_dir" "/workspace"; do + if [[ -z "$dir" ]]; then + continue + fi + seen=0 + for existing in "${search_dirs[@]}"; do + if [[ "$existing" == "$dir" ]]; then + seen=1 + break + fi + done + if [[ "$seen" -eq 0 ]]; then + search_dirs+=("$dir") + fi + done + + local trace_file="" + local wait_attempts=10 + for (( i=1; i<=wait_attempts; i++ )); do + trace_file="$(_find_latest_profile_trace "${search_dirs[@]}")" + if [[ -n "$trace_file" ]]; then + break + fi + sleep 10 + done + + if [[ -z "$trace_file" ]]; then + echo "[PROFILE] No trace found for relay under: ${search_dirs[*]}" >&2 + return 0 + fi + + local dest_trace="/workspace/profile_${RESULT_FILENAME}.trace.json.gz" + if [[ "$trace_file" == *.gz ]]; then + cp -f "$trace_file" "$dest_trace" + else + gzip -c "$trace_file" > "$dest_trace" + fi + + echo "[PROFILE] Relay trace prepared: $dest_trace (source: $trace_file)" +} + + +# ------------------------------ +# Eval (lm-eval-harness) helpers +# ------------------------------ + +_install_lm_eval_deps() { + # torchvision causes circular imports in ATOM; TRT-LLM/SGLang need it at module level. + if [[ "${IMAGE:-}" == *atom* ]]; then + python3 -m pip uninstall -y torchvision 2>/dev/null || true + fi + python3 -m pip install -q --no-cache-dir --break-system-packages "lm-eval[api]" || true + local lm_eval_ref="b315ef3b05176acc9732bb7fdec116abe1ecc476" + if command -v git >/dev/null 2>&1; then + if ! python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "git+https://github.com/EleutherAI/lm-evaluation-harness.git@${lm_eval_ref}"; then + python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true + fi + else + python3 -m pip install -q --no-cache-dir --no-deps --force-reinstall --break-system-packages \ + "https://github.com/EleutherAI/lm-evaluation-harness/archive/${lm_eval_ref}.tar.gz" || true + fi +} + +# Patch lm-eval filters to be robust to empty strings via sitecustomize +_patch_lm_eval() { + local patch_dir + patch_dir="$(mktemp -d)" + cat > "$patch_dir/sitecustomize.py" <<'PY' +# --- Patch LocalChatCompletion.parse_generations to handle empty content with reasoning_content --- +import re, sys, unicodedata, json +from lm_eval.filters import extraction as ex +from lm_eval.models.openai_completions import LocalChatCompletion as _LCC + +def _le_parse_generations(outputs, **kwargs): + res = [] + if not isinstance(outputs, list): + outputs = [outputs] + for out in (outputs or []): + try: + choices = out.get("choices", []) + tmp = ["" for _ in choices] + for choice in choices: + idx = choice.get("index", 0) + msg = (choice.get("message") or {}) + content = msg.get("content") + if content in (None, "", []): + content = msg.get("reasoning_content") or "" + tmp[idx] = content + except Exception: + tmp = [""] + res.extend(tmp) + return res + +# Keep staticmethod semantics +_LCC.parse_generations = staticmethod(_le_parse_generations) + +# --- Patch TemplateAPI.apply_chat_template to avoid injecting "type": "text" for TRT --- +try: + from lm_eval.models import api_models as _api_models + _TemplateAPI = _api_models.TemplateAPI + _JsonChatStr = _api_models.JsonChatStr +except Exception: + _TemplateAPI = None + _JsonChatStr = None + +if _TemplateAPI is not None and _JsonChatStr is not None: + _orig_apply_chat_template = _TemplateAPI.apply_chat_template + + def _patched_apply_chat_template( + self, + chat_history, + add_generation_prompt: bool = True, + ): + """Applies a chat template to a list of chat history between user and model.""" + if self.tokenizer_backend == "huggingface" and self.tokenized_requests: + return self.tokenizer.apply_chat_template( + chat_history, + tokenize=False, + add_generation_prompt=add_generation_prompt, + continue_final_message=not add_generation_prompt, + ) + elif self.tokenizer_backend == "remote" and self.tokenized_requests: + return chat_history + else: + # NOTE: we no longer inject `"type": "text"` when tokenizer is None / non-HF + return _JsonChatStr( + json.dumps( + [{**item} for item in chat_history], + ensure_ascii=False, + ) + ) + + _TemplateAPI.apply_chat_template = _patched_apply_chat_template +PY + export PYTHONPATH="${patch_dir}:${PYTHONPATH:-}" +} + +get_native_max_context_length() { + local model_path="$1" + # Prefer MODEL_PATH (local model directory) when available, since the + # argument may be a served-model name that is neither a valid HF repo + # ID nor a local path (e.g. "deepseek-r1-fp4" on the B300 cluster). + if [ -n "${MODEL_PATH:-}" ] && [ -d "${MODEL_PATH}" ]; then + model_path="${MODEL_PATH}" + fi + python3 -c " +try: + from transformers import AutoConfig + config = AutoConfig.from_pretrained('${model_path}', trust_remote_code=True) + for attr in ['max_position_embeddings', 'max_sequence_length', 'seq_length', 'n_positions']: + if hasattr(config, attr): + print(getattr(config, attr)) + break + else: + print(0) +except Exception: + print(0) +" +} + +# Compute the context length for eval-only mode. +# Uses the requested benchmark context capped at the model's native max. +# Sets EVAL_MAX_MODEL_LEN (needed by run_lm_eval). +# Echoes the computed value for scripts to capture. +# +# Usage: local ctx=$(compute_eval_context_length "$MODEL" "${current_ctx}") +compute_eval_context_length() { + local model="$1" + local benchmark_ctx="${2:-0}" + local native_max + native_max=$(get_native_max_context_length "$model") + native_max="${native_max:-0}" + + if [ "$benchmark_ctx" -eq 0 ] 2>/dev/null; then + benchmark_ctx="${native_max:-0}" + fi + local eval_ctx=$(( benchmark_ctx * 1 )) + if [ "$native_max" -gt 0 ] 2>/dev/null && [ "$eval_ctx" -gt "$native_max" ]; then + eval_ctx="$native_max" + fi + # If eval_ctx is still 0 (both benchmark_ctx and native_max were 0), fall back + if [ "$eval_ctx" -le 0 ] 2>/dev/null; then + echo "WARN: compute_eval_context_length could not determine context length for $model" >&2 + eval_ctx="${MAX_MODEL_LEN:-16384}" + fi + EVAL_MAX_MODEL_LEN="$eval_ctx" + echo "$eval_ctx" +} + +# Convenience wrapper: compute eval context from ISL/OSL and export EVAL_MAX_MODEL_LEN. +# Call directly (not in a subshell) so the export persists. +# Scripts then wire $EVAL_MAX_MODEL_LEN into whichever server variable they need. +setup_eval_context() { + EVAL_MAX_MODEL_LEN=$(compute_eval_context_length "$MODEL" "$((ISL + OSL + 256))") + export EVAL_MAX_MODEL_LEN +} + +run_lm_eval() { + local port="${PORT:-8888}" + local tasks_dir="${EVAL_TASKS_DIR:-utils/evals/gsm8k.yaml}" + local results_dir="${EVAL_RESULT_DIR:-$(mktemp -d /tmp/eval_out-XXXXXX)}" + local eval_context_len="${EVAL_MAX_MODEL_LEN:-16384}" + local temperature=0 + local top_p=1 + local concurrent_requests="${EVAL_CONCURRENT_REQUESTS:-${CONC:-64}}" + + while [[ $# -gt 0 ]]; do + case $1 in + --port) port="$2"; shift 2 ;; + --task) tasks_dir="$2"; shift 2 ;; + --results-dir) results_dir="$2"; shift 2 ;; + --gen-max-tokens) eval_context_len="$2"; shift 2 ;; + --temperature) temperature="$2"; shift 2 ;; + --top-p) top_p="$2"; shift 2 ;; + *) echo "Unknown parameter: $1"; return 1 ;; + esac + done + + # Anchor a relative task-yaml to the repo root. On the llmd-vllm path + # the eval runs inside the serving container, whose WORKDIR is + # /vllm-workspace, not the repo bind-mount (/workspace) - so a relative + # path like "utils/evals/gsm8k.yaml" resolves to a nonexistent file and + # lm_eval fails with "Tasks not found". benchmark_lib.sh always lives at + # /benchmarks/, so derive the repo root from BASH_SOURCE and + # relocate the path there. Only rewrites a relative *.yaml that is + # missing from cwd but present under the repo root; builtin lm_eval task + # names (no .yaml), absolute paths, and paths that already resolve from + # cwd (the dynamo/srt-slurm path) are left untouched. + local _repo_root + _repo_root="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)" + if [[ "$tasks_dir" == *.yaml && "$tasks_dir" != /* \ + && ! -f "$tasks_dir" && -f "$_repo_root/$tasks_dir" ]]; then + echo "run_lm_eval: anchoring relative task '$tasks_dir' to repo root -> $_repo_root/$tasks_dir" + tasks_dir="$_repo_root/$tasks_dir" + fi + + if [ "${INFERENCEX_LM_EVAL_RUNTIME_READY:-false}" != "true" ]; then + _install_lm_eval_deps + _patch_lm_eval + export INFERENCEX_LM_EVAL_RUNTIME_READY=true + fi + + local openai_server_base="http://0.0.0.0:${port}" + local openai_chat_base="${openai_server_base}/v1/chat/completions" + export OPENAI_API_KEY=${OPENAI_API_KEY:-EMPTY} + MODEL_NAME=${MODEL_NAME:-$MODEL} # Prefer MODEL_NAME, else MODEL + + # Cap output tokens: must fit within context window (leave room for input), + # and avoid excessive KV cache reservation per request on TRT. + local max_output_tokens=$(( eval_context_len > 4096 ? eval_context_len - 4096 : eval_context_len / 2 )) + if [ "$max_output_tokens" -gt 16384 ]; then + max_output_tokens=16384 + fi + echo "Eval budget: eval_context_len=${eval_context_len}, max_output_tokens=${max_output_tokens}" + + # Export for append_lm_eval_summary to pick up + export EVAL_RESULT_DIR="$results_dir" + set -x + python3 -m lm_eval --model local-chat-completions --apply_chat_template \ + --tasks "${tasks_dir}" \ + --output_path "${results_dir}" \ + --log_samples \ + --model_args "model=${MODEL_NAME},base_url=${openai_chat_base},api_key=${OPENAI_API_KEY},eos_string=,max_retries=5,num_concurrent=${concurrent_requests},timeout=1800,tokenized_requests=False,max_length=${eval_context_len}" \ + --gen_kwargs "max_tokens=${max_output_tokens},temperature=${temperature},top_p=${top_p}" + local eval_exit=$? + set +x + return $eval_exit +} + +_stage_lm_eval_artifacts() { + local results_dir="$1" + local eval_conc="$2" + local moved=0 + local failed=0 + local jf base stem extension target suffix + + if [ ! -d "${results_dir}" ]; then + echo "WARN: eval result directory '${results_dir}' does not exist" >&2 + return 1 + fi + + while IFS= read -r -d '' jf; do + base=$(basename "$jf") + case "$base" in + meta_env.json) + continue + ;; + *.jsonl) + stem="${base%.jsonl}" + extension=".jsonl" + ;; + *.json) + stem="${base%.json}" + extension=".json" + ;; + *) + continue + ;; + esac + + target="./${stem}_conc${eval_conc}${extension}" + suffix=2 + while [ -e "$target" ]; do + target="./${stem}_conc${eval_conc}_${suffix}${extension}" + suffix=$((suffix + 1)) + done + + if mv -f "$jf" "$target"; then + moved=1 + else + echo "WARN: failed to stage eval artifact ${jf}" >&2 + failed=1 + fi + done < <( + find "${results_dir}" -type f \ + \( -name "*.json" -o -name "*.jsonl" \) -print0 2>/dev/null + ) + + rm -rf --one-file-system "${results_dir}" 2>/dev/null \ + || rm -rf "${results_dir}" \ + || true + + if [ "$moved" -eq 0 ]; then + echo "WARN: no eval artifacts were produced for concurrency ${eval_conc}" >&2 + return 1 + fi + return "$failed" +} + +_eval_concs_to_json() { + local values="$1" + local value + local joined="" + + for value in $values; do + if [[ ! "$value" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: invalid eval concurrency '${value}'" >&2 + return 1 + fi + if [ -n "$joined" ]; then + joined="${joined}, " + fi + joined="${joined}${value}" + done + + printf '[%s]' "$joined" +} + +append_lm_eval_summary() { + local batch_concs="${EVAL_BATCHED_CONCS:-}" + local results_dir="${EVAL_RESULT_DIR:-}" + local out_dir="${results_dir}" + local meta_json + local metadata_conc="${CONC:-1}" + local batch_metadata="" + + if [ -n "$batch_concs" ]; then + meta_json="./meta_env.json" + metadata_conc="${batch_concs%% *}" + + local eval_concs_json completed_concs_json failed_concs_json + eval_concs_json=$(_eval_concs_to_json "$batch_concs") || return 1 + completed_concs_json=$( + _eval_concs_to_json "${EVAL_BATCHED_COMPLETED_CONCS:-}" + ) || return 1 + failed_concs_json=$( + _eval_concs_to_json "${EVAL_BATCHED_FAILED_CONCS:-}" + ) || return 1 + printf -v batch_metadata \ + ' "eval_concs": %s,\n "completed_eval_concs": %s,\n "failed_eval_concs": %s,\n' \ + "$eval_concs_json" \ + "$completed_concs_json" \ + "$failed_concs_json" + else + if [ -z "${results_dir}" ]; then + echo "WARN: EVAL_RESULT_DIR is empty; skipping artifact collection" >&2 + return 1 + fi + if [ ! -d "${out_dir}" ]; then + echo "WARN: EVAL_RESULT_DIR='${out_dir}' does not exist; skipping artifact collection" >&2 + return 1 + fi + meta_json="${out_dir}/meta_env.json" + fi + + # Write minimal meta for collectors that expect it + local model_name="${MODEL_NAME:-$MODEL}" + local is_multinode_json="false" + if [ "${IS_MULTINODE:-false}" = "true" ]; then + is_multinode_json="true" + fi + + local prefill_tp="${PREFILL_TP:-${TP:-1}}" + local prefill_pp="${PREFILL_PP_SIZE:-${PP_SIZE:-1}}" + local prefill_dcp_size="${PREFILL_DCP_SIZE:-${DCP_SIZE:-1}}" + local prefill_pcp_size="${PREFILL_PCP_SIZE:-${PCP_SIZE:-1}}" + local prefill_ep="${PREFILL_EP:-${EP_SIZE:-1}}" + local prefill_num_workers="${PREFILL_NUM_WORKERS:-1}" + local decode_tp="${DECODE_TP:-${TP:-1}}" + local decode_pp="${DECODE_PP_SIZE:-${PP_SIZE:-1}}" + local decode_dcp_size="${DECODE_DCP_SIZE:-${DCP_SIZE:-1}}" + local decode_pcp_size="${DECODE_PCP_SIZE:-${PCP_SIZE:-1}}" + local decode_ep="${DECODE_EP:-${EP_SIZE:-1}}" + local decode_num_workers="${DECODE_NUM_WORKERS:-1}" + + local dp_json="false" + if [ "${DP_ATTENTION:-false}" = "true" ]; then dp_json="true"; fi + local prefill_dp_json="$dp_json" + if [ "${PREFILL_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then + prefill_dp_json="true" + else + prefill_dp_json="false" + fi + local decode_dp_json="$dp_json" + if [ "${DECODE_DP_ATTENTION:-${DP_ATTENTION:-false}}" = "true" ]; then + decode_dp_json="true" + else + decode_dp_json="false" + fi + + # Derive framework/precision from env, fallback to parsing RESULT_FILENAME + # RESULT_FILENAME format (from workflow): + # ___tp<...>_ep<...>_dpa_<...>_conc<...>_ + local fw="${FRAMEWORK:-}" + local prec="${PRECISION:-}" + if [[ -z "$fw" || -z "$prec" ]]; then + if [[ -n "${RESULT_FILENAME}" ]]; then + # Extract the two fields immediately before "_tp" + # Handles arbitrary underscores in exp_name by matching from the end + local parsed + parsed=$(echo "${RESULT_FILENAME}" | sed -n 's/.*_\([^_][^_]*\)_\([^_][^_]*\)_tp.*/\1 \2/p') + local p1="${parsed%% *}" + local p2="${parsed#* }" + if [[ -z "$prec" && -n "$p1" && "$p1" != "$parsed" ]]; then + prec="$p1" + fi + if [[ -z "$fw" && -n "$p2" && "$p2" != "$parsed" ]]; then + fw="$p2" + fi + fi + fi + cat > "${meta_json}" <&2 + fi + if [ -d "${out_dir}" ]; then + while IFS= read -r -d '' jf; do + base=$(basename "$jf") + if [ "$base" != "meta_env.json" ]; then + mv -f "$jf" ./ || echo "WARN: failed to move ${jf}" >&2 + fi + done < <(find "${out_dir}" -type f -name "*.json*" -print0 2>/dev/null) + fi + + # Best-effort cleanup of the temp directory + if [ -n "${out_dir}" ] && [ -d "${out_dir}" ]; then + rm -rf --one-file-system "${out_dir}" || rm -rf "${out_dir}" || true + fi + + echo "Moved eval artifacts to: $(pwd)" +} + +# ------------------------------ +# Unified eval entrypoint +# ------------------------------ + +run_eval() { + local framework="${EVAL_FRAMEWORK:-lm-eval}" + local forwarded=() + + while [[ $# -gt 0 ]]; do + case "$1" in + --framework) framework="$2"; shift 2 ;; + *) forwarded+=("$1"); shift ;; + esac + done + + # Compute EVAL_MAX_MODEL_LEN if not already set by the calling script + if [ -z "${EVAL_MAX_MODEL_LEN:-}" ]; then + compute_eval_context_length "$MODEL" "${MAX_MODEL_LEN:-0}" > /dev/null + fi + + unset EVAL_BATCHED_CONCS + unset EVAL_BATCHED_COMPLETED_CONCS + unset EVAL_BATCHED_FAILED_CONCS + + local requested_concs="${EVAL_CONCURRENT_REQUESTS:-}" + local eval_concs=() + if [ -n "$requested_concs" ]; then + read -r -a eval_concs <<< "$requested_concs" + fi + + if [ "${#eval_concs[@]}" -gt 1 ]; then + if [[ "$framework" != "lm-eval" && "$framework" != "lm_eval" ]]; then + echo "ERROR: batched eval concurrency is only supported for lm-eval" >&2 + return 1 + fi + + local eval_conc results_dir eval_rc stage_rc + local completed_concs=() + local failed_concs=() + + for eval_conc in "${eval_concs[@]}"; do + if [[ ! "$eval_conc" =~ ^[1-9][0-9]*$ ]]; then + echo "ERROR: invalid eval concurrency '${eval_conc}'" >&2 + return 1 + fi + + if ! results_dir=$(mktemp -d /tmp/eval_out-conc"${eval_conc}"-XXXXXX); then + echo "ERROR: failed to create eval output directory for concurrency ${eval_conc}" >&2 + failed_concs+=("$eval_conc") + continue + fi + + echo "Running lm-eval at concurrency ${eval_conc} using the existing engine" + export EVAL_CONCURRENT_REQUESTS="$eval_conc" + export CONC="$eval_conc" + eval_rc=0 + stage_rc=0 + run_lm_eval "${forwarded[@]}" --results-dir "$results_dir" \ + || eval_rc=$? + _stage_lm_eval_artifacts "$results_dir" "$eval_conc" \ + || stage_rc=$? + + if [ "$eval_rc" -eq 0 ] && [ "$stage_rc" -eq 0 ]; then + completed_concs+=("$eval_conc") + else + echo "ERROR: lm-eval failed at concurrency ${eval_conc} (eval_rc=${eval_rc}, stage_rc=${stage_rc})" >&2 + failed_concs+=("$eval_conc") + fi + done + + export EVAL_CONCURRENT_REQUESTS="$requested_concs" + export EVAL_RESULT_DIR="" + export EVAL_BATCHED_CONCS="${eval_concs[*]}" + export EVAL_BATCHED_COMPLETED_CONCS="${completed_concs[*]}" + export EVAL_BATCHED_FAILED_CONCS="${failed_concs[*]}" + + if [ "${#failed_concs[@]}" -gt 0 ]; then + echo "ERROR: batched eval failed for concurrency: ${failed_concs[*]}" >&2 + echo "Deferring failure until post-upload score validation preserves all artifacts" >&2 + fi + return 0 + fi + + local eval_rc=0 + case "$framework" in + lm-eval|lm_eval) run_lm_eval "${forwarded[@]}" || eval_rc=$? ;; + *) echo "Unknown framework '${framework}'"; eval_rc=1 ;; + esac + + if [ "$eval_rc" -ne 0 ]; then + echo "ERROR: run_eval failed with exit code $eval_rc" >&2 + if [ "${EVAL_ONLY}" = "true" ]; then + echo "Eval-only mode: failing after artifact collection" >&2 + return "$eval_rc" + fi + fi + return $eval_rc +} + + +# -------------------------------- +# Agentic trace replay helpers (aiperf driver) +# -------------------------------- + +INFMAX_CONTAINER_WORKSPACE="${INFMAX_CONTAINER_WORKSPACE:-/workspace}" +AGENTIC_DIR="${AGENTIC_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/agentic-benchmark}" +AIPERF_DIR="${AIPERF_DIR:-${INFMAX_CONTAINER_WORKSPACE}/utils/aiperf}" +AIPERF_RUNTIME_DIR="${AIPERF_RUNTIME_DIR:-${TMPDIR:-/tmp}/inferencex-agentic-${SLURM_JOB_ID:-$$}}" +AIPERF_VENV="${AIPERF_VENV:-${AIPERF_RUNTIME_DIR}/venv}" +AIPERF_UV_INSTALL_DIR="${AIPERF_UV_INSTALL_DIR:-${AIPERF_RUNTIME_DIR}/uv/bin}" +AIPERF_UV_CACHE_DIR="${AIPERF_UV_CACHE_DIR:-${AIPERF_RUNTIME_DIR}/uv-cache}" +AIPERF_PYTHON="${AIPERF_VENV}/bin/python" +AIPERF_CLI="${AIPERF_VENV}/bin/aiperf" +AIPERF_HF_CLI="${AIPERF_VENV}/bin/hf" +AIPERF_DEPS_READY=0 +AIPERF_FAILED_REQUEST_THRESHOLD="${AIPERF_FAILED_REQUEST_THRESHOLD:-0.10}" + +agentic_pip_install() { + local pip_install=(python3 -m pip install) + if python3 -m pip install --help 2>/dev/null | grep -q -- "--break-system-packages"; then + pip_install+=(--break-system-packages) + fi + + "${pip_install[@]}" "$@" +} + +ensure_agentic_uv() { + if command -v uv >/dev/null 2>&1; then + AIPERF_UV_BIN="$(command -v uv)" + return + fi + + AIPERF_UV_BIN="${AIPERF_UV_INSTALL_DIR}/uv" + if [ ! -x "$AIPERF_UV_BIN" ]; then + mkdir -p "$AIPERF_UV_INSTALL_DIR" + curl -LsSf https://astral.sh/uv/install.sh | + UV_INSTALL_DIR="$AIPERF_UV_INSTALL_DIR" sh + fi + + if [ ! -x "$AIPERF_UV_BIN" ]; then + echo "ERROR: uv installation did not create $AIPERF_UV_BIN" >&2 + return 1 + fi +} + +install_agentic_deps() { + if [ "$AIPERF_DEPS_READY" = "1" ]; then + return + fi + + # AIPerf must not share site-packages with the inference server. Installing + # it into vLLM/SGLang's system Python can upgrade FastAPI, Starlette, + # transformers, or other packages while the server imports from that same + # environment. + if ! command -v git >/dev/null 2>&1; then + apt-get update && apt-get install -y git + fi + + ensure_agentic_uv + rm -rf "$AIPERF_VENV" + mkdir -p "$AIPERF_UV_CACHE_DIR" + + UV_CACHE_DIR="$AIPERF_UV_CACHE_DIR" \ + "$AIPERF_UV_BIN" venv --python "$(command -v python3)" "$AIPERF_VENV" + UV_CACHE_DIR="$AIPERF_UV_CACHE_DIR" \ + "$AIPERF_UV_BIN" pip install --python "$AIPERF_PYTHON" \ + -r "$AGENTIC_DIR/requirements.txt" \ + -e "$AIPERF_DIR" \ + "datasets>=4.7.0" \ + "huggingface_hub[cli]>=0.25.0" \ + urllib3 \ + requests + + if [ ! -x "$AIPERF_CLI" ] || [ ! -x "$AIPERF_HF_CLI" ]; then + echo "ERROR: isolated AIPerf environment is incomplete at $AIPERF_VENV" >&2 + return 1 + fi + AIPERF_DEPS_READY=1 +} + +ensure_hf_cli() { + install_agentic_deps +} + +resolve_trace_source() { + # Per-recipe override: set WEKA_LOADER_OVERRIDE to one of the aiperf + # public-dataset loader names allowed by the inferencex-agentx-mvp + # scenario. Used by recipes whose servers have non-default context + # caps (e.g. minimaxm2.5 at max_model_len ~256k can't replay the + # unfiltered corpus and switches to the 256k-capped variant), or + # by recipes that want to pin an older corpus generation. + # + # Default (no override): the 062126 v7 corpus, selected by the model + # family's native context length. Models with a 1M-token default context + # use the unfiltered corpus; shorter-context families use the 256k-capped + # variant. Any recipe can still pin a specific corpus via + # WEKA_LOADER_OVERRIDE. + local default_loader + case "${MODEL_PREFIX:-}" in + dsv4*|minimaxm3*) + default_loader="semianalysis_cc_traces_weka_062126" + ;; + *) + default_loader="semianalysis_cc_traces_weka_062126_256k" + ;; + esac + local loader="${WEKA_LOADER_OVERRIDE:-$default_loader}" + local dataset + case "$loader" in + semianalysis_cc_traces_weka_with_subagents) + dataset="semianalysisai/cc-traces-weka-061526" + ;; + semianalysis_cc_traces_weka_with_subagents_256k) + dataset="semianalysisai/cc-traces-weka-061526-256k" + ;; + semianalysis_cc_traces_weka_with_subagents_060226) + dataset="semianalysisai/cc-traces-weka-with-subagents-060226" + ;; + semianalysis_cc_traces_weka_with_subagents_060226_256k) + dataset="semianalysisai/cc-traces-weka-with-subagents-060226-256k" + ;; + semianalysis_cc_traces_weka_with_subagents_060526) + dataset="semianalysisai/cc-traces-weka-with-subagents-060526" + ;; + semianalysis_cc_traces_weka_with_subagents_060526_256k) + dataset="semianalysisai/cc-traces-weka-with-subagents-060526-256k" + ;; + semianalysis_cc_traces_weka_with_subagents_060826) + dataset="semianalysisai/cc-traces-weka-with-subagents-060826" + ;; + semianalysis_cc_traces_weka_with_subagents_060826_256k) + dataset="semianalysisai/cc-traces-weka-with-subagents-060826-256k" + ;; + semianalysis_cc_traces_weka_061326) + dataset="semianalysisai/cc-traces-weka-061326" + ;; + semianalysis_cc_traces_weka_061326_256k) + dataset="semianalysisai/cc-traces-weka-061326-256k" + ;; + semianalysis_cc_traces_weka_061526) + dataset="semianalysisai/cc-traces-weka-061526" + ;; + semianalysis_cc_traces_weka_061526_256k) + dataset="semianalysisai/cc-traces-weka-061526-256k" + ;; + semianalysis_cc_traces_weka_062126) + dataset="semianalysisai/cc-traces-weka-062126" + ;; + semianalysis_cc_traces_weka_062126_256k) + dataset="semianalysisai/cc-traces-weka-062126-256k" + ;; + *) + echo "Error: unknown WEKA_LOADER_OVERRIDE='$loader'. Allowed: semianalysis_cc_traces_weka_with_subagents, semianalysis_cc_traces_weka_with_subagents_256k, semianalysis_cc_traces_weka_with_subagents_060226, semianalysis_cc_traces_weka_with_subagents_060226_256k, semianalysis_cc_traces_weka_with_subagents_060526, semianalysis_cc_traces_weka_with_subagents_060526_256k, semianalysis_cc_traces_weka_with_subagents_060826, semianalysis_cc_traces_weka_with_subagents_060826_256k, semianalysis_cc_traces_weka_061326, semianalysis_cc_traces_weka_061326_256k, semianalysis_cc_traces_weka_061526, semianalysis_cc_traces_weka_061526_256k, semianalysis_cc_traces_weka_062126, semianalysis_cc_traces_weka_062126_256k" >&2 + exit 1 + ;; + esac + TRACE_SOURCE_FLAG="--public-dataset $loader" + echo "Loading traces via aiperf public-dataset: $loader ($dataset) [MODEL_PREFIX=${MODEL_PREFIX:-unset}]" + # Pre-download the dataset into the shared HF_HUB_CACHE (same mount used + # for model weights) so subsequent runs read from cache instead of + # re-downloading every job. + ensure_hf_cli + "$AIPERF_HF_CLI" download --repo-type dataset "$dataset" +} + +build_replay_cmd() { + # aiperf invocation for the inferencex-agentx-mvp scenario. + # + # Pre-canned assistant replay is the default: recorded assistant responses + # are used for future prompt construction, and live server responses are + # discarded. Set AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES=1 explicitly + # to use live-assistant mode, where the loader emits user-only deltas and + # the worker threads the server's live assistant response back into the + # session. + # + # The scenario plugin locks: --cache-bust first_turn_prefix and + # --trace-idle-gap-cap-seconds 10 (per-trace idle-gap compression + # against parent + subagent request-start timestamps; supersedes the + # legacy --use-think-time-only / --inter-turn-delay-cap-seconds path), + # and auto-injects them — so we do not pass them. See + # utils/aiperf/docs/tutorials/agentx-mvp.md. + local result_dir="$1" + local duration="$DURATION" + + export AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES="${AIPERF_DATASET_WEKA_LIVE_ASSISTANT_RESPONSES:-0}" + # Dataset configuration (load + reconstruct + inputs.json + mmap) + # routinely takes 4-5 min for the Weka corpus on fast /tmp + # (B300) but can stretch to 14 min on slower /tmp + parallel contention + # (observed on H200 where all 14 R3 jobs hit aiperf's 900s Configure + # Profiling timeout simultaneously). Bump to 1800s to absorb 3x + # worst-case slowdown — the post-setup measurement window is unaffected. + export AIPERF_DATASET_CONFIGURATION_TIMEOUT=1800 + # aiperf validates that SERVICE_PROFILE_CONFIGURE_TIMEOUT >= + # DATASET_CONFIGURATION_TIMEOUT at startup. Bump it in lockstep. + export AIPERF_SERVICE_PROFILE_CONFIGURE_TIMEOUT=1800 + REPLAY_CMD="$AIPERF_CLI profile --scenario inferencex-agentx-mvp" + REPLAY_CMD+=" --url http://localhost:$PORT" + REPLAY_CMD+=" --endpoint /v1/chat/completions" + REPLAY_CMD+=" --endpoint-type chat" + REPLAY_CMD+=" --streaming" + REPLAY_CMD+=" --model $MODEL" + REPLAY_CMD+=" --concurrency $CONC" + REPLAY_CMD+=" --benchmark-duration $duration" + REPLAY_CMD+=" --random-seed 42" + # Fail runs once more than 10% of requests error. This keeps known + # transient low-rate failures from killing long sweeps while still + # catching malformed payloads or server crashes before they get aggregated + # as benchmarkable data. + REPLAY_CMD+=" --failed-request-threshold $AIPERF_FAILED_REQUEST_THRESHOLD" + # Sample each trajectory's warmup start position uniformly from + # [25%, 75%] of the trace's turn count, clamped by AIPerf to leave at + # least one profile turn after warmup. + REPLAY_CMD+=" --trajectory-start-min-ratio 0.25" + REPLAY_CMD+=" --trajectory-start-max-ratio 0.75" + # After the normal t* snapshot warmup, continue those exact trajectories + # with one-token outputs and no idle delays for 10 minutes. Profiling begins + # only after those requests drain and resumes from the resulting live state. + REPLAY_CMD+=" --agentic-cache-warmup-duration 600" + # Give long-context warmup requests up to 10 minutes to drain before + # cancelling any remaining requests and starting profiling. + REPLAY_CMD+=" --warmup-grace-period 600" + # Use server-reported usage fields (prompt_tokens / completion_tokens) for + # ISL/OSL instead of client-side tokenizer.encode(). Auto-enables + # stream_options.include_usage on the OpenAI chat endpoint. Skips the + # heavy per-record tokenization in the records pipeline that was pinning + # CPU on minimax-m2.5 at high concurrency. Lossless for vLLM (server + # usage is authoritative). + REPLAY_CMD+=" --use-server-token-count" + # Dynamo's KV router needs an explicit conversation session binding to + # keep later turns on the prefill worker that owns their prefix blocks. + # X-Correlation-ID is useful tracing metadata but does not establish that + # binding by itself. AIPerf emits nvext.session_control bind/close actions + # keyed by the stable conversation correlation ID when this flag is set. + if [[ "${FRAMEWORK:-}" == dynamo-* ]]; then + REPLAY_CMD+=" --use-dynamo-conv-aware-routing" + # The upstream 300s affinity TTL is shorter than an overloaded + # high-concurrency agentic request. Keep bindings alive across long + # prefills, generation, and capped inter-turn delay. This controls the + # router's inactivity lease; it does not relax HTTP/request failures. + REPLAY_CMD+=" --dynamo-session-timeout-seconds ${AIPERF_DYNAMO_SESSION_TIMEOUT_SECONDS:-3600}" + fi + # Disable DCGM GPU telemetry collection. aiperf's GpuMetricTimeSeries + # freezes its metric schema on the first DCGM scrape, then KeyErrors when + # an optional field (xid_errors, power_violation, encoder_utilization) + # first appears mid-run. We don't consume the gpu_telemetry artifact in + # downstream processing, and the server-metrics path (Prometheus /metrics + # from vLLM) is unaffected by this flag and still gives us KV usage, + # prefix cache hit rate, etc. + REPLAY_CMD+=" --no-gpu-telemetry" + # aiperf's dataset manager (separate from the inference parser) loads + # the model's tokenizer for trace-prompt tokenization regardless of + # --use-server-token-count. Models like kimi (amd/Kimi-K2.5-MXFP4, + # moonshotai/Kimi-K2.5) ship a custom tokenizer in their HF repo and + # need trust_remote_code=True to load. Benign for models without + # custom tokenizer code, so we set it unconditionally. + REPLAY_CMD+=" --tokenizer-trust-remote-code" + # Keep replay inputs inside the same context window used to launch the + # server. The WEKA corpus contains a few very long parent/subagent traces; + # if we mmap and replay them against a smaller-context server they become + # deterministic 4xxs and can still pressure the engine while queued. + if [ -n "${MAX_MODEL_LEN:-}" ] && [ "$MAX_MODEL_LEN" != "0" ]; then + REPLAY_CMD+=" --max-context-length $MAX_MODEL_LEN" + fi + # Default --num-dataset-entries is 100; the with-subagents Weka corpus + # has 393. Cap at 393 so all unique traces are loaded (the loader treats + # this as a ``min(cap, available)`` ceiling, not a target — see + # semianalysis_cc_traces_weka.py). + REPLAY_CMD+=" --num-dataset-entries 393" + # 1-second timeslices on the server-metrics scrape so the post-run + # plotter has per-window time series (KV usage, cache hit rate, + # throughput, etc.). Matches kv-cache-tester's poll_interval=1.0 + # snapshot cadence so metrics_plots.png is visually comparable. + # Without this, aiperf only emits aggregate stats and the 6x2 panels + # collapse to flat lines. + REPLAY_CMD+=" --slice-duration 1.0" + # Multi-node launchers can provide the Prometheus endpoints for every + # inference worker as a comma-separated list. AIPerf accepts multiple + # values after one --server-metrics flag and preserves endpoint_url on + # every exported series. The inference frontend's automatically detected + # /metrics endpoint remains enabled as well. + if [ -n "${AIPERF_SERVER_METRICS_URLS:-}" ]; then + local metrics_url + local -a metrics_urls + IFS=',' read -r -a metrics_urls <<< "$AIPERF_SERVER_METRICS_URLS" + REPLAY_CMD+=" --server-metrics" + for metrics_url in "${metrics_urls[@]}"; do + if [ -z "$metrics_url" ] || [[ "$metrics_url" == *[[:space:]]* ]]; then + echo "ERROR: AIPERF_SERVER_METRICS_URLS must be a comma-separated list of non-empty URLs" >&2 + return 1 + fi + REPLAY_CMD+=" $metrics_url" + done + fi + REPLAY_CMD+=" --output-artifact-dir $result_dir/aiperf_artifacts" + # The inferencex-agentx-mvp scenario enforces a 900s minimum + # benchmark duration. For smoke tests with shorter durations, opt + # into --unsafe-override (the run's submission_valid will be flagged + # false; that's expected for non-canonical runs). + if [ "$duration" -lt 900 ] || [ "${AIPERF_UNSAFE_OVERRIDE:-false}" = "true" ]; then + REPLAY_CMD+=" --unsafe-override" + fi + REPLAY_CMD+=" $TRACE_SOURCE_FLAG" +} + +write_agentic_result_json() { + # Aggregate aiperf's profile_export.{json,jsonl} + server_metrics_export.json + # into $AGENTIC_OUTPUT_DIR/$RESULT_FILENAME.json. The workflow checks that + # this file exists; run_agentic_replay_and_write_outputs separately rejects + # aggregates whose request error rate exceeds the configured limit. + local result_dir="$1" + ( + cd "$INFMAX_CONTAINER_WORKSPACE" + RESULT_DIR="$result_dir" AGENTIC_OUTPUT_DIR="${AGENTIC_OUTPUT_DIR:-$INFMAX_CONTAINER_WORKSPACE}" \ + "$AIPERF_PYTHON" -m utils.agentic.aggregation.process_agentic_result + ) + + # Generate metrics_plots.png from the same aiperf artifacts. Best-effort: + # don't fail the launcher if plot generation has trouble (e.g. matplotlib + # missing in a stripped-down image). The agg JSON is the success gate. + "$AIPERF_PYTHON" "$INFMAX_CONTAINER_WORKSPACE/utils/generate_aiperf_plots.py" "$result_dir" 2>&1 || true +} + +run_agentic_replay_and_write_outputs() { + local result_dir="$1" + local replay_rc + local validation_rc + + echo "$REPLAY_CMD" > "$result_dir/benchmark_command.txt" + + set +e + set -x + $REPLAY_CMD 2>&1 | tee "$result_dir/benchmark.log" + replay_rc=${PIPESTATUS[0]} + set +x + set -e + + write_agentic_result_json "$result_dir" + + "$AIPERF_PYTHON" "$AGENTIC_DIR/scripts/analyze_benchmark_distributions.py" \ + "$result_dir/aiperf_artifacts" -o "$result_dir" 2>&1 || true + + set +e + ( + cd "$INFMAX_CONTAINER_WORKSPACE" + "$AIPERF_PYTHON" -m utils.agentic.validation.validate_agentic_result \ + "$result_dir/aiperf_artifacts" \ + --failed-request-threshold "$AIPERF_FAILED_REQUEST_THRESHOLD" + ) + validation_rc=$? + set -e + + if [ "$replay_rc" -ne 0 ]; then + echo "ERROR: agentic trace replay exited with code $replay_rc after writing available results" >&2 + return "$replay_rc" + fi + + if [ "$validation_rc" -ne 0 ]; then + echo "ERROR: agentic trace replay produced invalid results after writing available artifacts" >&2 + return "$validation_rc" + fi +} diff --git a/scripts/MiniMax-M3/amd_utils/bench.sh b/scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/bench.sh old mode 100644 new mode 100755 similarity index 100% rename from scripts/MiniMax-M3/amd_utils/bench.sh rename to scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/bench.sh diff --git a/scripts/MiniMax-M3/amd_utils/env_atom.sh b/scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/env_atom.sh old mode 100644 new mode 100755 similarity index 100% rename from scripts/MiniMax-M3/amd_utils/env_atom.sh rename to scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/env_atom.sh diff --git a/scripts/MiniMax-M3/amd_utils/job.slurm b/scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/job.slurm similarity index 100% rename from scripts/MiniMax-M3/amd_utils/job.slurm rename to scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/job.slurm diff --git a/scripts/MiniMax-M3/amd_utils/models_atom.yaml b/scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/models_atom.yaml similarity index 100% rename from scripts/MiniMax-M3/amd_utils/models_atom.yaml rename to scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/models_atom.yaml diff --git a/scripts/MiniMax-M3/amd_utils/server.sh b/scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/server.sh old mode 100644 new mode 100755 similarity index 100% rename from scripts/MiniMax-M3/amd_utils/server.sh rename to scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/server.sh diff --git a/scripts/MiniMax-M3/amd_utils/server_atom.sh b/scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/server_atom.sh old mode 100644 new mode 100755 similarity index 100% rename from scripts/MiniMax-M3/amd_utils/server_atom.sh rename to scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/server_atom.sh diff --git a/scripts/MiniMax-M3/amd_utils/setup_deps.sh b/scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/setup_deps.sh old mode 100644 new mode 100755 similarity index 100% rename from scripts/MiniMax-M3/amd_utils/setup_deps.sh rename to scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/setup_deps.sh diff --git a/scripts/MiniMax-M3/amd_utils/submit.sh b/scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/submit.sh old mode 100644 new mode 100755 similarity index 100% rename from scripts/MiniMax-M3/amd_utils/submit.sh rename to scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/submit.sh diff --git a/scripts/MiniMax-M3/amd_utils/sync.py b/scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/sync.py similarity index 100% rename from scripts/MiniMax-M3/amd_utils/sync.py rename to scripts/MiniMax-M3/benchmarks/multi_node/amd_utils/sync.py diff --git a/scripts/MiniMax-M3/cluster.yaml b/scripts/MiniMax-M3/cluster.yaml new file mode 100644 index 00000000..36f434f7 --- /dev/null +++ b/scripts/MiniMax-M3/cluster.yaml @@ -0,0 +1,25 @@ +# cluster.yaml — MiniMax-M3 MXFP4 ATOM disagg cluster definition (MAD wrapper input). +# Read by run.sh. The verbatim InferenceX code (benchmarks/multi_node/amd_utils/) is unchanged; +# this file + run.sh only supply the environment/paths that InferenceX's job.slurm normally would. + +reservation: # SLURM reservation name (or "" to omit --reservation) +gpus_per_node: 8 # gfx950 (MI355X), 8 GPUs/node + +# Model staging root: same absolute path on every node (local NVMe recommended). +# Stage amd/MiniMax-M3-MXFP4 to /MiniMax-M3-MXFP4 (MODEL_NAME must match the dir +# name AND the models_atom.yaml key: MiniMax-M3-MXFP4). +models_root: /path/to/models + +results_root: ./results + +# Ports (match InferenceX server_atom.sh defaults) +router_port: 8000 +prefill_port: 8010 +decode_port: 8020 +handshake_port: 6301 + +# Node pool: list at least 3 (2p1d_dpa needs 3; 1p1d needs 2). All-TP4 => 1 node/worker. +nodes: + - + - + - diff --git a/scripts/MiniMax-M3/run.sh b/scripts/MiniMax-M3/run.sh new file mode 100755 index 00000000..9ff11db9 --- /dev/null +++ b/scripts/MiniMax-M3/run.sh @@ -0,0 +1,100 @@ +#!/bin/bash +# run.sh — MAD launcher for the VERBATIM InferenceX MiniMax-M3 MXFP4 ATOM disagg benchmark. +# +# The engine logic under benchmarks/multi_node/amd_utils/ is copied UNCHANGED from +# SemiAnalysisAI/InferenceX. InferenceX normally drives it via job.slurm/submit.sh; this wrapper +# reproduces the same per-node env contract + container mounts so it runs standalone from MAD. +# +# All model flags (TP/DPA, kv fp8, online_quant, block_size, env, MTP) are read by server_atom.sh +# from amd_utils/models_atom.yaml — this wrapper sets NONE of them, so behavior matches InferenceX. +# +# Usage: +# TOPO=1p1d ISL=8192 OSL=1024 ./run.sh # 2 nodes, conc 1-256 +# TOPO=1p1d ISL=1024 OSL=1024 ./run.sh # 2 nodes, conc 1-256 +# TOPO=2p1d_dpa ISL=8192 OSL=1024 ./run.sh # 3 nodes, conc 256-1024 +# NODES=n1,n2,n3 TOPO=2p1d_dpa ./run.sh # pin nodes +# ACTION=dry ./run.sh # print server/bench cmds, don't launch +set -uo pipefail +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # scripts/MiniMax-M3 +AMD_UTILS="$HERE/benchmarks/multi_node/amd_utils" +CLUSTER="$HERE/cluster.yaml" +IMG="rocm/atom-dev:nightly_202607011530" # from models_atom.yaml / config +MODEL_NAME="MiniMax-M3-MXFP4" # must match models_atom.yaml key + dir + +y(){ python3 -c "import yaml,sys;print(yaml.safe_load(open('$CLUSTER')).get('$1',''))"; } +RESV=$(y reservation); GPN=$(y gpus_per_node); MROOT=$(y models_root); RROOT=$(y results_root) +RPORT=$(y router_port); PPORT=$(y prefill_port); DPORT=$(y decode_port); HPORT=$(y handshake_port) +mapfile -t CFG_NODES < <(python3 -c "import yaml;[print(n) for n in yaml.safe_load(open('$CLUSTER')).get('nodes',[])]") + +TOPO="${TOPO:-1p1d}"; ACTION="${ACTION:-bench}"; ISL="${ISL:-8192}"; OSL="${OSL:-1024}" +# Topologies mirror configs/minimaxm3-fp4-mi355x-atom-disagg.yaml (all TP4, STP: DECODE_MTP_SIZE=0). +case "$TOPO" in + 1p1d) xP=1; yD=1; DPA=false; DEF_CONC="1 2 4 8 16 32 64 128 256" ;; + 2p1d_dpa) xP=2; yD=1; DPA=true; DEF_CONC="256 512 768 1024" ;; + *) echo "ERROR: TOPO must be 1p1d|2p1d_dpa (got $TOPO)"; exit 1 ;; +esac +TP=4 +CONC_LIST="${CONC_LIST:-$DEF_CONC}" +CONC_X=$(echo "$CONC_LIST" | tr ' ' 'x') # server_atom.sh derives decode --max-num-seqs from this +NNODES=$(( xP + yD )) # 1 node/worker at TP4 on 8-GPU nodes + +# --- node selection -------------------------------------------------------------------------- +if [[ -n "${NODES:-}" ]]; then IFS=',' read -ra NODES <<< "$NODES"; else NODES=("${CFG_NODES[@]}"); fi +[[ ${#NODES[@]} -lt $NNODES ]] && { echo "ERROR: need $NNODES nodes, have ${#NODES[@]} (fill cluster.yaml nodes: or pass NODES=)"; exit 1; } +NODES=("${NODES[@]:0:$NNODES}") + +STAMP=$(date +%Y%m%d_%H%M%S) +RUN="$RROOT/minimaxm3-fp4_atom_${TOPO}_isl${ISL}_${STAMP}"; mkdir -p "$RUN"; chmod 777 "$RUN" +echo "[run] TOPO=$TOPO xP=$xP yD=$yD DPA=$DPA nodes=${NODES[*]} shape=${ISL}/${OSL} conc='$CONC_LIST'" +echo "[run] image=$IMG model=$MODEL_NAME run=$RUN" +RESV_ARG=""; [[ -n "$RESV" && "$RESV" != "" ]] && RESV_ARG="--reservation=$RESV" + +# --- resolve node IPs (IPADDRS = prefill nodes first, then decode; server_atom.sh slices it) -- +declare -a IPS=() +for n in "${NODES[@]}"; do + ip=$(srun $RESV_ARG --nodelist="$n" --nodes=1 --ntasks=1 --overlap --mem=1G \ + bash -c "ip route get 1.1.1.1 2>/dev/null | awk '/src/{print \$7}'" 2>/dev/null \ + | grep -oE '([0-9]+\.){3}[0-9]+' | head -1) + IPS+=("$ip"); echo " $n -> $ip" +done +IPADDRS=$(IFS=,; echo "${IPS[*]}"); NODE0_IP="${IPS[0]}" + +DRY=0; [[ "$ACTION" == "dry" ]] && DRY=1 +RDMA="--device=/dev/infiniband --ulimit memlock=-1 --ulimit stack=67108864 --init" + +# --- launch verbatim server_atom.sh on each node (NODE_RANK self-selects role) ---------------- +# Mounts: recipe root -> /workspace (so bench.sh ../../benchmark_lib.sh + REPO_ROOT/utils resolve); +# ATOM_WS_PATH -> /workspace/benchmarks/multi_node/amd_utils (server_atom.sh's own dir). +rank=0 +for n in "${NODES[@]}"; do + rdir="$RUN/node${rank}_${n}"; mkdir -p "$rdir"; chmod 777 "$rdir" + srun $RESV_ARG --nodelist="$n" --nodes=1 --ntasks=1 --exclusive bash -c " + docker rm -f mm3_atom_n${rank} >/dev/null 2>&1 || true + docker run --rm --name mm3_atom_n${rank} --entrypoint bash \ + --device=/dev/kfd --device=/dev/dri $RDMA --group-add video --ipc host --shm-size 128G \ + --network host --privileged --cap-add=SYS_PTRACE --security-opt seccomp=unconfined \ + -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ + -e ATOM_WS_PATH=/workspace/benchmarks/multi_node/amd_utils \ + -e ENGINE=atom-disagg \ + -e NODE_RANK=${rank} -e NODE0_ADDR=${NODE0_IP} -e IPADDRS='${IPADDRS}' \ + -e MODEL_DIR=${MROOT} -e MODEL_NAME=${MODEL_NAME} \ + -e xP=${xP} -e yD=${yD} -e GPUS_PER_NODE=${GPN} \ + -e PREFILL_TP_SIZE=${TP} -e DECODE_TP_SIZE=${TP} \ + -e PREFILL_ENABLE_EP=false -e DECODE_ENABLE_EP=false \ + -e PREFILL_ENABLE_DP=${DPA} -e DECODE_ENABLE_DP=${DPA} \ + -e DECODE_MTP_SIZE=0 \ + -e PREFILL_PORT=${PPORT} -e DECODE_PORT=${DPORT} -e ROUTER_PORT=${RPORT} -e HANDSHAKE_PORT=${HPORT} \ + -e BENCH_INPUT_LEN=${ISL} -e BENCH_OUTPUT_LEN=${OSL} -e BENCH_MAX_CONCURRENCY='${CONC_X}' \ + -e BENCH_REQUEST_RATE=inf -e BENCH_RANDOM_RANGE_RATIO=1 -e BENCH_NUM_PROMPTS_MULTIPLIER=10 \ + -e DRY_RUN=${DRY} -e SLURM_JOB_ID=${STAMP} \ + -v ${MROOT}:${MROOT} -v ${HERE}:/workspace -v ${rdir}:/run_logs/slurm_job-${STAMP} \ + ${IMG} /workspace/benchmarks/multi_node/amd_utils/server.sh + " > "$rdir/srun.log" 2>&1 & + echo " [launch] node${rank} ($n) dispatched (role self-selected by NODE_RANK)" + rank=$(( rank + 1 )) +done + +echo "[run] all $NNODES nodes dispatched. node0=${NODES[0]} hosts atomesh router :$RPORT + drives bench." +echo "[run] tail node0: $RUN/node0_${NODES[0]}/srun.log" +wait +echo "[run] done -> $RUN" From da4c6c633f24c5ffed2678ddbbe6c20421918149 Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Thu, 16 Jul 2026 14:04:46 +0000 Subject: [PATCH 08/12] MiniMax-M3: trim README to plain runnable doc --- scripts/MiniMax-M3/README.md | 78 ++++++++++++------------------------ 1 file changed, 26 insertions(+), 52 deletions(-) diff --git a/scripts/MiniMax-M3/README.md b/scripts/MiniMax-M3/README.md index c33294b5..efc4e497 100644 --- a/scripts/MiniMax-M3/README.md +++ b/scripts/MiniMax-M3/README.md @@ -1,66 +1,40 @@ # MiniMax-M3 — MXFP4 disaggregated (P/D) ATOM benchmark -The ATOM MXFP4 prefill/decode disaggregated benchmark for MiniMax-M3 on AMD MI355X (gfx950), -config **`minimaxm3-fp4-mi355x-atom-disagg`** (atomesh router + mooncake RDMA KV) — the -top-throughput MiniMax-M3 disagg config on InferenceX (~10,829 tok/s/GPU at 8k/1k, 2P1D+DPA). -Dashboard: +Prefill/decode disaggregated ATOM benchmark for MiniMax-M3 MXFP4 on AMD MI355X (gfx950), +config `minimaxm3-fp4-mi355x-atom-disagg` (atomesh router + mooncake RDMA KV). -## Layout -The engine code is copied **verbatim** from [SemiAnalysisAI/InferenceX](https://github.com/SemiAnalysisAI/InferenceX) -@ `a174f20` (unmodified — update by re-copying from the same paths). A thin MAD wrapper -(`run.sh` + `cluster.yaml`) supplies the per-node env + mounts that InferenceX's -`job.slurm`/`submit.sh` normally provide, so it runs standalone from MAD. - -| Path | Origin | -|------|--------| -| `benchmarks/multi_node/amd_utils/` | **verbatim** InferenceX (`server_atom.sh`, `models_atom.yaml`, `env_atom.sh`, `bench.sh`, `setup_deps.sh`, `server.sh`, `job.slurm`, `submit.sh`, `sync.py`) | -| `benchmarks/benchmark_lib.sh` | **verbatim** InferenceX (bench.sh sources `../../benchmark_lib.sh`) | -| `utils/bench_serving/` | **verbatim** InferenceX benchmark client | -| `configs/minimaxm3-fp4-mi355x-atom-disagg.yaml` | config entry from `amd-master.yaml` | -| `run.sh`, `cluster.yaml` | **MAD wrapper** (the only non-verbatim files) | +## Nodes required +- `1p1d` → **2 nodes** (1 prefill + 1 decode) +- `2p1d_dpa` → **3 nodes** (2 prefill + 1 decode) -Model flags (all-TP4, `--kv_cache_dtype fp8`, `online_quant_config` ptpc_fp8, `ATOM_FORCE_ATTN_TRITON=1`, -`AITER_QUICK_REDUCE_QUANTIZATION=INT4`, block 128, mem-frac 0.8, max-model-len 32768) live in -`models_atom.yaml` and are read by `server_atom.sh` — the wrapper sets none of them, so runtime -behavior matches InferenceX exactly. +8× MI355X per node, all workers TP4 (1 worker/node). 3 nodes covers both topologies. ## Prerequisites -1. **`amdgpu_peermem` kernel module loaded on every host** — required for GPUDirect RDMA so - mooncake can stream the KV cache GPU→NIC→GPU across nodes. Check: `lsmod | grep amdgpu_peermem`; - load: `sudo modprobe amdgpu_peermem`. Without it, mooncake KV transfer stalls / times out. -2. **RDMA fabric** reachable between nodes (the launcher passes `--device=/dev/infiniband --ulimit memlock=-1`). -3. **Weights staged** at `/MiniMax-M3-MXFP4` (from `amd/MiniMax-M3-MXFP4`), same path on every node. -4. **SLURM** with the nodes/reservation from `cluster.yaml`; Docker + ROCm on each node. -5. Image `rocm/atom-dev:nightly_202607011530` pullable on each node. +1. `amdgpu_peermem` kernel module loaded on every host (GPUDirect RDMA for mooncake KV): + `lsmod | grep amdgpu_peermem` / `sudo modprobe amdgpu_peermem`. +2. RDMA fabric between nodes. +3. Weights staged at `/MiniMax-M3-MXFP4` (from `amd/MiniMax-M3-MXFP4`), same path on every node. +4. SLURM with the nodes/reservation in `cluster.yaml`; Docker + ROCm on each node. +5. Image `rocm/atom-dev:nightly_202607011530` available on each node. ## Usage Edit `cluster.yaml` (reservation, `models_root`, `nodes`), then: ```bash -# 8k/1k balanced, 2 nodes (1 prefill + 1 decode), conc 1-256: -TOPO=1p1d ISL=8192 OSL=1024 ./run.sh - -# 8k/1k high-throughput, 3 nodes (2 prefill + 1 decode + DP-attn), conc 256-1024: -TOPO=2p1d_dpa ISL=8192 OSL=1024 ./run.sh +TOPO=1p1d ISL=8192 OSL=1024 ./run.sh # 2 nodes, conc 1-256 +TOPO=2p1d_dpa ISL=8192 OSL=1024 ./run.sh # 3 nodes, conc 256-1024 +TOPO=1p1d ISL=1024 OSL=1024 ./run.sh # 2 nodes, 1k/1k -# 1k/1k, 2 nodes: -TOPO=1p1d ISL=1024 OSL=1024 ./run.sh - -# pin nodes / dry-run (print commands, launch nothing): -NODES=n1,n2,n3 TOPO=2p1d_dpa ./run.sh -ACTION=dry TOPO=1p1d ./run.sh +NODES=n1,n2,n3 TOPO=2p1d_dpa ./run.sh # pin nodes +ACTION=dry TOPO=1p1d ./run.sh # print commands, launch nothing ``` -Bench drives the atomesh router (`:8000`, OpenAI-compatible) → true aggregate throughput (no ×N). -Result JSONs land under `results/minimaxm3-fp4_atom___/`. - -## Topologies (`configs/minimaxm3-fp4-mi355x-atom-disagg.yaml`) -| TOPO | shape | conc | layout | nodes | -|------|-------|------|--------|-------| -| `1p1d` | 8192/1024, 1024/1024 | 1–256 | 1 prefill TP4 + 1 decode TP4 (STP) | 2 | -| `2p1d_dpa` | 8192/1024 | 256–1024 | 2 prefill + 1 decode, TP4 + DP-attn | 3 | +Bench drives the atomesh router (`:8000`, OpenAI-compatible). Results land under +`results/minimaxm3-fp4_atom___/`. -All-TP4 ⇒ 1 node/worker on 8-GPU nodes; STP only (`DECODE_MTP_SIZE=0`). +## Layout +Engine code (`benchmarks/`, `utils/bench_serving/`) is copied unchanged from +SemiAnalysisAI/InferenceX @ `a174f20`; all model flags live in +`benchmarks/multi_node/amd_utils/models_atom.yaml`. `run.sh` + `cluster.yaml` are the MAD +launcher (supply the env + mounts InferenceX's job.slurm normally provides). -## Status -Engine code is verbatim InferenceX (byte-identical). The MAD wrapper is syntax-checked and its -topology/mount contract mirrors InferenceX `job.slurm`, but has **not** been dry-run on live -MI355X hardware — validate with `ACTION=dry` first, then a small `TOPO=1p1d` run. +Topologies are defined in `configs/minimaxm3-fp4-mi355x-atom-disagg.yaml` (all-TP4, STP, +`DECODE_MTP_SIZE=0`). Validate with `ACTION=dry` before a real run. From 478244ff9781c5b9e613d5d58fbc1c853ff4feb5 Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Thu, 23 Jul 2026 07:28:51 +0000 Subject: [PATCH 09/12] Add MiniMax-M3 single-node P/D over XGMI recipe (scripts/MiniMax-M3-single-pd) Single-node prefill/decode disaggregation for MiniMax-M3-MXFP4 on AMD MI355X via vLLM + MoRIIO, with GPU-to-GPU XGMI KV transfer (backend=xgmi): prefill TP4 on GPU0-3 + decode TP4 on GPU4-7, co-located on one node, fronted by vllm-router. Complements the cross-node ATOM recipe in scripts/MiniMax-M3/. Co-location works here because vLLM exposes --disable-custom-all-reduce (avoiding the AITER custom all-reduce IPC collision that blocks ATOM) and MoRIIO has a native XGMI backend. README documents the 7 non-obvious fixes, exact config, verification steps, validation results (1k/1k + 8k/1k throughput, GSM8K-style accuracy), and provenance (ROCm/mori #481, vllm-project/router #181). Files: run_pd_singlenode.sh (launcher), prep_model.sh (flat model + remote-code), bench_pd.sh (vllm bench serve sweep), README.md. Adds a root README blueprint row. Co-Authored-By: Claude --- README.md | 1 + scripts/MiniMax-M3-single-pd/README.md | 93 ++++++++++++++++++ scripts/MiniMax-M3-single-pd/bench_pd.sh | 38 +++++++ scripts/MiniMax-M3-single-pd/prep_model.sh | 25 +++++ .../MiniMax-M3-single-pd/run_pd_singlenode.sh | 98 +++++++++++++++++++ 5 files changed, 255 insertions(+) create mode 100644 scripts/MiniMax-M3-single-pd/README.md create mode 100755 scripts/MiniMax-M3-single-pd/bench_pd.sh create mode 100755 scripts/MiniMax-M3-single-pd/prep_model.sh create mode 100755 scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh diff --git a/README.md b/README.md index 047ca027..9bd997f4 100644 --- a/README.md +++ b/README.md @@ -27,6 +27,7 @@ Below are blueprints of supported models along with their documentation. | [**DeepSeek-V4-Pro disaggregated P/D inference**](scripts/DeepSeekV4/README.md) | Distributed Inference P/D disaggregation for DeepSeek-V4-Pro (ATOM+mooncake / SGLang+MoRI) | DeepSeek-V4-Pro | | [**Qwen3.5-35B-A3B aggregated inference**](scripts/Qwen3.5/README.md) | Single-node aggregated serving/benchmark (SGLang/vLLM); recommended default SGLang+MXFP4+TP1 (8 instances/node) | Qwen3.5-35B-A3B (MXFP4, FP8) | | [**MiniMax-M3 disaggregated P/D inference**](scripts/MiniMax-M3/README.md) | Distributed Inference P/D disaggregation for MiniMax-M3 on MI355X (ATOM+mooncake, MXFP4); top-throughput InferenceX config | MiniMax-M3 (MXFP4) | +| [**MiniMax-M3 single-node P/D over XGMI**](scripts/MiniMax-M3-single-pd/README.md) | Single-node prefill/decode disaggregation for MiniMax-M3 on MI355X via vLLM + MoRIIO with GPU-to-GPU XGMI KV transfer (TP4/TP4) | MiniMax-M3 (MXFP4) | | [**KVCache Transfer Bench**](scripts/kvcache_transfer_bench/README.md) | Inter-node Transfer Benchmark | no specific models | ## Table of Contents diff --git a/scripts/MiniMax-M3-single-pd/README.md b/scripts/MiniMax-M3-single-pd/README.md new file mode 100644 index 00000000..aaba6906 --- /dev/null +++ b/scripts/MiniMax-M3-single-pd/README.md @@ -0,0 +1,93 @@ +# MiniMax-M3 MXFP4 — single-node vLLM + MoRIIO P/D disaggregation over XGMI + +Simulates prefill/decode (P/D) disaggregation on a **single** AMD MI355X node (8× gfx950): +prefill (kv_producer) TP4 on GPU0-3 + decode (kv_consumer) TP4 on GPU4-7, fronted by +`vllm-router`, with the KV cache moving prefill→decode via **MoRIIO over the XGMI GPU fabric** +(`backend=xgmi`) — no NIC/RDMA hop for the intra-node transfer. + +Validated end-to-end: router-mediated P/D, XGMI backend engaged, correct inference, and a +throughput sweep at 1k/1k and 8k/1k (see **Results** below). + +> This is the **single-node XGMI** counterpart to the cross-node ATOM recipe in +> `scripts/MiniMax-M3/`. It uses the **vLLM + MoRIIO** stack (not ATOM/mooncake) — see "Why this +> stack" for the reason single-node co-location requires vLLM here. + +## Why this stack (vLLM+MoRIIO, not ATOM/mooncake) +Two co-located TP groups on one node collide in the ATOM/aiter custom all-reduce (a fixed-key IPC +region → `hipIpcGetMemHandle ... invalid argument` on the 2nd engine), and ATOM has no switch to +disable it (its all-gather asserts `ca_comm != None`). ATOM's mooncake transport also has **no +XGMI path** (RDMA/TCP only). vLLM provides a supported `--disable-custom-all-reduce` **and** MoRIIO +has a native XGMI backend — so single-node co-located P/D over XGMI is only achievable on this stack. + +## What was non-obvious (each was a hard failure until fixed) +1. **`--disable-custom-all-reduce` is mandatory** for two co-located TP groups (else the AITER + custom all-reduce IPC handles collide and the 2nd engine aborts). Falls back to RCCL. +2. **Distinct MoRIIO ports per engine.** Both containers share host network; each binds + `handshake_port` (+1/rank) and `notify_port`. Use different bases (prefill 6301/61005, + decode 6311/61015) or they collide with `Address already in use`. +3. **`--max-model-len 32768`.** The model default is 1,048,576 (1M) → ~44 GiB KV, which won't fit. +4. **GPU memory headroom.** `--gpu-memory-utilization 0.90` needs a mostly-free node; on a shared + node (another tenant holding VRAM) lower it (≈0.30) or you OOM at warmup / get near-zero KV. +5. **Per-container GPU isolation.** Set both `HIP_VISIBLE_DEVICES` and `CUDA_VISIBLE_DEVICES` + (0-3 / 4-7) so MoRIIO/QuickReduce map physical GPUs correctly. +6. **Flat model dir with remote-code** — the HF weights-only download lacks + `configuration_minimax_m3_vl.py` referenced by `config.json` `auto_map` (see `prep_model.sh`). +7. **XGMI backend** via `kv_connector_extra_config: {"backend": "xgmi"}` (default `"rdma"`). + Confirmed by `Using MoRIIO backend: XGMI` in every worker log. + +## Files +| File | Purpose | +|------|---------| +| `run_pd_singlenode.sh` | launch router + prefill + decode, wait for registration, smoke inference | +| `prep_model.sh` | stage `amd/MiniMax-M3-MXFP4` as a flat dir (weights symlinks + remote-code .py) | +| `bench_pd.sh` | `vllm bench serve` throughput sweep through the router (shapes/concurrency via env) | + +## Prerequisites +- AMD MI355X node, ROCm; Docker with `--device=/dev/kfd,/dev/dri,/dev/infiniband`. +- Images: `rocm/vllm-dev:vllm-0.23.1-rocm723-mi35x-mori-0625` (engine) and + `vllm/vllm-router:nightly-` (router; pin a currently-available nightly — Docker Hub keeps + only ~16). Override via `IMG` / `ROUTER_IMG`. +- Model staged flat (see `prep_model.sh`), reachable at `$MODEL` inside the container mount `$MOUNT`. +- Enough free VRAM (mostly-idle node for `GPU_MEM_UTIL=0.90`; otherwise lower it). + +## Run +```bash +# stage the model once (adjust paths): +SNAPSHOT=/path/to/models--amd--MiniMax-M3-MXFP4/snapshots/ \ +PYSRC=/path/to/minimax-m3-remote-code DEST=/models/MiniMax-M3-MXFP4 ./prep_model.sh + +# launch P/D over XGMI + smoke inference (run detached so containers persist): +MODEL=/models/MiniMax-M3-MXFP4 MOUNT=/models BACKEND=xgmi nohup ./run_pd_singlenode.sh & + +# throughput benchmark through the router (1k/1k and 8k/1k): +SHAPES="1024,1024 8192,1024" CONCS="1 8 32" ./bench_pd.sh +``` +`BACKEND=rdma` switches KV transfer to RDMA-over-NIC (the upstream default) for comparison. + +## Verifying it worked +- Router log shows `🔵Add Prefill [...]` and `🔵Add Decode [...]` (both registered). +- A chat request to `http://:30000/v1/chat/completions` returns a correct answer, and the + router log maps it to **both** `:2584` (prefill) and `:2585` (decode). +- Decode/prefill logs show `Using MoRIIO backend: XGMI`. + +## Results (validation, MI355X, backend=xgmi) +Functional/scaling smoke test through the router (`vllm bench serve`), not a tuned peak sweep: + +| ISL/OSL | Concurrency | Output tok/s | Total tok/s | Mean TPOT (ms) | +|---------|-------------|--------------|-------------|----------------| +| 1024/1024 | 1 | 72 | 157 | 9.5 | +| 1024/1024 | 8 | 564 | 1,225 | 13.8 | +| 1024/1024 | 32 | 1,822| 3,957 | 16.6 | +| 8192/1024 | 8 | 502 | 4,603 | 14.3 | + +Accuracy (temperature 0, through the router): GSM8K-style prompts answered correctly +(e.g. Natalia=72, robe=3, 3×24−17=55, "17×23"=391). `lm_eval` was not in the image; validation +used direct known-answer probes rather than the full harness. + +## Provenance +Derived from the InferenceX `minimaxm3-fp4-mi355x-vllm-disagg` config (which runs MoRIIO +**cross-node over RDMA**); adapted here to **single-node co-located P/D over XGMI**. +Related upstream fixes: [ROCm/mori #481](https://github.com/ROCm/mori/pull/481) (dma-buf MR +registration for peermem-less MI355X + Ionic clusters), +[vllm-project/router #181](https://github.com/vllm-project/router/pull/181) (discovery DP-rank +round-robin). diff --git a/scripts/MiniMax-M3-single-pd/bench_pd.sh b/scripts/MiniMax-M3-single-pd/bench_pd.sh new file mode 100755 index 00000000..ff059ed8 --- /dev/null +++ b/scripts/MiniMax-M3-single-pd/bench_pd.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# ============================================================================= +# Throughput benchmark for the single-node vLLM+MoRIIO P/D deployment. +# Runs `vllm bench serve` through the router across a concurrency sweep, for one +# or more ISL/OSL shapes. Assumes run_pd_singlenode.sh is already up (router :30000). +# +# Usage: +# ./bench_pd.sh # default: 1k/1k @ conc 1,8,32 +# SHAPES="8192,1024" CONCS="8 16" ./bench_pd.sh +# MODEL=/models/MiniMax-M3-MXFP4 ROUTER_PORT=30000 ./bench_pd.sh +# ============================================================================= +set -uo pipefail +IMG="${IMG:-rocm/vllm-dev:vllm-0.23.1-rocm723-mi35x-mori-0625}" +MODEL="${MODEL:-/models/MiniMax-M3-MXFP4}" +MOUNT="${MOUNT:-/models}" +LOG="${LOG:-./logs/bench}"; mkdir -p "$LOG" +ROUTER_PORT="${ROUTER_PORT:-30000}" +SHAPES="${SHAPES:-1024,1024}" # comma-separated "isl,osl"; space-separate multiple shapes +CONCS="${CONCS:-1 8 32}" +HOST_IP=$(ip route get 1.1.1.1 2>/dev/null | awk '/src/{print $7}') + +echo "[bench] router=$HOST_IP:$ROUTER_PORT shapes='$SHAPES' concs='$CONCS' $(date)" | tee "$LOG/bench.log" +for shape in $SHAPES; do + ISL="${shape%,*}"; OSL="${shape#*,}" + for c in $CONCS; do + echo "[bench] --- isl=$ISL osl=$OSL conc=$c ---" | tee -a "$LOG/bench.log" + docker run --rm --entrypoint bash --network host -v "$MOUNT":"$MOUNT" "$IMG" -c " + vllm bench serve --backend openai-chat \ + --base-url http://$HOST_IP:$ROUTER_PORT --endpoint /v1/chat/completions \ + --model $MODEL --served-model-name minimaxm3 --dataset-name random \ + --random-input-len $ISL --random-output-len $OSL \ + --num-prompts \$(( c > 4 ? c*2 : 8 )) --max-concurrency $c \ + --percentile-metrics ttft,tpot,e2el --save-result \ + --result-dir $LOG --result-filename bench_isl${ISL}_osl${OSL}_conc${c}.json 2>&1" \ + | grep -iE "Successful|throughput|Mean TTFT|Median TTFT|Mean TPOT" | tee -a "$LOG/bench.log" + done +done +echo "[bench] DONE $(date) results in $LOG" | tee -a "$LOG/bench.log" diff --git a/scripts/MiniMax-M3-single-pd/prep_model.sh b/scripts/MiniMax-M3-single-pd/prep_model.sh new file mode 100755 index 00000000..3b03fc18 --- /dev/null +++ b/scripts/MiniMax-M3-single-pd/prep_model.sh @@ -0,0 +1,25 @@ +#!/bin/bash +# prep_model.sh — stage MiniMax-M3-MXFP4 as a FLAT dir vLLM can load offline. +# +# The HF snapshot of amd/MiniMax-M3-MXFP4 ships weights + config + tokenizer, but the +# config.json auto_map references remote-code Python (configuration_minimax_m3_vl.py) that +# is NOT in the weights-only download. Without it, `--trust-remote-code` fails to import +# the config. This builds a flat dir: symlinks to the snapshot's weights + real copies of +# the remote-code .py files. +# +# Usage: SNAPSHOT= PYSRC= DEST= ./prep_model.sh +set -euo pipefail +SNAPSHOT="${SNAPSHOT:?path to models--amd--MiniMax-M3-MXFP4/snapshots/}" +PYSRC="${PYSRC:?dir containing configuration_minimax_m3_vl.py etc (e.g. an M3 repo with remote code)}" +DEST="${DEST:?flat output dir, e.g. /models/MiniMax-M3-MXFP4}" + +mkdir -p "$DEST" +# symlink every real file from the snapshot (weights, config, tokenizer) +for f in "$SNAPSHOT"/*; do ln -sfn "$(readlink -f "$f")" "$DEST/$(basename "$f")"; done +# overlay real copies of the remote-code .py files +for p in configuration_minimax_m3_vl.py image_processor.py processing_minimax.py video_processor.py; do + [ -f "$PYSRC/$p" ] && cp -f "$PYSRC/$p" "$DEST/$p" && echo "copied $p" +done +echo "flat model dir ready: $DEST" +ls -L "$DEST"/config.json "$DEST"/model.safetensors.index.json >/dev/null && echo "config + index resolve OK" +echo "shards: $(ls -L "$DEST"/model-*-of-*.safetensors 2>/dev/null | wc -l)" diff --git a/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh b/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh new file mode 100755 index 00000000..d61f8352 --- /dev/null +++ b/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh @@ -0,0 +1,98 @@ +#!/bin/bash +# ============================================================================= +# MiniMax-M3 MXFP4 — single-node vLLM + MoRIIO P/D disaggregation over XGMI +# ============================================================================= +# Simulates prefill/decode disaggregation on ONE MI355X node: +# prefill (kv_producer) on GPU0-3 TP4 + decode (kv_consumer) on GPU4-7 TP4 +# fronted by vllm-router; KV moves prefill->decode via MoRIIO. +# +# BACKEND=xgmi -> GPU-to-GPU XGMI KV transfer (intra-node, the goal) +# BACKEND=rdma -> RDMA-over-NIC KV transfer (default upstream mode) +# +# Usage: +# BACKEND=xgmi ./run_pd_singlenode.sh # launch P+D+router, run smoke inference +# BACKEND=rdma ./run_pd_singlenode.sh +# Override MODEL / LOG / UCXDEV via env as needed. +# Run under `nohup ... &` or via ssh so containers persist (they are `docker -d`). +# ============================================================================= +set -uo pipefail + +# ---- config (override via env) ---------------------------------------------- +IMG="${IMG:-rocm/vllm-dev:vllm-0.23.1-rocm723-mi35x-mori-0625}" +ROUTER_IMG="${ROUTER_IMG:-vllm/vllm-router:nightly-20260721-1fbcde7}" +MODEL="${MODEL:-/models/MiniMax-M3-MXFP4}" # flat dir w/ weights + remote-code .py (see README) +LOG="${LOG:-./logs}"; mkdir -p "$LOG" +BACKEND="${BACKEND:-xgmi}" # xgmi | rdma +UCXDEV="${UCXDEV:-ionic_0:1}" # RoCE NIC for UCX (rdma backend); harmless for xgmi +GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.90}" # lower to ~0.30 on shared/occupied nodes +MAX_MODEL_LEN="${MAX_MODEL_LEN:-32768}" +MOUNT="${MOUNT:-/models}" # host dir mounted into containers (must contain MODEL) + +HOST_IP=$(ip route get 1.1.1.1 2>/dev/null | awk '/src/{print $7}') +PING=36367; ROUTER_PORT=30000; PF_PORT=2584; DC_PORT=2585 + +# Per-worker MoRIIO ZMQ ports MUST differ between the two co-located engines +# (both share host network) or they collide on bind (Address already in use). +PF_HS=6301; PF_NT=61005 # prefill handshake / notify +DC_HS=6311; DC_NT=61015 # decode handshake / notify + +MENV="VLLM_USE_V1=1 VLLM_ROCM_USE_AITER=1 VLLM_ROCM_USE_AITER_MOE=1 VLLM_USE_BREAKABLE_CUDAGRAPH=0 VLLM_ENGINE_READY_TIMEOUT_S=3600" +# NOTE: --disable-custom-all-reduce is REQUIRED for two co-located TP groups on one node +# (otherwise AITER custom all-reduce IPC handles collide -> hipIpcGetMemHandle EINVAL). +SERVE_FLAGS="--tensor-parallel-size 4 --max-num-batched-tokens 32768 --max-num-seqs 512 \ +--block-size 128 --language-model-only --attention-backend TRITON_ATTN --moe-backend aiter \ +--no-enable-prefix-caching --gpu-memory-utilization $GPU_MEM_UTIL --tool-call-parser minimax_m3 \ +--reasoning-parser minimax_m3 --enable-auto-tool-choice --disable-custom-all-reduce \ +--max-model-len $MAX_MODEL_LEN" + +echo "[run] $(date) host=$HOST_IP backend=$BACKEND model=$MODEL mem_util=$GPU_MEM_UTIL" | tee "$LOG/run.log" +docker rm -f vm_prefill vm_decode vm_router >/dev/null 2>&1 || true + +serve(){ local name=$1 gpus=$2 role=$3 port=$4 hs=$5 nt=$6 + local envargs=""; for kv in $MENV; do envargs+=" -e $kv"; done + docker run -d --name "$name" --entrypoint bash \ + --device=/dev/kfd --device=/dev/dri --device=/dev/infiniband \ + --ulimit memlock=-1 --ulimit stack=67108864 --init --group-add video --ipc host \ + --shm-size 128G --network host --privileged --cap-add=SYS_PTRACE \ + --security-opt seccomp=unconfined -v "$MOUNT":"$MOUNT" -v "$LOG":/out \ + -e HIP_VISIBLE_DEVICES=$gpus -e CUDA_VISIBLE_DEVICES=$gpus \ + -e UCX_NET_DEVICES=$UCXDEV -e VLLM_NIXL_SIDE_CHANNEL_HOST=$HOST_IP \ + -e VLLM_NIXL_SIDE_CHANNEL_PORT=$((5600+port)) $envargs \ + "$IMG" -c "vllm serve $MODEL --served-model-name minimaxm3 --host 0.0.0.0 --port $port --trust-remote-code \ + --kv-transfer-config '{\"kv_connector\":\"MoRIIOConnector\",\"kv_role\":\"$role\",\"kv_connector_extra_config\":{\"proxy_ip\":\"$HOST_IP\",\"proxy_ping_port\":\"$PING\",\"http_port\":\"$port\",\"handshake_port\":\"$hs\",\"notify_port\":\"$nt\",\"read_mode\":true,\"backend\":\"$BACKEND\"}}' \ + $SERVE_FLAGS 2>&1 | tee /out/$name.log"; } + +# 1) router first (workers self-register to it via ZMQ discovery on $PING) +docker run -d --name vm_router --network host --ulimit nofile=1048576:1048576 -v "$LOG":/out "$ROUTER_IMG" \ + bash -lc "exec vllm-router --vllm-pd-disaggregation --kv-connector moriio --vllm-discovery-address 0.0.0.0:$PING \ + --port $ROUTER_PORT --host 0.0.0.0 --policy consistent_hash --prefill-policy consistent_hash \ + --decode-policy consistent_hash --log-level info 2>&1 | tee /out/vm_router.log" +echo "[run] router :$ROUTER_PORT (discovery :$PING)" | tee -a "$LOG/run.log"; sleep 8 + +# 2) both workers (distinct handshake/notify ports) +serve vm_prefill 0,1,2,3 kv_producer $PF_PORT $PF_HS $PF_NT +serve vm_decode 4,5,6,7 kv_consumer $DC_PORT $DC_HS $DC_NT +echo "[run] prefill($PF_PORT,GPU0-3)+decode($DC_PORT,GPU4-7) launched" | tee -a "$LOG/run.log" + +# 3) wait until BOTH register to the router (discovery mode: /v1/models stays empty; +# check the router log for the "Add Prefill"/"Add Decode" registration lines). +for i in $(seq 1 120); do + P=$(docker logs vm_router 2>&1 | grep -c "Add Prefill" || true) + D=$(docker logs vm_router 2>&1 | grep -c "Add Decode" || true) + ps=$(docker inspect -f '{{.State.Status}}' vm_prefill 2>/dev/null) + ds=$(docker inspect -f '{{.State.Status}}' vm_decode 2>/dev/null) + echo "[run] t=$((i*10))s router_prefill=$P router_decode=$D pf=$ps dc=$ds" | tee -a "$LOG/run.log" + [ "$P" -ge 1 ] && [ "$D" -ge 1 ] && { echo "[run] both workers registered" | tee -a "$LOG/run.log"; break; } + { [ "$ps" = exited ] || [ "$ds" = exited ]; } && { echo "[run] a worker exited early — see $LOG/vm_*.log" | tee -a "$LOG/run.log"; exit 1; } + sleep 10 +done +sleep 10 + +# 4) smoke inference through the router (exercises P->D KV transfer over $BACKEND) +echo "[run] === INFERENCE via router ($BACKEND) ===" | tee -a "$LOG/run.log" +curl -sf http://$HOST_IP:$ROUTER_PORT/v1/chat/completions -H 'Content-Type: application/json' -d '{ + "model":"minimaxm3","messages":[{"role":"user","content":"What is 17 times 23? Reply with only the final number."}], + "max_tokens":200,"temperature":0}' 2>&1 | tee "$LOG/infer.json" | tee -a "$LOG/run.log" +echo "" | tee -a "$LOG/run.log" +echo "[run] confirm backend: $(docker logs vm_decode 2>&1 | grep -oE 'Using MoRIIO backend: [A-Z]+' | tail -1)" | tee -a "$LOG/run.log" +echo "[run] DONE $(date)" | tee -a "$LOG/run.log" From 00a72483fd0490a16e23c53f7d451b580cf5b043 Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Thu, 23 Jul 2026 07:34:44 +0000 Subject: [PATCH 10/12] MiniMax-M3-single-pd: fix docker volume paths (absolute LOG, mount bench result-dir) - LOG defaulted to a relative path but is bind-mounted (docker -v rejects relative paths); resolve it to absolute in run_pd_singlenode.sh and bench_pd.sh. - bench_pd.sh mounted only $MOUNT, so --result-dir writes were lost inside the container; also mount $LOG so result JSONs persist. Co-Authored-By: Claude --- scripts/MiniMax-M3-single-pd/bench_pd.sh | 4 ++-- scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/scripts/MiniMax-M3-single-pd/bench_pd.sh b/scripts/MiniMax-M3-single-pd/bench_pd.sh index ff059ed8..f1adb1af 100755 --- a/scripts/MiniMax-M3-single-pd/bench_pd.sh +++ b/scripts/MiniMax-M3-single-pd/bench_pd.sh @@ -13,7 +13,7 @@ set -uo pipefail IMG="${IMG:-rocm/vllm-dev:vllm-0.23.1-rocm723-mi35x-mori-0625}" MODEL="${MODEL:-/models/MiniMax-M3-MXFP4}" MOUNT="${MOUNT:-/models}" -LOG="${LOG:-./logs/bench}"; mkdir -p "$LOG" +LOG="${LOG:-$PWD/logs/bench}"; mkdir -p "$LOG"; LOG="$(cd "$LOG" && pwd)" # absolute (docker -v needs it) ROUTER_PORT="${ROUTER_PORT:-30000}" SHAPES="${SHAPES:-1024,1024}" # comma-separated "isl,osl"; space-separate multiple shapes CONCS="${CONCS:-1 8 32}" @@ -24,7 +24,7 @@ for shape in $SHAPES; do ISL="${shape%,*}"; OSL="${shape#*,}" for c in $CONCS; do echo "[bench] --- isl=$ISL osl=$OSL conc=$c ---" | tee -a "$LOG/bench.log" - docker run --rm --entrypoint bash --network host -v "$MOUNT":"$MOUNT" "$IMG" -c " + docker run --rm --entrypoint bash --network host -v "$MOUNT":"$MOUNT" -v "$LOG":"$LOG" "$IMG" -c " vllm bench serve --backend openai-chat \ --base-url http://$HOST_IP:$ROUTER_PORT --endpoint /v1/chat/completions \ --model $MODEL --served-model-name minimaxm3 --dataset-name random \ diff --git a/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh b/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh index d61f8352..80d06ecb 100755 --- a/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh +++ b/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh @@ -21,7 +21,7 @@ set -uo pipefail IMG="${IMG:-rocm/vllm-dev:vllm-0.23.1-rocm723-mi35x-mori-0625}" ROUTER_IMG="${ROUTER_IMG:-vllm/vllm-router:nightly-20260721-1fbcde7}" MODEL="${MODEL:-/models/MiniMax-M3-MXFP4}" # flat dir w/ weights + remote-code .py (see README) -LOG="${LOG:-./logs}"; mkdir -p "$LOG" +LOG="${LOG:-$PWD/logs}"; mkdir -p "$LOG"; LOG="$(cd "$LOG" && pwd)" # absolute (docker -v needs it) BACKEND="${BACKEND:-xgmi}" # xgmi | rdma UCXDEV="${UCXDEV:-ionic_0:1}" # RoCE NIC for UCX (rdma backend); harmless for xgmi GPU_MEM_UTIL="${GPU_MEM_UTIL:-0.90}" # lower to ~0.30 on shared/occupied nodes From 48cfaaf35ec3052ee6fa49633b18386c9be4899b Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Thu, 23 Jul 2026 07:54:31 +0000 Subject: [PATCH 11/12] MiniMax-M3-single-pd: retry smoke inference (engine warmup after registration) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Dry-run showed the single curl right after worker registration can fail silently (curl -sf returns empty during first-token warmup), leaving infer.json empty even though the stack serves correctly. Replace with a retry loop (12×15s) that checks for a valid "choices" response and prints the answer. Verified on MI355X: manual request returns 391 (17x23). Co-Authored-By: Claude --- .../MiniMax-M3-single-pd/run_pd_singlenode.sh | 19 ++++++++++++++----- 1 file changed, 14 insertions(+), 5 deletions(-) diff --git a/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh b/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh index 80d06ecb..8354bd07 100755 --- a/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh +++ b/scripts/MiniMax-M3-single-pd/run_pd_singlenode.sh @@ -88,11 +88,20 @@ for i in $(seq 1 120); do done sleep 10 -# 4) smoke inference through the router (exercises P->D KV transfer over $BACKEND) +# 4) smoke inference through the router (exercises P->D KV transfer over $BACKEND). +# Retry: right after registration the first request can fail while the engine warms up. echo "[run] === INFERENCE via router ($BACKEND) ===" | tee -a "$LOG/run.log" -curl -sf http://$HOST_IP:$ROUTER_PORT/v1/chat/completions -H 'Content-Type: application/json' -d '{ - "model":"minimaxm3","messages":[{"role":"user","content":"What is 17 times 23? Reply with only the final number."}], - "max_tokens":200,"temperature":0}' 2>&1 | tee "$LOG/infer.json" | tee -a "$LOG/run.log" -echo "" | tee -a "$LOG/run.log" +REQ='{"model":"minimaxm3","messages":[{"role":"user","content":"What is 17 times 23? Reply with only the final number."}],"max_tokens":200,"temperature":0}' +for attempt in $(seq 1 12); do + resp=$(curl -s --max-time 120 http://$HOST_IP:$ROUTER_PORT/v1/chat/completions \ + -H 'Content-Type: application/json' -d "$REQ" 2>&1) + if echo "$resp" | grep -q '"choices"'; then + echo "$resp" | tee "$LOG/infer.json" >/dev/null + echo "[run] inference OK (attempt $attempt)" | tee -a "$LOG/run.log" + echo "$resp" | python3 -c "import sys,json;d=json.load(sys.stdin);print('[run] answer:',d['choices'][0]['message'].get('content'))" 2>/dev/null | tee -a "$LOG/run.log" + break + fi + echo "[run] inference attempt $attempt not ready, retrying in 15s..." | tee -a "$LOG/run.log"; sleep 15 +done echo "[run] confirm backend: $(docker logs vm_decode 2>&1 | grep -oE 'Using MoRIIO backend: [A-Z]+' | tail -1)" | tee -a "$LOG/run.log" echo "[run] DONE $(date)" | tee -a "$LOG/run.log" From 85905d52fac9cb1533e224af6555c33a934b05a1 Mon Sep 17 00:00:00 2001 From: Ravi Gupta Date: Mon, 27 Jul 2026 09:35:55 +0000 Subject: [PATCH 12/12] Add GLM-5.2 disaggregated P/D recipe (scripts/GLM-5.2-disagg) Two-node prefill/decode disaggregation for GLM-5.2-MXFP4 on MI355X: ATOM native engine + atomesh router, KV cache over mooncake TCP, TP8/TP8. Two non-obvious settings, both applied by serve_atom.sh: - MXFP4 online_quant_config excludes "*expert*" so MoE experts stay MXFP4 (FP8 per-block expert re-quant deadlocks the decode engine). - MC_FORCE_TCP=1 so mooncake moves KV over TCP; without it it attempts GPU-direct RDMA writes and every transfer returns "chunk error -1" (these nodes have no amdgpu_peermem). Uses public upstream image rocm/atom-dev; README notes the exact tested digest. --- README.md | 1 + scripts/GLM-5.2-disagg/.gitignore | 3 + scripts/GLM-5.2-disagg/README.md | 90 ++++++++++++++++++++++ scripts/GLM-5.2-disagg/bench_pd.sh | 71 ++++++++++++++++++ scripts/GLM-5.2-disagg/serve_atom.sh | 108 +++++++++++++++++++++++++++ 5 files changed, 273 insertions(+) create mode 100644 scripts/GLM-5.2-disagg/.gitignore create mode 100644 scripts/GLM-5.2-disagg/README.md create mode 100755 scripts/GLM-5.2-disagg/bench_pd.sh create mode 100755 scripts/GLM-5.2-disagg/serve_atom.sh diff --git a/README.md b/README.md index 9bd997f4..2e2b2217 100644 --- a/README.md +++ b/README.md @@ -28,6 +28,7 @@ Below are blueprints of supported models along with their documentation. | [**Qwen3.5-35B-A3B aggregated inference**](scripts/Qwen3.5/README.md) | Single-node aggregated serving/benchmark (SGLang/vLLM); recommended default SGLang+MXFP4+TP1 (8 instances/node) | Qwen3.5-35B-A3B (MXFP4, FP8) | | [**MiniMax-M3 disaggregated P/D inference**](scripts/MiniMax-M3/README.md) | Distributed Inference P/D disaggregation for MiniMax-M3 on MI355X (ATOM+mooncake, MXFP4); top-throughput InferenceX config | MiniMax-M3 (MXFP4) | | [**MiniMax-M3 single-node P/D over XGMI**](scripts/MiniMax-M3-single-pd/README.md) | Single-node prefill/decode disaggregation for MiniMax-M3 on MI355X via vLLM + MoRIIO with GPU-to-GPU XGMI KV transfer (TP4/TP4) | MiniMax-M3 (MXFP4) | +| [**GLM-5.2 disaggregated P/D inference**](scripts/GLM-5.2-disagg/README.md) | Two-node prefill/decode disaggregation for GLM-5.2 on MI355X (ATOM + atomesh + mooncake TCP KV transfer, TP8/TP8) | GLM-5.2 (MXFP4) | | [**KVCache Transfer Bench**](scripts/kvcache_transfer_bench/README.md) | Inter-node Transfer Benchmark | no specific models | ## Table of Contents diff --git a/scripts/GLM-5.2-disagg/.gitignore b/scripts/GLM-5.2-disagg/.gitignore new file mode 100644 index 00000000..71e0fbb8 --- /dev/null +++ b/scripts/GLM-5.2-disagg/.gitignore @@ -0,0 +1,3 @@ +# runtime artifacts created by serve_atom.sh / bench_pd.sh +logs/ +cache/ diff --git a/scripts/GLM-5.2-disagg/README.md b/scripts/GLM-5.2-disagg/README.md new file mode 100644 index 00000000..b86046d0 --- /dev/null +++ b/scripts/GLM-5.2-disagg/README.md @@ -0,0 +1,90 @@ +# GLM-5.2-MXFP4 — disaggregated (P/D) serving on MI355X (ATOM + mooncake TCP) + +Two-node **prefill/decode disaggregated** serving for GLM-5.2-MXFP4 on AMD Instinct MI355X (gfx950): +prefill (`kv_producer`) TP8 on node A + decode (`kv_consumer`) TP8 on node B, fronted by one `atomesh` +router, with the KV cache streamed prefill → decode over the **mooncake TransferEngine (TCP)**. + +## Why this stack + +- **ATOM native engine.** GLM-5.2 is `GlmMoeDsa` — a DeepSeek-v4-style MoE with multi-head latent + attention (MLA) and sparse (DSA) attention. ATOM's native engine implements this architecture directly. + (SGLang's GLM-5.2 DSA top-k kernels are CUDA-only and not yet ported to ROCm/hipcc.) +- **mooncake KV transfer over TCP.** These nodes do not load the `amdgpu_peermem` kernel module, so + GPU-direct RDMA memory registration is unavailable. mooncake's TCP transport moves the KV cache over + the host RoCE interface instead — no GPU-direct RDMA hop required. + +## The two settings that make it work + +Both are already applied by `serve_atom.sh`; they are the non-obvious part of this recipe. + +1. **MXFP4 quant config excludes `*expert*`.** `--online_quant_config` keeps the MoE experts in MXFP4 + (not re-quantized). The FP8 per-block expert re-quantization deadlocks the decode engine at + "Weight post-processing (online quantization)". +2. **`MC_FORCE_TCP=1`.** Forces mooncake's TCP data path. Setting only `protocol: tcp` in the + kv-transfer-config is not enough — mooncake still attempts GPU-direct RDMA writes and every transfer + returns `block RDMA chunk error -1`. + +## Prerequisites + +- **2 nodes**, 8× gfx950 (MI355X) each, on the same RoCE fabric. +- Image `rocm/atom-dev:latest` (ships `atom.entrypoints.openai_server`, `atomesh`, and mooncake). + It is public on Docker Hub (anonymously pullable). + > **FYI — tested image.** This recipe was validated on `rocm/atom-dev@sha256:280d2fe1a4d79db51cb37ae96112cecfcc208436b29ba0880bf04039606f5157` + > (the `:latest` tag at time of testing). `:latest` floats, so for a reproducible run pin that digest via + > `IMG=rocm/atom-dev@sha256:280d2fe1a4d7...`. +- GLM-5.2-MXFP4 weights staged to the same path on every node (default `/models/GLM-5.2-MXFP4`; + the host `/models` dir is mounted into the containers). +- ionic RDMA provider libs present on the host (bind-mounted automatically by the launcher). + +## Files + +| File | Role | +|------|------| +| `serve_atom.sh` | Per-node launcher (`ROLE=prefill\|decode\|router`); starts the ATOM engine / atomesh router. | +| `bench_pd.sh` | Concurrency-sweep driver against the router (OpenAI-compatible `/v1/chat/completions`). | + +## Usage + +Pick the prefill node IP (`PROXY_IP`) and decode node IP (`DECODE_IP`). Weights must be staged to +`$MODEL` on both nodes. + +```bash +# On node A (prefill): +PROXY_IP= DECODE_IP= ROLE=prefill ./serve_atom.sh + +# On node B (decode): +PROXY_IP= DECODE_IP= ROLE=decode ./serve_atom.sh + +# On node A, once prefill (:8010) and decode (:8020) are serving — start the router: +PROXY_IP= DECODE_IP= ROLE=router ./serve_atom.sh +``` + +The router exposes an OpenAI-compatible endpoint on `http://:30000`: + +```bash +curl -s http://:30000/v1/chat/completions -H 'Content-Type: application/json' -d '{ + "model":"/models/GLM-5.2-MXFP4", + "messages":[{"role":"user","content":"What is the capital of France? One word."}], + "max_tokens":32,"temperature":0}' +``` + +Throughput sweep through the router: + +```bash +ROUTER_IP= ./bench_pd.sh # 1k/1k @ conc 1,8,32 +SHAPES="8192,1024" CONCS="8 16" ROUTER_IP= ./bench_pd.sh +``` + +## Overridable env (both scripts) + +| Var | Default | Meaning | +|-----|---------|---------| +| `MODEL` | `/models/GLM-5.2-MXFP4` | model path (same on every node) | +| `MOUNT` | `/models` | host dir mounted into the container (must contain `MODEL`) | +| `TP` | `8` | tensor-parallel size per engine | +| `PF_PORT` / `DC_PORT` | `8010` / `8020` | prefill / decode server ports | +| `HS_PORT` | `6301` | mooncake handshake port | +| `ROUTER_PORT` | `30000` | atomesh router port | +| `SOCKET_IFNAME` | `ens3` | host interface for mooncake/NCCL sockets | +| `CACHE_ROOT` | `$PWD/cache` | host dir for persisted JIT caches | +| `IMG` | `rocm/atom-dev:latest` | container image | diff --git a/scripts/GLM-5.2-disagg/bench_pd.sh b/scripts/GLM-5.2-disagg/bench_pd.sh new file mode 100755 index 00000000..ac6e008c --- /dev/null +++ b/scripts/GLM-5.2-disagg/bench_pd.sh @@ -0,0 +1,71 @@ +#!/bin/bash +# ============================================================================= +# Throughput benchmark for the GLM-5.2 ATOM + mooncake-TCP P/D deployment. +# Drives the atomesh router (OpenAI-compatible) across an ISL/OSL x concurrency +# sweep. Assumes serve_atom.sh prefill+decode+router are already up. +# +# Usage: +# ./bench_pd.sh # default: 1k/1k @ conc 1,8,32 +# SHAPES="8192,1024" CONCS="8 16" ./bench_pd.sh +# ROUTER_IP=10.0.0.1 ROUTER_PORT=30000 MODEL=/models/GLM-5.2-MXFP4 ./bench_pd.sh +# ============================================================================= +set -uo pipefail +MODEL="${MODEL:-/models/GLM-5.2-MXFP4}" +ROUTER_IP="${ROUTER_IP:-$(ip route get 1.1.1.1 2>/dev/null | awk '/src/{print $7}')}" +ROUTER_PORT="${ROUTER_PORT:-30000}" +SHAPES="${SHAPES:-1024,1024}" # comma-separated "isl,osl"; space-separate multiple shapes +CONCS="${CONCS:-1 8 32}" +LOG="${LOG:-$PWD/logs/bench}"; mkdir -p "$LOG"; LOG="$(cd "$LOG" && pwd)" + +URL="http://$ROUTER_IP:$ROUTER_PORT/v1/chat/completions" +echo "[bench] router=$URL shapes='$SHAPES' concs='$CONCS' $(date)" | tee "$LOG/bench.log" + +for shape in $SHAPES; do + ISL="${shape%,*}"; OSL="${shape#*,}" + for c in $CONCS; do + echo "[bench] --- isl=$ISL osl=$OSL conc=$c ---" | tee -a "$LOG/bench.log" + URL="$URL" MODEL="$MODEL" ISL="$ISL" OSL="$OSL" CONC="$c" \ + RESULT="$LOG/bench_isl${ISL}_osl${OSL}_conc${c}.json" \ + python3 - <<'PY' 2>&1 | tee -a "$LOG/bench.log" +import json, os, time, threading, urllib.request, statistics +URL=os.environ["URL"]; MODEL=os.environ["MODEL"] +ISL=int(os.environ["ISL"]); OSL=int(os.environ["OSL"]); CONC=int(os.environ["CONC"]) +# ~ISL-token prompt from a repeated block (~44 tok/block); vary per request id. +BLOCK=("In distributed machine learning inference, disaggregated prefill and decode " +"serving separates the compute-bound prefill stage from the memory-bandwidth-bound " +"decode stage across distinct hardware pools. ") +reps=max(1, ISL//44) +def prompt(rid): return "".join(f"[doc {rid} seg {i}] "+BLOCK for i in range(reps))+f"\n\nSummarize (req {rid})." +res=[]; lock=threading.Lock() +def worker(rid): + body=json.dumps({"model":MODEL,"messages":[{"role":"user","content":prompt(rid)}], + "max_tokens":OSL,"temperature":0}).encode() + req=urllib.request.Request(URL,data=body,headers={"Content-Type":"application/json"}) + t0=time.time() + try: + with urllib.request.urlopen(req,timeout=1200) as r: d=json.loads(r.read()) + u=d.get("usage",{}) + with lock: res.append((time.time()-t0,u.get("prompt_tokens"),u.get("completion_tokens"),u.get("ttft_s"),None)) + except Exception as e: + with lock: res.append((time.time()-t0,None,None,None,str(e)[:120])) +T0=time.time() +ths=[threading.Thread(target=worker,args=(i,)) for i in range(CONC)] +for t in ths: t.start() +for t in ths: t.join() +WALL=time.time()-T0 +ok=[r for r in res if r[4] is None]; err=[r for r in res if r[4] is not None] +out={"isl":ISL,"osl":OSL,"conc":CONC,"wall_s":round(WALL,2),"success":len(ok),"errors":len(err)} +if ok: + cts=[r[2] for r in ok if r[2]]; pts=[r[1] for r in ok if r[1]]; lats=[r[0] for r in ok] + ttfts=[r[3] for r in ok if r[3]] + out.update(output_tok_s=round(sum(cts)/WALL,1), total_tok_s=round((sum(pts)+sum(cts))/WALL,1), + mean_lat_s=round(statistics.mean(lats),2), p50_lat_s=round(statistics.median(lats),2)) + if ttfts: out["mean_ttft_s"]=round(statistics.mean(ttfts),3) +print(f"Successful: {len(ok)}/{CONC} output_tok/s={out.get('output_tok_s')} " + f"total_tok/s={out.get('total_tok_s')} mean_ttft={out.get('mean_ttft_s')}s p50_lat={out.get('p50_lat_s')}s") +if err: print(f" errors={len(err)} (first: {err[0][4]})") +open(os.environ["RESULT"],"w").write(json.dumps(out,indent=2)) +PY + done +done +echo "[bench] DONE $(date) results in $LOG" | tee -a "$LOG/bench.log" diff --git a/scripts/GLM-5.2-disagg/serve_atom.sh b/scripts/GLM-5.2-disagg/serve_atom.sh new file mode 100755 index 00000000..8b05a423 --- /dev/null +++ b/scripts/GLM-5.2-disagg/serve_atom.sh @@ -0,0 +1,108 @@ +#!/bin/bash +# ============================================================================= +# GLM-5.2-MXFP4 — 1P1D prefill/decode disaggregation on 2x MI355X (8x gfx950) +# ATOM native engine (atom.entrypoints.openai_server) + atomesh router +# KV cache moves prefill -> decode over the mooncake TransferEngine (TCP). +# +# Per-NODE launcher. Run once per role: +# PROXY_IP= DECODE_IP= ROLE=prefill ./serve_atom.sh # node A +# PROXY_IP= DECODE_IP= ROLE=decode ./serve_atom.sh # node B +# PROXY_IP= DECODE_IP= ROLE=router ./serve_atom.sh # node A, after A+B are up +# +# Two non-obvious settings make this work on this class of node (see README): +# 1) MXFP4 online_quant_config excludes "*expert*" -> experts stay MXFP4 and +# are NOT re-quantized. The FP8 per-block re-quant deadlocks the decode. +# 2) MC_FORCE_TCP=1 -> mooncake moves KV over TCP. Without it mooncake attempts +# GPU-direct RDMA writes and every transfer returns "chunk error -1" (these +# nodes have no amdgpu_peermem, so GPU-direct RDMA registration is unavailable). +# ============================================================================= +set -uo pipefail +ROLE="${ROLE:?prefill|decode|router}" +IMG="${IMG:-rocm/atom-dev:latest}" +MODEL="${MODEL:-/models/GLM-5.2-MXFP4}" # weights staged here on every node +MOUNT="${MOUNT:-/models}" # host dir mounted into the container (must contain MODEL) +TP="${TP:-8}" +CONTAINER="${CONTAINER:-atom_${ROLE}}" +LOG="${LOG:-$PWD/logs}"; mkdir -p "$LOG"; LOG="$(cd "$LOG" && pwd)" # absolute (docker -v needs it) +CACHE_ROOT="${CACHE_ROOT:-$PWD/cache}"; mkdir -p "$CACHE_ROOT" 2>/dev/null || true +CACHE_ROOT="$(cd "$CACHE_ROOT" && pwd)" + +PROXY_IP="${PROXY_IP:?prefill/proxy node IP}" +PING_PORT="${PING_PORT:-36367}"; ROUTER_PORT="${ROUTER_PORT:-30000}" +PF_PORT="${PF_PORT:-8010}"; DC_PORT="${DC_PORT:-8020}"; HS_PORT="${HS_PORT:-6301}" + +# --- ionic RDMA provider libs (mooncake initializes an RDMA context even in TCP +# mode, so the host provider libs must be visible inside the container) ------ +IONIC_VER_SO=$(ls /usr/lib/x86_64-linux-gnu/libionic.so.1.0.* 2>/dev/null | head -1) +NICM=() +[ -f /usr/lib/x86_64-linux-gnu/libionic.so ] && NICM+=(-v /usr/lib/x86_64-linux-gnu/libionic.so:/usr/lib/x86_64-linux-gnu/libionic.so:ro) +[ -f /usr/lib/x86_64-linux-gnu/libionic.so.1 ] && NICM+=(-v /usr/lib/x86_64-linux-gnu/libionic.so.1:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro) +[ -n "$IONIC_VER_SO" ] && NICM+=(-v "${IONIC_VER_SO}:${IONIC_VER_SO}:ro") +[ -f /usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so ] && NICM+=(-v /usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:ro) +[ -f /etc/libibverbs.d/ionic.driver ] && NICM+=(-v /etc/libibverbs.d/ionic.driver:/etc/libibverbs.d/ionic.driver:ro) + +# --- engine + mooncake-TCP env ----------------------------------------------- +MENV=( + # ATOM: QuickReduce custom all-reduce (TP8) + GLM DSA QK-norm quant fusion. + -e AITER_QUICK_REDUCE_QUANTIZATION=INT4 -e AITER_USE_FLYDSL_MOE_SORTING=1 + -e ATOM_ENABLE_DS_QKNORM_QUANT_FUSION=1 + # mooncake KV over TCP (the transfer fix). MC_FORCE_TCP is load-bearing: + # without it mooncake tries GPU-direct RDMA writes -> "chunk error -1". + -e MC_FORCE_TCP="${MC_FORCE_TCP:-1}" -e MC_TCP_ENABLE_CONNECTION_POOL=true + -e MC_SOCKET_IFNAME="${SOCKET_IFNAME:-ens3}" -e MC_LOG_LEVEL="${MC_LOG_LEVEL:-INFO}" + -e MORI_SOCKET_IFNAME="${SOCKET_IFNAME:-ens3}" + -e RDMAV_FORK_SAFE=1 -e HSA_NO_SCRATCH_RECLAIM=1 -e PYTHONHASHSEED=0 + -e NCCL_SOCKET_IFNAME="${SOCKET_IFNAME:-ens3}" -e GLOO_SOCKET_IFNAME="${SOCKET_IFNAME:-ens3}" + -e NCCL_IB_RETRY_CNT=15 -e NCCL_IB_TIMEOUT=22 + # First-run AITER JIT compiles (tens of seconds each, serialized) can exceed + # NCCL's watchdog/heartbeat -> rank killed. Disable the heartbeat kill + extend + # timeouts so the cold compiles survive. + -e TORCH_NCCL_ENABLE_MONITORING=0 -e TORCH_NCCL_HEARTBEAT_TIMEOUT_SEC=3600 + -e TORCH_NCCL_ASYNC_ERROR_HANDLING=0 -e NCCL_TIMEOUT=3600 + -e ATOM_ENGINE_READY_TIMEOUT_S=3600 +) +DEV=(--device /dev/kfd --device /dev/dri --device /dev/infiniband) +COMMON=(--network host --ipc host --privileged --group-add video + --cap-add IPC_LOCK --cap-add NET_ADMIN --ulimit memlock=-1 --ulimit stack=67108864 + --ulimit nofile=1048576:1048576 --shm-size 128G + -v "$MOUNT":"$MOUNT":ro -v /sys:/sys + -v /sys/class/infiniband:/sys/class/infiniband:ro -v "$LOG":/out) + +# per-worker isolated JIT caches (host-persisted so recompiles survive restarts) +CACHE="$CACHE_ROOT/atomcache_${ROLE}"; mkdir -p "$CACHE" 2>/dev/null || true +CENV=(-e HOME=/out -e TORCHINDUCTOR_CACHE_DIR=/cache/inductor -e TRITON_CACHE_DIR=/cache/triton + -e AITER_JIT_DIR=/cache/aiter -v "$CACHE":/cache) + +# GLM-5.2 MXFP4 online-quant config. exclude_layer "*expert*" keeps the MoE +# experts in MXFP4 (not re-quantized) -> avoids the decode online-quant deadlock. +QCFG='{"global_quant_config":"ptpc_fp8","exclude_layer":["lm_head","model.embed_tokens","*.mlp.gate","*expert*"]}' + +serve(){ local port=$1 kvrole=$2 + docker rm -f "$CONTAINER" >/dev/null 2>&1 || true + docker run -d --name "$CONTAINER" --entrypoint bash "${DEV[@]}" "${COMMON[@]}" "${NICM[@]}" "${MENV[@]}" "${CENV[@]}" \ + -e HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 "$IMG" -c " + # mooncake native libs live in the venv site-packages; put them + rocm on the loader path. + export LD_LIBRARY_PATH=\$(python3 -c 'import sysconfig; print(sysconfig.get_path(\"purelib\"))')/mooncake:/opt/rocm/lib:\${LD_LIBRARY_PATH:-} + python -m atom.entrypoints.openai_server --model $MODEL --host 0.0.0.0 --server-port $port \ + -tp $TP --kv_cache_dtype fp8 --no-enable_prefix_caching \ + --online_quant_config '$QCFG' \ + --kv-transfer-config '{\"kv_role\":\"$kvrole\",\"kv_connector\":\"mooncake\",\"protocol\":\"tcp\",\"proxy_ip\":\"$PROXY_IP\",\"proxy_ping_port\":$PING_PORT,\"http_port\":$port,\"handshake_port\":$HS_PORT}' \ + 2>&1 | tee /out/atom_$ROLE.log" + echo "[atom-$ROLE] launched on :$port (mooncake-tcp $kvrole, proxy=$PROXY_IP)" | tee "$LOG/serve_$ROLE.log" +} + +case "$ROLE" in + prefill) serve "$PF_PORT" kv_producer ;; + decode) serve "$DC_PORT" kv_consumer ;; + router) + docker rm -f "$CONTAINER" >/dev/null 2>&1 || true + docker run -d --name "$CONTAINER" --network host -v "$LOG":/out \ + -v "$MOUNT":"$MOUNT":ro "$IMG" bash -lc " + exec atomesh launch --host 0.0.0.0 --port $ROUTER_PORT --pd-disaggregation \ + --prefill http://$PROXY_IP:$PF_PORT --decode http://${DECODE_IP:?}:$DC_PORT \ + --policy random --backend atom --disable-health-check --disable-circuit-breaker \ + --log-level info 2>&1 | tee /out/atom_router.log" + echo "[atom-router] launched :$ROUTER_PORT (prefill=$PROXY_IP:$PF_PORT decode=$DECODE_IP:$DC_PORT)" | tee "$LOG/serve_router.log" + ;; + *) echo "ROLE must be prefill|decode|router"; exit 1 ;; +esac