From f2f2783d806c67a4e82d70117569c3ac19940f80 Mon Sep 17 00:00:00 2001 From: Basem Barakat Date: Tue, 18 Aug 2026 15:14:52 -0500 Subject: [PATCH] Fix 4P/4D GLM NIAH: Wei combine() topk_ids and pair warmup with score. EP32 was gathering with dispatched ids; NIAH now warms each length then scores it so a later EngineCore 500 cannot zero earlier sizes. Serve argv unchanged. --- scripts/vllm_dissag/README.MD | 5 +- .../apply_mori_combine_original_topk_fix.py | 120 ++++++++++++++++++ scripts/vllm_dissag/benchmark_niah.py | 76 ++++++++--- scripts/vllm_dissag/benchmark_niah.sh | 7 +- scripts/vllm_dissag/connectors/moriio.sh | 25 ++++ scripts/vllm_dissag/run_xPyD_models.slurm | 12 +- 6 files changed, 219 insertions(+), 26 deletions(-) create mode 100644 scripts/vllm_dissag/apply_mori_combine_original_topk_fix.py diff --git a/scripts/vllm_dissag/README.MD b/scripts/vllm_dissag/README.MD index 17660abe..92a796c9 100644 --- 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 new file mode 100644 index 00000000..b4ea1915 --- /dev/null +++ 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 +""" +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]} ", 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 index 43d81850..0bc9ab3f 100755 --- 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 index d366b5e4..cdccc735 100755 --- 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 index e1c7a51a..c71d48fc 100644 --- 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 index 365883b8..73d2a146 100755 --- 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} \