Enable DeepSeek-V3 Wide-EP disaggregated inference + fix DP=32 correctness

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 <registry>/<tag> .
 #
-# 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}

scripts/vllm_dissag/benchmark_niah.py

@@ -1,21 +1,37 @@
 #!/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:
 #   NIAH_URL     endpoint (default http://127.0.0.1:30000/v1/chat/completions)
 #   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")
 MODEL = os.environ.get("NIAH_MODEL", "")
 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__":

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}"

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 \

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"

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} \

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
 )"