raviguptaamd · basemam · Aug 18, 2026
diff --git a/scripts/vllm_dissag/README.MD b/scripts/vllm_dissag/README.MD
@@ -234,8 +234,9 @@ lengths and the model is scored on how many it retrieves. Select it in place of
 sweep via the launcher's benchmark hook:
 
 ```bash
-export BENCHMARK_SCRIPT_FILE=benchmark_niah.sh
-export NIAH_WORDS="2000,8000,20000,35000"     # context sizes (words); optional
+export BENCHMARK_SCRIPT=niah
+export NIAH_WORDS="2000,8000,16000,20000,28000,35000"
+export NIAH_SEEDS=0
 ```
 
 Or run `benchmark_niah.py` standalone against any live OpenAI-compatible endpoint

diff --git a/scripts/vllm_dissag/apply_mori_combine_original_topk_fix.py b/scripts/vllm_dissag/apply_mori_combine_original_topk_fix.py
@@ -0,0 +1,120 @@
+#!/usr/bin/env python3
+"""Fix MoRI combine() using dispatched topk_ids (vLLM EP32 garbage).
+
+Source: amd-weisun/vllm @ fix/mori-combine-original-topk-ids
+  commit e25bf1826f9a8dc66d7312f046375e8c9491d638 (2026-08-17)
+
+MoriPrepareAndFinalize.finalize() was calling mori_op.combine() with the
+topk_ids modular_kernel already replaced with DISPATCHED (post-prepare) ids.
+Those ids describe OTHER ranks' tokens that landed on this rank's local
+experts, not the original routing of THIS rank's tokens. combine() then
+gathers from the wrong nodes. Worse at EP32 (4 nodes) than EP16 (2 nodes).
+
+Qwen3-30B-A3B (Li / Wei Sun, OCI): EP32 probe 3/10 and GSM8K 18.95% before;
+10/10 and 89.31% after. EP8/EP16 were already fine. Hunyuan needs more than
+this patch; GLM 4P/2D 216534 rank-0 Who isWho matches the EP32 signature.
+
+Matches acb0f1dc mori.py. Idempotent. Missing file -> skip. Found-old that
+fails to apply is a hard error.
+
+Usage: apply_mori_combine_original_topk_fix.py <vllm_install_dir>
+"""
+import os
+import sys
+
+REL = "model_executor/layers/fused_moe/prepare_finalize/mori.py"
+
+OLD_INIT = """        self.max_tokens_per_rank = max_tokens_per_rank
+        self.use_fp8_dispatch = use_fp8_dispatch
+"""
+NEW_INIT = """        self.max_tokens_per_rank = max_tokens_per_rank
+        self.use_fp8_dispatch = use_fp8_dispatch
+        # Original (pre-dispatch) topk_ids, stashed in prepare() for use in
+        # finalize() -- see the comment in prepare() for why.
+        self._original_topk_ids: torch.Tensor | None = None
+"""
+
+OLD_PREPARE = """        assert not apply_router_weight_on_input, (
+            "mori does not support apply_router_weight_on_input=True now."
+        )
+        scale = None
+"""
+NEW_PREPARE = """        assert not apply_router_weight_on_input, (
+            "mori does not support apply_router_weight_on_input=True now."
+        )
+        # combine() needs the ORIGINAL (pre-dispatch) topk_ids, not the
+        # dispatched ids that modular_kernel.py's forward() substitutes in
+        # for the expert GEMM stage. MoRI's combine kernel uses these ids
+        # (via tokenIndices) to decide which remote nodes' partial results
+        # to gather for each of THIS rank's own tokens -- the dispatched
+        # ids describe OTHER ranks' tokens that were routed here, not the
+        # original routing of this rank's own tokens, so passing them to
+        # combine() causes it to gather from the wrong nodes. Stash the
+        # original ids here; finalize() below uses this instead of the
+        # topk_ids it's given (which is also the post-dispatch value, for
+        # the same reason).
+        self._original_topk_ids = topk_ids
+        scale = None
+"""
+
+OLD_COMBINE = """        result = self.mori_op.combine(
+            fused_expert_output,
+            None,
+            topk_ids,
+        )[0]
+"""
+NEW_COMBINE = """        result = self.mori_op.combine(
+            fused_expert_output,
+            None,
+            self._original_topk_ids,
+        )[0]
+"""
+
+
+def main() -> int:
+    if len(sys.argv) != 2:
+        print(f"usage: {sys.argv[0]} <vllm_install_dir>", file=sys.stderr)
+        return 2
+    path = os.path.join(sys.argv[1], REL)
+    if not os.path.isfile(path):
+        print(f"[mori-combine] {REL} not found under {sys.argv[1]} -- skipping.")
+        return 0
+
+    src = open(path).read()
+    if "self._original_topk_ids" in src and "self._original_topk_ids," in src:
+        print(f"[mori-combine] already patched in {path} -- no-op.")
+        return 0
+
+    if OLD_COMBINE not in src:
+        print(
+            f"[mori-combine] ERROR: combine(topk_ids) anchor missing in {path}.",
+            file=sys.stderr,
+        )
+        return 1
+    if OLD_PREPARE not in src:
+        print(
+            f"[mori-combine] ERROR: prepare() stash anchor missing in {path}.",
+            file=sys.stderr,
+        )
+        return 1
+    if OLD_INIT not in src:
+        print(
+            f"[mori-combine] ERROR: __init__ anchor missing in {path}.",
+            file=sys.stderr,
+        )
+        return 1
+
+    src = src.replace(OLD_INIT, NEW_INIT, 1)
+    src = src.replace(OLD_PREPARE, NEW_PREPARE, 1)
+    src = src.replace(OLD_COMBINE, NEW_COMBINE, 1)
+    if "self._original_topk_ids," not in src:
+        print(f"[mori-combine] ERROR: post-write combine() still uses topk_ids in {path}.",
+              file=sys.stderr)
+        return 1
+    open(path, "w").write(src)
+    print(f"[mori-combine] patched: combine() uses original topk_ids (Wei Sun e25bf182) in {path}.")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/scripts/vllm_dissag/benchmark_niah.py b/scripts/vllm_dissag/benchmark_niah.py
@@ -11,13 +11,16 @@
 #   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
+#   NIAH_WARMUP  1 (default) = for each length, one throwaway request THEN the scored
+#                request (warmup N → score N). Do not warm 16k–35k before scoring 2k:
+#                a later EngineCore 500 makes earlier lengths look like found=0/10.
+#                Set 0 to disable warmup. NIAH_HALT_ON_FAIL=1 stops the ladder after a
+#                warmup/score timeout so a dead engine does not keep printing 0/10.
+#   NIAH_PAIR_SLEEP_S          seconds after finishing a length before the next pair
+#                              (default 30). 0 disables.
+#   NIAH_WARMUP_SCORE_SLEEP_S  seconds between warmup ok and the scored request
+#                              (default 5). 0 disables.
+import os, sys, json, random, time, urllib.request
 
 URL = os.environ.get("NIAH_URL", "http://127.0.0.1:30000/v1/chat/completions")
 MODEL = os.environ.get("NIAH_MODEL", "")
@@ -29,8 +32,10 @@
 # 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.
+HALT_ON_FAIL = os.environ.get("NIAH_HALT_ON_FAIL", "1") == "1"
+PAIR_SLEEP = float(os.environ.get("NIAH_PAIR_SLEEP_S", "30"))
+WARMUP_SCORE_SLEEP = float(os.environ.get("NIAH_WARMUP_SCORE_SLEEP_S", "5"))
+# Warmup uses a generous timeout (cold compile of a long-context shape can take minutes).
 WARMUP_TIMEOUT = max(TIMEOUT, 1800.0)
 
 FILLER = (
@@ -83,12 +88,11 @@ def _request(n_words, seed, max_tokens, timeout):
 
 
 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)."""
+    """Throwaway request for this length only. Returns True on HTTP success."""
     _, 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)
+    return err is None
 
 
 def run(n_words, seed=0):
@@ -112,16 +116,46 @@ 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  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:
+    print("url=%s  model=%s  sizes=%s  seeds=%s  warmup=%s halt_on_fail=%s "
+          "pair_sleep=%.0fs warmup_score_sleep=%.0fs"
+          % (URL, MODEL, WORDS, SEEDS, WARMUP, HALT_ON_FAIL, PAIR_SLEEP, WARMUP_SCORE_SLEEP),
+          flush=True)
+    # Pair each length: warmup N then score N. Warm-all-then-score-all lets a later
+    # EngineCore 500 make earlier lengths look like found=0/10.
+    results = {}  # n_words -> list of scores across seeds (None = timeout/error)
+    halted = False
+    for i, n in enumerate(WORDS):
+        if halted:
+            results[n] = [None] * len(SEEDS)
+            print("words=%6d  SKIP (ladder halted)" % n, flush=True)
+            continue
+        if i > 0 and PAIR_SLEEP > 0:
+            print(
+                "[niah] sleep %.0fs between pairs (after words=%d, before words=%d)"
+                % (PAIR_SLEEP, WORDS[i - 1], n),
+                flush=True,
+            )
+            time.sleep(PAIR_SLEEP)
+        if WARMUP:
+            print("=== pair words=%d (warmup then score) ===" % n, flush=True)
+            if not warmup(n):
+                results[n] = [None] * len(SEEDS)
+                print("words=%6d  SKIP score (warmup failed)" % n, flush=True)
+                if HALT_ON_FAIL:
+                    print("NIAH_HALT_ON_FAIL=1 — stopping remaining lengths", flush=True)
+                    halted = True
+                continue
+            if WARMUP_SCORE_SLEEP > 0:
+                print(
+                    "[niah] sleep %.0fs between warmup and score words=%d"
+                    % (WARMUP_SCORE_SLEEP, n),
+                    flush=True,
+                )
+                time.sleep(WARMUP_SCORE_SLEEP)
         results[n] = [run(n, s) for s in SEEDS]
+        if HALT_ON_FAIL and all(v is None for v in results[n]):
+            print("NIAH_HALT_ON_FAIL=1 — stopping remaining lengths (score failed)", flush=True)
+            halted = True
     print("=== NIAH summary (mean/min/max across %d seed(s)) ===" % len(SEEDS), flush=True)
     for n in WORDS:
         scored = results[n]

diff --git a/scripts/vllm_dissag/benchmark_niah.sh b/scripts/vllm_dissag/benchmark_niah.sh
@@ -29,14 +29,17 @@ 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.
+# Pairing (warmup N → score N) lives in benchmark_niah.py. Chat POSTs unchanged.
 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_SEEDS="${NIAH_SEEDS:-0}" \
 NIAH_MAXTOK="${NIAH_MAXTOK:-2048}" \
 NIAH_TIMEOUT="${NIAH_TIMEOUT:-1800}" \
 NIAH_WARMUP="${NIAH_WARMUP:-1}" \
+NIAH_HALT_ON_FAIL="${NIAH_HALT_ON_FAIL:-0}" \
+NIAH_PAIR_SLEEP_S="${NIAH_PAIR_SLEEP_S:-30}" \
+NIAH_WARMUP_SCORE_SLEEP_S="${NIAH_WARMUP_SCORE_SLEEP_S:-5}" \
   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
@@ -141,9 +141,34 @@ connector_runtime_patch() {
     # large-transfer notify/mapping fixes #424/#436/#432 baked in); if a newer MoRI is
     # needed, update MORI_REF and rebuild the image — no runtime library swap here.
     [ "${MODEL_NAME:-}" = "GLM-5.1-FP8" ] || return 0
+    # Wei Sun original topk_ids (e25bf182). EP32 combine() was gathering with
+    # dispatched ids -> 4P/4D garbage. Not in the baked DSA image, so this runs
+    # even when GLM_SKIP_PATCHERS=1. Does not change vllm serve argv.
+    _mori_combine_original_topk_fix
     _glm_dsa_runtime_patch
 }
 
+# MoRI combine() original topk_ids — Wei Sun e25bf182. All GLM MoE EP.
+_mori_combine_original_topk_fix() {
+    local _patch_dir="${SCRIPT_DIR:-$(cd "$(dirname "${BASH_SOURCE[0]:-$0}")/.." && pwd)}"
+    local _py="${_patch_dir}/apply_mori_combine_original_topk_fix.py"
+    if [ ! -f "${_py}" ]; then
+        echo "Error: [mori-combine] ${_py} not found. Aborting." >&2
+        exit 1
+    fi
+    local _vllm_dir
+    _vllm_dir="$(python3 -c 'import vllm, os; print(os.path.dirname(vllm.__file__))' 2>/dev/null || true)"
+    if [ -z "${_vllm_dir}" ] || [ ! -d "${_vllm_dir}" ]; then
+        echo "Error: [mori-combine] cannot locate vLLM install dir. Aborting." >&2
+        exit 1
+    fi
+    echo "[mori-combine] applying ${_py} against ${_vllm_dir}"
+    python3 "${_py}" "${_vllm_dir}" 2>&1 || {
+        echo "Error: [mori-combine] patch failed — EP32 would emit garbage. Aborting." >&2
+        exit 1
+    }
+}
+
 # GLM-5.1 DSA patchers (see connector_runtime_patch). Ported from MAD-private #338.
 # Resolves the vLLM install dir, then applies the 4 required patchers in order,
 # aborting on a hard failure (a real failure means GLM emits garbage or stalls, so

diff --git a/scripts/vllm_dissag/run_xPyD_models.slurm b/scripts/vllm_dissag/run_xPyD_models.slurm
@@ -457,12 +457,14 @@ BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-}"
 #   long_context           -> benchmark_long_context.sh (per-shape warmup, c=1-first)
 #   keepalive              -> keepalive_bench.sh        (hold server up KEEPALIVE_MINS
 #                                                        for external accuracy probes)
+#   niah                   -> benchmark_niah.sh         (pairing NIAH client; servers unchanged)
 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" ;;
     keepalive)    BENCHMARK_SCRIPT_FILE="keepalive_bench.sh" ;;
-    *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, keepalive)" >&2; exit 1 ;;
+    niah)         BENCHMARK_SCRIPT_FILE="benchmark_niah.sh" ;;
+    *) echo "Error: invalid BENCHMARK_SCRIPT='$BENCHMARK_SCRIPT' (valid: sweep, long_context, keepalive, niah)" >&2; exit 1 ;;
 esac
 if [[ ! -f "$BENCHMARK_SCRIPT_FILE" ]]; then
     echo "Error: selected benchmark script '$BENCHMARK_SCRIPT_FILE' not found in $(pwd)." >&2
@@ -631,6 +633,14 @@ docker run --rm \
     ${KV_CACHE_DTYPE:+-e KV_CACHE_DTYPE=$KV_CACHE_DTYPE} \
     ${MORIIO_TOY_PROXY:+-e MORIIO_TOY_PROXY=$MORIIO_TOY_PROXY} \
     ${BENCHMARK_SCRIPT_FILE:+-e BENCHMARK_SCRIPT_FILE=$BENCHMARK_SCRIPT_FILE} \
+    ${NIAH_WORDS:+-e NIAH_WORDS=$NIAH_WORDS} \
+    ${NIAH_SEEDS:+-e NIAH_SEEDS=$NIAH_SEEDS} \
+    ${NIAH_MAXTOK:+-e NIAH_MAXTOK=$NIAH_MAXTOK} \
+    ${NIAH_TIMEOUT:+-e NIAH_TIMEOUT=$NIAH_TIMEOUT} \
+    ${NIAH_WARMUP:+-e NIAH_WARMUP=$NIAH_WARMUP} \
+    ${NIAH_HALT_ON_FAIL:+-e NIAH_HALT_ON_FAIL=$NIAH_HALT_ON_FAIL} \
+    ${NIAH_PAIR_SLEEP_S:+-e NIAH_PAIR_SLEEP_S=$NIAH_PAIR_SLEEP_S} \
+    ${NIAH_WARMUP_SCORE_SLEEP_S:+-e NIAH_WARMUP_SCORE_SLEEP_S=$NIAH_WARMUP_SCORE_SLEEP_S} \
     ${KEEPALIVE_MINS:+-e KEEPALIVE_MINS=$KEEPALIVE_MINS} \
     ${PREFILL_CUDAGRAPH_MODE:+-e PREFILL_CUDAGRAPH_MODE=$PREFILL_CUDAGRAPH_MODE} \
     ${DECODE_CUDAGRAPH_MODE:+-e DECODE_CUDAGRAPH_MODE=$DECODE_CUDAGRAPH_MODE} \