diff --git a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile index 06dd18aa..8d1e84b7 100644 --- a/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile +++ b/docker/vllm_disagg_inference.ubuntu.amd.Dockerfile @@ -68,12 +68,12 @@ # Build context = repo root: # docker build -f docker/vllm_disagg_inference.ubuntu.amd.Dockerfile -t / . # -# BASE_IMAGE is the open rocm/vllm-dev ci_base pinned by the validated recipe -# (dist-inf-cookbook Dockerfile.vllm.mori121_shareable). Override --build-arg -# BASE_IMAGE=... to build on a different ROCm base. vLLM compile is long (~30-60 min). +# BASE_IMAGE is the shikpate ROCm/vLLM/MoRI-1.2.3 base carrying the validated +# Wide-EP MoRI stack. Override --build-arg BASE_IMAGE=... to build on a different +# ROCm base. vLLM compile is long (~30-60 min). # ============================================================================= -ARG BASE_IMAGE=rocm/vllm-dev:ci_base-0fcd9b99cc9d63202da4c858d8ebc6582c9e2491 +ARG BASE_IMAGE=rocmshared/pytorch-private:vllm-rocm_07_22_2026_shikpate_mori1.2.3 FROM ${BASE_IMAGE} ENTRYPOINT [] @@ -156,9 +156,11 @@ RUN rm -rf /opt/vllm_cache/aiter_jit /root/.aiter && echo "cleared stale AITER J # committed commits (no working-tree edits). # ----------------------------------------------------------------------------- # VLLM_REPO/REF are a PUBLIC GitHub repo + branch (the Wide-EP WRITE-mode vLLM the -# dist-inf-cookbook mori121 image builds from). Override to your own vLLM fork/branch. -ARG VLLM_REPO=https://github.com/shikamd123/vllm.git -ARG VLLM_REF=vllm_2p2d_wide-ep_write_shikpate_test_06_29_customer +# disagg image builds from). dsv3-wideep-clean = the shikpate/Shiksha Wide-EP MoRI-IO +# base + a generic crash-safe rmsnorm probe fix, with all GLM sparse-DSA commits +# stripped (DeepSeek-V3 uses standard MLA, not DSA). Override to your own fork/branch. +ARG VLLM_REPO=https://github.com/raviguptaamd/vllm.git +ARG VLLM_REF=dsv3-wideep-clean ENV VLLM_TARGET_DEVICE=rocm \ PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} \ MAX_JOBS=${MAX_JOBS} diff --git a/scripts/vllm_dissag/benchmark_niah.py b/scripts/vllm_dissag/benchmark_niah.py index 0cdd027e..51454882 100755 --- a/scripts/vllm_dissag/benchmark_niah.py +++ b/scripts/vllm_dissag/benchmark_niah.py @@ -1,6 +1,6 @@ #!/usr/bin/env python3 # Needle-in-a-haystack long-context retrieval test. -# Adapted from vllm-project/vllm issue #47042 (GLM-5.2 sparse-MLA decode collapse), +# Adapted from vllm-project/vllm issue #47042 (long-context sparse-MLA decode collapse), # generalized to run against any OpenAI-compatible endpoint / model. # # Env: @@ -8,7 +8,15 @@ # NIAH_MODEL model name/tag the server serves (required — the served path) # NIAH_WORDS comma list of context sizes in words (default 2000,8000,20000,35000) # NIAH_MAXTOK max_tokens for the answer (default 2048) +# NIAH_SEEDS comma list of needle-layout seeds (default 0,1,2); summary reports +# mean/min/max across seeds to separate real accuracy from variance # NIAH_TIMEOUT per-request timeout seconds (default 1800) +# NIAH_WARMUP 1 (default) = send one throwaway request per context length BEFORE +# scoring, so the first-hit JIT/kernel-autotune compile happens outside +# the scored/gated window. On a freshly-booted node the first request of +# a shape can take minutes to compile; without warmup that lands on the +# first scored request -> false 0/10 or timeout. Warmup failures are +# tolerated (logged, not fatal). Set 0 to disable. import os, sys, json, random, urllib.request URL = os.environ.get("NIAH_URL", "http://127.0.0.1:30000/v1/chat/completions") @@ -16,6 +24,14 @@ WORDS = [int(x) for x in os.environ.get("NIAH_WORDS", "2000,8000,20000,35000").split(",") if x.strip()] MAXTOK = int(os.environ.get("NIAH_MAXTOK", "2048")) TIMEOUT = float(os.environ.get("NIAH_TIMEOUT", "1800")) +# Needle layout is seeded, so a single run is deterministic (bit-exact repro on the +# same stack). Run multiple seeds to distinguish real accuracy from single-needle +# variance; the summary reports mean/min/max across seeds. Default 0,1,2. +SEEDS = [int(x) for x in os.environ.get("NIAH_SEEDS", "0,1,2").split(",") if x.strip()] +WARMUP = os.environ.get("NIAH_WARMUP", "1") == "1" +# Warmup uses a generous timeout (cold compile of a long-context shape can take minutes) +# and never fails the run — its only job is to trigger compilation before scoring. +WARMUP_TIMEOUT = max(TIMEOUT, 1800.0) FILLER = ( "table chair window bottle pencil garden river mountain coffee planet " @@ -41,27 +57,53 @@ def make_haystack(n_words, seed=0): return " ".join(words) -def run(n_words): +def _request(n_words, seed, max_tokens, timeout): + """POST one NIAH request; return (message_dict, error_str). Exactly one is non-None.""" body = { "model": MODEL, "messages": [ {"role": "system", "content": SYSTEM}, - {"role": "user", "content": "Find the animals in this list:\n\n" + make_haystack(n_words)}, + {"role": "user", "content": "Find the animals in this list:\n\n" + make_haystack(n_words, seed)}, ], "temperature": 0.0, - "max_tokens": MAXTOK, + "max_tokens": max_tokens, + # Thinking models emit chain-of-thought into a separate reasoning field + # and leave `content` empty until the final answer; with a + # small max_tokens the answer never appears in `content` and the score is a + # false 0/10. Disable thinking so the answer lands in `content` directly. + "chat_template_kwargs": {"enable_thinking": False}, } data = json.dumps(body).encode() req = urllib.request.Request(URL, data=data, headers={"Content-Type": "application/json"}) try: - with urllib.request.urlopen(req, timeout=TIMEOUT) as r: - msg = json.loads(r.read())["choices"][0]["message"] + with urllib.request.urlopen(req, timeout=timeout) as r: + return json.loads(r.read())["choices"][0]["message"], None except Exception as e: - print("words=%6d ERROR %s" % (n_words, e), flush=True) + return None, str(e) + + +def warmup(n_words): + """One throwaway request per length so first-hit compile happens off the scored path. + Never fatal: a warmup timeout just means the shape is still compiling; the scored + request will pay whatever remains (bounded by NIAH_TIMEOUT).""" + _, err = _request(n_words, seed=0, max_tokens=8, timeout=WARMUP_TIMEOUT) + status = "ok" if err is None else ("timeout/err: %s" % err) + print("words=%6d [warmup] %s" % (n_words, status), flush=True) + + +def run(n_words, seed=0): + # Sentinel: None = timeout/transport error (NOT a wrong answer); int = score 0..10. + msg, err = _request(n_words, seed, MAXTOK, TIMEOUT) + if err is not None: + print("words=%6d seed=%d TIMEOUT/ERROR %s" % (n_words, seed, err), flush=True) return None - text = ((msg.get("content") or "") + " " + (msg.get("reasoning_content") or "")).lower() + # Score content plus any reasoning field (some servers surface CoT as + # `reasoning` or `reasoning_content`) so a thinking model is never mis-scored. + text = ((msg.get("content") or "") + " " + + (msg.get("reasoning_content") or "") + " " + + (msg.get("reasoning") or "")).lower() found = sorted(a for a in ANIMALS if a in text) - print("words=%6d found=%2d/10 %s" % (n_words, len(found), found), flush=True) + print("words=%6d seed=%d found=%2d/10 %s" % (n_words, seed, len(found), found), flush=True) return len(found) @@ -70,14 +112,29 @@ def main(): print("NIAH_MODEL must be set (the served model path/name)", file=sys.stderr) sys.exit(2) print("=== NIAH retrieval test ===", flush=True) - print("url=%s model=%s sizes=%s" % (URL, MODEL, WORDS), flush=True) - results = {} + print("url=%s model=%s sizes=%s seeds=%s warmup=%s" % (URL, MODEL, WORDS, SEEDS, WARMUP), flush=True) + # Warmup pass: compile every shape once before scoring, so cold JIT never lands on a + # scored/gated request (the common cause of false 0/10 or timeout on a fresh boot). + if WARMUP: + print("=== NIAH warmup (one throwaway request per length) ===", flush=True) + for n in WORDS: + warmup(n) + results = {} # n_words -> list of scores across seeds (None = timeout/error, not a wrong answer) for n in WORDS: - results[n] = run(n) - print("=== NIAH summary ===", flush=True) + results[n] = [run(n, s) for s in SEEDS] + print("=== NIAH summary (mean/min/max across %d seed(s)) ===" % len(SEEDS), flush=True) for n in WORDS: - v = results[n] - print(" words=%6d found=%s/10" % (n, "ERR" if v is None else v), flush=True) + scored = results[n] + vals = [v for v in scored if v is not None] + n_to = sum(1 for v in scored if v is None) # timeouts/errors, excluded from mean + if not vals: + print(" words=%6d NO-RESULT (%d/%d timed out or errored — likely cold compile; " + "raise NIAH_TIMEOUT or keep NIAH_WARMUP=1)" % (n, n_to, len(scored)), flush=True) + continue + mean = sum(vals) / len(vals) + extra = (" [%d timeout/err excluded]" % n_to) if n_to else "" + print(" words=%6d mean=%.1f/10 min=%d max=%d (n=%d)%s" + % (n, mean, min(vals), max(vals), len(vals), extra), flush=True) if __name__ == "__main__": diff --git a/scripts/vllm_dissag/benchmark_niah.sh b/scripts/vllm_dissag/benchmark_niah.sh index ba49a359..d366b5e4 100755 --- a/scripts/vllm_dissag/benchmark_niah.sh +++ b/scripts/vllm_dissag/benchmark_niah.sh @@ -14,15 +14,29 @@ LOG="/run_logs/${SLURM_JOB_ID}/niah_${SLURM_JOB_ID}_${timestamp}_xP${xP}_yD${yD} echo "==== NIAH long-context retrieval test ====" echo "port=${BENCHMARK_PORT} model=${MODEL_PATH} sizes=${NIAH_WORDS:-2000,8000,20000,35000}" -# Give the router a moment to be fully ready for chat completions. -sleep 10 +# Wait until the router actually serves before starting (replaces a blind sleep). On a +# fresh boot the router may register a few seconds after the workers report ready; poll +# /v1/models until it answers, up to ~5 min. Non-fatal: fall through if the probe can't +# confirm (the harness's own warmup + timeout still protect the run). +_ready=0 +for _i in $(seq 1 60); do + if curl -s -o /dev/null -w '%{http_code}' --max-time 5 \ + "http://127.0.0.1:${BENCHMARK_PORT}/v1/models" 2>/dev/null | grep -q '^200$'; then + _ready=1; echo "[niah] router ready after ~$((_i*5))s"; break + fi + sleep 5 +done +[ "$_ready" = 1 ] || echo "[niah] WARN: router readiness not confirmed in 300s; proceeding (warmup will absorb)" # The server registers the model under its path (served_model_name = MODEL_PATH). +# NIAH_WARMUP=1 (harness default): first-hit JIT compiles off the scored path so a cold +# boot does not produce false 0/10 or timeouts on the first scored request. NIAH_URL="http://127.0.0.1:${BENCHMARK_PORT}/v1/chat/completions" \ NIAH_MODEL="${MODEL_PATH}" \ NIAH_WORDS="${NIAH_WORDS:-2000,8000,20000,35000}" \ NIAH_MAXTOK="${NIAH_MAXTOK:-2048}" \ NIAH_TIMEOUT="${NIAH_TIMEOUT:-1800}" \ +NIAH_WARMUP="${NIAH_WARMUP:-1}" \ python3 "${DIR}/benchmark_niah.py" 2>&1 | tee -a "${LOG}" echo "NIAH results -> ${LOG}" diff --git a/scripts/vllm_dissag/connectors/moriio.sh b/scripts/vllm_dissag/connectors/moriio.sh index 44c64d68..886add47 100644 --- a/scripts/vllm_dissag/connectors/moriio.sh +++ b/scripts/vllm_dissag/connectors/moriio.sh @@ -40,7 +40,19 @@ connector_init() { # decode=low_latency (InterNodeV1LL) and REJECT the bare "mori" alias. Default # to the per-role names; override via PREFILL_MORI_BACKEND/DECODE_MORI_BACKEND # (or VLLM_ALL2ALL_BACKEND for the prefill side). - PREFILL_MORI_BACKEND="${PREFILL_MORI_BACKEND:-${VLLM_ALL2ALL_BACKEND:-mori_high_throughput}}" + # + # DP-aware prefill default: the high_throughput InterNodeV1 dispatch silently + # mis-routes MoE tokens at >=32-way expert parallelism (validated on DeepSeek-V3 + # NIAH: PREFILL_DP=32 collapses to ~4/10 with empty answers; decode's low_latency + # path is clean at DP=32). Fall back to low_latency for prefill when PREFILL_DP>=32 + # so wide topologies (4P+/xPyD) are correct out of the box. Any explicit + # PREFILL_MORI_BACKEND/VLLM_ALL2ALL_BACKEND override still wins. + local _prefill_default="mori_high_throughput" + if [[ "${PREFILL_DP_SIZE:-0}" -ge 32 ]]; then + _prefill_default="mori_low_latency" + echo "[moriio] PREFILL_DP_SIZE=${PREFILL_DP_SIZE} >= 32: defaulting prefill backend to mori_low_latency (high_throughput InterNodeV1 mis-routes at EP>=32)." + fi + PREFILL_MORI_BACKEND="${PREFILL_MORI_BACKEND:-${VLLM_ALL2ALL_BACKEND:-${_prefill_default}}}" DECODE_MORI_BACKEND="${DECODE_MORI_BACKEND:-mori_low_latency}" } @@ -200,6 +212,12 @@ connector_launch_worker() { local _kvdtype="${KV_CACHE_DTYPE:-fp8}" local mem_args=() [[ -n "${KV_CACHE_MEMORY_BYTES:-}" ]] && mem_args+=(--kv-cache-memory-bytes "${KV_CACHE_MEMORY_BYTES}") + # Bound the attention workspace: without --max-model-len vLLM sizes it for the + # model's full context (DeepSeek-V3 = 163840), producing a ~128 GiB workspace + # alloc that OOMs alongside the EP-sharded weights. MAX_MODEL_LEN (models.yaml) + # caps it to what the workload needs. + local mml_args=() + [[ -n "${MAX_MODEL_LEN:-}" ]] && mml_args+=(--max-model-len "${MAX_MODEL_LEN}") if [[ "${DRY_RUN:-0}" == "1" ]]; then _dryrun_emit "moriio" "${log_prefix}" "${role}" \ @@ -213,6 +231,7 @@ connector_launch_worker() { --port "${SERVE_PORT}" \ --gpu-memory-utilization "${GPU_MEMORY_UTILIZATION:-0.8}" \ "${mem_args[@]}" \ + "${mml_args[@]}" \ --kv-cache-dtype "${_kvdtype}" \ --block-size "${_block}" \ --no-enable-prefix-caching \ @@ -233,6 +252,7 @@ connector_launch_worker() { --port ${SERVE_PORT} \ --gpu-memory-utilization ${GPU_MEMORY_UTILIZATION:-0.8} \ "${mem_args[@]}" \ + "${mml_args[@]}" \ --kv-cache-dtype "${_kvdtype}" \ --block-size "${_block}" \ --no-enable-prefix-caching \ diff --git a/scripts/vllm_dissag/models.yaml b/scripts/vllm_dissag/models.yaml index 23d66059..583ec85d 100644 --- a/scripts/vllm_dissag/models.yaml +++ b/scripts/vllm_dissag/models.yaml @@ -144,6 +144,10 @@ _deepseek_recipe_env: &deepseek_recipe_env KV_CACHE_DTYPE: "fp8" KV_CACHE_MEMORY_BYTES: "20000000000" GPU_MEMORY_UTILIZATION: "0.80" + # Cap context so the attention workspace fits alongside EP-sharded weights. + # Without this vLLM sizes the workspace for the model's full 163840 ctx (~128 GiB + # alloc -> OOM). 65536 covers the NIAH sweep (longest 35k words ~= 50k tokens). + MAX_MODEL_LEN: "65536" VLLM_CUDAGRAPH_MODE: "PIECEWISE" PREFILL_CUDAGRAPH_MODE: "NONE" DECODE_CUDAGRAPH_MODE: "PIECEWISE" diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm index c71fc7e8..79330c5f 100755 --- a/scripts/vllm_dissag/run_xPyD_models.slurm +++ b/scripts/vllm_dissag/run_xPyD_models.slurm @@ -68,6 +68,22 @@ for f in "${REQUIRED_FILES[@]}"; do done echo "Running from: $(pwd)" +# ------------------------------------------------------------------------------ +# models.yaml env precedence: capture which recipe knobs the USER explicitly set +# at submit time. The driver (vllm_disagg.sh) uses this to let models.yaml `env:` +# OVERRIDE image-baked ENV defaults (e.g. a DeepSeek-tuned image bakes +# KV_BLOCK_SIZE=16 / VLLM_ROCM_USE_AITER_MLA=0, which would otherwise shadow a +# model's own recipe defined in models.yaml env:), while a genuine +# submit-time `-e VAR=...` still wins. Precedence: image-baked < models.yaml < submit -e. +# Captured HERE (before the slurm sets any defaults) so it reflects user intent only. +_RECIPE_ENV_KEYS="VLLM_USE_V1 VLLM_ROCM_USE_AITER VLLM_ROCM_USE_AITER_RMSNORM VLLM_ROCM_USE_AITER_MLA KV_BLOCK_SIZE KV_CACHE_DTYPE KV_CACHE_MEMORY_BYTES GPU_MEMORY_UTILIZATION VLLM_CUDAGRAPH_MODE PREFILL_CUDAGRAPH_MODE DECODE_CUDAGRAPH_MODE CUDAGRAPH_CAPTURE_SIZES VLLM_ALL2ALL_BACKEND PREFILL_MORI_BACKEND DECODE_MORI_BACKEND MORI_SHMEM_HEAP_SIZE" +MODELS_YAML_PROTECT="" +for _k in $_RECIPE_ENV_KEYS; do + [ -n "${!_k+x}" ] && MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT} ${_k}" +done +export MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT# }" +echo "models.yaml protect-list (submit-time overrides): '${MODELS_YAML_PROTECT}'" + # ------------------------ # Print current time in UTC and PST formats # ------------------------ @@ -427,11 +443,13 @@ BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-}" # Benchmark script selector: BENCHMARK_SCRIPT tag -> file run by the launcher. # sweep (default) -> benchmark_xPyD.sh (general concurrency sweep) # long_context -> benchmark_long_context.sh (per-shape warmup, c=1-first) +# niah -> benchmark_niah.sh (needle-in-haystack accuracy) BENCHMARK_SCRIPT="${BENCHMARK_SCRIPT:-sweep}" case "$BENCHMARK_SCRIPT" in sweep) BENCHMARK_SCRIPT_FILE="benchmark_xPyD.sh" ;; long_context) BENCHMARK_SCRIPT_FILE="benchmark_long_context.sh" ;; - *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context)" >&2; exit 1 ;; + niah) BENCHMARK_SCRIPT_FILE="benchmark_niah.sh" ;; + *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, niah)" >&2; exit 1 ;; esac if [[ ! -f "$BENCHMARK_SCRIPT_FILE" ]]; then echo "Error: selected benchmark script '$BENCHMARK_SCRIPT_FILE' not found in $(pwd)." >&2 @@ -583,6 +601,7 @@ docker run --rm \ ${VLLM_ALL2ALL_BACKEND:+-e VLLM_ALL2ALL_BACKEND=$VLLM_ALL2ALL_BACKEND} \ ${PREFILL_MORI_BACKEND:+-e PREFILL_MORI_BACKEND=$PREFILL_MORI_BACKEND} \ ${DECODE_MORI_BACKEND:+-e DECODE_MORI_BACKEND=$DECODE_MORI_BACKEND} \ + -e MODELS_YAML_PROTECT="${MODELS_YAML_PROTECT:-}" \ ${KV_BLOCK_SIZE:+-e KV_BLOCK_SIZE=$KV_BLOCK_SIZE} \ ${KV_CACHE_MEMORY_BYTES:+-e KV_CACHE_MEMORY_BYTES=$KV_CACHE_MEMORY_BYTES} \ ${VLLM_ROCM_USE_AITER_MLA:+-e VLLM_ROCM_USE_AITER_MLA=$VLLM_ROCM_USE_AITER_MLA} \ diff --git a/scripts/vllm_dissag/vllm_disagg.sh b/scripts/vllm_dissag/vllm_disagg.sh index 06fbf84f..bf98647a 100755 --- a/scripts/vllm_dissag/vllm_disagg.sh +++ b/scripts/vllm_dissag/vllm_disagg.sh @@ -160,17 +160,33 @@ MODEL_CONFIG_DECODE="" if [[ -n "$MODEL_NAME" && -f "$MODELS_YAML" ]]; then export MODELS_YAML MODEL_NAME PARALLEL_MODE # 1) Export per-model env: block FIRST (so connector ${VAR:-default} yields to it). - # Only set a var that is NOT already in the environment, so a submit-time - # `docker -e VAR=...` (already exported) WINS over the yaml value. Precedence: - # connector default < models.yaml env: < submit-time -e. + # Precedence: image-baked ENV < models.yaml env: < submit-time -e. + # models.yaml MUST override image-baked ENV: a DeepSeek-tuned disagg image + # bakes KV_BLOCK_SIZE=16 / VLLM_ROCM_USE_AITER_MLA=0 / VLLM_CUDAGRAPH_MODE= + # PIECEWISE etc. as container ENV, which would otherwise shadow a model's own + # recipe defined in models.yaml env:. But a genuine + # submit-time `-e VAR=...` must still win. The slurm can tell the two apart + # (it runs on the host) and passes MODELS_YAML_PROTECT = the space-separated + # list of keys the USER set at submit; the driver protects only those. When + # MODELS_YAML_PROTECT is unset (script run directly, no slurm), fall back to + # the old "skip if in env" behavior so nothing regresses. _yaml_env="$(python3 - <<'PY' import os, yaml, shlex m = yaml.safe_load(open(os.environ["MODELS_YAML"])) or {} cfg = m.get(os.environ["MODEL_NAME"]) or {} +protect_raw = os.environ.get("MODELS_YAML_PROTECT") +have_protect = protect_raw is not None +protect = set((protect_raw or "").split()) for k, v in (cfg.get("env") or {}).items(): - # skip if already present in the environment (submit-time -e override wins) - if k in os.environ: - continue + if have_protect: + # 3-tier: yaml overrides baked ENV; only a user submit-time -e (in the + # protect-list) wins over yaml. + if k in protect: + continue + else: + # No protect-list (direct run): legacy behavior — any existing env wins. + if k in os.environ: + continue print(f'export {k}={shlex.quote(str(v))}') PY )"