From 153afef35496492c4655d9fd3b3283996ff6131f Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Thu, 2 Jul 2026 22:21:15 -0700 Subject: [PATCH 01/13] feat: make moe_ep (EP) part of the default install; drop nccl submodule Previously the EP transport backends required an opt-in install: `BUILD_NVEP=1 pip install -e ".[nvep]"`. Now a plain `pip install .` enables them by default: - Move the [nvep] extra's runtime deps (cuda-python>=13.0, nccl4py>=0.3.1, nvidia-nccl-cu13>=2.30.7) into the base dependencies (requirements.txt). The [nvep] extra remains as an empty deprecated alias. - build_backend.py: NCCL-EP and NIXL-EP default ON with tri-state env flags. Unset -> build best-effort (missing build deps skip the backend with a warning instead of failing the install). BUILD_NIXL_EP=1 -> strict (missing deps abort). BUILD_NVEP=0 / BUILD_NIXL_EP=0 / BUILD_NCCL_EP=0 -> opt out. - New _ensure_nixl_wheel(): the build hook pre-installs the nixl-cu13 wheel (--no-deps) that the default NIXL-EP build links against, so no manual pre-install step is needed. - Remove the 3rdparty/nccl submodule: NCCL-EP has been provided by the released nccl4py wheel since nccl-ep-v0.1.0; nothing builds from the submodule anymore. - validate_arch_for_backend(): raise a clear error when torch is built for CUDA < 13 (the EP wheels ship CUDA-13 binaries only). - Update Dockerfiles, install scripts, docs, test markers, and rebuild hints for the new default. Verified in nvcr.io/nvidia/pytorch:26.05-py3 (CUDA 13.2, torch 2.12): - Stock image: `pip install .` succeeds; NIXL-EP is skipped gracefully (image UCX 1.20 lacks the 1.21 device API); backends = ['nccl_ep']. - With DOCA 3.2 + UCX v1.21.x from source: backends = ['nccl_ep', 'nixl_ep'], nixl_ep_cpp.so staged. Smoke pass. AI-assisted. Co-Authored-By: Claude Fable 5 --- .dockerignore | 2 +- .gitignore | 2 +- .gitmodules | 3 - 3rdparty/nccl | 1 - benchmarks/MoE_benchmarks.md | 5 +- benchmarks/bench_moe_ep.py | 2 +- build_backend.py | 126 +++++++++++++----- docker/Dockerfile.flashinfer-ep-pytorch | 12 +- docker/Dockerfile.flashinfer-nvep | 48 +++---- docker/install/build_flashinfer_ep_pytorch.sh | 4 +- flashinfer/moe_ep/__init__.py | 29 ++-- flashinfer/moe_ep/_validators.py | 14 +- flashinfer/moe_ep/nccl_ep/__init__.py | 6 +- flashinfer/moe_ep/nccl_ep/fleet.py | 6 +- flashinfer/moe_ep/nixl_ep/__init__.py | 12 +- flashinfer/moe_ep/nixl_ep/fleet.py | 6 +- pyproject.toml | 32 ++--- requirements.txt | 10 ++ scripts/build_in_container.sh | 9 +- tests/conftest.py | 7 +- tests/moe_ep/smoke_nccl_ep.py | 4 +- 21 files changed, 208 insertions(+), 132 deletions(-) delete mode 160000 3rdparty/nccl diff --git a/.dockerignore b/.dockerignore index a60c884cb2a..e6f0200f1e4 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,5 +1,5 @@ # Exclude bulky / host-specific paths from the docker build context. -# Submodules are still copied (they're needed for BUILD_NVEP=1). +# Submodules are still copied (needed for the default NIXL-EP build). # Python venvs / build trees .venv/ diff --git a/.gitignore b/.gitignore index 4ab7544987c..18b7fb0fcf7 100644 --- a/.gitignore +++ b/.gitignore @@ -202,7 +202,7 @@ cython_debug/ docs/tutorials/generated/ docs/sg_execution_times.rst -# moe_ep build artifacts (BUILD_NVEP=1 in-tree build of NIXL-EP + NCCL-EP) +# moe_ep build artifacts (default in-tree build of NIXL-EP) build_nvep/ flashinfer/moe_ep/*/_libs/ flashinfer/moe_ep/*/_vendored/ diff --git a/.gitmodules b/.gitmodules index 50d7a9bfd07..7a12c990d78 100644 --- a/.gitmodules +++ b/.gitmodules @@ -10,6 +10,3 @@ [submodule "3rdparty/nixl"] path = 3rdparty/nixl url = https://github.com/ai-dynamo/nixl.git -[submodule "3rdparty/nccl"] - path = 3rdparty/nccl - url = https://github.com/NVIDIA/nccl.git diff --git a/3rdparty/nccl b/3rdparty/nccl deleted file mode 160000 index 63cf786b015..00000000000 --- a/3rdparty/nccl +++ /dev/null @@ -1 +0,0 @@ -Subproject commit 63cf786b015b2b6bff6cf263461621acf584bd18 diff --git a/benchmarks/MoE_benchmarks.md b/benchmarks/MoE_benchmarks.md index 15889231418..296f3235859 100644 --- a/benchmarks/MoE_benchmarks.md +++ b/benchmarks/MoE_benchmarks.md @@ -28,7 +28,8 @@ DOCA / UCX-from-source / GDRCopy layers of the NIXL image are unnecessary for NC Build (`docker/install/build_flashinfer_ep_pytorch.sh` does the install): it pins the verified set over the base image's constraints — `nvidia-nccl-cu13==2.30.7` (via `PIP_CONSTRAINT=` to beat torch's 2.30.4 pin), `nccl4py[cu13]==0.3.1`, `cuda-core==1.0.1`, -`cuda-bindings==13.2.0` — then `BUILD_NCCL_EP=1 pip install -e ".[nvep]"`. +`cuda-bindings==13.2.0` — then `BUILD_NCCL_EP=1 BUILD_NIXL_EP=0 pip install -e .` +(the moe_ep deps are base dependencies now; no extra needed). ```bash # local docker @@ -53,7 +54,7 @@ Smoke: `python -c "import nccl.ep; from flashinfer.moe_ep import available_backe ## 2. How to run ### 2a. Comm matrix vs ep_bench (`bench_ep_matrix.py`) -Standalone — needs only FlashInfer (`.[nvep]`), torch, and a multi-rank launcher; it does +Standalone — needs only FlashInfer (EP is in the default install), torch, and a multi-rank launcher; it does **not** call `ep_bench` (that's a separate C++ reference). It emits ep_bench-compatible text so `scripts/parse_results.py` parses both. The 28-case driver issues one `srun` per config: diff --git a/benchmarks/bench_moe_ep.py b/benchmarks/bench_moe_ep.py index 2a81980b0d4..486edfdd8a3 100644 --- a/benchmarks/bench_moe_ep.py +++ b/benchmarks/bench_moe_ep.py @@ -1,7 +1,7 @@ """MoE Expert-Parallel benchmark: dispatch → compute → combine. Canonical cases mirror the NCCL-EP ``ep_bench`` reference -(``3rdparty/nccl/contrib/nccl_ep/README.md``): **128 tokens/rank, hidden 7168, +(``contrib/nccl_ep/README.md`` in the NCCL repo): **128 tokens/rank, hidden 7168, top-k 8, 256 experts, BF16**, swept over **8/16/32/64 GPUs** and over the two EP **algorithms — Low-Latency (LL) and High-Throughput (HT)** (one table each). Select that geometry with ``--reference`` and the algorithm with diff --git a/build_backend.py b/build_backend.py index 86c4c11aabe..350e4c35cd1 100644 --- a/build_backend.py +++ b/build_backend.py @@ -29,21 +29,34 @@ _data_dir = _root / "flashinfer" / "data" -# moe_ep build infra. Three opt-in switches, all `0` by default: -# BUILD_NCCL_EP=1 → enable NCCL-EP (provided by the `nccl4py>=0.3.1` wheel; -# NO in-tree build — see the `[nvep]` extra in pyproject) -# BUILD_NIXL_EP=1 → build NIXL-EP from 3rdparty/nixl (meson) -# BUILD_NVEP=1 → legacy alias: turns BOTH on (back-compat with earlier docs) +# moe_ep build infra. Both EP backends are ON BY DEFAULT since the moe_ep +# runtime deps moved into the base dependencies (`pip install .` is enough): +# NCCL-EP — provided by the `nccl4py>=0.3.1` wheel (a base dep now); NO +# in-tree build. +# NIXL-EP — built in-tree from 3rdparty/nixl (meson). Missing build deps +# (meson/ninja/nvcc/UCX/...) skip the backend with a warning +# instead of failing the install (best-effort). # -# Only NIXL-EP is built in-tree (it needs DOCA gpunetio + UCX 1.21.x). NCCL-EP is -# a pure pip dependency (`nccl4py`, which ships the `nccl.ep` API + bundled -# libnccl_ep.so). Hosts that only have one backend's deps should opt in with the -# matching flag instead of BUILD_NVEP. +# Env switches (tri-state; unset means "default on, best-effort"): +# BUILD_NIXL_EP=0 → skip the NIXL-EP submodule build +# BUILD_NIXL_EP=1 → strict: a missing build dep FAILS the install +# BUILD_NCCL_EP=0/1 → same idea for NCCL-EP (no build step; only affects +# the informational logging) +# BUILD_NVEP=0 → legacy alias: turns BOTH off +# BUILD_NVEP=1 → legacy alias: both on, best-effort (back-compat) def _flag(name: str) -> bool: v = os.environ.get(name, "") return v == "1" or v.lower() in ("true", "yes", "on") +def _tri_flag(name: str) -> bool | None: + """Tri-state env flag: True / False when set, None when unset/empty.""" + v = os.environ.get(name) + if v is None or v.strip() == "": + return None + return v == "1" or v.lower() in ("true", "yes", "on") + + @contextmanager def _time_phase(label: str): """Emit a wall-clock duration line for a build phase. @@ -61,18 +74,29 @@ def _time_phase(label: str): print(f"[BUILD_NVEP] {label}: done in {dt:.1f}s", flush=True) -_BUILD_NVEP = _flag("BUILD_NVEP") -_BUILD_NCCL_EP = _flag("BUILD_NCCL_EP") or _BUILD_NVEP -_BUILD_NIXL_EP = _flag("BUILD_NIXL_EP") or _BUILD_NVEP +# Resolution order per backend: explicit BUILD_{NIXL,NCCL}_EP, then the +# legacy BUILD_NVEP alias, then the default (ON). +_BUILD_NVEP = _tri_flag("BUILD_NVEP") + + +def _backend_enabled(name: str) -> bool: + explicit = _tri_flag(name) + if explicit is not None: + return explicit + if _BUILD_NVEP is not None: + return _BUILD_NVEP + return True + -# Was the user opting in via the legacy "give me everything possible" alias -# (BUILD_NVEP=1) AND NOT explicitly via the per-backend switches? If yes, -# treat a missing build-time dep as "skip that backend with a warning" -# instead of "abort the entire install". When the user explicitly asks for -# BUILD_NCCL_EP=1 / BUILD_NIXL_EP=1, a missing dep is a hard error — they -# asked for that backend specifically. -_BUILD_NVEP_BEST_EFFORT = _BUILD_NVEP and not ( - _flag("BUILD_NCCL_EP") or _flag("BUILD_NIXL_EP") +_BUILD_NCCL_EP = _backend_enabled("BUILD_NCCL_EP") +_BUILD_NIXL_EP = _backend_enabled("BUILD_NIXL_EP") + +# Missing build-time deps skip the backend with a warning instead of aborting +# the install — EXCEPT when the user explicitly asked for a backend with +# BUILD_NCCL_EP=1 / BUILD_NIXL_EP=1; then a missing dep is a hard error. The +# default-on install and the legacy BUILD_NVEP=1 alias are both best-effort. +_BUILD_NVEP_BEST_EFFORT = not ( + _tri_flag("BUILD_NCCL_EP") is True or _tri_flag("BUILD_NIXL_EP") is True ) _nvep_build_root = _root / "build_nvep" @@ -201,10 +225,9 @@ def _build_nixl_ep() -> None: wheel_lib_dir = _find_nixl_wheel_lib_dir() if wheel_lib_dir is None: raise RuntimeError( - "BUILD_NIXL_EP requires nixl-cu13 to be pre-installed.\n" + "The NIXL-EP build requires the nixl-cu13 wheel (the build " + "hook normally pre-installs it; see _ensure_nixl_wheel).\n" "Run: uv pip install --no-deps 'nixl-cu13>=1.0.1'\n" - "(the FlashInfer Dockerfile does this automatically; bare-host\n" - "installs need to do it before `pip install -e .[nvep]`).\n" "Or set BUILD_NIXL_EP_HERMETIC=1 to build the full NIXL tree." ) setup_args += [ @@ -307,6 +330,38 @@ def _fix_rpaths() -> None: print(f"[BUILD_NVEP] WARNING: patchelf failed on {so.name}: {err}") +def _ensure_nixl_wheel() -> None: + """Pre-install the nixl-cu* wheel the NIXL-EP build links against. + + The default (non-hermetic) NIXL-EP build links nixl_ep_cpp.so against the + libnixl.so shipped by the `nixl-cu13` pip wheel. Since the EP build now + runs by default on `pip install .`, install that wheel up front instead of + requiring users to pre-install it. `--no-deps` for the same reason as + _install_nvep_runtime_wheels: the wheel's transitive constraints downgrade + torch. Best-effort: on failure the _nixl_buildable probe reports the + missing wheel and the backend is gated as usual (skip or hard error). + """ + if _find_nixl_wheel_lib_dir() is not None: + return + cuda_major = _detect_cuda_major() + wheel = f"nixl-cu{cuda_major}>=1.0.1" + print(f"[BUILD_NVEP] pre-installing NIXL wheel --no-deps: {wheel}") + + uv_bin = shutil.which("uv") + if uv_bin: + cmd = [uv_bin, "pip", "install", "--python", sys.executable, "--no-deps", wheel] + else: + cmd = [sys.executable, "-m", "pip", "install", "--no-deps", wheel] + print(f"[BUILD_NVEP] $ {' '.join(cmd)}") + try: + subprocess.run(cmd, check=True) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + print( + f"[BUILD_NVEP] WARNING: could not pre-install {wheel} ({e}); " + "the NIXL-EP pre-flight probe will decide whether to skip or fail." + ) + + def _nixl_buildable() -> tuple[bool, str]: """Probe for hard NIXL-EP build-time deps. Returns (ok, reason_if_not). @@ -524,9 +579,9 @@ def _gate_backend(name: str, requested: bool, probe) -> bool: # User opted in explicitly (BUILD_NCCL_EP=1 or BUILD_NIXL_EP=1) — fail hard. raise RuntimeError( f"{name} build requested but a hard dep is missing: {reason}. " - "Either install the missing dependency, or use BUILD_NVEP=1 " - "(best-effort mode) to skip this backend with a warning instead " - "of failing the install." + "Either install the missing dependency, unset the BUILD_*_EP flag " + "(the default build is best-effort and skips this backend with a " + "warning), or set it to 0 to skip the backend entirely." ) @@ -542,17 +597,22 @@ def _build_nvep_if_enabled() -> None: mode = "best-effort" if _BUILD_NVEP_BEST_EFFORT else "strict" print(f"[BUILD_NVEP] requested: {', '.join(requested)} (mode: {mode})") - # NCCL-EP is no longer built from the submodule — it is provided by the - # released `nccl4py` wheel (>=0.3.1, the `nccl.ep` API + bundled - # libnccl_ep.so), declared in the `[nvep]` extra. So BUILD_NCCL_EP requires - # no in-tree build step; we only note it here. + # NCCL-EP is not built from source — it is provided by the released + # `nccl4py` wheel (>=0.3.1, the `nccl.ep` API + bundled libnccl_ep.so), + # which is a base dependency now. So BUILD_NCCL_EP requires no in-tree + # build step; we only note it here. if _BUILD_NCCL_EP: print( - "[BUILD_NVEP] NCCL-EP is provided by the nccl4py wheel (>=0.3.1); " - "no submodule build. Ensure it is installed (e.g. `pip install " - "\".[nvep]\"` or `pip install 'nccl4py>=0.3.1'`)." + "[BUILD_NVEP] NCCL-EP is provided by the nccl4py wheel (>=0.3.1), " + "a base dependency of flashinfer-python; no in-tree build." ) + # The default (non-hermetic) NIXL-EP build links against the nixl-cu13 + # wheel's libnixl.so — install it up front so plain `pip install .` works + # without a manual pre-install step. + if _BUILD_NIXL_EP and not _flag("BUILD_NIXL_EP_HERMETIC"): + _ensure_nixl_wheel() + # Pre-flight gating — probe the NIXL-EP build-time deps (NCCL-EP needs none). will_build_nixl = _gate_backend("NIXL-EP", _BUILD_NIXL_EP, _nixl_buildable) diff --git a/docker/Dockerfile.flashinfer-ep-pytorch b/docker/Dockerfile.flashinfer-ep-pytorch index c662d0e61a5..7d7b80bea0b 100644 --- a/docker/Dockerfile.flashinfer-ep-pytorch +++ b/docker/Dockerfile.flashinfer-ep-pytorch @@ -83,20 +83,18 @@ print('nccl.ep OK; libnccl', v.value); assert v.value>=23007" # Build & install FlashInfer with the NCCL-EP backend only. # --no-build-isolation uses the base image's python/torch for the build hook -# (so meson/torch are found); BUILD_NCCL_EP=1 + BUILD_NIXL_EP=0 selects the -# NCCL-EP path, which links contrib's Makefile against the nvidia-nccl-cu13 -# wheel (no in-tree NCCL build). The [nvep] extra pulls nccl4py (already pinned -# above; the >= constraint no-ops). +# (so meson/torch are found); BUILD_NIXL_EP=0 opts out of the (default-on) +# NIXL-EP submodule build. NCCL-EP needs no build step — the nccl4py wheel is +# a base dependency (already pinned above; the >= constraint no-ops). ARG BUILD_NCCL_EP=1 ARG BUILD_NIXL_EP=0 ARG FLASHINFER_SRC=/workspace/flashinfer COPY . ${FLASHINFER_SRC} WORKDIR ${FLASHINFER_SRC} # Submodule trees are copied as-is (see .dockerignore); no `git submodule update`. -RUN BUILD_NVEP=0 \ - BUILD_NCCL_EP=${BUILD_NCCL_EP} \ +RUN BUILD_NCCL_EP=${BUILD_NCCL_EP} \ BUILD_NIXL_EP=${BUILD_NIXL_EP} \ - pip install --no-cache-dir --no-build-isolation -e ".[nvep]" + pip install --no-cache-dir --no-build-isolation -e . # Smoke probe — fail the build if NCCL-EP didn't actually come up. RUN python -c "\ diff --git a/docker/Dockerfile.flashinfer-nvep b/docker/Dockerfile.flashinfer-nvep index d059564aa17..53cf9f155ac 100644 --- a/docker/Dockerfile.flashinfer-nvep +++ b/docker/Dockerfile.flashinfer-nvep @@ -1,9 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # # Reference Dockerfile for building FlashInfer with the moe_ep transport -# backends enabled: NIXL-EP is compiled from the in-tree git submodule; -# NCCL-EP is the released `nccl4py>=0.3.1` wheel (the `[nvep]` extra), which -# ships the `nccl.ep` API + bundled libnccl_ep.so (no in-tree NCCL build). +# backends enabled (the default): NIXL-EP is compiled from the in-tree git +# submodule; NCCL-EP is the released `nccl4py>=0.3.1` wheel (a base +# dependency), which ships the `nccl.ep` API + bundled libnccl_ep.so (no +# in-tree NCCL build). # # Usage: # cd /path/to/flashinfer @@ -12,7 +13,7 @@ # # Build args: # BUILD_NVEP — 0 to skip the moe_ep submodule builds (still installs -# FlashInfer). Defaults to 1. +# FlashInfer). Defaults to 1 (also the build hook's default). # CUDA_IMAGE — base image. Defaults to a CUDA 13 devel image. ARG CUDA_IMAGE=nvcr.io/nvidia/cuda:13.0.0-cudnn-devel-ubuntu24.04 @@ -145,22 +146,21 @@ RUN uv pip install --python ${VENV}/bin/python --no-deps \ # sys.executable inside build_backend.py points at uv's isolated env, meson # can't find torch, and NIXL EP is silently skipped. # -# Backend selection — three knobs, all honored by build_backend.py: -# BUILD_NVEP=1 → legacy alias, turns both on (default for this image) -# BUILD_NCCL_EP=1 → build only the NCCL-EP backend -# BUILD_NIXL_EP=1 → build only the NIXL-EP backend -# Override at build time, e.g.: `docker build --build-arg BUILD_NVEP=0 -# --build-arg BUILD_NCCL_EP=1 ...` for an NCCL-only image. +# Backend selection — three knobs, all honored by build_backend.py (all +# default ON in the hook; unset means "on, best-effort"): +# BUILD_NVEP=0 → legacy alias, turns both off +# BUILD_NCCL_EP=0/1 → disable / strictly require the NCCL-EP backend +# BUILD_NIXL_EP=0/1 → disable / strictly require the NIXL-EP backend +# Override at build time, e.g.: `docker build --build-arg BUILD_NIXL_EP=0 ...` +# for an NCCL-only image. # -# Both backends default to a wheel-linked build: NCCL-EP points contrib's -# Makefile BUILDDIR at the nvidia-nccl-cu13 wheel (skipping `make src.build`); -# NIXL-EP applies a meson overlay (3rdparty_patches/nixl/0002-ep-only-build.patch) -# that skips `subdir('src')` and links nixl_ep_cpp.so against the nixl-cu13 -# wheel's libnixl.so. The runtime loaders in flashinfer/moe_ep/{nccl,nixl}_ep -# ctypes-preload the base libs from the same wheels. Set -# BUILD_NCCL_EP_HERMETIC=1 / BUILD_NIXL_EP_HERMETIC=1 to opt back into the -# full from-source build for either backend (for hosts without PyPI access -# or when investigating ABI mismatches). +# NIXL-EP defaults to a wheel-linked build: a meson overlay +# (3rdparty_patches/nixl/0002-ep-only-build.patch) skips `subdir('src')` and +# links nixl_ep_cpp.so against the nixl-cu13 wheel's libnixl.so. The runtime +# loader in flashinfer/moe_ep/nixl_ep ctypes-preloads the base libs from the +# same wheel. Set BUILD_NIXL_EP_HERMETIC=1 to opt back into the full +# from-source build (for hosts without PyPI access or when investigating ABI +# mismatches). NCCL-EP has no build step — it ships in the nccl4py wheel. ARG BUILD_NVEP=1 ARG BUILD_NCCL_EP= ARG BUILD_NIXL_EP= @@ -172,17 +172,17 @@ WORKDIR ${FLASHINFER_SRC} # the submodule .git pointers reference the host superproject's .git/modules/ # tree which isn't copied, and `git apply` in _apply_patches works on a plain # directory without a git repo. -# The `[nvep]` extra pulls the released `nccl4py>=0.3.1` wheel (NCCL-EP: the -# `nccl.ep` API + bundled libnccl_ep.so) — no in-tree NCCL build. BUILD_NIXL_EP -# still builds NIXL-EP from the submodule via the build hook. +# The base deps pull the released `nccl4py>=0.3.1` wheel (NCCL-EP: the +# `nccl.ep` API + bundled libnccl_ep.so) — no in-tree NCCL build. The build +# hook builds NIXL-EP from the submodule by default. RUN BUILD_NVEP=${BUILD_NVEP} \ BUILD_NCCL_EP=${BUILD_NCCL_EP} \ BUILD_NIXL_EP=${BUILD_NIXL_EP} \ uv pip install --python ${VENV}/bin/python \ - --no-build-isolation -e ".[nvep]" + --no-build-isolation -e . # Smoke probe. Assertions fail the build if the EP backends weren't -# actually produced — without these, a silent skip in BUILD_NVEP=1's +# actually produced — without these, a silent skip in the default # best-effort mode (e.g. a probe miss on UCX, DOCA, or a wheel) would # ship a working FlashInfer image with no EP support and the failure # wouldn't surface until first use. diff --git a/docker/install/build_flashinfer_ep_pytorch.sh b/docker/install/build_flashinfer_ep_pytorch.sh index 76923cd63ff..6a7103c3935 100644 --- a/docker/install/build_flashinfer_ep_pytorch.sh +++ b/docker/install/build_flashinfer_ep_pytorch.sh @@ -41,8 +41,8 @@ python -c "import nccl.ep; from nccl.core import Communicator; print('nccl.ep + echo "== build & install FlashInfer (NCCL-EP only) ==" cd "${FI_SRC}" -BUILD_NVEP=0 BUILD_NCCL_EP=1 BUILD_NIXL_EP=0 \ - pip install --no-cache-dir --no-build-isolation -e ".[nvep]" +BUILD_NCCL_EP=1 BUILD_NIXL_EP=0 \ + pip install --no-cache-dir --no-build-isolation -e . echo "== smoke probe ==" python -c "\ diff --git a/flashinfer/moe_ep/__init__.py b/flashinfer/moe_ep/__init__.py index 0e25e00e019..cab04feb521 100644 --- a/flashinfer/moe_ep/__init__.py +++ b/flashinfer/moe_ep/__init__.py @@ -3,16 +3,18 @@ This package is a thin Python wrapper over two transport backends: - ``flashinfer.moe_ep.nccl_ep`` — primary backend, wraps NVIDIA's nccl4py - ``nccl.ep`` API (nccl-ep-v0.1.0; built from ``3rdparty/nccl/bindings/nccl4py``). + ``nccl.ep`` API (nccl-ep-v0.1.0; provided by the released ``nccl4py`` wheel, + a base dependency of flashinfer-python). - ``flashinfer.moe_ep.nixl_ep`` — alternate backend, wraps ai-dynamo's ``nixl_ep`` (built in-tree from ``3rdparty/nixl/examples/device/ep``). NCCL-EP availability is the importability of ``nccl.ep`` (no in-tree -``libnccl_ep.so`` as of v0.1.0); NIXL-EP still ships a staged ``nixl_ep_cpp*.so``. -Both are produced by the FlashInfer build only when ``BUILD_NVEP=1`` (or the -per-backend ``BUILD_NCCL_EP`` / ``BUILD_NIXL_EP``) is set at install time: - - BUILD_NVEP=1 pip install -e ".[nvep]" +``libnccl_ep.so`` as of v0.1.0); NIXL-EP ships a staged ``nixl_ep_cpp*.so``. +Both are enabled by default: plain ``pip install .`` pulls the nccl4py wheel +and builds NIXL-EP best-effort (skipped with a warning when its build deps — +meson, UCX, nvcc, ... — are missing). Opt out with ``BUILD_NVEP=0`` (or the +per-backend ``BUILD_NCCL_EP=0`` / ``BUILD_NIXL_EP=0``); force a hard error on +missing NIXL-EP build deps with ``BUILD_NIXL_EP=1``. Without a built backend the package imports succeed but calling :func:`create_fleet` raises :class:`MoEEpNotBuiltError` with rebuild @@ -101,10 +103,11 @@ _pkg_dir = Path(__file__).parent _REBUILD_HINT = ( - "flashinfer.moe_ep is not built. Rebuild with:\n" - ' BUILD_NVEP=1 pip install -e ".[nvep]"\n' - "from the FlashInfer source tree. See " - "flashinfer/moe_ep/README.md for required system dependencies." + "flashinfer.moe_ep is not built. It builds by default; rebuild with:\n" + " pip install -e .\n" + "from the FlashInfer source tree (use BUILD_NIXL_EP=1 to turn missing\n" + "build deps into hard errors instead of skip-with-warning). See the\n" + "moe_ep section of build_backend.py for required system dependencies." ) @@ -176,9 +179,9 @@ def _require_built(backend: str) -> None: # Quiet diagnostic at import time when a build flag was set but the libs # are absent — most likely cause is a partial build (probe failure -# swallowed in BUILD_NVEP=1 best-effort mode). Helpful for first-time -# users. Covers all three opt-in flags: the legacy BUILD_NVEP alias plus -# the per-backend BUILD_NCCL_EP / BUILD_NIXL_EP. +# swallowed in the default best-effort mode). Helpful for first-time +# users. Covers all three flags: the legacy BUILD_NVEP alias plus the +# per-backend BUILD_NCCL_EP / BUILD_NIXL_EP. _set_build_flags = [ name for name in ("BUILD_NVEP", "BUILD_NCCL_EP", "BUILD_NIXL_EP") diff --git a/flashinfer/moe_ep/_validators.py b/flashinfer/moe_ep/_validators.py index e54e3a825e1..76aac6971fc 100644 --- a/flashinfer/moe_ep/_validators.py +++ b/flashinfer/moe_ep/_validators.py @@ -36,9 +36,21 @@ class MoEEpArchError(MoEEpConfigError): def validate_arch_for_backend(backend: str) -> None: - """Check ``torch.cuda.get_device_capability(0)`` is supported by `backend`.""" + """Check the GPU arch and CUDA version are supported by `backend`.""" import torch + # The EP runtime wheels (nccl4py, nvidia-nccl-cu13, nixl-cu13) are + # CUDA-13-only, so a torch built for CUDA 12 can't drive either backend — + # fail here with a clear message instead of a cryptic dlopen error later. + cuda_ver = torch.version.cuda + if cuda_ver is not None and int(cuda_ver.split(".")[0]) < 13: + raise MoEEpConfigError( + f"{backend} requires CUDA 13: the EP runtime wheels (nccl4py, " + f"nvidia-nccl-cu13, nixl-cu13) ship CUDA-13 binaries only, but " + f"the installed torch was built for CUDA {cuda_ver}. Install a " + "CUDA-13 torch build to use flashinfer.moe_ep." + ) + if not torch.cuda.is_available(): return # Mock/test path — let backend probes catch missing libs instead. cc = torch.cuda.get_device_capability(0) diff --git a/flashinfer/moe_ep/nccl_ep/__init__.py b/flashinfer/moe_ep/nccl_ep/__init__.py index e8f0ae2e956..7f8ca1a63ec 100644 --- a/flashinfer/moe_ep/nccl_ep/__init__.py +++ b/flashinfer/moe_ep/nccl_ep/__init__.py @@ -3,9 +3,9 @@ As of ``nccl-ep-v0.1.0`` the backend is driven entirely by the **nccl4py** Python package's ``nccl.ep`` API — there is no longer an in-tree ``libnccl_ep.so`` to dlopen or a flat ``nccl_ep`` ctypes module to import. -The ``nccl`` package (built with ``BUILD_NCCL4PY`` and shipped as a wheel, or -installed editable from ``3rdparty/nccl/bindings/nccl4py``) self-loads its -native library; we just import ``nccl.ep`` lazily in :mod:`.fleet` / :mod:`.handle`. +The ``nccl`` package (the released ``nccl4py`` wheel, a base dependency of +flashinfer-python) self-loads its native library; we just import ``nccl.ep`` +lazily in :mod:`.fleet` / :mod:`.handle`. Availability is probed via :func:`flashinfer.moe_ep._probe_nccl_ep`, which checks that ``nccl.ep`` is importable. diff --git a/flashinfer/moe_ep/nccl_ep/fleet.py b/flashinfer/moe_ep/nccl_ep/fleet.py index 7ad7d9ead8b..67ba134d7aa 100644 --- a/flashinfer/moe_ep/nccl_ep/fleet.py +++ b/flashinfer/moe_ep/nccl_ep/fleet.py @@ -46,9 +46,9 @@ def _import_nccl_ep(): return nccl_ep except ImportError as e: # pragma: no cover - exercised only without build raise MoEEpNotBuiltError( - "nccl.ep (nccl-ep-v0.1.0) python package unavailable. Rebuild with " - "BUILD_NCCL_EP=1 (which builds the nccl4py bindings), or install the " - "nccl4py wheel that ships nccl.ep." + "nccl.ep (nccl-ep-v0.1.0) python package unavailable. It ships in " + "the nccl4py wheel, a base dependency of flashinfer-python — " + "install with `pip install 'nccl4py>=0.3.1'`." ) from e diff --git a/flashinfer/moe_ep/nixl_ep/__init__.py b/flashinfer/moe_ep/nixl_ep/__init__.py index eba06db9754..0c26606a770 100644 --- a/flashinfer/moe_ep/nixl_ep/__init__.py +++ b/flashinfer/moe_ep/nixl_ep/__init__.py @@ -5,8 +5,8 @@ 1. The base NIXL runtime libraries (``libnixl.so``, ``libnixl_capi.so``, ``libnixl_common.so``, ``libserdes.so``, etc.) — *not* shipped inside this package. They're expected to come from the ``nixl-cu13`` pip wheel, - installed automatically when the user runs ``BUILD_NVEP=1 pip install ...`` - (see ``build_backend._install_nvep_runtime_wheels``). + installed automatically by the default ``pip install .`` build (see + ``build_backend._ensure_nixl_wheel`` / ``_install_nvep_runtime_wheels``). 2. The EP torch extension, ``nixl_ep_cpp*.so`` — built in-tree from ``3rdparty/nixl/examples/device/ep`` and staged into ``_libs/`` here. @@ -150,9 +150,11 @@ def _load_nixl_ep_cpp() -> ctypes.CDLL: so_files = list(_libs_dir.glob("nixl_ep_cpp*.so")) if not so_files: raise MoEEpNotBuiltError( - f"nixl_ep_cpp*.so is not staged under {_libs_dir}. Rebuild with:\n" - ' BUILD_NVEP=1 pip install -e ".[nvep]"\n' - "or BUILD_NIXL_EP=1 for a NIXL-EP-only build." + f"nixl_ep_cpp*.so is not staged under {_libs_dir}. It builds by " + "default; rebuild with:\n" + " pip install -e .\n" + "(BUILD_NIXL_EP=1 makes missing build deps a hard error instead " + "of skip-with-warning)." ) _preload_libnixl() try: diff --git a/flashinfer/moe_ep/nixl_ep/fleet.py b/flashinfer/moe_ep/nixl_ep/fleet.py index ed41c778d53..4a54a5946e7 100644 --- a/flashinfer/moe_ep/nixl_ep/fleet.py +++ b/flashinfer/moe_ep/nixl_ep/fleet.py @@ -43,7 +43,8 @@ def _load_nixl_ep(): from . import _load_nixl_ep_cpp # noqa: F401 except ImportError as e: raise MoEEpNotBuiltError( - "nixl_ep loaders not staged; rebuild with BUILD_NIXL_EP=1" + "nixl_ep loaders not staged; rebuild with `pip install -e .` " + "(BUILD_NIXL_EP=1 makes missing build deps a hard error)" ) from e _load_nixl_ep_cpp() try: @@ -59,7 +60,8 @@ def _load_nixl_ep(): import nixl_ep # type: ignore[import-not-found] except ImportError as e: raise MoEEpNotBuiltError( - "nixl_ep python module not importable; rebuild with BUILD_NIXL_EP=1" + "nixl_ep python module not importable; rebuild with `pip install -e .` " + "(BUILD_NIXL_EP=1 makes missing build deps a hard error)" ) from e return nixl_ep diff --git a/pyproject.toml b/pyproject.toml index 04aebc10623..00907e99aa1 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,25 +26,13 @@ license-files = ["LICENSE", "LICENSE*.txt"] [project.optional-dependencies] cu12 = ["nvidia-cutlass-dsl>=4.5.0"] cu13 = ["nvidia-cutlass-dsl[cu13]>=4.5.0"] -# Runtime deps for the moe_ep transport backends. Pair with the BUILD_NVEP=1 -# env var at install time: `BUILD_NVEP=1 pip install -e ".[nvep]"`. The env var -# triggers the NIXL-EP submodule build (meson); the build_backend.py hook also -# pip-installs `nixl-cu13>=1.0.1` with --no-deps (matching SGLang's Dockerfile -# pattern) so it doesn't drag transitive constraints that downgrade torch. -# -# NCCL-EP is NO LONGER built from the in-tree submodule: it is provided by the -# released `nccl4py` wheel (>=0.3.1), which ships the `nccl.ep` Pythonic API and -# bundles libnccl_ep.so (loaded via cuda-pathfinder, so no LD_LIBRARY_PATH hack). -# `flashinfer.moe_ep.nccl_ep` imports `nccl.ep`; availability is probed via -# `find_spec("nccl.ep")`. -nvep = [ - "cuda-python>=13.0", - "nccl4py>=0.3.1", - # NCCL-EP group-create fails on B200 with older NCCL (2.27.x/2.29.x); >=2.30.7 - # carries the B200 EP support. Ensure this wheel's libnccl is loaded (first on - # LD_LIBRARY_PATH) rather than a base-image system NCCL. - "nvidia-nccl-cu13>=2.30.7", -] +# DEPRECATED alias, kept so existing `pip install ".[nvep]"` commands keep +# working. The moe_ep runtime deps (cuda-python, nccl4py, nvidia-nccl-cu13) +# are now part of the BASE dependencies (requirements.txt) and the NIXL-EP +# submodule build runs by default (best-effort) on `pip install .` — see the +# moe_ep section at the top of build_backend.py. Opt out of the native build +# with BUILD_NVEP=0 (or per-backend BUILD_NIXL_EP=0 / BUILD_NCCL_EP=0). +nvep = [] [project.scripts] flashinfer = "flashinfer.__main__:cli" @@ -94,9 +82,9 @@ exclude = ["flashinfer-jit-cache*", "flashinfer-cubin*"] "flashinfer.data.cutlass" = ["include/**", "tools/util/include/**"] "flashinfer.data.spdlog" = ["include/**"] "flashinfer.data.cccl" = ["cub/cub/**", "libcudacxx/include/**", "thrust/thrust/**"] -# EP backend shared libraries built in-tree from 3rdparty/{nixl,nccl} when -# BUILD_NVEP=1 is set during `pip install`. The .so files live under -# flashinfer/moe_ep/{nixl_ep,nccl_ep}/_libs/ (gitignored; populated by +# EP backend shared libraries built in-tree from 3rdparty/nixl by default +# during `pip install` (opt out with BUILD_NIXL_EP=0). The .so files live +# under flashinfer/moe_ep/{nixl_ep,nccl_ep}/_libs/ (gitignored; populated by # build_backend._build_nvep_if_enabled). "flashinfer.moe_ep.nixl_ep" = ["_libs/**"] "flashinfer.moe_ep.nccl_ep" = ["_libs/*.so*"] diff --git a/requirements.txt b/requirements.txt index 7dc363300c4..83ee20d6c15 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,12 +1,22 @@ apache-tvm-ffi>=0.1.6,!=0.1.8,!=0.1.8.post0,<0.2 click +# cuda-python + nccl4py + nvidia-nccl-cu13: runtime deps for the moe_ep EP +# transport backends (default since the EP install became opt-out; CUDA-13 +# wheels only — see build_backend.py). NCCL-EP is the released nccl4py wheel +# (`nccl.ep` API + bundled libnccl_ep.so); no in-tree NCCL build. +cuda-python>=13.0 cuda-tile>=1.4.0 einops +nccl4py>=0.3.1 ninja numpy nvidia-cudnn-frontend>=1.13.0 nvidia-cutlass-dsl>=4.5.0 nvidia-ml-py +# NCCL-EP group-create fails on B200 with older NCCL (2.27.x/2.29.x); >=2.30.7 +# carries the B200 EP support. Ensure this wheel's libnccl is loaded (first on +# LD_LIBRARY_PATH) rather than a base-image system NCCL. +nvidia-nccl-cu13>=2.30.7 packaging>=24.2 requests tabulate diff --git a/scripts/build_in_container.sh b/scripts/build_in_container.sh index 0b1ddb73cee..3fab2b87909 100755 --- a/scripts/build_in_container.sh +++ b/scripts/build_in_container.sh @@ -6,8 +6,8 @@ # --container-image=nvcr.io/nvidia/cuda:13.0.0-cudnn-devel-ubuntu24.04 # --container-writable` session; it installs system deps, builds UCX # v1.21.x + GDRCopy v2.5.1 from source, creates a venv with FlashInfer -# pinned, and finally runs `BUILD_NCCL_EP=1 BUILD_NIXL_EP=1 pip install -# -e ".[nvep]"`. +# pinned, and finally runs `BUILD_NIXL_EP=1 pip install -e .` (the EP +# backends build by default; the explicit flag makes missing deps fatal). # # Env knobs: # REPO_ROOT path to the flashinfer checkout (defaults to PWD) @@ -118,8 +118,9 @@ uv pip install --python "${VENV}/bin/python" \ uv pip install --python "${VENV}/bin/python" --no-deps \ "nixl-cu13>=1.0.1" -# FlashInfer runtime deps + the [nvep] extra, installed explicitly here (WITH -# their own deps) so the editable flashinfer install below can use --no-deps. +# FlashInfer runtime deps (incl. the moe_ep deps, now part of the base +# dependencies), installed explicitly here (WITH their own deps) so the +# editable flashinfer install below can use --no-deps. # Why: torch 2.12's `cuda-toolkit[nvjitlink]` metapackage pin trips uv's # resolver during the editable `-e .` resolution (nvidia-nvjitlink METADATA # mismatch). Installing the leaf deps first + `--no-deps -e .` sidesteps that. diff --git a/tests/conftest.py b/tests/conftest.py index c2e7f2a251e..57664e423ee 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -156,7 +156,9 @@ def pytest_configure(config): for fn in TORCH_COMPILE_FNS: _monkeypatch_add_torch_compile(fn) # moe_ep markers (Part B of the EP API design integration). - config.addinivalue_line("markers", "nvep: requires BUILD_NVEP=1 install") + config.addinivalue_line( + "markers", "nvep: requires a moe_ep-enabled install (default)" + ) config.addinivalue_line("markers", "gpu_2: requires >=2 GPUs") config.addinivalue_line("markers", "gpu_4: requires >=4 GPUs") config.addinivalue_line("markers", "gpu_8: requires >=8 GPUs") @@ -191,7 +193,8 @@ def pytest_collection_modifyitems(config, items): if "nvep" in item.keywords and not nvep_built: item.add_marker( pytest.mark.skip( - reason="needs BUILD_NCCL_EP=1 / BUILD_NIXL_EP=1 install" + reason="no moe_ep backend built (EP builds by default; " + "check install log for skipped-backend warnings)" ) ) for mk, req in (("gpu_2", 2), ("gpu_4", 4), ("gpu_8", 8)): diff --git a/tests/moe_ep/smoke_nccl_ep.py b/tests/moe_ep/smoke_nccl_ep.py index 5a41d543e6c..ffc707889d7 100644 --- a/tests/moe_ep/smoke_nccl_ep.py +++ b/tests/moe_ep/smoke_nccl_ep.py @@ -9,8 +9,8 @@ topk_weights, the output approximates the input within bf16 tolerance. Designed for the Phase 4 on-cluster validation step. On the dev box this -also exits 0 with ``--nproc_per_node=1`` provided the EP backends were -built (``BUILD_NCCL_EP=1``). +also exits 0 with ``--nproc_per_node=1`` provided the EP backends are +available (they are by default: nccl4py is a base dependency). """ from __future__ import annotations From 155b66bfeeca27a11d8018b71bc02dbafbbed964 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Thu, 2 Jul 2026 22:42:33 -0700 Subject: [PATCH 02/13] fix: address review comments on EP-default build - Key NIXL-EP strictness solely off BUILD_NIXL_EP: only NIXL-EP goes through _gate_backend (NCCL-EP has no build step), so an explicit BUILD_NCCL_EP=1 no longer forces NIXL-EP into strict mode (coderabbit). - Detect PEP 517 isolated build envs (pip-build-env-* / uv builds-v0) and print a prominent warning that the NIXL-EP build needs --no-build-isolation, since hook-installed wheels don't persist into the target environment (gemini-code-assist). - Parse torch.version.cuda defensively in validate_arch_for_backend so custom/nightly version strings can't crash the CUDA-13 check (gemini-code-assist). AI-assisted. Co-Authored-By: Claude Fable 5 --- build_backend.py | 42 +++++++++++++++++++++++++++----- flashinfer/moe_ep/_validators.py | 8 +++++- 2 files changed, 43 insertions(+), 7 deletions(-) diff --git a/build_backend.py b/build_backend.py index 350e4c35cd1..f827ce5e998 100644 --- a/build_backend.py +++ b/build_backend.py @@ -92,17 +92,33 @@ def _backend_enabled(name: str) -> bool: _BUILD_NIXL_EP = _backend_enabled("BUILD_NIXL_EP") # Missing build-time deps skip the backend with a warning instead of aborting -# the install — EXCEPT when the user explicitly asked for a backend with -# BUILD_NCCL_EP=1 / BUILD_NIXL_EP=1; then a missing dep is a hard error. The -# default-on install and the legacy BUILD_NVEP=1 alias are both best-effort. -_BUILD_NVEP_BEST_EFFORT = not ( - _tri_flag("BUILD_NCCL_EP") is True or _tri_flag("BUILD_NIXL_EP") is True -) +# the install — EXCEPT when the user explicitly asked for the NIXL-EP build +# with BUILD_NIXL_EP=1; then a missing dep is a hard error. Only NIXL-EP goes +# through _gate_backend (NCCL-EP has no build step), so strictness is keyed +# solely off the NIXL-EP flag — an explicit BUILD_NCCL_EP=1 must not force +# NIXL-EP into strict mode. The default-on install and the legacy +# BUILD_NVEP=1 alias are both best-effort. +_BUILD_NVEP_BEST_EFFORT = _tri_flag("BUILD_NIXL_EP") is not True _nvep_build_root = _root / "build_nvep" _moe_ep_pkg = _root / "flashinfer" / "moe_ep" +def _in_isolated_build_env() -> bool: + """Heuristic: are we running inside a PEP 517 isolated build env? + + pip's isolated build envs live in a ``pip-build-env-*`` temp dir injected + on sys.path; uv's ephemeral build envs live under a ``builds-v0`` cache + dir. In such an env, wheels installed by this hook (nixl-cu13) vanish + when the build finishes and never reach the user's target environment — + and the env usually has no ``pip`` module at all, so the installs fail + outright. The moe_ep build path therefore needs --no-build-isolation. + """ + markers = ("pip-build-env-", f"{os.sep}builds-v0{os.sep}") + paths = [sys.prefix, *sys.path] + return any(m in p for m in markers for p in paths) + + def _detect_cuda_major() -> int: """Best-effort detection of the CUDA major version on the host.""" try: @@ -597,6 +613,20 @@ def _build_nvep_if_enabled() -> None: mode = "best-effort" if _BUILD_NVEP_BEST_EFFORT else "strict" print(f"[BUILD_NVEP] requested: {', '.join(requested)} (mode: {mode})") + if _BUILD_NIXL_EP and _in_isolated_build_env(): + print( + "[BUILD_NVEP] WARNING: PEP 517 build isolation detected. Wheels " + "installed by this hook (nixl-cu13) land in the throwaway build " + "env — the NIXL-EP build will most likely be skipped, and even " + "if it succeeds its runtime wheel will NOT persist into the " + "target environment. To enable NIXL-EP when installing from " + "source, disable isolation:\n" + " pip install --no-build-isolation .\n" + "If NIXL-EP libs were still staged, install the runtime wheel " + "manually afterwards: pip install --no-deps 'nixl-cu13>=1.0.1'.", + flush=True, + ) + # NCCL-EP is not built from source — it is provided by the released # `nccl4py` wheel (>=0.3.1, the `nccl.ep` API + bundled libnccl_ep.so), # which is a base dependency now. So BUILD_NCCL_EP requires no in-tree diff --git a/flashinfer/moe_ep/_validators.py b/flashinfer/moe_ep/_validators.py index 76aac6971fc..af80db7d680 100644 --- a/flashinfer/moe_ep/_validators.py +++ b/flashinfer/moe_ep/_validators.py @@ -42,8 +42,14 @@ def validate_arch_for_backend(backend: str) -> None: # The EP runtime wheels (nccl4py, nvidia-nccl-cu13, nixl-cu13) are # CUDA-13-only, so a torch built for CUDA 12 can't drive either backend — # fail here with a clear message instead of a cryptic dlopen error later. + # Parse defensively: custom/nightly torch builds can carry version + # strings this check shouldn't crash on; skip it when unparseable. cuda_ver = torch.version.cuda - if cuda_ver is not None and int(cuda_ver.split(".")[0]) < 13: + try: + cuda_major = int(cuda_ver.split(".")[0]) if cuda_ver else None + except ValueError: + cuda_major = None + if cuda_major is not None and cuda_major < 13: raise MoEEpConfigError( f"{backend} requires CUDA 13: the EP runtime wheels (nccl4py, " f"nvidia-nccl-cu13, nixl-cu13) ship CUDA-13 binaries only, but " From d12a6b5b3aa9ed99ec8f76bb0d464ad7a09117ef Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sat, 4 Jul 2026 10:20:13 -0700 Subject: [PATCH 03/13] test: skip nixl_ep mock tests on CUDA < 13 torch builds create_fleet now runs the CUDA-13 check in validate_arch_for_backend (the EP runtime wheels ship CUDA-13 binaries only), so on CI runners with a CUDA-12 torch these mocked tests failed in validation before reaching the fake Buffer. Replace the bare torch.cuda.is_available() skips with a _skip_unless_ep_capable() helper that also skips when torch.version.cuda < 13. Verified in nvcr.io/nvidia/pytorch:26.05-py3 (CUDA 13.2): all three tests run and pass; helper matrix-tested against stubbed torch versions (12.8 -> skip, 13.2/None/unparseable -> run, no CUDA -> skip). AI-assisted. Co-Authored-By: Claude Fable 5 --- tests/moe_ep/nixl_ep/test_fleet_mock.py | 34 +++++++++++++++++-------- 1 file changed, 24 insertions(+), 10 deletions(-) diff --git a/tests/moe_ep/nixl_ep/test_fleet_mock.py b/tests/moe_ep/nixl_ep/test_fleet_mock.py index 378467adb4a..d5d270c659c 100644 --- a/tests/moe_ep/nixl_ep/test_fleet_mock.py +++ b/tests/moe_ep/nixl_ep/test_fleet_mock.py @@ -21,6 +21,27 @@ import pytest +def _skip_unless_ep_capable(): + """Skip on hosts that can't construct an EP Fleet even with mocks. + + ``create_fleet`` runs ``validate_arch_for_backend``, which requires a + CUDA device and a CUDA-13 torch build (the EP runtime wheels ship + CUDA-13 binaries only), so on older stacks these tests would fail in + validation before reaching the mocked Buffer. + """ + import torch + + if not torch.cuda.is_available(): + pytest.skip("needs CUDA") + cuda_ver = torch.version.cuda + try: + cuda_major = int(cuda_ver.split(".")[0]) if cuda_ver else None + except ValueError: + cuda_major = None + if cuda_major is not None and cuda_major < 13: + pytest.skip(f"moe_ep requires a CUDA-13 torch build (got CUDA {cuda_ver})") + + @pytest.fixture def fake_buffer_cls(): """Build a `Buffer` class that records ctor + method calls.""" @@ -119,10 +140,7 @@ def patched_loader(fake_nixl_ep_module): def test_fleet_init_calls_update_memory_and_connect(patched_loader, fake_buffer_cls): - import torch - - if not torch.cuda.is_available(): - pytest.skip("needs CUDA") + _skip_unless_ep_capable() from flashinfer.moe_ep import ( BootstrapConfig, @@ -163,8 +181,7 @@ def test_fleet_init_calls_update_memory_and_connect(patched_loader, fake_buffer_ def test_handle_combine_requires_topk_weights(patched_loader, fake_buffer_cls): import torch - if not torch.cuda.is_available(): - pytest.skip("needs CUDA") + _skip_unless_ep_capable() from flashinfer.moe_ep import ( BootstrapConfig, @@ -198,10 +215,7 @@ def test_handle_combine_requires_topk_weights(patched_loader, fake_buffer_cls): def test_update_topology_diffs_ranks(patched_loader, fake_buffer_cls): - import torch - - if not torch.cuda.is_available(): - pytest.skip("needs CUDA") + _skip_unless_ep_capable() from flashinfer.moe_ep import ( BootstrapConfig, From e2bfadcb31814f9842aa43811d0f41e45df2bf73 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sat, 4 Jul 2026 11:28:40 -0700 Subject: [PATCH 04/13] fix: drop nvidia-nccl-cu13 floor from base deps; enforce at runtime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The nvidia-nccl-cu13>=2.30.7 base dependency broke the aarch64 AOT CI job: torch's cu13 wheels pin nvidia-nccl-cu13 EXACTLY (2.11.0 ==2.28.9, 2.12.1 ==2.29.7), so the floor made pip's resolver evict the installed CUDA torch and backtrack to the CPU-only aarch64 torch-2.10.0 wheel ("Torch not compiled with CUDA enabled"). torch's own cu13 pin supplies libnccl, so the base dep is only needed for the B200 EP floor — which can't be expressed in metadata without fighting torch's exact pins. Instead: - requirements.txt: drop nvidia-nccl-cu13 (keep nccl4py, cuda-python) with a comment explaining why a floor there is a footgun. - build_backend.py: new _ensure_nccl_floor() installs nvidia-nccl-cu13>=2.30.7 with --no-deps on source installs (mirrors the nixl-cu13 pattern; never enters the resolver). Best-effort: failures warn and defer to the runtime check. - moe_ep/_validators.py: enforce the floor where it actually matters — NCCL-EP Fleet construction on Blackwell (sm_100+), where group-create fails with NCCL < 2.30.7. _installed_nccl_version() probes the nvidia-nccl-cu13 wheel metadata, falling back to ncclGetVersion via ctypes; undeterminable versions never block. Verified in nvcr.io/nvidia/pytorch:26.05-py3 on a sm_100 GPU: plain `pip install .` leaves torch untouched, the hook installs nvidia-nccl-cu13 2.30.7 out-of-band, available_backends() has nccl_ep, the validator passes on real Blackwell, and the moe_ep mock tests pass. Floor logic matrix-tested (old NCCL on sm_100 raises with actionable message; Hopper and unknown-version paths don't block). AI-assisted. Co-Authored-By: Claude Fable 5 --- build_backend.py | 72 ++++++++++++++++++++++++-------- flashinfer/moe_ep/_validators.py | 57 +++++++++++++++++++++++++ pyproject.toml | 12 +++--- requirements.txt | 19 +++++---- 4 files changed, 130 insertions(+), 30 deletions(-) diff --git a/build_backend.py b/build_backend.py index f827ce5e998..3dafd349955 100644 --- a/build_backend.py +++ b/build_backend.py @@ -378,6 +378,43 @@ def _ensure_nixl_wheel() -> None: ) +def _ensure_nccl_floor() -> None: + """Best-effort upgrade of nvidia-nccl-cu13 to the B200 EP floor (>=2.30.7). + + Deliberately NOT a base dependency: torch's cu13 wheels pin + nvidia-nccl-cu13 EXACTLY (e.g. ==2.29.7), so declaring a >=2.30.7 floor + in package metadata makes pip's resolver evict torch — on aarch64 it + backtracks to the CPU-only torch wheel. Installing here with --no-deps + (mirroring the nixl-cu13 pattern) upgrades the wheel without ever + entering the resolver. Failures only warn: torch's own NCCL is + sufficient everywhere except NCCL-EP group-create on B200, and + moe_ep/_validators.py enforces the floor at runtime with an actionable + error where it actually matters. + """ + cuda_major = _detect_cuda_major() + if cuda_major < 13: + return # EP is CUDA-13-only; nothing to upgrade on cu12 hosts. + wheel = "nvidia-nccl-cu13>=2.30.7" + print(f"[BUILD_NVEP] ensuring NCCL-EP floor --no-deps: {wheel}") + + uv_bin = shutil.which("uv") + if uv_bin: + cmd = [uv_bin, "pip", "install", "--python", sys.executable, "--no-deps", wheel] + else: + cmd = [sys.executable, "-m", "pip", "install", "--no-deps", wheel] + print(f"[BUILD_NVEP] $ {' '.join(cmd)}") + try: + subprocess.run(cmd, check=True) + except (subprocess.CalledProcessError, FileNotFoundError) as e: + print( + f"[BUILD_NVEP] WARNING: could not install {wheel} ({e}). " + "NCCL-EP on B200 needs NCCL >= 2.30.7 (group-create fails on " + "older releases); the runtime validator will raise there. " + "Install manually if needed: pip install --no-deps " + f"'{wheel}'" + ) + + def _nixl_buildable() -> tuple[bool, str]: """Probe for hard NIXL-EP build-time deps. Returns (ok, reason_if_not). @@ -420,15 +457,16 @@ def _nixl_buildable() -> tuple[bool, str]: return True, "" -def _install_nvep_runtime_wheels(built_nixl: bool, built_nccl: bool) -> None: - """Install the EP-related runtime wheels with --no-deps, gated per backend. +def _install_nvep_runtime_wheels(built_nixl: bool) -> None: + """Install the NIXL runtime wheel with --no-deps when NIXL-EP was built. - These wheels supply the BASE libraries (libnccl.so.2, libnixl.so) that the - EP plugins (libnccl_ep.so, nixl_ep_cpp.so) dynamically load at runtime. - We do NOT stage the base libs into the FlashInfer package tree — relying - on these pip wheels keeps the wheel small and avoids the duplication. + The wheel supplies the BASE libraries (libnixl.so + siblings) that the + nixl_ep_cpp.so plugin dynamically loads at runtime. We do NOT stage the + base libs into the FlashInfer package tree — relying on the pip wheel + keeps the wheel small and avoids the duplication. (NCCL-EP's libnccl + comes from torch's own nvidia-nccl-cu13 pin; see _ensure_nccl_floor.) - The wheels carry transitive constraints (e.g. an nvidia-nccl-cu12 pin via + The wheel carries transitive constraints (e.g. an nvidia-nccl-cu12 pin via the `nixl` meta-package) that conflict with a recent torch and force a downgrade when resolved normally. SGLang's Dockerfile mirrors this with `pip install nixl nixl-cu13 --no-deps`; we do the same. @@ -438,21 +476,18 @@ def _install_nvep_runtime_wheels(built_nixl: bool, built_nccl: bool) -> None: no pip module). This is the path most users hit. 2. `python -m pip install` — for venvs with pip seeded. - Each wheel is gated on what was ACTUALLY built (not what was requested), - so `pip list` stays honest when a backend was skipped due to missing - build-time deps in best-effort mode. + Gated on what was ACTUALLY built (not what was requested), so `pip list` + stays honest when the backend was skipped due to missing build-time deps + in best-effort mode. - This step is now FATAL on failure. Since we no longer stage the base - libs, a half-installed env where the wheels failed to install would - leave the EP plugins unable to load at runtime. Better to fail loudly - at install time. + This step is FATAL on failure. Since we no longer stage the base libs, a + half-installed env where the wheel failed to install would leave the EP + plugin unable to load at runtime. Better to fail loudly at install time. """ cuda_major = _detect_cuda_major() wheels: list[str] = [] if built_nixl: wheels.append(f"nixl-cu{cuda_major}>=1.0.1") - if built_nccl: - wheels.append(f"nvidia-nccl-cu{cuda_major}>=2.30.4") if not wheels: return @@ -636,6 +671,9 @@ def _build_nvep_if_enabled() -> None: "[BUILD_NVEP] NCCL-EP is provided by the nccl4py wheel (>=0.3.1), " "a base dependency of flashinfer-python; no in-tree build." ) + # torch's cu13 wheels pin nvidia-nccl-cu13 exactly (< the B200 EP + # floor), so upgrade it out-of-band; best-effort by design. + _ensure_nccl_floor() # The default (non-hermetic) NIXL-EP build links against the nixl-cu13 # wheel's libnixl.so — install it up front so plain `pip install .` works @@ -688,7 +726,7 @@ def _build_nvep_if_enabled() -> None: with _time_phase("_fix_rpaths"): _fix_rpaths() with _time_phase("_install_nvep_runtime_wheels"): - _install_nvep_runtime_wheels(built_nixl=built_nixl, built_nccl=False) + _install_nvep_runtime_wheels(built_nixl=built_nixl) print( f"[BUILD_NVEP] total build phase wall time: " diff --git a/flashinfer/moe_ep/_validators.py b/flashinfer/moe_ep/_validators.py index af80db7d680..152aaa5113f 100644 --- a/flashinfer/moe_ep/_validators.py +++ b/flashinfer/moe_ep/_validators.py @@ -26,6 +26,14 @@ # NIXL EP's `FINISHED_SUM_TAG` is hard-coded to 1024 in the kernel. _NIXL_EP_MAX_TOKENS_PER_RANK = 1024 +# NCCL-EP group-create fails on Blackwell (B200) with older NCCL +# (2.27.x/2.29.x, at nccl_ep.cc:1438); >=2.30.7 carries the B200 EP support. +# This floor is enforced HERE rather than as a base-dependency pin because +# torch's cu13 wheels pin nvidia-nccl-cu13 exactly (e.g. ==2.29.7) — a +# metadata floor makes pip evict torch (see requirements.txt). The build hook +# upgrades the wheel --no-deps on source installs (build_backend.py). +_NCCL_EP_BLACKWELL_MIN_NCCL = (2, 30, 7) + class MoEEpConfigError(ValueError): """Raised when an EP config field is out-of-range for the chosen backend.""" @@ -35,6 +43,37 @@ class MoEEpArchError(MoEEpConfigError): """Raised when the GPU arch doesn't support the chosen backend.""" +def _installed_nccl_version() -> "tuple[int, int, int] | None": + """Best-effort probe of the NCCL version the EP backend will load. + + Prefers the nvidia-nccl-cu13 pip wheel's metadata (cuda-pathfinder loads + that wheel's libnccl first when present); falls back to ncclGetVersion on + the dynamic linker's default search path (covers NGC-style images with a + system NCCL and no pip wheel). Returns None when undeterminable — callers + must not block in that case. + """ + try: + from importlib.metadata import version + + parts = version("nvidia-nccl-cu13").split(".")[:3] + return tuple(int(p) for p in parts) # type: ignore[return-value] + except Exception: + pass + try: + import ctypes + + lib = ctypes.CDLL("libnccl.so.2") + out = ctypes.c_int() + if lib.ncclGetVersion(ctypes.byref(out)) == 0: + # NCCL_VERSION_CODE encoding: major*10000 + minor*100 + patch + # (e.g. 2.30.7 -> 23007). + code = out.value + return (code // 10000, (code // 100) % 100, code % 100) + except Exception: + pass + return None + + def validate_arch_for_backend(backend: str) -> None: """Check the GPU arch and CUDA version are supported by `backend`.""" import torch @@ -64,6 +103,24 @@ def validate_arch_for_backend(backend: str) -> None: if cc < (9, 0): raise MoEEpArchError(f"{backend} requires sm_90+, host has sm_{cc[0]}{cc[1]}") + # NCCL-EP group-create fails on Blackwell with NCCL < 2.30.7 — catch it + # here (Fleet construction) with an actionable message instead of the + # cryptic nccl_ep.cc:1438 failure. Skipped when the version can't be + # determined (no pip wheel + no loadable libnccl.so.2). + if backend == "nccl_ep" and cc >= (10, 0): + nccl_ver = _installed_nccl_version() + if nccl_ver is not None and nccl_ver < _NCCL_EP_BLACKWELL_MIN_NCCL: + floor = ".".join(map(str, _NCCL_EP_BLACKWELL_MIN_NCCL)) + found = ".".join(map(str, nccl_ver)) + raise MoEEpConfigError( + f"nccl_ep on Blackwell (sm_{cc[0]}{cc[1]}) requires NCCL >= " + f"{floor} (group-create fails with older releases); found " + f"{found}. Upgrade with:\n" + f" pip install --no-deps 'nvidia-nccl-cu13>={floor}'\n" + "and ensure that wheel's libnccl is the one loaded (first on " + "LD_LIBRARY_PATH) rather than a base-image system NCCL." + ) + def validate_fleet_params( params: FleetParams, diff --git a/pyproject.toml b/pyproject.toml index 00907e99aa1..ffb137ca373 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -27,11 +27,13 @@ license-files = ["LICENSE", "LICENSE*.txt"] cu12 = ["nvidia-cutlass-dsl>=4.5.0"] cu13 = ["nvidia-cutlass-dsl[cu13]>=4.5.0"] # DEPRECATED alias, kept so existing `pip install ".[nvep]"` commands keep -# working. The moe_ep runtime deps (cuda-python, nccl4py, nvidia-nccl-cu13) -# are now part of the BASE dependencies (requirements.txt) and the NIXL-EP -# submodule build runs by default (best-effort) on `pip install .` — see the -# moe_ep section at the top of build_backend.py. Opt out of the native build -# with BUILD_NVEP=0 (or per-backend BUILD_NIXL_EP=0 / BUILD_NCCL_EP=0). +# working. The moe_ep runtime deps (cuda-python, nccl4py) are now part of the +# BASE dependencies (requirements.txt) and the NIXL-EP submodule build runs by +# default (best-effort) on `pip install .` — see the moe_ep section at the top +# of build_backend.py. libnccl comes from torch's own nvidia-nccl-cu13 pin; +# the >=2.30.7 B200 floor is enforced at runtime (moe_ep/_validators.py) and +# upgraded --no-deps by the build hook on source installs. Opt out of the +# native build with BUILD_NVEP=0 (or BUILD_NIXL_EP=0 / BUILD_NCCL_EP=0). nvep = [] [project.scripts] diff --git a/requirements.txt b/requirements.txt index 83ee20d6c15..6aa592684a8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,9 +1,16 @@ apache-tvm-ffi>=0.1.6,!=0.1.8,!=0.1.8.post0,<0.2 click -# cuda-python + nccl4py + nvidia-nccl-cu13: runtime deps for the moe_ep EP -# transport backends (default since the EP install became opt-out; CUDA-13 -# wheels only — see build_backend.py). NCCL-EP is the released nccl4py wheel -# (`nccl.ep` API + bundled libnccl_ep.so); no in-tree NCCL build. +# cuda-python + nccl4py: runtime deps for the moe_ep EP transport backends +# (default since the EP install became opt-out; CUDA-13 wheels only — see +# build_backend.py). NCCL-EP is the released nccl4py wheel (`nccl.ep` API + +# bundled libnccl_ep.so); no in-tree NCCL build. +# +# nvidia-nccl-cu13 is deliberately NOT a base dep: torch's cu13 wheels pin it +# EXACTLY (e.g. ==2.29.7), so any floor here (>=2.30.7 for B200 EP) makes the +# resolver evict torch — on aarch64 it backtracks all the way to the CPU-only +# torch 2.10.0 wheel. torch supplies libnccl; the >=2.30.7 B200 floor is +# enforced at runtime in flashinfer/moe_ep/_validators.py, and the build hook +# installs the newer wheel --no-deps on source installs (build_backend.py). cuda-python>=13.0 cuda-tile>=1.4.0 einops @@ -13,10 +20,6 @@ numpy nvidia-cudnn-frontend>=1.13.0 nvidia-cutlass-dsl>=4.5.0 nvidia-ml-py -# NCCL-EP group-create fails on B200 with older NCCL (2.27.x/2.29.x); >=2.30.7 -# carries the B200 EP support. Ensure this wheel's libnccl is loaded (first on -# LD_LIBRARY_PATH) rather than a base-image system NCCL. -nvidia-nccl-cu13>=2.30.7 packaging>=24.2 requests tabulate From 127d539cdfce704b118960f1e5a2dd3ddd1ca021 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sun, 5 Jul 2026 16:07:15 -0700 Subject: [PATCH 05/13] fix(moe_ep): clamp nccl_ep HT max_tokens_per_rank to MAX_SUPPORTED_TOKENS_PER_RANK MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit ncclEpCreateGroup asserts (SIGABRT, nccl_ep.cc:1253) when a HIGH_THROUGHPUT group is created with max_dispatch_tokens_per_rank > MAX_SUPPORTED_TOKENS_PER_RANK (build-time 8192 in the nccl4py wheel). vLLM sizes the HT fleet from moe.max_num_tokens = scheduler max_num_batched_tokens (e.g. 16384), which tripped the assert and aborted all EP ranks the moment the modular EP (DP-EP) path was actually exercised — masked until now because the monolithic (TP-only) path never builds the HT prepare/finalize. LL is unbounded and unaffected. Clamp the HT fleet's max_tokens_per_rank to the cap (warn once) in NcclEpFleet so group creation succeeds; clamp the stored FleetParams (not just GroupConfig) so the handle's recv-buffer sizing agrees. Add a clear MoEEpConfigError guard in _dispatch_ht for the case where a single forward genuinely dispatches more than the cap per rank (previously the C++ abort / buffer overflow), pointing at --max-num-batched-tokens. AI-assisted (Claude Code): root-caused from an 8-rank per-rank stderr capture of the SIGABRT under DP-EP on Pre-Nyx. Co-Authored-By: Claude Opus 4.8 --- flashinfer/moe_ep/nccl_ep/fleet.py | 47 +++++++++++++++++++++++++++++ flashinfer/moe_ep/nccl_ep/handle.py | 16 ++++++++++ 2 files changed, 63 insertions(+) diff --git a/flashinfer/moe_ep/nccl_ep/fleet.py b/flashinfer/moe_ep/nccl_ep/fleet.py index 67ba134d7aa..60f8e5e5f21 100644 --- a/flashinfer/moe_ep/nccl_ep/fleet.py +++ b/flashinfer/moe_ep/nccl_ep/fleet.py @@ -11,6 +11,8 @@ from __future__ import annotations import contextlib +import dataclasses +import logging from typing import TYPE_CHECKING, Sequence from .. import MoEEpNotBuiltError, _require_built @@ -34,9 +36,50 @@ from ..handle import Handle +logger = logging.getLogger(__name__) + # ``GroupConfig`` fields left at 0 forward as NCCL_EP_AUTO. NCCL_EP_AUTO = 0 +# nccl_ep HT hard limit: ``ncclEpCreateGroup`` *asserts* (SIGABRT, nccl_ep.cc:1253) +# when a HIGH_THROUGHPUT group is created with +# ``max_dispatch_tokens_per_rank > MAX_SUPPORTED_TOKENS_PER_RANK``. The constant is +# a build-time template bound in the wheel +# (``nccl/ep/include/nccl_ep/common.hpp``: ``#define MAX_SUPPORTED_TOKENS_PER_RANK +# 8192``). We mirror it here to *clamp* the HT dispatch budget (graceful) rather +# than let a large caller value (e.g. vLLM ``max_num_batched_tokens``) hit the C++ +# assert. LL has no such cap. Kept in sync with the nccl4py wheel. +_HT_MAX_SUPPORTED_TOKENS_PER_RANK = 8192 + + +def _clamp_ht_max_tokens(params: FleetParams) -> FleetParams: + """Clamp a HT fleet's ``max_tokens_per_rank`` to the nccl_ep build-time cap. + + HT's ``ncclEpCreateGroup`` aborts when ``max_dispatch_tokens_per_rank`` exceeds + ``MAX_SUPPORTED_TOKENS_PER_RANK`` (8192). We return a clamped copy so group + creation succeeds; a single forward that actually dispatches more than the cap + per rank is caught with a clear error at dispatch (see ``NcclEpHandle._dispatch_ht``) + rather than silently truncated. No-op for LL (unbounded) or when already within cap. + """ + if ( + params.algorithm is EpAlgorithm.HIGH_THROUGHPUT + and params.max_tokens_per_rank > _HT_MAX_SUPPORTED_TOKENS_PER_RANK + ): + logger.warning( + "nccl_ep HT caps max_dispatch_tokens_per_rank at %d " + "(MAX_SUPPORTED_TOKENS_PER_RANK); requested %d — clamping to avoid the " + "ncclEpCreateGroup abort. Ensure the per-forward token count per rank " + "stays <= %d (e.g. vLLM --max-num-batched-tokens); a larger dispatch " + "will raise at forward time.", + _HT_MAX_SUPPORTED_TOKENS_PER_RANK, + params.max_tokens_per_rank, + _HT_MAX_SUPPORTED_TOKENS_PER_RANK, + ) + return dataclasses.replace( + params, max_tokens_per_rank=_HT_MAX_SUPPORTED_TOKENS_PER_RANK + ) + return params + def _import_nccl_ep(): """Import the ``nccl.ep`` package or raise an actionable build error.""" @@ -108,6 +151,10 @@ def __init__( _require_built("nccl_ep") validate_arch_for_backend("nccl_ep") + # HT: clamp the per-rank dispatch budget to the library's build-time cap so + # ncclEpCreateGroup doesn't abort; must clamp the stored params (not just the + # GroupConfig) so the handle's recv-buffer sizing agrees. + params = _clamp_ht_max_tokens(params) self._params = params self._fleet_knobs = _index_knobs(algo_knobs) validate_fleet_params( diff --git a/flashinfer/moe_ep/nccl_ep/handle.py b/flashinfer/moe_ep/nccl_ep/handle.py index 750aa9190c2..3c741c3ffc3 100644 --- a/flashinfer/moe_ep/nccl_ep/handle.py +++ b/flashinfer/moe_ep/nccl_ep/handle.py @@ -44,6 +44,7 @@ HandleAlgoKnobUserStream, _index_knobs, ) +from .._validators import MoEEpConfigError from ..config import ( CombineInputParams, CombineOutput, @@ -414,6 +415,21 @@ def _dispatch_ht(self, x) -> DispatchOutput: world = self._fleet.params.num_experts // self._num_local_experts num_recv = max_per_rank * world + # The HT staging buffers (and this recv buffer) are sized to max_per_rank, + # which the fleet clamps to the library's MAX_SUPPORTED_TOKENS_PER_RANK. A + # forward that dispatches more than that per rank would overflow the staging + # buffers (and previously hit a C++ abort at group-create for the un-clamped + # value). Fail with an actionable error instead of corrupting memory. + n_tokens = x.shape[0] + if n_tokens > max_per_rank: + raise MoEEpConfigError( + f"nccl_ep HT dispatch received {n_tokens} tokens on this rank, " + f"exceeding max_tokens_per_rank ({max_per_rank} = the library's " + "MAX_SUPPORTED_TOKENS_PER_RANK). Reduce the per-forward token count " + "per rank (e.g. vLLM --max-num-batched-tokens <= " + f"{max_per_rank}), or use the low-latency algorithm." + ) + tw = self._handle_knobs.get(HandleAlgoKnobTopKWeights) if tw is None: raise ValueError( From 96ab9620bb0ae7120562f8da7c392c57f87db546 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sun, 5 Jul 2026 20:06:26 -0700 Subject: [PATCH 06/13] perf(moe_ep): fleet-level host-path caches for the per-forward handle flow vLLM creates a fresh NcclEpHandle every MoE layer x step (routing binds at create_handle), so the existing NV_FI_EP_FAST_PATH per-handle caches never hit and every forward paid ~149us of host time (EP_PROFILE_HOST, LL decode: FFI descriptor builds 32.6+31.3us, handle setup 35.9us, allocs, plus ~45us of C calls). At decode the GPU is host-paced, and nccl.ep's fused send+recv dispatch kernel absorbs the resulting inter-rank lag as in-kernel spin (median 256us/launch, 33% of GPU time in the capped LL profile). Anchor the caches on the long-lived Fleet instead: - recv buffers (LL + HT) and the recv-count tensor (not re-zeroed across forwards; the dispatch metadata fully overwrites it -- the same contract the per-handle _FAST reuse relied on), - static FFI descriptor tuples (DispatchOutputs/LayoutInfo/configs), - a (data_ptr, dtype, shape)-keyed memo for per-call Tensor wraps, restricted to tensors <= 2 MiB: the nccl.ep Tensor wrapper keeps the torch tensor alive, so memoizing large prefill activations pinned GBs across allocator addresses and OOM'd at --gpu-memory-utilization 0.9 (small tensors are exactly the host-bound decode path this cache targets). Also adds EP_PROFILE_HOST timers around handle create/destroy. Host path 149 -> ~119us/layer/step; measured DP-EP throughput (Qwen3-30B-A3B, 8xB200, eager): LL 128/2048 5416 -> 5654 tok/s, HT 128/2048 3755 -> 3926, HT 2048/128 44467 (stable). GSM8K through the DP-EP transport re-validated: LL 0.8560/0.8976, HT 0.8567/0.8984 (flex/strict). AI-assisted (Claude Code). Co-Authored-By: Claude Opus 4.8 --- flashinfer/moe_ep/nccl_ep/handle.py | 205 ++++++++++++++++++---------- 1 file changed, 136 insertions(+), 69 deletions(-) diff --git a/flashinfer/moe_ep/nccl_ep/handle.py b/flashinfer/moe_ep/nccl_ep/handle.py index 3c741c3ffc3..8a26a3dea60 100644 --- a/flashinfer/moe_ep/nccl_ep/handle.py +++ b/flashinfer/moe_ep/nccl_ep/handle.py @@ -130,8 +130,22 @@ def __init__( from ..config import EpAlgorithm, EpLayout + _t = _pc() if _HP else None self._fleet = fleet self._ep = fleet.nccl_ep + # Cross-handle host-path cache. vLLM creates a fresh Handle every MoE + # layer x step (routing binds at create_handle), so per-handle caches + # never hit; anchoring them on the long-lived Fleet makes the recv + # buffers, counter tensors and FFI descriptor objects reusable across + # forwards. Tensor wrappers are memoized by (data_ptr, dtype, shape), + # so an entry can only ever describe the same memory layout it was + # built for; the dict is cleared when it grows past a bound (entries + # are then rebuilt, which is always safe — each handle only needs + # address stability within its own lifetime). + hot = getattr(fleet, "_hot_cache", None) + if hot is None: + hot = fleet._hot_cache = {} + self._hot = hot self._handle_knobs = _index_knobs(algo_knobs) self._stream = self._knob_stream() self._staged = HandleAlgoKnobSplitOperation in self._handle_knobs @@ -163,15 +177,22 @@ def __init__( self._topk_idx = topk_idx # keepalive self._num_tokens_in = topk_idx.shape[0] self._top_k = topk_idx.shape[1] - self._topk_idx_t = self._ep.Tensor(topk_idx) + self._topk_idx_t = self._wrap(topk_idx) # Per-source counter the library writes at dispatch (LL): EXPERT_MAJOR # gets per-local-expert recv counts [num_local_experts]; RANK_MAJOR gets - # per-source-rank token counts [world]. + # per-source-rank token counts [world]. Fleet-cached; NOT re-zeroed + # across forwards — the dispatch metadata fully overwrites every entry + # (the same contract the NV_FI_EP_FAST_PATH per-handle reuse relies on). recv_count_len = world_size if self._is_rank_major else self._num_local_experts - self._recv_count_t = torch.zeros( - recv_count_len, dtype=torch.int32, device=topk_idx.device - ) + ck = ("recv_count", recv_count_len, topk_idx.device) + self._recv_count_t = self._hot.get(ck) + if self._recv_count_t is None: + self._recv_count_t = torch.zeros( + recv_count_len, dtype=torch.int32, device=topk_idx.device + ) + self._hot[ck] = self._recv_count_t + _t = _hp("hinit.setup", _t) if self._is_ht: layout = self._ep.Layout.FLAT @@ -217,6 +238,7 @@ def __init__( config=None, stream=self._stream, ) + _t = _hp("hinit.create_handle_c", _t) # ----------------------------------------------------------------- knobs @@ -224,6 +246,38 @@ def _knob_stream(self) -> int: k = self._handle_knobs.get(HandleAlgoKnobUserStream) return int(k.stream) if k is not None else self._fleet.stream # type: ignore[attr-defined] + # Only memoize wrappers of SMALL tensors: the wrapper keeps the torch tensor + # alive, so caching wraps of large activations (e.g. 8k-token prefill inputs, + # the [num_recv, hidden] combine views) pins GBs across allocator addresses + # and OOMs at high --gpu-memory-utilization. Small tensors (weights, topk, + # counters, decode-sized activations) are exactly the host-bound decode path + # this cache exists for. 2 MiB * 256 entries caps pinning at 512 MiB worst + # case (steady-state decode reuses a handful of addresses). + _WRAP_MEMO_MAX_BYTES = 2 << 20 + _WRAP_MEMO_MAX_ENTRIES = 256 + + def _wrap(self, t): + """Memoized ``nccl.ep.Tensor`` wrapper (fleet-level, address-keyed). + + Building an FFI Tensor descriptor costs ~10us of host time; vLLM's + allocator recycles workspace addresses across decode steps, so keying + by (data_ptr, dtype, shape) hits almost always after warmup. A hit can + never alias the wrong layout — a reused address with a different + shape/dtype misses and builds a fresh wrapper. Large tensors are + wrapped per call (see _WRAP_MEMO_MAX_BYTES). + """ + if t.numel() * t.element_size() > self._WRAP_MEMO_MAX_BYTES: + return self._ep.Tensor(t) + hot = self._hot + key = (t.data_ptr(), t.dtype, tuple(t.shape)) + w = hot.get(key) + if w is None: + if len(hot) > self._WRAP_MEMO_MAX_ENTRIES: + hot.clear() + w = self._ep.Tensor(t) + hot[key] = w + return w + # ----------------------------------------------------------------- dispatch # @flashinfer_api # disabled per PR #3453 review @@ -244,25 +298,30 @@ def _dispatch_ll(self, x) -> DispatchOutput: max_per_rank = self._fleet.params.max_tokens_per_rank hidden = self._fleet.params.token_hidden_size - # (3) cache the recv buffer instead of torch.empty() every dispatch. - out_t = getattr(self, "_ll_recv_buf", None) if _FAST else None - if out_t is None: - out_t = torch.empty( - self._num_local_experts, - max_per_rank * world_size, - hidden, - dtype=x.dtype, - device=x.device, - ) - if _FAST: - self._ll_recv_buf = out_t + # Fleet-cached recv buffer (a fresh Handle is created every forward, so + # per-handle caching never hits; the fleet persists). + shape = (self._num_local_experts, max_per_rank * world_size, hidden) + out_t = self._hot.get("ll_recv_buf") + if ( + out_t is None + or out_t.shape != shape + or out_t.dtype != x.dtype + or out_t.device != x.device + ): + out_t = torch.empty(*shape, dtype=x.dtype, device=x.device) + self._hot["ll_recv_buf"] = out_t _t = _hp("ll_disp.alloc", _t) - # (2) cache the FFI wrapper objects over STABLE tensors (out_t / recv_count / - # config). Only the input-token wrap is rebuilt each call (x may alias a new - # tensor). On the slow path everything is rebuilt as before. - cache = getattr(self, "_ll_disp_cache", None) if _FAST else None - if cache is None: + # Fleet-cached FFI descriptor objects over the STABLE tensors (recv + # buffer / counters / config). Only the input-token wrap varies per call + # (memoized by address in _wrap). + cache = self._hot.get("ll_disp_ffi") + if ( + cache is None + or cache[0] is not out_t + or cache[1] is not self._recv_count_t + or cache[2] != self._staged + ): outputs = self._ep.DispatchOutputs(tokens=self._ep.Tensor(out_t)) layout_info = self._ep.LayoutInfo( expert_counters=self._ep.Tensor(self._recv_count_t) @@ -270,11 +329,17 @@ def _dispatch_ll(self, x) -> DispatchOutput: config = self._ep.DispatchConfig( send_only=int(self._staged), round_scales=0 ) - if _FAST: - self._ll_disp_cache = (outputs, layout_info, config) + self._hot["ll_disp_ffi"] = ( + out_t, + self._recv_count_t, + self._staged, + outputs, + layout_info, + config, + ) else: - outputs, layout_info, config = cache - inputs = self._ep.DispatchInputs(tokens=self._ep.Tensor(x)) + outputs, layout_info, config = cache[3], cache[4], cache[5] + inputs = self._ep.DispatchInputs(tokens=self._wrap(x)) _t = _hp("ll_disp.build_ffi_objs", _t) self._handle.dispatch( @@ -451,8 +516,14 @@ def _dispatch_ht(self, x) -> DispatchOutput: # first dispatch. Fresh torch.empty buffers each call gave the cached # dispatch new addresses and deadlocked the next collective. _t = _pc() if _HP else None - cached = getattr(self, "_ht_recv_bufs", None) - if cached is None or cached[0].shape[0] != num_recv: + cached = self._hot.get("ht_recv_bufs") + if ( + cached is None + or cached[0].shape[0] != num_recv + or cached[1].shape[1] != self._top_k + or cached[0].dtype != x.dtype + or cached[0].device != x.device + ): out_t = torch.empty(num_recv, hidden, dtype=x.dtype, device=x.device) out_w = torch.empty( num_recv, self._top_k, dtype=torch.float32, device=x.device @@ -460,15 +531,15 @@ def _dispatch_ht(self, x) -> DispatchOutput: out_idx = torch.empty( num_recv, self._top_k, dtype=torch.int64, device=x.device ) - self._ht_recv_bufs = (out_t, out_w, out_idx) + self._hot["ht_recv_bufs"] = (out_t, out_w, out_idx) else: out_t, out_w, out_idx = cached _t = _hp("ht_disp.alloc_cached", _t) - # (2) cache the output wraps (over cached recv bufs) + weights wrap + config; - # rebuild only the per-call input-token wrap. - cache = getattr(self, "_ht_disp_cache", None) if _FAST else None - if cache is None: + # Fleet-cached output wraps (over the cached recv bufs) + config; the + # per-call input-token and weights wraps go through the _wrap memo. + cache = self._hot.get("ht_disp_ffi") + if cache is None or cache[0] is not out_t or cache[1] != self._staged: outputs = self._ep.DispatchOutputs( tokens=self._ep.Tensor(out_t), topk_weights=self._ep.Tensor(out_w), @@ -477,13 +548,11 @@ def _dispatch_ht(self, x) -> DispatchOutput: config = self._ep.DispatchConfig( send_only=int(self._staged), round_scales=0 ) - weights_t = self._ep.Tensor(weights) - if _FAST: - self._ht_disp_cache = (outputs, config, weights_t) + self._hot["ht_disp_ffi"] = (out_t, self._staged, outputs, config) else: - outputs, config, weights_t = cache + outputs, config = cache[2], cache[3] inputs = self._ep.DispatchInputs( - tokens=self._ep.Tensor(x), topk_weights=weights_t + tokens=self._wrap(x), topk_weights=self._wrap(weights) ) _t = _hp("ht_disp.build_ffi_objs", _t) @@ -544,15 +613,13 @@ def combine(self, params: CombineInputParams) -> CombineOutput: x2d = x.reshape(-1, hidden) # (2) cache output wrap + config (guarded by out_t identity); rebuild # only the per-call input wrap (x2d is a fresh view each call). - cache = getattr(self, "_ht_comb_cache", None) if _FAST else None - if cache is None or cache[2] is not out_t: - outputs = self._ep.CombineOutputs(tokens=self._ep.Tensor(out_t)) + ck = ("ht_comb_cfg", self._staged) + config = self._hot.get(ck) + if config is None: config = self._ep.CombineConfig(send_only=int(self._staged)) - if _FAST: - self._ht_comb_cache = (outputs, config, out_t) - else: - outputs, config, _ = cache - inputs = self._ep.CombineInputs(tokens=self._ep.Tensor(x2d)) + self._hot[ck] = config + outputs = self._ep.CombineOutputs(tokens=self._wrap(out_t)) + inputs = self._ep.CombineInputs(tokens=self._wrap(x2d)) _t = _hp("ht_comb.build_ffi_objs", _t) self._handle.combine(inputs, outputs, config=config, stream=self._stream) _t = _hp("ht_comb.ffi_combine", _t) @@ -584,31 +651,29 @@ def combine(self, params: CombineInputParams) -> CombineOutput: self._combine_outputs = outputs return CombineOutput(x=out_t) - # LL EXPERT_MAJOR combine: weights applied on the receive side. - # (2) cache the stable weights wrap + config; rebuild only the per-call - # token wraps (x / out_t may alias new tensors). - cache = getattr(self, "_ll_comb_cache", None) if _FAST else None - if cache is None: - tw = self._handle_knobs.get(HandleAlgoKnobTopKWeights) - if tw is None: - raise ValueError( - "NcclEpHandle.combine requires HandleAlgoKnobTopKWeights set " - "at handle creation; NCCL EP LL needs per-token weights to " - "reweight on combine." - ) - weights = tw.weights # type: ignore[attr-defined] - if weights.dtype != torch.float32: - weights = weights.to(torch.float32) - weights_t = self._ep.Tensor(weights) + # LL EXPERT_MAJOR combine: weights applied on the receive side. The + # weights tensor changes every forward (per-step routing), but its + # allocator address recycles across decode steps — the _wrap memo makes + # the descriptor build ~free. The config is static per staged-mode. + tw = self._handle_knobs.get(HandleAlgoKnobTopKWeights) + if tw is None: + raise ValueError( + "NcclEpHandle.combine requires HandleAlgoKnobTopKWeights set " + "at handle creation; NCCL EP LL needs per-token weights to " + "reweight on combine." + ) + weights = tw.weights # type: ignore[attr-defined] + if weights.dtype != torch.float32: + weights = weights.to(torch.float32) + weights_t = self._wrap(weights) + ck = ("ll_comb_cfg", self._staged) + config = self._hot.get(ck) + if config is None: config = self._ep.CombineConfig(send_only=int(self._staged)) - if _FAST: - self._ll_comb_cache = (weights, weights_t, config) - else: - weights, weights_t, config = cache - - inputs = self._ep.CombineInputs(tokens=self._ep.Tensor(x)) + self._hot[ck] = config + inputs = self._ep.CombineInputs(tokens=self._wrap(x)) outputs = self._ep.CombineOutputs( - tokens=self._ep.Tensor(out_t), + tokens=self._wrap(out_t), topk_weights=weights_t, ) _t = _hp("ll_comb.build_ffi_objs", _t) @@ -637,9 +702,11 @@ def complete(self) -> None: def destroy(self) -> None: if not self._destroyed: + _t = _pc() if _HP else None with contextlib.suppress(Exception): self._handle.destroy() self._destroyed = True + _hp("hdestroy.destroy_c", _t) def __del__(self) -> None: self.destroy() From 34512db1880957e61a00a0639e89f0ac657993fa Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sun, 5 Jul 2026 20:49:45 -0700 Subject: [PATCH 07/13] docs(moe_ep): update runbook + results for the DP-EP perf iteration Runbook is now a faithful manual-reproduction script for the final numbers: - pin clone refs (flashinfer feat/nvep-default >= fa09bc46, vLLM feat/flashinfer-ep-all2all >= ab1415e) with merge-base sanity checks; note that editable installs resolve /host clones at runtime (no image rebuild for the Python-only perf fixes) - per-backend --max-num-batched-tokens handling everywhere (HT needs 8192, LL must leave it unset so the batched-DP 256 auto-cap engages) - 3b' GSM8K-through-the-transport step with final expected scores - 3c fixed to NP=256 (what the reference matrix used) + per-run sum helper - reference-numbers section replaced with the final measured matrix (FI-HT ahead of DeepEP-HT 19-27% on 2 of 3 shapes; FI-LL within 4-12%) Results doc: final same-day 4-backend matrix (1.1e) + the perf-iteration log (1.1f: batched-DP cap membership, HT recv-trim, fleet host-path caches, the ep.Tensor pinning OOM lesson, remaining decode levers). Integration doc: 0 headline updated to the transport-exercised results. AI-assisted (Claude Code). Co-Authored-By: Claude Opus 4.8 --- docs/design_docs/vllm_moe_ep_integration.md | 48 +- .../design_docs/vllm_moe_ep_results_prenyx.md | 627 ++++++++++++++++++ docs/design_docs/vllm_moe_ep_runbook.md | 566 ++++++++++++++++ 3 files changed, 1231 insertions(+), 10 deletions(-) create mode 100644 docs/design_docs/vllm_moe_ep_results_prenyx.md create mode 100644 docs/design_docs/vllm_moe_ep_runbook.md diff --git a/docs/design_docs/vllm_moe_ep_integration.md b/docs/design_docs/vllm_moe_ep_integration.md index 03e82788066..028d09cfe23 100644 --- a/docs/design_docs/vllm_moe_ep_integration.md +++ b/docs/design_docs/vllm_moe_ep_integration.md @@ -31,10 +31,29 @@ FlashInfer run from the branch. All checks below **pass**: | vLLM e2e smoke (OLMoE, coherent output) | ✅ | ✅ | | **GSM8K 5-shot, Qwen3-30B-A3B** (flex / strict) | **0.852 / 0.894** | **0.858 / 0.897** | -Both backends clear the GSM8K ≥ 0.80 gate (reference ~0.88). - -**Throughput vs DeepEP** (`vllm bench throughput --dataset-name random`, Qwen3-30B-A3B, 8-GPU -EP, 1000 prompts; total tok/s): +Both backends clear the GSM8K ≥ 0.80 gate (reference ~0.88). Correctness above (GAP tests + +`--validate` transport round-trip) directly exercises the dispatch/combine path. + +**Transport-exercised results (DP-EP, the numbers that matter — results doc §1.1d–f):** +GSM8K through a real DP-EP server: **LL 0.856/0.898, HT 0.857/0.898** (flex/strict). +DP-EP eager throughput vs DeepEP after the perf iteration (Qwen3-30B-A3B, 8×GPU, total tok/s, +128/128 · 2048/128 · 128/2048): **FI-LL 9,088/23,106/5,825 (0.90/0.96/0.88× of DeepEP-LL)**; +**FI-HT 6,797/45,224/3,795 (1.19/1.27/0.84× of DeepEP-HT — ahead on 2 of 3 shapes)**. +The initial DP-EP pass was 2–6× behind; the closure came from three root-cause fixes +(batched-DP scheduler cap membership, HT recv-count trim, fleet-level host-path caches) — +full iteration log in results doc §1.1f. + +> 🛑 **The throughput/GSM8K/memory numbers below are historical and do NOT compare the two +> transports.** They were run with `--tensor-parallel-size 8` (`dp_size=1`), so vLLM took the +> `MoEPrepareAndFinalizeNoDPEPMonolithic` path — experts computed locally, reconciled by TP +> all-reduce — and **`--all2all-backend` was a no-op** (confirmed by nsys: identical kernels, only +> TP all-reduce, no dispatch/combine, for both FI-EP and DeepEP). The all2all transport is only +> selected when `dp_size > 1` (`fused_moe/config.py::use_all2all_kernels`). The tables above / +> in results doc §1.1e use `--data-parallel-size 8 --enable-expert-parallel` (verify the log says +> `Using FlashInferEPLL/HT…PrepareAndFinalize`, not `…Monolithic`). See runbook §3.0. + +**Throughput vs DeepEP** *(provisional — monolithic path, transport not exercised)* +(`vllm bench throughput --dataset-name random`, Qwen3-30B-A3B, 8-GPU, 1000 prompts; total tok/s): | ISL/OSL | FI-EP LL | FI-EP HT | DeepEP LL | DeepEP HT | |---|---|---|---|---| @@ -42,10 +61,10 @@ EP, 1000 prompts; total tok/s): | 2048 / 128 | 140,823 | 141,744 | 141,786 | 143,238 | | 128 / 2048 | 18,515 | 18,461 | 18,891 | 18,764 | -GSM8K accuracy is within noise across all four backends; **throughput is within ~1–2% of -DeepEP** across all three shapes. **Memory footprint is identical** across all four -(150.45 GiB / 6.57M-token KV cache at `--gpu-memory-utilization 0.9`) — EP backend choice is -memory-neutral. See +All four numbers land within ~1–2% — but that is because all four ran the *same* monolithic +TP-all-reduce path, not because the transports are equivalent. GSM8K accuracy likewise within +noise (monolithic path — end-to-end accuracy, not transport). Memory identical across all four +(150.45 GiB / 6.57M-token KV cache at `--gpu-memory-utilization 0.9`; monolithic path). See [`vllm_moe_ep_results_prenyx.md`](vllm_moe_ep_results_prenyx.md) for the full method, per-backend req/s, GSM8K-vs-DeepEP table, multi-node (2-node/16-GPU), and reproduction. **Not measured:** raw NCCL-EP (N/A upstream), TTFT/TPOT via `bench serve`. **2-node/16-GPU:** @@ -216,10 +235,16 @@ curl -s localhost:8000/v1/completions -H 'Content-Type: application/json' \ Pass criterion **≥ 0.80** (reference ~0.88). Run for each backend, LL and HT: ```bash lm_eval --model vllm \ - --model_args "pretrained=Qwen/Qwen3-30B-A3B,data_parallel_size=8,enable_expert_parallel=True,all2all_backend=flashinfer_ep_low_latency,trust_remote_code=True" \ + --model_args "pretrained=Qwen/Qwen3-30B-A3B,tensor_parallel_size=8,enable_expert_parallel=True,all2all_backend=flashinfer_ep_low_latency,trust_remote_code=True" \ --tasks gsm8k --num_fewshot 5 --batch_size auto # temperature 0, seed 42 (harness defaults for gsm8k are greedy) ``` +> **Accuracy gate only — does not exercise the transport.** lm_eval's `data_parallel_size` spawns +> independent replica engines (each `dp_size=1` ⇒ monolithic path), so it can't drive a unified EP +> group; keep `tensor_parallel_size=8`. The dispatch/combine transport is validated by §5.2 +> (`--validate`) and the nsys capture (runbook §3e). To exercise the transport end-to-end in vLLM, +> use the **server** path (§5.3, `vllm serve --data-parallel-size 8 --enable-expert-parallel`, +> which *does* build a real DP-EP deployment) or offline `torchrun … external_launcher` (runbook §3.0). ### 5.5 Multi-node (2 nodes, 16 GPU) Repeat 5.2–5.4 across 2 nodes. FlashInfer tests: `srun --nodes=2 --ntasks-per-node=1 @@ -263,7 +288,10 @@ torchrun --nproc_per_node=8 benchmarks/bench_moe_ep.py \ ### 7.2 End-to-end serving perf (the comparison matrix) Fixed load: **Qwen3-30B-A3B BF16, ISL/OSL 128/128, `max_concurrency=32`, -`NUM_PROMPTS=1000`**. For each cell, start `vllm serve` with the backend, then: +`NUM_PROMPTS=1000`**. Launch each `vllm serve` with **`--data-parallel-size 8 +--enable-expert-parallel`** (NOT TP-only) so the all2all transport is actually on the path — +otherwise every cell collapses to the identical monolithic path (§0 caveat / runbook §3.0). For +each cell, start `vllm serve` with the backend, then: ```bash vllm bench serve \ --model Qwen/Qwen3-30B-A3B \ diff --git a/docs/design_docs/vllm_moe_ep_results_prenyx.md b/docs/design_docs/vllm_moe_ep_results_prenyx.md new file mode 100644 index 00000000000..62afc035d18 --- /dev/null +++ b/docs/design_docs/vllm_moe_ep_results_prenyx.md @@ -0,0 +1,627 @@ +# vLLM ⇄ `flashinfer.moe_ep` — validated results & reproduction (Pre-Nyx) + +Measured results for the two vLLM all2all backends backed by `flashinfer.moe_ep` +(`flashinfer_ep_low_latency`, `flashinfer_ep_high_throughput`) and the exact steps to +reproduce them on the **Pre-Nyx** cluster (SLURM + pyxis/enroot, B200-class GPUs, CUDA 13.2). + +For the design/architecture and the full test/bench catalog see +[`vllm_moe_ep_integration.md`](vllm_moe_ep_integration.md). + +Code under test: +- FlashInfer: branch `feat/vllm-moe-ep-api` (`github.com/Anerudhan/flashinfer`, `cfc93a9c`). +- vLLM: branch `feat/flashinfer-ep-all2all` (`github.com/Anerudhan/vllm`, `f4e2618`). + +--- + +## 1. Results + +> 🛑 **IMPORTANT CAVEAT — the throughput/GSM8K/memory numbers below did NOT exercise the all-to-all +> transport.** They were collected with `--tensor-parallel-size 8 --enable-expert-parallel` +> (`dp_size = 1`). vLLM only routes MoE through the modular EP dispatch/combine path — the only +> path that uses `--all2all-backend` — when **`dp_size > 1`** (`config.py::use_all2all_kernels = +> dp_size > 1 and use_ep`). With `dp_size = 1` it falls back to +> `MoEPrepareAndFinalizeNoDPEPMonolithic`, where experts run locally and are reconciled by the +> ordinary **TP all-reduce**. nsys confirmed this: for *both* `flashinfer_ep_low_latency` and +> `deepep_low_latency` the kernel summaries were dominated by `multimem_all_reduce_kernel` / +> `vllm::cross_device_reduce_*` with **no** dispatch/combine kernels, and both logged +> `Using MoEPrepareAndFinalizeNoDPEPMonolithic`. **Consequence:** the ~1–2% FI-EP↔DeepEP closeness +> below reflects the *shared monolithic path*, NOT a comparison of the two transports. +> **Correctness (§1.1: GAP tests + `--validate`) is unaffected** — those exercise the transport +> directly. The transport-level comparison is being re-collected with `--data-parallel-size 8 +> --enable-expert-parallel` (see runbook §3.0); tables below are marked *provisional (monolithic)* +> until then. + +All runs: **single node, 8×GPU, 8-way expert parallel, CUDA 13.2**, base image +`nvcr.io/nvidia/pytorch:26.05-py3`, vLLM built from source (upstream `torch==2.11.0` pin +stripped via `use_existing_torch.py`; builds cleanly against NGC 26.05 torch). + +### 1.1 Correctness + +| Check | LL (`flashinfer_ep_low_latency`) | HT (`flashinfer_ep_high_throughput`) | +|---|---|---| +| `moe_ep` GAP 1/2/3 unit tests (mocked nccl, host-only) | 14 / 14 passed | — | +| EP dispatch+combine `--validate` @ world=8 | ✅ `ll/em per_rank=128 OK` | ✅ `ht/fl per_rank=4096 OK` | +| vLLM e2e smoke (OLMoE-1B-7B, greedy) | ✅ coherent | ✅ coherent (identical text) | +| **GSM8K 5-shot, Qwen3-30B-A3B** — flexible-extract | **0.8522** ± 0.0098 | **0.8575** ± 0.0096 | +| **GSM8K 5-shot, Qwen3-30B-A3B** — strict-match | **0.8939** ± 0.0085 | **0.8969** ± 0.0084 | + +Both clear the **≥ 0.80** gate and match the ~0.88 reference for Qwen3-30B-A3B. + +### 1.1b DP-EP transport-verified nsys capture (the all2all transport IS exercised) + +To exercise the all2all transport in vLLM (not just `--validate`), the MoE layer must take the +**modular EP** path, which requires `dp_size > 1` (§1 caveat). Offline `vllm bench throughput` +can't take `--data-parallel-size` directly — it must be launched under `torchrun --nproc_per_node=8` +with `--distributed-executor-backend external_launcher` (see runbook §3.0). Config: Qwen3-30B-A3B, +8×GPU, `--data-parallel-size 8 --enable-expert-parallel`, `--enforce-eager`, 8/64 prompts, nsys +`-t cuda,nvtx,nccl --trace-fork-before-exec=true`. + +With DP-EP the oracle switches the expert backend to a *batched* one and selects the EP +prepare/finalize — the log now prints (per backend): + +| Backend | expert backend | prepare/finalize (log) | all2all transport engaged? | +|---|---|---|---| +| `flashinfer_ep_low_latency` | `BATCHED_TRITON` | `Using FlashInferEPLLPrepareAndFinalize` | ✅ | +| `deepep_low_latency` | `BATCHED_TRITON` | `Using DeepEPLLPrepareAndFinalize` | ✅ | +| `deepep_high_throughput` | `TRITON` | `Using DeepEPHTPrepareAndFinalize` | ✅ | +| `flashinfer_ep_high_throughput` | `TRITON` | `Using FlashInferEPHTPrepareAndFinalize` | ✅ (after the 3 HT fixes below) | + +**Actual dispatch/combine kernels now on the GPU** (`nsys stats cuda_gpu_kern_sum`), which were +**absent** in the earlier `--tensor-parallel-size 8` (monolithic) capture: + +| Backend | dispatch kernel | combine kernel | extra | +|---|---|---|---| +| `flashinfer_ep_low_latency` (nccl.ep) | `nccl_ep::internode_ll::dispatch<…,(ncclEpLayout_t)1,…>` | `nccl_ep::internode_ll::combine<…,(ncclEpLayout_t)1>` | — | +| `deepep_low_latency` (DeepEP+NVSHMEM) | `deep_ep::legacy::internode_ll::dispatch<…>` | `deep_ep::legacy::internode_ll::combine<…>` | `nvshmemi_init_array_kernel<…>` | +| `deepep_high_throughput` (DeepEP intranode) | `deep_ep::legacy::intranode::notify_dispatch<8>` + `dispatch<8,768,8192>` | `intranode::cached_notify_combine<8>` + `combine` | `intranode::layout::get_dispatch_layout` | + +So the FlashInfer-EP LL vs DeepEP LL **transport difference is real and now visible**: +FlashInfer-EP LL runs the **`nccl_ep` GIN internode-LL** dispatch/combine; DeepEP LL runs the +**`deep_ep` NVSHMEM internode-LL** dispatch/combine (+ an NVSHMEM init kernel). Both LL backends +picked the *internode* LL path even on a single node. + +> ⚠ **Do NOT read these as a perf comparison.** The runs are tiny (8/64 prompts, decode-heavy, no +> CUDA graph), and the LL dispatch/combine kernels **busy-wait on the network**, so their +> `Total/Max` times are dominated by wait/sync (e.g. multi-second `Max` on an 8-prompt run) and are +> not throughput. These captures **prove the transport is exercised and identify the exact kernels**; +> a real perf comparison needs the DP-EP throughput sweep (runbook §3c) with CUDA graphs + larger +> batches. Raw dumps: `$RW/logs/kern_dpep_.txt`, `kern_a2a_.txt`, +> `nsys_dpep_.nsys-rep`. + +### 1.1c FlashInfer-EP HT under DP-EP — 3 bugs found & fixed (was SIGABRT, now works) + +Enabling the real DP-EP path surfaced three sequential HT-only bugs (all masked by the monolithic +path, which never builds the HT prepare/finalize). Each was root-caused from per-rank +`CUDA_LAUNCH_BLOCKING=1` stderr on Pre-Nyx and fixed; **HT now runs end-to-end and is GSM8K-validated** +(§1.1d): + +1. **SIGABRT at `ncclEpCreateGroup` (nccl_ep.cc:1253).** HT asserts + `max_dispatch_tokens_per_rank ≤ MAX_SUPPORTED_TOKENS_PER_RANK` (build-time `8192` in the nccl4py + wheel, `nccl/ep/include/nccl_ep/common.hpp`). vLLM sizes the HT fleet from + `moe.max_num_tokens = scheduler max_num_batched_tokens` (16384 here) → abort on all ranks. LL is + uncapped, so only HT hit it. **Fix (flashinfer `feat/nvep-default`, `60ff0fc1`):** clamp the HT + fleet's `max_tokens_per_rank` to the cap in `NcclEpFleet` (+ a clear `MoEEpConfigError` guard in + `_dispatch_ht`). Run HT with `--max-num-batched-tokens ≤ 8192`. +2. **Triton illegal-memory-access at `moe_align_block_size.py:101` (`expert_map[expert_ids]`).** + FlashInfer's FLAT recv gives **local** expert ids with `-1` for non-local/padding picks, but + vLLM's Standard experts feed `topk_ids` through `moe_align` + `expert_map` expecting **global** + ids (the skip is applied via `expert_map`, never as `-1` in `topk_ids`). **Fix (vLLM + `flashinfer_ep_ht.py`):** rebuild global ids from `expert_map`, remap `-1`→a non-owned global id + (which `expert_map` re-tags skipped) — matching the DeepEP HT contract. +3. **`finalize` over-strict assert.** Standard Triton experts declare `TopKWeightAndReduceNoOP` + (they already applied the dispatched routing weights and reduced their local picks), but finalize + only accepted `TopKWeightAndReduceDelegate`. FlashInfer HT combine applies **no** weights + (captured at dispatch) and only reduces per-rank partials across ranks → no double-weighting. + **Fix (vLLM `flashinfer_ep_ht.py`):** accept `NoOP` too. + +`deepep_high_throughput` with the same `TRITON` experts was unaffected (it returns global ids), which +is why only FlashInfer-EP HT hit bugs 2–3. + +### 1.1d GSM8K over a REAL DP-EP deployment (transport-exercised accuracy) + +Run via `vllm serve --data-parallel-size 8 --enable-expert-parallel --all2all-backend ` +(the online server path builds a genuine DP-EP engine, unlike lm_eval's own `data_parallel_size` +which spins up independent monolithic replicas) + `lm_eval --model local-completions` (needs +`pip install lm-eval[api]`). Both backends log the modular EP prepare/finalize and clear ≥0.80: + +| Backend | prepare/finalize | flexible-extract | strict-match | +|---|---|---|---| +| `flashinfer_ep_low_latency` | `FlashInferEPLLPrepareAndFinalize` | **0.8582** | **0.8976** | +| `flashinfer_ep_high_throughput` | `FlashInferEPHTPrepareAndFinalize` | **0.8461** | **0.8946** | + +On par with each other, with the monolithic-path GSM8K (§1.1), and with the ~0.88 Qwen3-30B-A3B +reference — i.e. the FlashInfer-EP dispatch/combine transport (both LL and HT) is numerically +correct end-to-end, not just at `--validate`. (This *replaces* the earlier "GSM8K is TP/accuracy-only" +caveat: it is now measured through the actual all2all transport.) + +### 1.1e DP-EP transport-exercised throughput (all 4 backends) + +Offline `vllm bench throughput` launched under `torchrun --nproc_per_node=8 +--distributed-executor-backend external_launcher --data-parallel-size 8 --enable-expert-parallel` +(runbook §3.0), Qwen3-30B-A3B, `--enforce-eager`, NP=256, `--dataset-name random`. HT backends +add `--max-num-batched-tokens 8192` (the nccl_ep HT cap); **LL backends must leave the flag +unset** (§1.1f fix #1 — the batched-DP 256 auto-cap must engage; DeepEP-LL also rejects 8192). +Deployment total = sum of the 8 DP ranks' `Throughput:` lines. Every cell verified +`Using PrepareAndFinalize` in the log. + +**Final (after the §1.1f optimizations; single same-day pass, total tok/s):** + +| Backend | 128/128 | 2048/128 | 128/2048 | +|---|---|---|---| +| `flashinfer_ep_low_latency` | **9,088** | **23,106** | **5,825** | +| `deepep_low_latency` | 10,116 | 24,013 | 6,595 | +| *FI-LL / DeepEP-LL* | *0.90×* | *0.96×* | *0.88×* | +| `flashinfer_ep_high_throughput` | **6,797** | **45,224** | 3,795 | +| `deepep_high_throughput` | 5,736 | 35,623 | 4,539 | +| *FI-HT / DeepEP-HT* | ***1.19×*** | ***1.27×*** | *0.84×* | + +**FlashInfer-EP HT is now ahead of DeepEP-HT by 19–27%** on the balanced and prefill-heavy +shapes; LL is within 4–12% of DeepEP-LL; the decode-heavy shape is within 12–16% for both modes. + +For the record, the **first** transport-exercised pass (before the §1.1f fixes) was +FI-LL 1,926/12,975/1,038 and FI-HT 2,761/21,011/1,485 — i.e. **2–6× behind DeepEP** — so the +gap closure came from the three root-cause fixes below, each GSM8K-gated. + +> ⚠ **Interpretation.** Eager-mode (no CUDA graphs) DP-EP runs over the Triton/batched-Triton +> expert backends with a small NP — chosen so the all2all transport is genuinely on the critical +> path and comparable across backends, **not** a production throughput number (the CUDA-graph +> monolithic numbers in §1.2 are much higher). CUDA-graph capture of dispatch/combine is the +> remaining big lever for both backends. + +### 1.1f Closing the 2–6× gap: root causes & fixes (perf iteration log) + +All found by diffing `nsys cuda_gpu_kern_sum` per-kernel medians between the FI-EP and DeepEP +runs, then `EP_PROFILE_HOST=1` host-phase timing. Each fix validated by GSM8K over the real +DP-EP server (§1.1d method). + +1. **[vLLM, 1 line] `flashinfer_ep_low_latency` was missing from `use_batched_dp_moe`** + (`vllm/config/parallel.py`). That property auto-caps the scheduler to the 256-token + batched-DP budget for BatchedExperts-format backends (`deepep_low_latency`, `nixl_ep`) — + without it FI-LL ran with the 8192-token offline default, so the padded + `[local_experts, max_tokens×world, N]` workspaces were **32× larger** than DeepEP-LL's: + every `fill_(0)` (394µs vs 13µs), full-workspace `act_and_mul` (2.07ms vs 36µs), the padded + batched-GEMM grid (1.29ms vs 256µs) and the LL transport slot buffers all paid it. + *Effect: FI-LL 1,926/12,975/1,038 → 8,854/22,939/5,416.* After the fix the compute kernels + are **byte-identical** to DeepEP-LL's (280µs/35.8µs/13.2µs medians on both). +2. **[vLLM adapter] HT recv-count trim.** FI-HT ran the whole Standard MoE stack (`moe_align` + 241µs, `count_and_sort` 344µs, 2×`fused_moe` ~1ms, `act_and_mul` 1.06ms, `moe_sum` 344µs) + over the **static 65,536-row** recv buffer every forward. The GAP-3 `recv_total` counter + (written by the HT metadata step at create_handle) now trims the compute view to + `round_up(actual,128)` rows (`.item()` host sync — eager-only); finalize copies the trimmed + expert output into a persistent full-size buffer (nccl.ep combine needs the address-stable + static staging; padding rows carry no routing state and are never sent). Also cached the + static local→global expert-id remap (was 2 `nonzero()` device syncs/layer/step). + *Effect: FI-HT 2,761/21,011/1,485 → 6,636/44,238/3,755 — ahead of DeepEP-HT on 2 of 3 shapes.* +3. **[flashinfer] Fleet-level host-path caches** (`nccl_ep/handle.py`). vLLM creates a fresh + `NcclEpHandle` every MoE layer×step (routing binds at `create_handle`), so the per-handle + `NV_FI_EP_FAST_PATH` caches never hit and each forward paid **~149µs host** (measured: + FFI descriptor builds 32.6+31.3µs, handle setup 35.9µs, create/destroy/dispatch/combine C + calls ~45µs). At decode the GPU is host-paced and nccl.ep's **fused send+recv** dispatch + kernel absorbs the inter-rank lag as in-kernel spin (median 256µs/launch, 33% of GPU time). + Fix: anchor recv buffers, counter tensors, static FFI tuples and a + `(data_ptr,dtype,shape)`-keyed wrap memo on the long-lived Fleet — **restricted to tensors + ≤2 MiB**, because the nccl.ep Tensor wrapper pins the torch tensor (memoizing large prefill + activations pinned GBs → OOM at `--gpu-memory-utilization 0.9`; found the hard way). + *Effect: host 149→119µs/layer; decode shapes +4–7%; final matrix above.* + +**Kernel-level notes for the remaining decode-heavy delta (~0.85×):** +- FI-LL fused `internode_ll::dispatch` (send+recv in one kernel) median 256µs vs DeepEP's + split send (14µs) + deferred recv hook (13µs): the fused kernel spins for the slowest peer, + absorbing per-layer host-path lag; the residual ≈ the remaining ~119µs host path (of which + ~45µs is create/destroy/dispatch/combine C calls — an nccl.ep handle-reuse/update API would + remove most of it). +- Remaining levers, in expected-impact order: CUDA-graph capture of dispatch/combine (removes + host pacing entirely), an nccl.ep API to reuse/update a handle instead of per-forward + create/destroy, splitting send/recv (staged mode) to overlap like DeepEP's hook, and + trimming the ~28µs combine-side FFI container builds. + +### 1.2 Throughput sweep — FlashInfer-EP vs DeepEP *(provisional — monolithic path, see §1 caveat)* + +`vllm bench throughput`, Qwen3-30B-A3B, 8-GPU EP, 1000 prompts, `--dataset-name random`, +`--enforce-eager`, `--max-model-len 4096`. DeepEP built into `vllm-fi-ep-deepep.sqsh` (§3.3). +Three ISL/OSL shapes: balanced (128/128), prefill-heavy (2048/128), decode-heavy (128/2048). + +| ISL / OSL | Backend | total tok/s | output tok/s | req/s | +|---|---|---|---|---| +| **128 / 128** | `flashinfer_ep_low_latency` | 32,506 | 16,253 | 126.98 | +| | `flashinfer_ep_high_throughput` | 32,050 | 16,025 | 125.20 | +| | `deepep_low_latency` | 32,535 | 16,267 | 127.09 | +| | `deepep_high_throughput` | 32,050 | 16,025 | 125.19 | +| **2048 / 128** (prefill-heavy) | `flashinfer_ep_low_latency` | 140,823 | 8,284 | 64.72 | +| | `flashinfer_ep_high_throughput` | 141,744 | 8,338 | 65.14 | +| | `deepep_low_latency` | 141,786 | 8,340 | 65.16 | +| | `deepep_high_throughput` | 143,238 | 8,426 | 65.83 | +| **128 / 2048** (decode-heavy) | `flashinfer_ep_low_latency` | 18,515 | 17,426 | 8.51 | +| | `flashinfer_ep_high_throughput` | 18,461 | 17,376 | 8.48 | +| | `deepep_low_latency` | 18,891 | 17,780 | 8.68 | +| | `deepep_high_throughput` | 18,764 | 17,660 | 8.62 | + +**Relative throughput (best-of-each-backend total tok/s, FlashInfer-EP ÷ DeepEP):** + +| ISL / OSL | FlashInfer-EP best | DeepEP best | FI-EP / DeepEP | +|---|---|---|---| +| 128 / 128 | 32,506 (LL) | 32,535 (LL) | **0.999** (−0.1%) | +| 2048 / 128 | 141,744 (HT) | 143,238 (HT) | **0.990** (−1.0%) | +| 128 / 2048 | 18,515 (LL) | 18,891 (LL) | **0.980** (−2.0%) | + +**Takeaway (provisional):** across all three shapes the four configs land within ~1–2% of each +other — but note (§1 caveat) that with `dp_size=1` all four ran the **same monolithic TP-all-reduce +path**, so this closeness is largely an artifact of the transport not being on the critical path, +not evidence that the FlashInfer-EP and DeepEP dispatch/combine transports perform equivalently. +The genuine transport comparison requires the `--data-parallel-size 8` re-run. (LL vs HT ordering +within a backend is likewise not meaningful here.) + +> ⚠ **Correction:** earlier throughput numbers reported for this work (~99–103k tok/s +> "128/128") were actually the vLLM **`sonnet` default dataset (~1024 in / 128 out)** — +> `--input-len`/`--output-len` are ignored unless `--dataset-name random` is passed. The table +> above is the corrected sweep (token counts verified: 128/128→128k+128k, 2048/128→2048k+128k, +> 128/2048→128k+2048k). GSM8K numbers were unaffected (lm_eval uses its own data). + +### 1.3 GSM8K — FlashInfer-EP vs DeepEP (accuracy on par) *(monolithic path — accuracy only)* + +| Backend | flexible-extract | strict-match | +|---|---|---| +| `flashinfer_ep_low_latency` | 0.8522 | 0.8939 | +| `flashinfer_ep_high_throughput` | 0.8575 | 0.8969 | +| `deepep_low_latency` | 0.8514 | 0.8931 | +| `deepep_high_throughput` | 0.8544 | 0.8946 | + +All four within run-to-run noise. Note this run used the monolithic path (§1 caveat), so it +validates end-to-end model accuracy but does **not** compare the dispatch/combine transports; the +transport is validated separately by the `--validate` correctness checks in §1.1. + +### 1.3b Memory footprint *(monolithic path, see §1 caveat)* + +From vLLM's engine-init memory profiling (Qwen3-30B-A3B, 8-GPU EP, `--gpu-memory-utilization +0.9`, `--max-model-len 4096`), **all four backends are identical**: + +| Backend | Available KV cache | KV cache size | Max concurrency @ 4096 tok | +|---|---|---|---| +| `flashinfer_ep_low_latency` | 150.45 GiB | 6,573,312 tokens | 1604.8× | +| `flashinfer_ep_high_throughput` | 150.45 GiB | 6,573,312 tokens | 1604.8× | +| `deepep_low_latency` | 150.45 GiB | 6,573,312 tokens | 1604.8× | +| `deepep_high_throughput` | 150.45 GiB | 6,573,312 tokens | 1604.8× | + +**EP backend choice is memory-neutral** here — the dispatch/combine transport buffers do not +measurably reduce the usable KV-cache budget on a B200 (≈180 GiB HBM); the model weights + +activation + transport all fit in the non-KV reservation identically for every backend, leaving +the same 150.45 GiB for KV cache. (vLLM in this build does not emit a grep-able +weights/non-torch/activation split; the identical KV-cache size is the operative footprint +metric.) + +### 1.4 Multi-node (2-node / 16-GPU) — plumbing works, cross-node init blocked (environmental) + +The 2-node path (`benchmarks/_perf_2node.sh`, §4.7) **stands up correctly**: a Ray cluster forms +across both containers (`cluster GPUs so far=16`) and `vllm bench throughput +--data-parallel-size 16 --enable-expert-parallel --distributed-executor-backend ray` with the EP +backend launches. But the vLLM **engine-core init then stalls** (>40 min, no throughput; no error +surfaced in the driver or Ray worker logs). (Original runs used `--tensor-parallel-size 16`, which +would additionally take the monolithic path per §1 caveat; the DP-EP form is the correct config +and stalls at the same cross-node init step below regardless.) + +**Isolation test — this is NOT the EP integration.** A plain 2-node run with +`--tensor-parallel-size 16` and **no** `--enable-expert-parallel` / all2all backend **stalls +identically** at engine init. So the blocker is the cluster's **cross-node vLLM/NCCL bring-up** +(fabric / NCCL cross-node env / ray-executor init), independent of FlashInfer-EP or DeepEP. +Resolving it is cluster-config work (e.g. `NCCL_SOCKET_IFNAME` / `NCCL_IB_HCA` tuning, verifying +IB reachability between the allocated nodes) — out of scope for this integration, which is fully +validated at 8-GPU single node. + +### 1.5 Not measured + +- **Raw NCCL-EP backend** — exists only in the GitLab `vllm-nccl-moe-integration` fork, not in + upstream vLLM; N/A here. +- **`vllm bench serve` TTFT/TPOT** — the server path hit a >20-min per-worker FlashInfer + cubin-download + JIT startup (plus an shm-broadcast wait) on this image; a serving-startup + issue, not the EP backend (GSM8K drives the same path). Offline `vllm bench throughput` was + used for the headline numbers. + +--- + +## 2. Environment + +- Login: `ssh prenyx` (→ `login-prenyx`, MFA/GSSAPI, `ProxyJump`; a persistent ControlMaster + socket lets subsequent non-interactive `ssh prenyx ''` reuse the session). +- SLURM: `--account=coreai_libraries_cudnn --partition=batch`; **whole-node allocations — do + NOT pass `--gres`** (`Invalid generic resource` otherwise). +- Containers: pyxis/enroot. `--container-name` does **not** persist across separate `srun` + jobs — pass `--container-image=<...>.sqsh` every time (squashfs mounts fast). +- Work dir (shared lustre, mounted `/host` in-container): + `RW=/lustre/fsw/coreai_libraries_cudnn/agopal-moe-ep`. + +```bash +RW=/lustre/fsw/coreai_libraries_cudnn/agopal-moe-ep +``` + +--- + +## 3. Build the container images (one-time, **no Docker** — pyxis/enroot) + +Pre-Nyx has no Docker daemon; images are `.sqsh` files built with **pyxis/enroot** via +`srun --container-save`. Conventions used throughout: +- Enroot registry syntax **`nvcr.io#nvidia/`** (not the docker-style `nvcr.io/nvidia/`). +- **`--container-writable`** — required so in-container `apt`/`pip` installs are captured by + `--container-save`. +- **Whole-node** allocation (no `--gres`); `--container-mounts=$RW:/host` (shared lustre work dir). +- `docker/Dockerfile.*` are the canonical spec (usable only on a machine with Docker). + +### 3.1 FlashInfer-EP base (`flashinfer-ep-pt2605.sqsh`) + +> For the results in this doc the base `.sqsh` was **pre-built and reused** (not rebuilt this +> run). The recipe below is how to (re)create it. Mirrors the real `build.sbatch` artifact. + +Clone FlashInfer at `feat/vllm-moe-ep-api` into `$RW/flashinfer`, then (the base build is long +— run it as an sbatch or a plain `srun`): + +```bash +srun --account=coreai_libraries_cudnn --partition=batch -N1 --ntasks-per-node=1 --time=03:00:00 \ + --container-image="nvcr.io#nvidia/pytorch:26.05-py3" --container-writable \ + --container-save=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host \ + bash -lc 'cd /host/flashinfer && bash docker/install/build_flashinfer_ep_pytorch.sh' +``` +As an sbatch wrapper (mirrors `$RW/build.sbatch`): +```bash +#!/bin/bash +#SBATCH -A coreai_libraries_cudnn -p batch -N1 --time=03:00:00 -J fi_ep_build +RW=/lustre/fsw/coreai_libraries_cudnn/agopal-moe-ep +srun --container-image="nvcr.io#nvidia/pytorch:26.05-py3" --container-writable \ + --container-save=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host \ + bash -lc 'cd /host/flashinfer && bash docker/install/build_flashinfer_ep_pytorch.sh' +``` +The install script pins `nvidia-nccl-cu13==2.30.7`, `nccl4py[cu13]==0.3.1`, `cuda-core==1.0.1`, +`cuda-bindings==13.2.0` and runs `BUILD_NCCL_EP=1 pip install -e ".[nvep]"` (editable, from +`/host/flashinfer`). See `docker/install/build_flashinfer_ep_pytorch.sh` / +`docker/Dockerfile.flashinfer-ep-pytorch`. + +### 3.2 vLLM-from-source image (`vllm-flashinfer-ep.sqsh`) + +Clone vLLM at `feat/flashinfer-ep-all2all` into `$RW/vllm`. Build script `build_vllm.sh` +(strip torch pins, install build-deps under `--no-build-isolation`, add a Rust toolchain since +vLLM bundles `rust/Cargo.toml`, pin torch to the base version): + +```bash +#!/bin/bash +set -eo pipefail +cd /host/vllm +python use_existing_torch.py +TORCH_VER=$(python -c 'import torch;print(torch.__version__.split("+")[0])') +echo "torch==$TORCH_VER" > /tmp/tc.txt +command -v cargo >/dev/null 2>&1 || \ + curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal +export PATH="$HOME/.cargo/bin:$PATH" +PIP_CONSTRAINT=/tmp/tc.txt pip install --no-cache-dir -r requirements/build/cuda.txt +MAX_JOBS=32 VLLM_USE_PRECOMPILED=0 \ + PIP_CONSTRAINT=/tmp/tc.txt pip install --no-cache-dir --no-build-isolation -e . -v +# vLLM's deps pull `flashinfer-python` from PyPI and shadow the branch editable (its moe_ep +# lacks EpLayout/FleetAlgoKnobAllocator) — restore the branch editable. +pip install --no-cache-dir --no-build-isolation --no-deps -e /host/flashinfer +python -c 'import vllm, flashinfer; from flashinfer.moe_ep import EpLayout, FleetAlgoKnobAllocator; \ + print("vllm", vllm.__version__, "| flashinfer", flashinfer.__file__)' +``` +```bash +srun --account=coreai_libraries_cudnn --partition=batch -N1 --ntasks-per-node=1 --time=03:00:00 \ + --container-image=$RW/flashinfer-ep-pt2605.sqsh --container-writable \ + --container-save=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ + bash /host/build_vllm.sh +``` +(`--container-writable` is required so the vLLM install is captured by `--container-save`; the +base image here is a local `.sqsh` file, so no `nvcr.io#` registry prefix.) Equivalent one-shot +on a Docker host: `docker/Dockerfile.vllm-flashinfer-ep`. + +### 3.3 DeepEP image (`vllm-fi-ep-deepep.sqsh`) — for the comparison + +Layer DeepEP + NVSHMEM on the vLLM image via vLLM's own installer. **Two fixes are required +on CUDA 13.2:** (a) `UV_BREAK_SYSTEM_PACKAGES=1` (the installer's `uv pip install --system` +trips PEP-668), and (b) **`TORCH_CUDA_ARCH_LIST=10.0a`** — otherwise DeepEP compiles for +`sm_75` and `ptxas` fails (`Feature 'elect'/'mbarrier'/'cp.async.bulk' requires .target +sm_90 or higher`) because its kernels are Hopper/Blackwell-only. + +```bash +#!/bin/bash # build_deepep.sh +export UV_BREAK_SYSTEM_PACKAGES=1 PIP_BREAK_SYSTEM_PACKAGES=1 UV_SYSTEM_PYTHON=1 +set -eo pipefail +export TORCH_CUDA_ARCH_LIST="10.0a" # B200; sm_90+ features +command -v uv >/dev/null 2>&1 || pip install -q uv # installer uses `uv pip install --system` +export PATH="$HOME/.local/bin:/usr/local/bin:$PATH" # ensure uv is on PATH +TORCH_VER=$(python -c 'import torch;print(torch.__version__.split("+")[0])') +echo "torch==$TORCH_VER" > /tmp/tc.txt +PIP_CONSTRAINT=/tmp/tc.txt UV_CONSTRAINT=/tmp/tc.txt \ + bash /host/vllm/tools/ep_kernels/install_python_libraries.sh --workspace /host/ep_kernels_workspace +# DeepEP asserts torch nccl == the nvidia-nccl wheel (2.30.7); the NGC image also ships a system +# libnccl 2.30.4 that torch loads by default. Force the wheel + bake into /etc/profile.d so +# `bash -lc` runtime runs load 2.30.7 too. +NCCL_LIB=$(python -c 'import nvidia.nccl,os;print(os.path.join(list(nvidia.nccl.__path__)[0],"lib"))') +echo "export LD_LIBRARY_PATH=\"$NCCL_LIB:\${LD_LIBRARY_PATH:-}\"" > /etc/profile.d/zz_nccl_wheel.sh +export LD_LIBRARY_PATH="$NCCL_LIB:${LD_LIBRARY_PATH:-}" +# `deep_ep OK` = DeepEP's assert (loaded libnccl == wheel) passed. Don't judge by +# torch.cuda.nccl.version() — that's torch's BUILD-time NCCL (cosmetic), not the loaded .so. +python -c 'import ctypes, torch; torch.cuda.init(); import deep_ep; \ + m=sorted({l.split()[-1] for l in open("/proc/self/maps") if "libnccl.so" in l}); \ + lib=ctypes.CDLL(m[0]); v=ctypes.c_int(); lib.ncclGetVersion(ctypes.byref(v)); \ + print("deep_ep OK; loaded", m, "ncclGetVersion", v.value)' +``` +```bash +srun --account=coreai_libraries_cudnn --partition=batch -N1 --ntasks-per-node=1 --time=02:00:00 \ + --container-image=$RW/vllm-flashinfer-ep.sqsh --container-writable \ + --container-save=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ + bash /host/build_deepep.sh +``` +Installs NVSHMEM 3.3.24 + `DeepEP@d4f41e4e93` (with the installer's CUDA-13 cccl patch). +The DeepEP runs in §4.4/§4.5 use this image with `BACKEND=deepep_low_latency` / +`deepep_high_throughput`. + +> **Getting the FlashInfer branch code into the vLLM image at runtime.** `moe_ep` is pure +> Python, editable-installed from `/host/flashinfer`. To run branch code without rebuilding, +> mount a checkout of `feat/vllm-moe-ep-api` and prepend it to `PYTHONPATH`. In this +> validation a **git worktree** `$RW/fi-vllmep` was used (to avoid touching a dirty primary +> checkout); it needs `3rdparty/{cutlass,spdlog,cccl}` present (init or symlink to a populated +> checkout), `flashinfer/_build_meta.py`, and the `flashinfer/data/*` symlinks. Simplest for a +> fresh repro: check the branch out directly in `$RW/flashinfer` before building 3.1 so the +> editable install already points at branch code, and skip `PYTHONPATH`. + +--- + +## 4. Reproduce each result + +Common env inside every run: `export PYTHONPATH=/host/fi-vllmep` (only if using the worktree), +`HF_HOME=/host/hf_cache`, `FLASHINFER_WORKSPACE_BASE=/host/fi_cache`, +`FLASHINFER_CUBIN_DIR=/host/fi_cubins` (persist JIT + cubins across runs), `NCCL_GIN_TYPE=3`. + +### 4.1 GAP unit tests (host-only, ~1 min) +```bash +srun ... --container-image=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host bash -lc ' + export PYTHONPATH=/host/fi-vllmep; python -m pip install -q pytest + cd /host/fi-vllmep && python -m pytest \ + tests/moe_ep/nccl_ep/test_gaps_mock.py tests/moe_ep/nccl_ep/test_fleet_mock.py -q' +``` + +### 4.2 EP dispatch+combine correctness @ world=8 (`--validate`) +Per-rank runner `benchmarks/run_ep_matrix_one_pt.sh` (points `cd /host/flashinfer`; for the +worktree use a copy that `cd`s to `/host/fi-vllmep` + sets `PYTHONPATH`). Launch **8 tasks/node** +with a `file://` rendezvous: +```bash +srun ... --ntasks-per-node=8 --container-image=$RW/flashinfer-ep-pt2605.sqsh \ + --container-mounts=$RW:/host bash -lc \ + 'EP_SYNC=/host/sync_ht NCCL_GIN_TYPE=3 bash /host//benchmarks/run_ep_matrix_one_pt.sh \ + --algorithm ht --layout fl --tokens 4096 --hidden 7168 --top-k 8 --experts 256 \ + --warmup 5 --iters 10 --validate' +# LL: --algorithm ll --layout em --tokens 128 (same COMMON args) +``` +Expect: `[validate] ht/fl world=8 per_rank=4096 dispatch+combine OK` (and `ll/em ... OK`). + +> The pytest `tests/moe_ep/test_moe_ep_ht_correctness.py` (launched via `torchrun`) hangs on a +> default-PG collective on this image — use the comm-matrix `file://` `--validate` path above. + +### 4.3 vLLM e2e smoke (both backends) +Run a **file-based** program (vLLM uses `spawn`; a heredoc/`stdin` program fails the workers): +```python +# _vllm_smoke_prog.py +import os +from vllm import LLM, SamplingParams +def main(): + llm = LLM(model=os.environ["MODEL"], tensor_parallel_size=8, + enable_expert_parallel=True, all2all_backend=os.environ["BACKEND"], + trust_remote_code=True, enforce_eager=True, max_model_len=2048) + out = llm.generate(["San Francisco is a"], SamplingParams(max_tokens=32, temperature=0.0)) + print(repr(out[0].outputs[0].text)) +if __name__ == "__main__": + main() +``` +```bash +srun ... --container-image=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ + --export=ALL,MODEL=allenai/OLMoE-1B-7B-0924,BACKEND=flashinfer_ep_low_latency bash -lc ' + export PYTHONPATH=/host/fi-vllmep HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 \ + FLASHINFER_WORKSPACE_BASE=/host/fi_cache + python -u /host/fi-vllmep/benchmarks/_vllm_smoke_prog.py' +# repeat with BACKEND=flashinfer_ep_high_throughput +``` + +### 4.4 GSM8K 5-shot (accuracy gate) +```bash +srun ... --container-image=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ + --export=ALL,BACKEND=flashinfer_ep_low_latency bash -lc ' + export PYTHONPATH=/host/fi-vllmep HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 \ + FLASHINFER_WORKSPACE_BASE=/host/fi_cache + python -m pip install -q lm_eval + lm_eval --model vllm --tasks gsm8k --num_fewshot 5 --batch_size auto \ + --model_args pretrained=Qwen/Qwen3-30B-A3B,tensor_parallel_size=8,enable_expert_parallel=True,all2all_backend=$BACKEND,trust_remote_code=True,max_model_len=4096,enforce_eager=True' +# repeat with BACKEND=flashinfer_ep_high_throughput +``` +> ⚠ **GSM8K is an accuracy gate only — it does not exercise the all2all transport.** lm_eval's own +> `data_parallel_size` launches independent replica engines (each `dp_size=1` ⇒ monolithic path), +> so there is no unified EP group to dispatch through; keep `tensor_parallel_size=8`. The transport +> is validated separately by the §4.2 `--validate` round-trip and the nsys dispatch/combine capture +> (runbook §3e). + +### 4.5 Throughput sweep +**Use `--dataset-name random` with `--random-input-len`/`--random-output-len`** — otherwise +`vllm bench throughput` falls back to the `sonnet` dataset (~1024/128) and silently ignores +`--input-len`/`--output-len`. +```bash +# ISL/OSL ∈ {128/128, 2048/128, 128/2048}; backend ∈ {flashinfer_ep_low_latency, ...high_throughput} +srun ... --container-image=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ + --export=ALL,BACKEND=flashinfer_ep_low_latency,ISL=2048,OSL=128 bash -lc ' + export PYTHONPATH=/host/fi-vllmep HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 \ + FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins + cd /tmp; torchrun --nproc_per_node=8 /host/dprun/driver.py --model Qwen/Qwen3-30B-A3B \ + --dataset-name random --random-input-len $ISL --random-output-len $OSL --num-prompts 1000 \ + --data-parallel-size 8 --distributed-executor-backend external_launcher \ + --enable-expert-parallel --all2all-backend $BACKEND \ + --trust-remote-code --max-model-len 4096 --enforce-eager' +``` +> ⚠ **DP-EP via `torchrun` + `external_launcher`** (needs `$RW/dprun/driver.py`, runbook §3.0) — +> plain `vllm bench throughput --data-parallel-size 8` errors offline. With `--tensor-parallel-size +> 8` you get the monolithic path (§1 caveat). Each of the 8 ranks prints its own `Throughput:`; +> the deployment total ≈ their sum. Confirm the log says `Using FlashInferEP…/DeepEP…`, not +> `…Monolithic`. + +### 4.6 DeepEP comparison (§1.2/§1.3) +Identical to §4.4/§4.5 but use the DeepEP image and DeepEP backend names: +`--container-image=$RW/vllm-fi-ep-deepep.sqsh` and +`BACKEND=deepep_low_latency` / `deepep_high_throughput`. + +### 4.7 Multi-node (2-node / 16-GPU) +`benchmarks/_perf_2node.sh` stands up a Ray cluster across the two nodes' containers (rank-0 +`ray start --head`, rank-1 `ray start --address=` via a shared `/host/ray_head_ip.$JOBID` +file), waits for 16 GPUs, then runs `vllm bench throughput --data-parallel-size 16 +--enable-expert-parallel --distributed-executor-backend ray` with `NCCL_MNNVL_ENABLE=1` for the +cross-node EP fabric. (Data-parallel EP, not TP-16 — see §1 caveat / runbook §3.0.) +```bash +srun --account=coreai_libraries_cudnn --partition=batch -N2 --ntasks-per-node=1 \ + --container-image=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ + --export=ALL,BACKEND=flashinfer_ep_low_latency,ISL=128,OSL=128 \ + bash /host/fi-vllmep/benchmarks/_perf_2node.sh +``` +Warm `/host/fi_cubins` first (from the single-node runs) — a cold cubin cache makes the 16-way +init stall for tens of minutes while every worker downloads FlashInfer cubins. + +### 4.8 Memory footprint (§1.3b) +Same as §4.5 but `--num-prompts 8` (init does the memory profiling regardless) and grep the +KV-cache line from the **full** stream (don't `tail`-truncate): +```bash +cd /tmp; torchrun --nproc_per_node=8 /host/dprun/driver.py --model Qwen/Qwen3-30B-A3B \ + --dataset-name random --random-input-len 128 --random-output-len 128 --num-prompts 8 \ + --data-parallel-size 8 --distributed-executor-backend external_launcher \ + --enable-expert-parallel --all2all-backend $BACKEND \ + --gpu-memory-utilization 0.9 --trust-remote-code --max-model-len 4096 --enforce-eager 2>&1 \ + | grep -iE "Available KV cache|GPU KV cache size|Maximum concurrency|Using .*PrepareAndFinalize" +``` + +--- + +## 5. Gotchas (learned during this validation) + +- **⚠ TP-only EP silently disables the all2all backend (biggest gotcha).** vLLM only takes the + modular EP dispatch/combine path when `dp_size > 1` + (`fused_moe/config.py::use_all2all_kernels = dp_size > 1 and use_ep`). Running + `--tensor-parallel-size 8 --enable-expert-parallel` (dp_size=1) picks + `MoEPrepareAndFinalizeNoDPEPMonolithic` — experts run locally, reconciled by TP all-reduce, and + `--all2all-backend` is a **no-op**. Use `--data-parallel-size 8 --enable-expert-parallel` + instead. **Always confirm** the log prints `Using FlashInferEPLL/HT…` or `DeepEPLL/HT…` + `PrepareAndFinalize`, never `…Monolithic`. (This is why the first §1.2/§1.3 sweep showed + FI-EP ≈ DeepEP and identical nsys kernels — the transport was never on the GPU.) +- **No `--gres`** on this cluster (whole-node); `--container-name` doesn't persist across jobs. +- **`NCCL_GIN_TYPE=3`** for the EP GIN transport; multi-node also needs `NCCL_MNNVL_ENABLE=1`. +- **vLLM `spawn`** re-imports the main module → run a real `.py` file, never a heredoc/stdin. +- **`pytest` / `lm_eval`** are not in the image → `pip install` them in the job. +- Don't recursive-glob `/usr/**` inside the container (pathologically slow → looks like a hang). +- Persist JIT to `/host` (`FLASHINFER_WORKSPACE_BASE`) and cubins + (`FLASHINFER_CUBIN_DIR`) so the ~267-unit CUTLASS MoE compile happens once. +- First vLLM forward JIT-compiles the FlashInfer CUTLASS MoE expert kernel — needs the `cccl` + submodule present in the checkout used at runtime. +- **DeepEP build on CUDA 13.2:** the vLLM installer calls `uv pip install --system`, so `uv` + must be present — `pip install uv` first (the NGC/vLLM image has `pip`, not `uv`). Also set + `UV_BREAK_SYSTEM_PACKAGES=1` (PEP-668) **and** `TORCH_CUDA_ARCH_LIST=10.0a` — without the arch + it builds for `sm_75` and `ptxas` rejects the Hopper/Blackwell-only features (`elect`, + `mbarrier`, `cp.async.bulk`). +- **vLLM shadows the branch flashinfer:** `pip install -e .` (vLLM) pulls `flashinfer-python` + from PyPI and uninstalls the branch editable → `ImportError: cannot import name 'EpLayout'` + (and the `flashinfer_ep_*` backend breaks: no `FleetAlgoKnobAllocator`). Fix: re-run + `pip install --no-build-isolation --no-deps -e /host/flashinfer` after the vLLM install (or + set `PYTHONPATH=/host/flashinfer` on every run). Verify `flashinfer.__file__` → + `/host/flashinfer/...`. +- **DeepEP NCCL-version assert:** DeepEP requires torch's loaded NCCL to equal the nvidia-nccl + wheel (2.30.7), but the NGC image also has a system `libnccl.so.2.30.4` that torch loads by + default → `AssertionError: Invalid NCCL versions: ...2.30.4 (loaded) v.s. ...wheel...`. Fix: + prepend the wheel's `nvidia/nccl/lib` to `LD_LIBRARY_PATH` (2.30.x is ABI-compatible) and bake + it into `/etc/profile.d` so `bash -lc` runs inherit it. diff --git a/docs/design_docs/vllm_moe_ep_runbook.md b/docs/design_docs/vllm_moe_ep_runbook.md new file mode 100644 index 00000000000..3d3eee06e02 --- /dev/null +++ b/docs/design_docs/vllm_moe_ep_runbook.md @@ -0,0 +1,566 @@ +# Runbook — NCCL-EP (FlashInfer) vs DeepEP on Pre-Nyx (pyxis/enroot, no Docker) + +Linear, copy-paste steps to build the container images **with pyxis (no Docker daemon)** and run +the **NCCL-EP vs DeepEP** comparison in vLLM, on **1 node (8 GPU)** and **2 nodes (16 GPU)**. + +- **NCCL-EP** = the `nccl.ep`-backed vLLM all2all backends `flashinfer_ep_low_latency` / + `flashinfer_ep_high_throughput` (FlashInfer branch **`feat/nvep-default`** ≥ `fa09bc46` + + vLLM branch **`feat/flashinfer-ep-all2all`** ≥ `ab1415e` — these commits carry the HT + token-cap clamp, the HT recv-trim, the batched-DP cap membership and the fleet host-path + caches; older refs reproduce the pre-optimization 2–6× gap). The standalone raw-`nccl.ep` + backend is **not** in upstream vLLM (GitLab-fork only), so it's out of scope here. +- **DeepEP** = `deepep_low_latency` / `deepep_high_throughput`. +- Model: `Qwen/Qwen3-30B-A3B` (128 experts, bf16). Base: `nvcr.io/nvidia/pytorch:26.05-py3` + (CUDA 13.2 — required; older stacks abort cross-node HT at `nccl_ep.cc:2884`). + +Everything runs from a shared-FS work dir mounted `/host` inside the container. + +```bash +# --- run once per shell --- +RW=/lustre/fsw/coreai_libraries_cudnn/agopal/agopal-moe-ep-verif # a FRESH shared-FS work dir +ACCT=coreai_libraries_cudnn ; PART=batch +mkdir -p $RW/logs +``` + +Starting **from scratch** in a fresh `$RW`? Run §1 (clone) → §2 (build all three images, +nothing is reused) → §3 (single-node runs) → §4 (multi-node). Everything below is written to be +copy-pasted top-to-bottom into a shell that has already `export`ed `RW`/`ACCT`/`PART` above. + +--- + +## 0. Why pyxis, not Docker + +Pre-Nyx login/compute nodes have **no Docker daemon**. Container images are **enroot squashfs +(`.sqsh`) files**, built by running the install steps *inside* a container under `srun` and +snapshotting it with `--container-save`: + +- `--container-image="nvcr.io#nvidia/"` — enroot registry syntax (note the **`#`**, not + `/`). A local image is just a path: `--container-image=$RW/foo.sqsh`. +- `--container-writable` — **required** so `apt`/`pip`/build outputs are captured by the save. +- `--container-save=$RW/out.sqsh` — writes the resulting image. +- `--container-mounts=$RW:/host` — the shared FS shows up at `/host` in the container. +- Whole-node allocation only — **do not pass `--gres`** (rejected on this cluster). +- A `--container-name` does **not** persist across separate `srun` jobs; always pass + `--container-image=<...>.sqsh`. + +--- + +## 1. Prerequisites — create the fresh dir and clone the repos + +```bash +# (RW/ACCT/PART already exported above — e.g. RW=.../agopal-moe-ep-verif) +mkdir -p $RW/logs +# Pinned refs — these carry ALL the DP-EP fixes/optimizations the reference numbers were +# measured with (HT clamp+trim, batched-DP cap membership, fleet host-path caches): +git clone -b feat/nvep-default https://github.com/Anerudhan/flashinfer.git $RW/flashinfer +git clone -b feat/flashinfer-ep-all2all https://github.com/Anerudhan/vllm.git $RW/vllm +git -C $RW/flashinfer submodule update --init --recursive # cutlass, cccl, spdlog, nccl +# sanity: the perf-critical commits must be present +git -C $RW/flashinfer merge-base --is-ancestor fa09bc46 HEAD && echo "flashinfer ref OK" +git -C $RW/vllm merge-base --is-ancestor ab1415e HEAD && echo "vllm ref OK" +``` + +--- + +## 2. Build the three images with pyxis + +> **Rebuilds vs `git pull`:** the images install flashinfer and vLLM as *editable* installs +> pointing at `/host/flashinfer` / `/host/vllm` — i.e. at **your `$RW` clones, resolved at +> runtime**. Python-only changes (all the perf fixes in the pinned refs are Python) take effect +> by just updating the clones (§1); **no image rebuild needed**. Rebuild only for dependency/ +> native changes (e.g. a new nccl4py/nvidia-nccl pin). + +### 2a. FlashInfer-EP base — `flashinfer-ep-pt2605.sqsh` +Runs `docker/install/build_flashinfer_ep_pytorch.sh` (pins nvidia-nccl-cu13 2.30.7 / nccl4py +0.3.1 / cuda-core 1.0.1 / cuda-bindings 13.2.0, then `BUILD_NCCL_EP=1 pip install -e .[nvep]`). + +```bash +srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=03:00:00 \ + --container-image="nvcr.io#nvidia/pytorch:26.05-py3" --container-writable \ + --container-save=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host \ + bash -lc 'cd /host/flashinfer && bash docker/install/build_flashinfer_ep_pytorch.sh' +# sanity: +srun -A $ACCT -p $PART -N1 --container-image=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host \ + bash -lc "python -c \"from flashinfer.moe_ep import available_backends; print(available_backends())\"" +# expect: ['nccl_ep'] +``` + +### 2b. vLLM from source — `vllm-flashinfer-ep.sqsh` +Create `$RW/build_vllm.sh`: +```bash +cat > $RW/build_vllm.sh <<'EOS' +#!/bin/bash +set -eo pipefail +cd /host/vllm +python use_existing_torch.py # strip torch==2.11 pin -> use NGC torch +TORCH_VER=$(python -c 'import torch;print(torch.__version__.split("+")[0])'); echo "torch==$TORCH_VER" > /tmp/tc.txt +command -v cargo >/dev/null 2>&1 || curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal +export PATH="$HOME/.cargo/bin:$PATH" # vLLM bundles a Rust crate +PIP_CONSTRAINT=/tmp/tc.txt pip install --no-cache-dir -r requirements/build/cuda.txt +MAX_JOBS=32 VLLM_USE_PRECOMPILED=0 PIP_CONSTRAINT=/tmp/tc.txt \ + pip install --no-cache-dir --no-build-isolation -e . -v +# vLLM's deps pull `flashinfer-python` from PyPI, which UNINSTALLS/shadows our branch editable +# (the PyPI moe_ep lacks EpLayout / FleetAlgoKnobAllocator). Restore the branch editable so +# `import flashinfer` resolves to /host/flashinfer at runtime. +pip install --no-cache-dir --no-build-isolation --no-deps -e /host/flashinfer +python -c 'import vllm, flashinfer; from flashinfer.moe_ep import EpLayout, FleetAlgoKnobAllocator; \ + print("vllm", vllm.__version__, "| flashinfer", flashinfer.__file__)' +EOS +``` +Build: +```bash +srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=03:00:00 \ + --container-image=$RW/flashinfer-ep-pt2605.sqsh --container-writable \ + --container-save=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ + bash /host/build_vllm.sh +``` + +### 2c. DeepEP — `vllm-fi-ep-deepep.sqsh` +Four things to get right on CUDA 13.2: (1) `uv` must be installed (the vLLM installer runs +`uv pip install --system`); (2) `UV_BREAK_SYSTEM_PACKAGES=1` (PEP-668); (3) +**`TORCH_CUDA_ARCH_LIST=10.0a`** (else DeepEP builds `sm_75` and `ptxas` rejects its +`elect`/`mbarrier`/`cp.async.bulk` kernels); and (4) **torch must load the 2.30.7 wheel +libnccl** — DeepEP asserts *torch nccl == the nvidia-nccl wheel*, but the NGC image also has a +system libnccl 2.30.4 that torch loads by default, so the script forces the wheel onto +`LD_LIBRARY_PATH` and bakes it into `/etc/profile.d`. (Symptom if skipped: +`AssertionError: Invalid NCCL versions: ...2.30.4 (loaded) v.s. ...nvidia/nccl/lib/libnccl.so.2 +(expected)`.) Create `$RW/build_deepep.sh`: +```bash +cat > $RW/build_deepep.sh <<'EOS' +#!/bin/bash +export UV_BREAK_SYSTEM_PACKAGES=1 PIP_BREAK_SYSTEM_PACKAGES=1 UV_SYSTEM_PYTHON=1 +set -eo pipefail +export TORCH_CUDA_ARCH_LIST="10.0a" # B200 / sm_100 +command -v uv >/dev/null 2>&1 || pip install -q uv # installer calls `uv pip install --system` +export PATH="$HOME/.local/bin:/usr/local/bin:$PATH" # ensure `uv` is on PATH +TORCH_VER=$(python -c 'import torch;print(torch.__version__.split("+")[0])'); echo "torch==$TORCH_VER" > /tmp/tc.txt +PIP_CONSTRAINT=/tmp/tc.txt UV_CONSTRAINT=/tmp/tc.txt \ + bash /host/vllm/tools/ep_kernels/install_python_libraries.sh --workspace /host/ep_kernels_workspace +# guard: ensure the branch flashinfer editable is still the active install (nothing shadowed it) +pip install --no-cache-dir --no-build-isolation --no-deps -e /host/flashinfer +# DeepEP asserts torch's loaded NCCL == the nvidia-nccl wheel (2.30.7). The NGC image also has a +# system libnccl 2.30.4 that torch loads by default -> mismatch. Force the wheel's libnccl first, +# and bake it into /etc/profile.d so every `bash -lc` runtime run picks it up. +NCCL_LIB=$(python -c 'import nvidia.nccl,os;print(os.path.join(list(nvidia.nccl.__path__)[0],"lib"))') +echo "export LD_LIBRARY_PATH=\"$NCCL_LIB:\${LD_LIBRARY_PATH:-}\"" > /etc/profile.d/zz_nccl_wheel.sh +export LD_LIBRARY_PATH="$NCCL_LIB:${LD_LIBRARY_PATH:-}" +# `import deep_ep` runs DeepEP's NCCL assert (loaded libnccl must == the wheel); if it prints OK +# the correct lib is loaded. NOTE: torch.cuda.nccl.version() reports torch's BUILD-time NCCL +# (cosmetic) — to see the actually-loaded runtime lib use ncclGetVersion on /proc/self/maps. +python -c 'import ctypes, torch; torch.cuda.init(); import deep_ep; \ + m=sorted({l.split()[-1] for l in open("/proc/self/maps") if "libnccl.so" in l}); \ + lib=ctypes.CDLL(m[0]); v=ctypes.c_int(); lib.ncclGetVersion(ctypes.byref(v)); \ + print("deep_ep OK; loaded libnccl", m, "runtime ncclGetVersion", v.value)' +EOS +``` +Build: +```bash +srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=02:00:00 \ + --container-image=$RW/vllm-flashinfer-ep.sqsh --container-writable \ + --container-save=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ + bash /host/build_deepep.sh +``` + +> The DeepEP image contains vLLM + FlashInfer-EP + DeepEP, so you can run **all four backends +> from `vllm-fi-ep-deepep.sqsh`**. (The FI-EP image lacks DeepEP.) + +### Warm the caches once (avoids a 30-min first-run JIT/cubin storm, esp. multi-node) +The very first vLLM forward JIT-compiles the CUTLASS MoE kernels (~267 units) and downloads +cubins. Persist them to `/host` so every later run (and every worker) reuses them: +``` +export FLASHINFER_WORKSPACE_BASE=/host/fi_cache # JIT cache +export FLASHINFER_CUBIN_DIR=/host/fi_cubins # cubin cache +``` +(These are already in the run commands below. Run one single-node throughput first to populate.) + +--- + +## 3. Single node (8 GPU) — the comparison + +Common per-run env (put at the top of each `bash -lc '...'`): +``` +export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 \ + FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins +``` + +### ⚠ 3.0 CRITICAL — you MUST use **data-parallel EP** or the all-to-all backend is a no-op + +The `--all2all-backend` flag ONLY takes effect when vLLM selects the **modular EP** dispatch/combine +path. That selection is gated in `vllm/model_executor/layers/fused_moe/config.py`: + +```python +@property +def use_all2all_kernels(self): + return self.dp_size > 1 and self.use_ep # ← dp_size MUST be > 1 +``` + +and every backend predicate (`use_flashinfer_ep_ll_kernels`, `use_deepep_ll_kernels`, …) is +`use_all2all_kernels and all2all_backend == ""`. With **pure `--tensor-parallel-size 8` +(dp_size = 1)** `use_all2all_kernels` is `False`, so `maybe_make_prepare_finalize()` +(`all2all_utils.py:148`) returns the **monolithic** `MoEPrepareAndFinalizeNoDPEPMonolithic` — the +experts run locally per rank and are reconciled with the ordinary **TP all-reduce**. The dispatch/ +combine transport (nccl.ep for FlashInfer-EP, NVSHMEM for DeepEP) is **never launched**, and the +two backends produce byte-identical communication kernels. This is exactly what our first nsys +capture showed (both logged `Using MoEPrepareAndFinalizeNoDPEPMonolithic`). + +**Correct config for a single 8-GPU node** (EP=8 across the 8 ranks, all2all engaged): +`--data-parallel-size 8 --enable-expert-parallel` (with TP=1, so DP×TP = 8 GPUs). Qwen3-30B-A3B in +bf16 fits per-GPU on B200/GB200 (non-expert weights replicated per DP rank; the 128 experts shard +16-per-rank). When this engages, the log shows the expert backend flip to a *batched* one +(`Using BATCHED_TRITON …` or `FlashInfer CUTLASS`) and `Using FlashInferEPLLPrepareAndFinalize`. + +**⚠ Offline `vllm bench throughput` cannot take `--data-parallel-size` directly** — it errors +`Data parallel is only supported with external launcher mode with synchronous engine in offline +benchmark` (`benchmarks/throughput.py:914`). You must launch it under **`torchrun` (one process per +DP rank) with `--distributed-executor-backend external_launcher`**. Create a tiny driver once +(the `sys.path` scrub is required — torchrun prepends the script dir, and a sibling `vllm/` repo +dir would otherwise shadow the installed `vllm` package → `ModuleNotFoundError: +vllm.benchmarks.throughput`): +```bash +mkdir -p $RW/dprun +cat > $RW/dprun/driver.py <<'PY' +import sys +sys.path = [p for p in sys.path if p not in ("", "/host", "/host/dprun")] +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.benchmarks.throughput import add_cli_args, main +p = FlexibleArgumentParser(); add_cli_args(p); main(p.parse_args()) +PY +``` +Then every offline throughput/nsys run below uses `cd /tmp && torchrun --nproc_per_node=8 +/host/dprun/driver.py --data-parallel-size 8 --distributed-executor-backend +external_launcher …` (each rank prints its own `Throughput:`; the deployment total ≈ sum of the 8). +`lm_eval` (GSM8K, §3b) does its own dataset-sharding DP and does **not** build a unified EP group, +so it stays an *accuracy* check on the monolithic path — the transport is validated by §3a +`--validate` and the §3e nsys dispatch/combine capture, not by GSM8K. To exercise the transport +end-to-end with GSM8K, use the **server** path instead (`vllm serve --data-parallel-size 8 +--enable-expert-parallel --all2all-backend ` + `lm_eval --model local-completions +--model_args base_url=http://127.0.0.1:8000/v1/completions,...`; needs `pip install lm-eval[api]`). + +**HT backends (`*_high_throughput`) require `--max-num-batched-tokens 8192`.** nccl_ep HT hard-caps +`max_dispatch_tokens_per_rank` at `MAX_SUPPORTED_TOKENS_PER_RANK=8192`; the flashinfer fleet clamps +to it, but a single forward with more than 8192 tokens per rank raises at dispatch — so cap the +scheduler to match. (DeepEP-LL, conversely, *rejects* `--max-num-batched-tokens 8192`, so only pass +it for the HT runs.) + +**LL backends: do NOT pass `--max-num-batched-tokens` at all.** `flashinfer_ep_low_latency` is in +vLLM's `use_batched_dp_moe` set (with `deepep_low_latency`/`nixl_ep`), so when the flag is unset the +scheduler auto-caps to the 256-token batched-DP budget the BatchedExperts format needs — the padded +`[local_experts, max_tokens×world, N]` workspaces (and the LL transport slot buffers) are sized from +it, and an explicit large value silently makes every fill/activation/GEMM pad 32× (this was the +original 2–5× perf gap vs DeepEP; see results doc §1.1f). + +**Always verify the path** after every run — the log must NOT say `Monolithic`: +```bash +grep -h "Using .*PrepareAndFinalize" $RW/logs/ # expect one of: +# Using FlashInferEPLLPrepareAndFinalize (flashinfer_ep_low_latency) +# Using FlashInferEPHTPrepareAndFinalize (flashinfer_ep_high_throughput) +# Using DeepEPLLPrepareAndFinalize (deepep_low_latency) +# Using DeepEPHTPrepareAndFinalize (deepep_high_throughput) +# ✗ BUG (transport NOT exercised): Using MoEPrepareAndFinalizeNoDPEPMonolithic +``` +(vLLM logs this via the oracle at `oracle/unquantized.py:332` — `logger.info_once("Using %s", …)`.) + +### 3a. (optional) Correctness — EP dispatch/combine `--validate` @ world=8 +```bash +srun -A $ACCT -p $PART -N1 --ntasks-per-node=8 --time=00:25:00 \ + --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ + bash -lc 'EP_SYNC=/host/sync_ht NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ + bash /host/flashinfer/benchmarks/run_ep_matrix_one_pt.sh \ + --algorithm ht --layout fl --tokens 4096 --hidden 7168 --top-k 8 --experts 256 --validate' +# LL: --algorithm ll --layout em --tokens 128 --validate (expect "... dispatch+combine OK") +``` + +### 3b. GSM8K accuracy (5-shot) — loop the 4 backends +```bash +for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do + srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=01:00:00 \ + --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ + bash -lc "export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ + FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins; \ + python -m pip install -q lm_eval; \ + lm_eval --model vllm --tasks gsm8k --num_fewshot 5 --batch_size auto \ + --model_args pretrained=Qwen/Qwen3-30B-A3B,tensor_parallel_size=8,enable_expert_parallel=True,all2all_backend=$B,trust_remote_code=True,max_model_len=4096,enforce_eager=True" \ + > $RW/logs/gsm8k_${B}.log 2>&1 & +done; wait +grep -H "flexible-extract\|strict-match" $RW/logs/gsm8k_*.log +``` +> **GSM8K is an accuracy gate only — it does NOT exercise the all2all transport.** lm_eval's own +> `data_parallel_size` spins up independent replica engines (each `dp_size=1` internally ⇒ +> monolithic path), so there is no single cross-rank EP group to dispatch through. Keep +> `tensor_parallel_size=8` here; it confirms end-to-end model accuracy is correct with the backend +> selected. The **transport** is validated by §3a (`--validate`), §3b′ below, and the §3e nsys capture. + +### 3b′. GSM8K THROUGH the transport (real DP-EP server) — the one that exercises dispatch/combine +Serve a genuine DP-EP engine, then point lm_eval's OpenAI-compatible client at it. Unlike §3b this +runs GSM8K over the actual all2all transport. One srun starts the server, waits for `/health`, runs +the eval, and tears down. HT backends need `--max-num-batched-tokens 8192` (see §3.0). +```bash +for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput; do + case $B in *high_throughput*) CAP="--max-num-batched-tokens 8192";; *) CAP="";; esac + srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=01:00:00 \ + --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ + bash -lc "export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ + FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins; \ + python -m pip install -q lm_eval tenacity; \ + vllm serve Qwen/Qwen3-30B-A3B --port 8000 --data-parallel-size 8 --enable-expert-parallel \ + --all2all-backend $B $CAP --trust-remote-code --max-model-len 4096 --enforce-eager \ + > /host/logs/serve_${B}.log 2>&1 & SP=\$!; \ + for i in \$(seq 1 120); do curl -sf http://127.0.0.1:8000/health && break; \ + kill -0 \$SP || { echo SERVER_DIED; tail -30 /host/logs/serve_${B}.log; exit 1; }; sleep 10; done; \ + lm_eval --model local-completions --tasks gsm8k --num_fewshot 5 --batch_size 1 \ + --model_args model=Qwen/Qwen3-30B-A3B,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=128,tokenized_requests=False; \ + kill \$SP" \ + > $RW/logs/gsm8k_dpep_${B}.log 2>&1 + grep -H "flexible-extract\|strict-match\|Using .*PrepareAndFinalize" $RW/logs/gsm8k_dpep_${B}.log +done +# Expect (transport-exercised, at the pinned refs): LL 0.856/0.898, HT 0.857/0.898 +# (flex/strict; ±0.01 run-to-run). Log says FlashInferEPLL/HTPrepareAndFinalize (NOT +# Monolithic). Needs lm-eval[api] (tenacity). +``` + +### 3c. Throughput — 4 backends × 3 ISL/OSL shapes +**Must pass `--dataset-name random`** or vLLM silently uses the `sonnet` default (~1024/128) and +ignores `--input-len`/`--output-len`. +Launched under `torchrun` + `external_launcher` (see §3.0) so DP-EP engages. Each of the 8 ranks +prints its own `Throughput:`; sum them (or ×8 the mean) for the deployment total. +```bash +for SHAPE in "128 128" "2048 128" "128 2048"; do set -- $SHAPE; ISL=$1; OSL=$2 + for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do + # HT backends need the 8192 cap (nccl_ep HT limit); LL backends must NOT get it + # (deepep_low_latency rejects --max-num-batched-tokens 8192). See §3.0. + case $B in *high_throughput*) CAP="--max-num-batched-tokens 8192";; *) CAP="";; esac + srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=01:15:00 \ + --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ + bash -lc "export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ + FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins; cd /tmp; \ + torchrun --nproc_per_node=8 /host/dprun/driver.py --model Qwen/Qwen3-30B-A3B \ + --dataset-name random --random-input-len $ISL --random-output-len $OSL --num-prompts 256 \ + --data-parallel-size 8 --distributed-executor-backend external_launcher \ + --enable-expert-parallel --all2all-backend $B $CAP \ + --trust-remote-code --max-model-len 4096 --enforce-eager" \ + > $RW/logs/tp_${B}_${ISL}x${OSL}.log 2>&1 & + done; wait # (or drop `wait` to run shapes concurrently if you have the nodes) +done +# Deployment total = SUM of the 8 per-rank Throughput lines; this prints it per run: +for f in $RW/logs/tp_*.log; do + tot=$(grep -a "Throughput:" "$f" | grep -oE "[0-9.]+ total tokens/s" | grep -oE "^[0-9.]+" \ + | awk '{s+=$1} END{printf "%.0f", s}') + pf=$(grep -a "Using .*PrepareAndFinalize" "$f" | head -1 | grep -oE "Using \w+") + echo "$(basename $f) sum_total_tok/s=$tot [$pf]" # [..] must be FlashInferEP…/DeepEP…, NOT Monolithic +done +``` +> `--num-prompts 256` matches the reference matrix below exactly; larger NP (e.g. 1000) runs +> longer/steadier but shifts absolute numbers — keep it fixed when comparing backends. +> Requires the `$RW/dprun/driver.py` from §3.0. Plain `vllm bench throughput --data-parallel-size 8` +> does **not** work offline (it errors and tells you to use external launcher / serving). +> These are **eager-mode DP-EP** numbers (transport on the critical path); they are far below the +> CUDA-graph monolithic numbers — see the measured values + interpretation in +> `vllm_moe_ep_results_prenyx.md` §1.1e. + +### 3d. (optional) Memory footprint +Same as 3c but `--num-prompts 8`, and grep the KV-cache line from the **full** stream (add +`--max-num-batched-tokens 8192` for HT backends, see §3.0): +```bash +srun ... bash -lc "... cd /tmp; torchrun --nproc_per_node=8 /host/dprun/driver.py \ + --model Qwen/Qwen3-30B-A3B --dataset-name random \ + --random-input-len 128 --random-output-len 128 --num-prompts 8 --data-parallel-size 8 \ + --distributed-executor-backend external_launcher \ + --enable-expert-parallel --all2all-backend $B $CAP --gpu-memory-utilization 0.9 \ + --trust-remote-code --max-model-len 4096 --enforce-eager 2>&1 \ + | grep -iE 'Available KV cache|GPU KV cache size|Maximum concurrency|Using .*PrepareAndFinalize'" +``` + +### 3e. Capture the launched kernels with nsys, and list the all-to-all kernels + +Profile a **short** run under Nsight Systems, then dump the GPU kernel summary and filter for +the EP dispatch/combine (all-to-all) kernels. Use `--enforce-eager` (already set) so kernels are +launched individually (not hidden inside CUDA graphs), and a tiny `--num-prompts` so the report +is small. `nsys profile` follows the vLLM worker child processes, so all ranks' GPU kernels +land in one `.nsys-rep`. + +> **Use the DP-EP `torchrun` form** (see §3.0). With `--tensor-parallel-size 8` the run takes the +> monolithic path and the `.nsys-rep` will contain **no** dispatch/combine kernels — which defeats +> the purpose. `nsys` wraps `torchrun`; `--trace-fork-before-exec=true` makes it follow the 8 +> external-launcher rank processes so all ranks' kernels land in one `.nsys-rep`. + +```bash +for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do + case $B in *high_throughput*) CAP="--max-num-batched-tokens 8192";; *) CAP="";; esac # see §3.0 + srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=00:45:00 \ + --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ + bash -lc "export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ + FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins; cd /tmp; \ + nsys profile -t cuda,nvtx,nccl --force-overwrite true --sample=none --cpuctxsw=none \ + --trace-fork-before-exec=true -o /host/logs/nsys_dpep_${B} \ + torchrun --nproc_per_node=8 /host/dprun/driver.py --model Qwen/Qwen3-30B-A3B \ + --dataset-name random --random-input-len 128 --random-output-len 128 --num-prompts 64 \ + --data-parallel-size 8 --distributed-executor-backend external_launcher \ + --enable-expert-parallel --all2all-backend $B $CAP \ + --trust-remote-code --max-model-len 4096 --enforce-eager" \ + > $RW/logs/nsys_dpep_${B}.log 2>&1 +done +# GATE: confirm every run actually took the modular EP path before trusting the kernel dump. +grep -H "Using .*PrepareAndFinalize" $RW/logs/nsys_dpep_*.log +# want: FlashInferEPLL/HT... or DeepEPLL/HT...PrepareAndFinalize ✗ reject: ...Monolithic +``` + +Dump the per-kernel GPU-time summary and **filter to the all-to-all / EP kernels**. `nsys` is +only inside the container (not on the login node), so run the parsing under `srun` too — it +needs no GPU, just reads the `.nsys-rep`. The `kern_*.txt` land in `$RW/logs/` (via `/host`), so +you can grep them afterward on the login node. +```bash +srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=00:15:00 \ + --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ + bash -lc 'for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do + echo "############ $B — GPU kernel summary ############" + nsys stats --report cuda_gpu_kern_sum --format table /host/logs/nsys_dpep_${B}.nsys-rep \ + | tee /host/logs/kern_dpep_${B}.txt | head -40 + echo "---- all-to-all / EP dispatch-combine kernels only ----" + grep -iE "nccl.?ep|gin|gdaki|hybridep|deep_?ep|nvshmem|intranode|internode|dispatch|combine|moe.?ep" \ + /host/logs/kern_dpep_${B}.txt | tee /host/logs/kern_a2a_${B}.txt + [ -s /host/logs/kern_a2a_${B}.txt ] || echo " (NONE — check §3.0: run took the monolithic path?)" + done' +# If nsys is not on PATH in the image: NSYS=$(ls /opt/nvidia/nsight-systems/*/bin/nsys | head -1); use "$NSYS" stats ... +``` + +**What to expect** — the exact symbols observed on this stack (Qwen3-30B-A3B, 8×DP-EP, 64 prompts, +eager): +- **FlashInfer-EP LL** (`nccl.ep`): `nccl_ep::internode_ll::dispatch` + `nccl_ep::internode_ll::combine`. +- **FlashInfer-EP HT** (`nccl.ep`, JIT): `nccl_ep_jit_ht_dispatch_kernel` + `nccl_ep_jit_ht_combine_kernel` + + `nccl_ep_jit_ht_scan_kernel` (FLAT metadata) + `nccl_ep::hybridep::{dense_to_sparse_prob, + convert_topk_to_routing_map,sparse_to_dense_prob}` helpers. +- **DeepEP LL** (`deep_ep`+NVSHMEM): `deep_ep::legacy::internode_ll::dispatch` + + `…::internode_ll::combine` + one-time `nvshmemi_init_array_kernel`. +- **DeepEP HT** (`deep_ep` intranode): `deep_ep::legacy::intranode::{notify_dispatch,dispatch, + combine,cached_notify_combine}` + `…::layout::get_dispatch_layout`. +- Both also show the shared **expert GEMM** and attention/norm kernels — those are *not* + all-to-all; the grep above narrows to transport. + +> **Reading the times (important):** the LL dispatch/combine kernels **busy-wait on the network**, +> so `Total`/`Avg`/`Max` are spin-dominated (multi-second `Max`, huge `StdDev`) — use the **median** +> per-launch. And **launch counts differ by design**: DeepEP LL issues ~2× the launches of +> FlashInfer-EP LL because DeepEP low-latency splits each dispatch/combine into a **send kernel + +> a deferred receive "hook"** (`low_latency_dispatch(return_recv_hook=True)` then `hook()`, for +> compute/comm overlap), whereas nccl.ep issues one fused kernel per call (`handle.dispatch()` + +> a `handle.complete()` stream-sync, not a 2nd launch). So compare DeepEP's (send+recv) sum vs +> FlashInfer's single launch, not launch-for-launch. HT is 1 launch each → directly comparable. +> Measured medians and the full table are in `vllm_moe_ep_results_prenyx.md` §1.1b; for a +> spin-free per-op latency use the standalone comm benchmark below. +- **If `kern_a2a_.txt` is empty for a backend, the transport did not run** — the summary will + instead be dominated by `multimem_all_reduce_kernel` / `vllm::cross_device_reduce_*` (TP + all-reduce). That means the run fell back to the monolithic path (§3.0); fix the DP flag and + re-capture. **This is the exact failure our first capture hit** (both backends identical, only + TP all-reduce, no dispatch/combine). + +Notes: +- **`nsys` lives only inside the container**, not on the Pre-Nyx login node — run *both* + `nsys profile` and `nsys stats` under `srun --container-image=...`. `nsys stats` needs no GPU + (it only reads the `.nsys-rep`), so it's a cheap short job. +- **Cleaner isolation (NCCL-EP only):** to see *just* the dispatch/combine kernels with no model + noise, profile the standalone comm benchmark instead of vLLM: + ```bash + srun ... --ntasks-per-node=8 bash -lc 'EP_SYNC=/host/sync_ns NCCL_GIN_TYPE=3 \ + nsys profile -t cuda,nvtx,nccl -o /host/logs/nsys_epcomm_r${SLURM_PROCID} \ + bash /host/flashinfer/benchmarks/run_ep_matrix_one_pt.sh \ + --algorithm ht --layout fl --tokens 4096 --hidden 7168 --top-k 8 --experts 256 --iters 20' + # then: nsys stats --report cuda_gpu_kern_sum /host/logs/nsys_epcomm_r0.nsys-rep + ``` +- If the `.nsys-rep` has no GPU kernels, nsys didn't follow the workers — re-run adding + `--trace-fork-before-exec=true`, or use the per-rank standalone form above (`-o ..._r${SLURM_PROCID}`). +- Copy `$RW/logs/*.nsys-rep` locally to open the timeline in the Nsight Systems GUI if you want + to see the dispatch→expert-GEMM→combine sequence visually. + +--- + +## 4. Multi-node (2 nodes / 16 GPU) + +Multi-node needs a **Ray cluster across both nodes' containers** + `--data-parallel-size 16 +--enable-expert-parallel --distributed-executor-backend ray` + `NCCL_MNNVL_ENABLE=1` for the +cross-node EP fabric. **Use data-parallel EP, not `--tensor-parallel-size 16`** — same reason as +§3.0: TP-only ⇒ `dp_size=1` ⇒ monolithic path ⇒ `all2all_backend` ignored (and the cross-node +all-to-all you're trying to measure never runs). +Create `$RW/run_2node.sh` (rank-0 starts the Ray head + runs the bench; rank-1 joins and blocks): + +```bash +cat > $RW/run_2node.sh <<'EOS' +#!/bin/bash +set -o pipefail +export HF_HOME=/host/hf_cache FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins +export NCCL_GIN_TYPE=3 NCCL_MNNVL_ENABLE=1 +# force the 2.30.7 wheel libnccl (DeepEP needs torch nccl == wheel; runs via bash, not -lc) +NCCL_LIB=$(python -c 'import nvidia.nccl,os;print(os.path.join(list(nvidia.nccl.__path__)[0],"lib"))') +export LD_LIBRARY_PATH="$NCCL_LIB:${LD_LIBRARY_PATH:-}" +B=${BACKEND:-flashinfer_ep_low_latency}; ISL=${ISL:-128}; OSL=${OSL:-128} +python -m pip install -q ray 2>/dev/null || true +HEADF=/host/ray_head.$SLURM_JOB_ID; TMP=/host/raylog_${SLURM_JOB_ID}_${SLURM_NODEID} +if [ "${SLURM_NODEID:-0}" = "0" ]; then + hostname -I | awk '{print $1}' > $HEADF + ray start --head --port=6379 --num-gpus=8 --disable-usage-stats --temp-dir=$TMP + for i in $(seq 1 60); do + n=$(python -c 'import ray;ray.init(address="auto");print(int(ray.cluster_resources().get("GPU",0)))' 2>/dev/null|tail -1) + [ "$n" = "16" ] && break; sleep 5; done; echo "cluster GPUs=$n" + vllm bench throughput --model Qwen/Qwen3-30B-A3B --dataset-name random \ + --random-input-len $ISL --random-output-len $OSL --num-prompts 1000 \ + --data-parallel-size 16 --enable-expert-parallel --all2all-backend $B \ + --distributed-executor-backend ray --trust-remote-code --max-model-len 4096 --enforce-eager + ray stop; rm -f $HEADF +else + for i in $(seq 1 60); do [ -f $HEADF ] && break; sleep 3; done; sleep 8 + ray start --address=$(cat $HEADF):6379 --num-gpus=8 --disable-usage-stats --temp-dir=$TMP --block +fi +EOS +``` +Run (one task per node; warm `/host/fi_cubins` from a single-node run first): +```bash +for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do + srun -A $ACCT -p $PART -N2 --ntasks-per-node=1 --time=01:00:00 \ + --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ + --export=ALL,BACKEND=$B,ISL=128,OSL=128 \ + bash /host/run_2node.sh > $RW/logs/tp2n_${B}.log 2>&1 + grep -H "Throughput:" $RW/logs/tp2n_${B}.log +done +``` + +> ⚠ **Known issue (environmental, not the EP code):** in our runs the Ray 16-GPU cluster forms +> and vLLM launches, but engine-core init **stalls** — and a plain `--tensor-parallel-size +> 16` run **without** EP stalls identically. So it's the cluster's cross-node vLLM/NCCL bring-up, +> not FlashInfer-EP/DeepEP. If you hit it, debug the cross-node fabric first, e.g. export before +> the bench: `NCCL_DEBUG=INFO`, and set `NCCL_SOCKET_IFNAME` / `NCCL_IB_HCA` to the node's IB +> interfaces (`ibdev2netdev` / `ibv_devices`); verify a plain 2-node NCCL all-reduce works. + +> ⚠ **Offline-DP caveat also applies here.** `vllm bench throughput --data-parallel-size 16` with +> the Ray backend hits the same offline-DP guard as §3.0 (it needs `external_launcher`, not `ray`). +> For a working 2-node DP-EP throughput run once the fabric is up, either launch the §3.0 driver +> under `torchrun --nnodes=2 --nproc_per_node=8 --distributed-executor-backend external_launcher`, +> or use the **server** path (`vllm serve --data-parallel-size 16 --enable-expert-parallel` + +> `vllm bench serve`). The Ray script above is kept as the cluster-bring-up reference. + +--- + +## Expected reference numbers (8-GPU, Qwen3-30B-A3B; full detail in `vllm_moe_ep_results_prenyx.md`) + +**DP-EP, transport-exercised** (the numbers to trust for a backend comparison; §1.1c–f). +Measured at the pinned refs (flashinfer `fa09bc46`, vLLM `ab1415e`), §3c command verbatim +(NP=256, eager, HT capped at 8192, LL uncapped): + +| Backend | 128/128 | 2048/128 | 128/2048 | +|---|---|---|---| +| `flashinfer_ep_low_latency` | **9,088** | **23,106** | **5,825** | +| `deepep_low_latency` | 10,116 | 24,013 | 6,595 | +| `flashinfer_ep_high_throughput` | **6,797** | **45,224** | 3,795 | +| `deepep_high_throughput` | 5,736 | 35,623 | 4,539 | + +(total tok/s, sum of 8 ranks; expect ±3–5% run-to-run.) **FI-HT is ahead of DeepEP-HT by +19–27%** on 128/128 and 2048/128; FI-LL within 4–12% of DeepEP-LL; decode-heavy within +12–16% for both modes. +- **GSM8K over a real DP-EP server** (§3b′, flex/strict): FI-EP **LL 0.856/0.898**, + **HT 0.857/0.898** — both ≥0.80, on par with the ~0.88 reference. Transport genuinely exercised. +- ⚠ These numbers **require the pinned refs** (§1) — the HT token clamp + recv-trim, the + batched-DP cap membership and the fleet host-path caches. On older refs FI-EP lands 2–6× + behind DeepEP and HT SIGABRTs at group-create (root causes + fixes: results doc §1.1c/§1.1f). +- Also required: HT runs get `--max-num-batched-tokens 8192`; LL runs get **no** such flag. + +**Monolithic (`--tensor-parallel-size 8`, dp_size=1) — NOT a transport comparison** (§1.2, retained +for reference only): GSM8K strict ~0.89 all four; throughput 32.5k · 141k · 18.5k (FI) ≈ DeepEP +(both ran the identical TP-all-reduce path, so the ~1–2% closeness is an artifact); memory identical +150.45 GiB KV cache. The all2all backend had **no effect** here. From 0ce698052c85914f2af00d44630db5f0aeb1ce9e Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Mon, 6 Jul 2026 01:52:25 -0700 Subject: [PATCH 08/13] fix(moe_ep): declare _hot_cache on NcclEpFleet (mypy) CI pre-commit (full-repo mypy) fails with 'NcclEpFleet has no attribute "_hot_cache"' -- the cross-handle host-path cache was injected dynamically from NcclEpHandle.__init__ via getattr/setattr. Declare it in NcclEpFleet.__init__ instead (where it belongs), clear it in update_topology() (world-size change invalidates cached buffer shapes), and have the handle reference it directly. No behavior change on the hot path; mypy on both files now passes with the pre-commit config. Co-Authored-By: Claude Opus 4.8 --- flashinfer/moe_ep/nccl_ep/fleet.py | 9 +++++++++ flashinfer/moe_ep/nccl_ep/handle.py | 24 +++++++++++------------- 2 files changed, 20 insertions(+), 13 deletions(-) diff --git a/flashinfer/moe_ep/nccl_ep/fleet.py b/flashinfer/moe_ep/nccl_ep/fleet.py index 60f8e5e5f21..2ab975af382 100644 --- a/flashinfer/moe_ep/nccl_ep/fleet.py +++ b/flashinfer/moe_ep/nccl_ep/fleet.py @@ -168,6 +168,12 @@ def __init__( self._nccl_ep = _import_nccl_ep() self._comm = _resolve_comm(bootstrap) # keepalive: Group borrows it + # Cross-handle host-path cache (recv buffers, counter tensors, FFI + # descriptor memos), populated and consumed by NcclEpHandle. Anchored on + # the Fleet because callers (e.g. vLLM) create a fresh Handle every + # forward while the Fleet persists — per-handle caches never hit. + self._hot_cache: dict = {} + self._group = self._nccl_ep.Group.create(self._comm, self._build_group_config()) self._destroyed = False @@ -288,6 +294,9 @@ def update_topology( self._bootstrap = bootstrap self._stream = bootstrap.stream self._comm = _resolve_comm(bootstrap) + # Topology (world size) changed — drop the cross-handle host caches so + # recv buffers / counters / FFI descriptors are rebuilt at the new sizes. + self._hot_cache.clear() self._group = self._nccl_ep.Group.create(self._comm, self._build_group_config()) self._destroyed = False diff --git a/flashinfer/moe_ep/nccl_ep/handle.py b/flashinfer/moe_ep/nccl_ep/handle.py index 8a26a3dea60..7e284511475 100644 --- a/flashinfer/moe_ep/nccl_ep/handle.py +++ b/flashinfer/moe_ep/nccl_ep/handle.py @@ -133,19 +133,17 @@ def __init__( _t = _pc() if _HP else None self._fleet = fleet self._ep = fleet.nccl_ep - # Cross-handle host-path cache. vLLM creates a fresh Handle every MoE - # layer x step (routing binds at create_handle), so per-handle caches - # never hit; anchoring them on the long-lived Fleet makes the recv - # buffers, counter tensors and FFI descriptor objects reusable across - # forwards. Tensor wrappers are memoized by (data_ptr, dtype, shape), - # so an entry can only ever describe the same memory layout it was - # built for; the dict is cleared when it grows past a bound (entries - # are then rebuilt, which is always safe — each handle only needs - # address stability within its own lifetime). - hot = getattr(fleet, "_hot_cache", None) - if hot is None: - hot = fleet._hot_cache = {} - self._hot = hot + # Cross-handle host-path cache (declared on NcclEpFleet). vLLM creates + # a fresh Handle every MoE layer x step (routing binds at + # create_handle), so per-handle caches never hit; anchoring them on the + # long-lived Fleet makes the recv buffers, counter tensors and FFI + # descriptor objects reusable across forwards. Tensor wrappers are + # memoized by (data_ptr, dtype, shape), so an entry can only ever + # describe the same memory layout it was built for; the dict is cleared + # when it grows past a bound (entries are then rebuilt, which is always + # safe — each handle only needs address stability within its own + # lifetime). + self._hot = fleet._hot_cache self._handle_knobs = _index_knobs(algo_knobs) self._stream = self._knob_stream() self._staged = HandleAlgoKnobSplitOperation in self._handle_knobs From c9dcc805b3f43f160f60a6a94ec3891e912a380b Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Mon, 6 Jul 2026 12:50:34 -0700 Subject: [PATCH 09/13] fix(deps): relax cuda-python floor to >=12.0 (don't force a CUDA major) The cuda-python>=13.0 base-dep floor bulldozes CUDA-12 environments. On the H100 CI runner (cu12 container) pip reports nvshmem4py-cu12 requires cuda-python<=12.9, but you have cuda-python 13.3.1 and then re-resolves torch to the PyPI-default cu13 build, whose cuda-toolkit[cudart] dependency drops libcudart.so.13 into an env that already carries libcudart.so.12 (nvidia-cutlass-dsl-libs-cu12). cudnn-frontend's cuda-pathfinder loader (pathfinder itself newly present via cuda-python 13) then fails every graph.check_support() with "RuntimeError: Multiple libcudart libraries found" -- the JIT Unittest (H100) failures on PR #3821. Base deps must not force a CUDA major (the same rule requirements.txt already documents for nvidia-nccl-cu13). With >=12.0 a cu12 env keeps cuda-python 12.x -- exactly the configuration main's CI already passes with -- while cu13 envs resolve 13.x and get the full nccl.ep stack. moe_ep itself stays CUDA-13-only, enforced with a clear error at runtime (moe_ep/_validators.py); availability probing is find_spec-based so nothing imports nccl.ep on cu12. scripts/build_in_container.sh keeps its cuda-python>=13.0 pin deliberately: that is the single-CUDA cu13 dev-container flow (installs nixl-cu13 / nvidia-nccl-cu13 explicitly). AI-assisted (Claude Code): root-caused by diffing the installed-package sets of the failing PR job vs main's passing H100 job. Co-Authored-By: Claude Opus 4.8 --- requirements.txt | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 6aa592684a8..985081ec2b8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -11,7 +11,17 @@ click # torch 2.10.0 wheel. torch supplies libnccl; the >=2.30.7 B200 floor is # enforced at runtime in flashinfer/moe_ep/_validators.py, and the build hook # installs the newer wheel --no-deps on source installs (build_backend.py). -cuda-python>=13.0 +# +# The same base-deps-must-not-force-a-CUDA-major rule applies to cuda-python: +# a >=13.0 floor bulldozes CUDA-12 environments (pip reports e.g. +# "nvshmem4py-cu12 requires cuda-python<=12.9" and then re-resolves torch to +# the cu13 default build, whose cuda-toolkit[cudart] dep drops libcudart.so.13 +# next to the env's libcudart.so.12 — cudnn-frontend's cuda-pathfinder loader +# then aborts with "Multiple libcudart libraries found"). With a >=12.0 floor +# a cu12 env keeps cuda-python 12.x (moe_ep is CUDA-13-only and its runtime +# validator raises a clear error there), while cu13 envs resolve 13.x and get +# the full nccl.ep stack. +cuda-python>=12.0 cuda-tile>=1.4.0 einops nccl4py>=0.3.1 From 566e8db47b4557ebcbbccaf92a9d739a10974d5d Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Mon, 6 Jul 2026 12:56:55 -0700 Subject: [PATCH 10/13] docs(moe_ep): drop internal validation docs from the PR vllm_moe_ep_results_prenyx.md and vllm_moe_ep_runbook.md are internal cluster-specific validation notes (Pre-Nyx runbook + measurement log); keep them out of the upstream PR. Files remain in local working trees, untracked. Co-Authored-By: Claude Opus 4.8 --- .../design_docs/vllm_moe_ep_results_prenyx.md | 627 ------------------ docs/design_docs/vllm_moe_ep_runbook.md | 566 ---------------- 2 files changed, 1193 deletions(-) delete mode 100644 docs/design_docs/vllm_moe_ep_results_prenyx.md delete mode 100644 docs/design_docs/vllm_moe_ep_runbook.md diff --git a/docs/design_docs/vllm_moe_ep_results_prenyx.md b/docs/design_docs/vllm_moe_ep_results_prenyx.md deleted file mode 100644 index 62afc035d18..00000000000 --- a/docs/design_docs/vllm_moe_ep_results_prenyx.md +++ /dev/null @@ -1,627 +0,0 @@ -# vLLM ⇄ `flashinfer.moe_ep` — validated results & reproduction (Pre-Nyx) - -Measured results for the two vLLM all2all backends backed by `flashinfer.moe_ep` -(`flashinfer_ep_low_latency`, `flashinfer_ep_high_throughput`) and the exact steps to -reproduce them on the **Pre-Nyx** cluster (SLURM + pyxis/enroot, B200-class GPUs, CUDA 13.2). - -For the design/architecture and the full test/bench catalog see -[`vllm_moe_ep_integration.md`](vllm_moe_ep_integration.md). - -Code under test: -- FlashInfer: branch `feat/vllm-moe-ep-api` (`github.com/Anerudhan/flashinfer`, `cfc93a9c`). -- vLLM: branch `feat/flashinfer-ep-all2all` (`github.com/Anerudhan/vllm`, `f4e2618`). - ---- - -## 1. Results - -> 🛑 **IMPORTANT CAVEAT — the throughput/GSM8K/memory numbers below did NOT exercise the all-to-all -> transport.** They were collected with `--tensor-parallel-size 8 --enable-expert-parallel` -> (`dp_size = 1`). vLLM only routes MoE through the modular EP dispatch/combine path — the only -> path that uses `--all2all-backend` — when **`dp_size > 1`** (`config.py::use_all2all_kernels = -> dp_size > 1 and use_ep`). With `dp_size = 1` it falls back to -> `MoEPrepareAndFinalizeNoDPEPMonolithic`, where experts run locally and are reconciled by the -> ordinary **TP all-reduce**. nsys confirmed this: for *both* `flashinfer_ep_low_latency` and -> `deepep_low_latency` the kernel summaries were dominated by `multimem_all_reduce_kernel` / -> `vllm::cross_device_reduce_*` with **no** dispatch/combine kernels, and both logged -> `Using MoEPrepareAndFinalizeNoDPEPMonolithic`. **Consequence:** the ~1–2% FI-EP↔DeepEP closeness -> below reflects the *shared monolithic path*, NOT a comparison of the two transports. -> **Correctness (§1.1: GAP tests + `--validate`) is unaffected** — those exercise the transport -> directly. The transport-level comparison is being re-collected with `--data-parallel-size 8 -> --enable-expert-parallel` (see runbook §3.0); tables below are marked *provisional (monolithic)* -> until then. - -All runs: **single node, 8×GPU, 8-way expert parallel, CUDA 13.2**, base image -`nvcr.io/nvidia/pytorch:26.05-py3`, vLLM built from source (upstream `torch==2.11.0` pin -stripped via `use_existing_torch.py`; builds cleanly against NGC 26.05 torch). - -### 1.1 Correctness - -| Check | LL (`flashinfer_ep_low_latency`) | HT (`flashinfer_ep_high_throughput`) | -|---|---|---| -| `moe_ep` GAP 1/2/3 unit tests (mocked nccl, host-only) | 14 / 14 passed | — | -| EP dispatch+combine `--validate` @ world=8 | ✅ `ll/em per_rank=128 OK` | ✅ `ht/fl per_rank=4096 OK` | -| vLLM e2e smoke (OLMoE-1B-7B, greedy) | ✅ coherent | ✅ coherent (identical text) | -| **GSM8K 5-shot, Qwen3-30B-A3B** — flexible-extract | **0.8522** ± 0.0098 | **0.8575** ± 0.0096 | -| **GSM8K 5-shot, Qwen3-30B-A3B** — strict-match | **0.8939** ± 0.0085 | **0.8969** ± 0.0084 | - -Both clear the **≥ 0.80** gate and match the ~0.88 reference for Qwen3-30B-A3B. - -### 1.1b DP-EP transport-verified nsys capture (the all2all transport IS exercised) - -To exercise the all2all transport in vLLM (not just `--validate`), the MoE layer must take the -**modular EP** path, which requires `dp_size > 1` (§1 caveat). Offline `vllm bench throughput` -can't take `--data-parallel-size` directly — it must be launched under `torchrun --nproc_per_node=8` -with `--distributed-executor-backend external_launcher` (see runbook §3.0). Config: Qwen3-30B-A3B, -8×GPU, `--data-parallel-size 8 --enable-expert-parallel`, `--enforce-eager`, 8/64 prompts, nsys -`-t cuda,nvtx,nccl --trace-fork-before-exec=true`. - -With DP-EP the oracle switches the expert backend to a *batched* one and selects the EP -prepare/finalize — the log now prints (per backend): - -| Backend | expert backend | prepare/finalize (log) | all2all transport engaged? | -|---|---|---|---| -| `flashinfer_ep_low_latency` | `BATCHED_TRITON` | `Using FlashInferEPLLPrepareAndFinalize` | ✅ | -| `deepep_low_latency` | `BATCHED_TRITON` | `Using DeepEPLLPrepareAndFinalize` | ✅ | -| `deepep_high_throughput` | `TRITON` | `Using DeepEPHTPrepareAndFinalize` | ✅ | -| `flashinfer_ep_high_throughput` | `TRITON` | `Using FlashInferEPHTPrepareAndFinalize` | ✅ (after the 3 HT fixes below) | - -**Actual dispatch/combine kernels now on the GPU** (`nsys stats cuda_gpu_kern_sum`), which were -**absent** in the earlier `--tensor-parallel-size 8` (monolithic) capture: - -| Backend | dispatch kernel | combine kernel | extra | -|---|---|---|---| -| `flashinfer_ep_low_latency` (nccl.ep) | `nccl_ep::internode_ll::dispatch<…,(ncclEpLayout_t)1,…>` | `nccl_ep::internode_ll::combine<…,(ncclEpLayout_t)1>` | — | -| `deepep_low_latency` (DeepEP+NVSHMEM) | `deep_ep::legacy::internode_ll::dispatch<…>` | `deep_ep::legacy::internode_ll::combine<…>` | `nvshmemi_init_array_kernel<…>` | -| `deepep_high_throughput` (DeepEP intranode) | `deep_ep::legacy::intranode::notify_dispatch<8>` + `dispatch<8,768,8192>` | `intranode::cached_notify_combine<8>` + `combine` | `intranode::layout::get_dispatch_layout` | - -So the FlashInfer-EP LL vs DeepEP LL **transport difference is real and now visible**: -FlashInfer-EP LL runs the **`nccl_ep` GIN internode-LL** dispatch/combine; DeepEP LL runs the -**`deep_ep` NVSHMEM internode-LL** dispatch/combine (+ an NVSHMEM init kernel). Both LL backends -picked the *internode* LL path even on a single node. - -> ⚠ **Do NOT read these as a perf comparison.** The runs are tiny (8/64 prompts, decode-heavy, no -> CUDA graph), and the LL dispatch/combine kernels **busy-wait on the network**, so their -> `Total/Max` times are dominated by wait/sync (e.g. multi-second `Max` on an 8-prompt run) and are -> not throughput. These captures **prove the transport is exercised and identify the exact kernels**; -> a real perf comparison needs the DP-EP throughput sweep (runbook §3c) with CUDA graphs + larger -> batches. Raw dumps: `$RW/logs/kern_dpep_.txt`, `kern_a2a_.txt`, -> `nsys_dpep_.nsys-rep`. - -### 1.1c FlashInfer-EP HT under DP-EP — 3 bugs found & fixed (was SIGABRT, now works) - -Enabling the real DP-EP path surfaced three sequential HT-only bugs (all masked by the monolithic -path, which never builds the HT prepare/finalize). Each was root-caused from per-rank -`CUDA_LAUNCH_BLOCKING=1` stderr on Pre-Nyx and fixed; **HT now runs end-to-end and is GSM8K-validated** -(§1.1d): - -1. **SIGABRT at `ncclEpCreateGroup` (nccl_ep.cc:1253).** HT asserts - `max_dispatch_tokens_per_rank ≤ MAX_SUPPORTED_TOKENS_PER_RANK` (build-time `8192` in the nccl4py - wheel, `nccl/ep/include/nccl_ep/common.hpp`). vLLM sizes the HT fleet from - `moe.max_num_tokens = scheduler max_num_batched_tokens` (16384 here) → abort on all ranks. LL is - uncapped, so only HT hit it. **Fix (flashinfer `feat/nvep-default`, `60ff0fc1`):** clamp the HT - fleet's `max_tokens_per_rank` to the cap in `NcclEpFleet` (+ a clear `MoEEpConfigError` guard in - `_dispatch_ht`). Run HT with `--max-num-batched-tokens ≤ 8192`. -2. **Triton illegal-memory-access at `moe_align_block_size.py:101` (`expert_map[expert_ids]`).** - FlashInfer's FLAT recv gives **local** expert ids with `-1` for non-local/padding picks, but - vLLM's Standard experts feed `topk_ids` through `moe_align` + `expert_map` expecting **global** - ids (the skip is applied via `expert_map`, never as `-1` in `topk_ids`). **Fix (vLLM - `flashinfer_ep_ht.py`):** rebuild global ids from `expert_map`, remap `-1`→a non-owned global id - (which `expert_map` re-tags skipped) — matching the DeepEP HT contract. -3. **`finalize` over-strict assert.** Standard Triton experts declare `TopKWeightAndReduceNoOP` - (they already applied the dispatched routing weights and reduced their local picks), but finalize - only accepted `TopKWeightAndReduceDelegate`. FlashInfer HT combine applies **no** weights - (captured at dispatch) and only reduces per-rank partials across ranks → no double-weighting. - **Fix (vLLM `flashinfer_ep_ht.py`):** accept `NoOP` too. - -`deepep_high_throughput` with the same `TRITON` experts was unaffected (it returns global ids), which -is why only FlashInfer-EP HT hit bugs 2–3. - -### 1.1d GSM8K over a REAL DP-EP deployment (transport-exercised accuracy) - -Run via `vllm serve --data-parallel-size 8 --enable-expert-parallel --all2all-backend ` -(the online server path builds a genuine DP-EP engine, unlike lm_eval's own `data_parallel_size` -which spins up independent monolithic replicas) + `lm_eval --model local-completions` (needs -`pip install lm-eval[api]`). Both backends log the modular EP prepare/finalize and clear ≥0.80: - -| Backend | prepare/finalize | flexible-extract | strict-match | -|---|---|---|---| -| `flashinfer_ep_low_latency` | `FlashInferEPLLPrepareAndFinalize` | **0.8582** | **0.8976** | -| `flashinfer_ep_high_throughput` | `FlashInferEPHTPrepareAndFinalize` | **0.8461** | **0.8946** | - -On par with each other, with the monolithic-path GSM8K (§1.1), and with the ~0.88 Qwen3-30B-A3B -reference — i.e. the FlashInfer-EP dispatch/combine transport (both LL and HT) is numerically -correct end-to-end, not just at `--validate`. (This *replaces* the earlier "GSM8K is TP/accuracy-only" -caveat: it is now measured through the actual all2all transport.) - -### 1.1e DP-EP transport-exercised throughput (all 4 backends) - -Offline `vllm bench throughput` launched under `torchrun --nproc_per_node=8 ---distributed-executor-backend external_launcher --data-parallel-size 8 --enable-expert-parallel` -(runbook §3.0), Qwen3-30B-A3B, `--enforce-eager`, NP=256, `--dataset-name random`. HT backends -add `--max-num-batched-tokens 8192` (the nccl_ep HT cap); **LL backends must leave the flag -unset** (§1.1f fix #1 — the batched-DP 256 auto-cap must engage; DeepEP-LL also rejects 8192). -Deployment total = sum of the 8 DP ranks' `Throughput:` lines. Every cell verified -`Using PrepareAndFinalize` in the log. - -**Final (after the §1.1f optimizations; single same-day pass, total tok/s):** - -| Backend | 128/128 | 2048/128 | 128/2048 | -|---|---|---|---| -| `flashinfer_ep_low_latency` | **9,088** | **23,106** | **5,825** | -| `deepep_low_latency` | 10,116 | 24,013 | 6,595 | -| *FI-LL / DeepEP-LL* | *0.90×* | *0.96×* | *0.88×* | -| `flashinfer_ep_high_throughput` | **6,797** | **45,224** | 3,795 | -| `deepep_high_throughput` | 5,736 | 35,623 | 4,539 | -| *FI-HT / DeepEP-HT* | ***1.19×*** | ***1.27×*** | *0.84×* | - -**FlashInfer-EP HT is now ahead of DeepEP-HT by 19–27%** on the balanced and prefill-heavy -shapes; LL is within 4–12% of DeepEP-LL; the decode-heavy shape is within 12–16% for both modes. - -For the record, the **first** transport-exercised pass (before the §1.1f fixes) was -FI-LL 1,926/12,975/1,038 and FI-HT 2,761/21,011/1,485 — i.e. **2–6× behind DeepEP** — so the -gap closure came from the three root-cause fixes below, each GSM8K-gated. - -> ⚠ **Interpretation.** Eager-mode (no CUDA graphs) DP-EP runs over the Triton/batched-Triton -> expert backends with a small NP — chosen so the all2all transport is genuinely on the critical -> path and comparable across backends, **not** a production throughput number (the CUDA-graph -> monolithic numbers in §1.2 are much higher). CUDA-graph capture of dispatch/combine is the -> remaining big lever for both backends. - -### 1.1f Closing the 2–6× gap: root causes & fixes (perf iteration log) - -All found by diffing `nsys cuda_gpu_kern_sum` per-kernel medians between the FI-EP and DeepEP -runs, then `EP_PROFILE_HOST=1` host-phase timing. Each fix validated by GSM8K over the real -DP-EP server (§1.1d method). - -1. **[vLLM, 1 line] `flashinfer_ep_low_latency` was missing from `use_batched_dp_moe`** - (`vllm/config/parallel.py`). That property auto-caps the scheduler to the 256-token - batched-DP budget for BatchedExperts-format backends (`deepep_low_latency`, `nixl_ep`) — - without it FI-LL ran with the 8192-token offline default, so the padded - `[local_experts, max_tokens×world, N]` workspaces were **32× larger** than DeepEP-LL's: - every `fill_(0)` (394µs vs 13µs), full-workspace `act_and_mul` (2.07ms vs 36µs), the padded - batched-GEMM grid (1.29ms vs 256µs) and the LL transport slot buffers all paid it. - *Effect: FI-LL 1,926/12,975/1,038 → 8,854/22,939/5,416.* After the fix the compute kernels - are **byte-identical** to DeepEP-LL's (280µs/35.8µs/13.2µs medians on both). -2. **[vLLM adapter] HT recv-count trim.** FI-HT ran the whole Standard MoE stack (`moe_align` - 241µs, `count_and_sort` 344µs, 2×`fused_moe` ~1ms, `act_and_mul` 1.06ms, `moe_sum` 344µs) - over the **static 65,536-row** recv buffer every forward. The GAP-3 `recv_total` counter - (written by the HT metadata step at create_handle) now trims the compute view to - `round_up(actual,128)` rows (`.item()` host sync — eager-only); finalize copies the trimmed - expert output into a persistent full-size buffer (nccl.ep combine needs the address-stable - static staging; padding rows carry no routing state and are never sent). Also cached the - static local→global expert-id remap (was 2 `nonzero()` device syncs/layer/step). - *Effect: FI-HT 2,761/21,011/1,485 → 6,636/44,238/3,755 — ahead of DeepEP-HT on 2 of 3 shapes.* -3. **[flashinfer] Fleet-level host-path caches** (`nccl_ep/handle.py`). vLLM creates a fresh - `NcclEpHandle` every MoE layer×step (routing binds at `create_handle`), so the per-handle - `NV_FI_EP_FAST_PATH` caches never hit and each forward paid **~149µs host** (measured: - FFI descriptor builds 32.6+31.3µs, handle setup 35.9µs, create/destroy/dispatch/combine C - calls ~45µs). At decode the GPU is host-paced and nccl.ep's **fused send+recv** dispatch - kernel absorbs the inter-rank lag as in-kernel spin (median 256µs/launch, 33% of GPU time). - Fix: anchor recv buffers, counter tensors, static FFI tuples and a - `(data_ptr,dtype,shape)`-keyed wrap memo on the long-lived Fleet — **restricted to tensors - ≤2 MiB**, because the nccl.ep Tensor wrapper pins the torch tensor (memoizing large prefill - activations pinned GBs → OOM at `--gpu-memory-utilization 0.9`; found the hard way). - *Effect: host 149→119µs/layer; decode shapes +4–7%; final matrix above.* - -**Kernel-level notes for the remaining decode-heavy delta (~0.85×):** -- FI-LL fused `internode_ll::dispatch` (send+recv in one kernel) median 256µs vs DeepEP's - split send (14µs) + deferred recv hook (13µs): the fused kernel spins for the slowest peer, - absorbing per-layer host-path lag; the residual ≈ the remaining ~119µs host path (of which - ~45µs is create/destroy/dispatch/combine C calls — an nccl.ep handle-reuse/update API would - remove most of it). -- Remaining levers, in expected-impact order: CUDA-graph capture of dispatch/combine (removes - host pacing entirely), an nccl.ep API to reuse/update a handle instead of per-forward - create/destroy, splitting send/recv (staged mode) to overlap like DeepEP's hook, and - trimming the ~28µs combine-side FFI container builds. - -### 1.2 Throughput sweep — FlashInfer-EP vs DeepEP *(provisional — monolithic path, see §1 caveat)* - -`vllm bench throughput`, Qwen3-30B-A3B, 8-GPU EP, 1000 prompts, `--dataset-name random`, -`--enforce-eager`, `--max-model-len 4096`. DeepEP built into `vllm-fi-ep-deepep.sqsh` (§3.3). -Three ISL/OSL shapes: balanced (128/128), prefill-heavy (2048/128), decode-heavy (128/2048). - -| ISL / OSL | Backend | total tok/s | output tok/s | req/s | -|---|---|---|---|---| -| **128 / 128** | `flashinfer_ep_low_latency` | 32,506 | 16,253 | 126.98 | -| | `flashinfer_ep_high_throughput` | 32,050 | 16,025 | 125.20 | -| | `deepep_low_latency` | 32,535 | 16,267 | 127.09 | -| | `deepep_high_throughput` | 32,050 | 16,025 | 125.19 | -| **2048 / 128** (prefill-heavy) | `flashinfer_ep_low_latency` | 140,823 | 8,284 | 64.72 | -| | `flashinfer_ep_high_throughput` | 141,744 | 8,338 | 65.14 | -| | `deepep_low_latency` | 141,786 | 8,340 | 65.16 | -| | `deepep_high_throughput` | 143,238 | 8,426 | 65.83 | -| **128 / 2048** (decode-heavy) | `flashinfer_ep_low_latency` | 18,515 | 17,426 | 8.51 | -| | `flashinfer_ep_high_throughput` | 18,461 | 17,376 | 8.48 | -| | `deepep_low_latency` | 18,891 | 17,780 | 8.68 | -| | `deepep_high_throughput` | 18,764 | 17,660 | 8.62 | - -**Relative throughput (best-of-each-backend total tok/s, FlashInfer-EP ÷ DeepEP):** - -| ISL / OSL | FlashInfer-EP best | DeepEP best | FI-EP / DeepEP | -|---|---|---|---| -| 128 / 128 | 32,506 (LL) | 32,535 (LL) | **0.999** (−0.1%) | -| 2048 / 128 | 141,744 (HT) | 143,238 (HT) | **0.990** (−1.0%) | -| 128 / 2048 | 18,515 (LL) | 18,891 (LL) | **0.980** (−2.0%) | - -**Takeaway (provisional):** across all three shapes the four configs land within ~1–2% of each -other — but note (§1 caveat) that with `dp_size=1` all four ran the **same monolithic TP-all-reduce -path**, so this closeness is largely an artifact of the transport not being on the critical path, -not evidence that the FlashInfer-EP and DeepEP dispatch/combine transports perform equivalently. -The genuine transport comparison requires the `--data-parallel-size 8` re-run. (LL vs HT ordering -within a backend is likewise not meaningful here.) - -> ⚠ **Correction:** earlier throughput numbers reported for this work (~99–103k tok/s -> "128/128") were actually the vLLM **`sonnet` default dataset (~1024 in / 128 out)** — -> `--input-len`/`--output-len` are ignored unless `--dataset-name random` is passed. The table -> above is the corrected sweep (token counts verified: 128/128→128k+128k, 2048/128→2048k+128k, -> 128/2048→128k+2048k). GSM8K numbers were unaffected (lm_eval uses its own data). - -### 1.3 GSM8K — FlashInfer-EP vs DeepEP (accuracy on par) *(monolithic path — accuracy only)* - -| Backend | flexible-extract | strict-match | -|---|---|---| -| `flashinfer_ep_low_latency` | 0.8522 | 0.8939 | -| `flashinfer_ep_high_throughput` | 0.8575 | 0.8969 | -| `deepep_low_latency` | 0.8514 | 0.8931 | -| `deepep_high_throughput` | 0.8544 | 0.8946 | - -All four within run-to-run noise. Note this run used the monolithic path (§1 caveat), so it -validates end-to-end model accuracy but does **not** compare the dispatch/combine transports; the -transport is validated separately by the `--validate` correctness checks in §1.1. - -### 1.3b Memory footprint *(monolithic path, see §1 caveat)* - -From vLLM's engine-init memory profiling (Qwen3-30B-A3B, 8-GPU EP, `--gpu-memory-utilization -0.9`, `--max-model-len 4096`), **all four backends are identical**: - -| Backend | Available KV cache | KV cache size | Max concurrency @ 4096 tok | -|---|---|---|---| -| `flashinfer_ep_low_latency` | 150.45 GiB | 6,573,312 tokens | 1604.8× | -| `flashinfer_ep_high_throughput` | 150.45 GiB | 6,573,312 tokens | 1604.8× | -| `deepep_low_latency` | 150.45 GiB | 6,573,312 tokens | 1604.8× | -| `deepep_high_throughput` | 150.45 GiB | 6,573,312 tokens | 1604.8× | - -**EP backend choice is memory-neutral** here — the dispatch/combine transport buffers do not -measurably reduce the usable KV-cache budget on a B200 (≈180 GiB HBM); the model weights + -activation + transport all fit in the non-KV reservation identically for every backend, leaving -the same 150.45 GiB for KV cache. (vLLM in this build does not emit a grep-able -weights/non-torch/activation split; the identical KV-cache size is the operative footprint -metric.) - -### 1.4 Multi-node (2-node / 16-GPU) — plumbing works, cross-node init blocked (environmental) - -The 2-node path (`benchmarks/_perf_2node.sh`, §4.7) **stands up correctly**: a Ray cluster forms -across both containers (`cluster GPUs so far=16`) and `vllm bench throughput ---data-parallel-size 16 --enable-expert-parallel --distributed-executor-backend ray` with the EP -backend launches. But the vLLM **engine-core init then stalls** (>40 min, no throughput; no error -surfaced in the driver or Ray worker logs). (Original runs used `--tensor-parallel-size 16`, which -would additionally take the monolithic path per §1 caveat; the DP-EP form is the correct config -and stalls at the same cross-node init step below regardless.) - -**Isolation test — this is NOT the EP integration.** A plain 2-node run with -`--tensor-parallel-size 16` and **no** `--enable-expert-parallel` / all2all backend **stalls -identically** at engine init. So the blocker is the cluster's **cross-node vLLM/NCCL bring-up** -(fabric / NCCL cross-node env / ray-executor init), independent of FlashInfer-EP or DeepEP. -Resolving it is cluster-config work (e.g. `NCCL_SOCKET_IFNAME` / `NCCL_IB_HCA` tuning, verifying -IB reachability between the allocated nodes) — out of scope for this integration, which is fully -validated at 8-GPU single node. - -### 1.5 Not measured - -- **Raw NCCL-EP backend** — exists only in the GitLab `vllm-nccl-moe-integration` fork, not in - upstream vLLM; N/A here. -- **`vllm bench serve` TTFT/TPOT** — the server path hit a >20-min per-worker FlashInfer - cubin-download + JIT startup (plus an shm-broadcast wait) on this image; a serving-startup - issue, not the EP backend (GSM8K drives the same path). Offline `vllm bench throughput` was - used for the headline numbers. - ---- - -## 2. Environment - -- Login: `ssh prenyx` (→ `login-prenyx`, MFA/GSSAPI, `ProxyJump`; a persistent ControlMaster - socket lets subsequent non-interactive `ssh prenyx ''` reuse the session). -- SLURM: `--account=coreai_libraries_cudnn --partition=batch`; **whole-node allocations — do - NOT pass `--gres`** (`Invalid generic resource` otherwise). -- Containers: pyxis/enroot. `--container-name` does **not** persist across separate `srun` - jobs — pass `--container-image=<...>.sqsh` every time (squashfs mounts fast). -- Work dir (shared lustre, mounted `/host` in-container): - `RW=/lustre/fsw/coreai_libraries_cudnn/agopal-moe-ep`. - -```bash -RW=/lustre/fsw/coreai_libraries_cudnn/agopal-moe-ep -``` - ---- - -## 3. Build the container images (one-time, **no Docker** — pyxis/enroot) - -Pre-Nyx has no Docker daemon; images are `.sqsh` files built with **pyxis/enroot** via -`srun --container-save`. Conventions used throughout: -- Enroot registry syntax **`nvcr.io#nvidia/`** (not the docker-style `nvcr.io/nvidia/`). -- **`--container-writable`** — required so in-container `apt`/`pip` installs are captured by - `--container-save`. -- **Whole-node** allocation (no `--gres`); `--container-mounts=$RW:/host` (shared lustre work dir). -- `docker/Dockerfile.*` are the canonical spec (usable only on a machine with Docker). - -### 3.1 FlashInfer-EP base (`flashinfer-ep-pt2605.sqsh`) - -> For the results in this doc the base `.sqsh` was **pre-built and reused** (not rebuilt this -> run). The recipe below is how to (re)create it. Mirrors the real `build.sbatch` artifact. - -Clone FlashInfer at `feat/vllm-moe-ep-api` into `$RW/flashinfer`, then (the base build is long -— run it as an sbatch or a plain `srun`): - -```bash -srun --account=coreai_libraries_cudnn --partition=batch -N1 --ntasks-per-node=1 --time=03:00:00 \ - --container-image="nvcr.io#nvidia/pytorch:26.05-py3" --container-writable \ - --container-save=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host \ - bash -lc 'cd /host/flashinfer && bash docker/install/build_flashinfer_ep_pytorch.sh' -``` -As an sbatch wrapper (mirrors `$RW/build.sbatch`): -```bash -#!/bin/bash -#SBATCH -A coreai_libraries_cudnn -p batch -N1 --time=03:00:00 -J fi_ep_build -RW=/lustre/fsw/coreai_libraries_cudnn/agopal-moe-ep -srun --container-image="nvcr.io#nvidia/pytorch:26.05-py3" --container-writable \ - --container-save=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host \ - bash -lc 'cd /host/flashinfer && bash docker/install/build_flashinfer_ep_pytorch.sh' -``` -The install script pins `nvidia-nccl-cu13==2.30.7`, `nccl4py[cu13]==0.3.1`, `cuda-core==1.0.1`, -`cuda-bindings==13.2.0` and runs `BUILD_NCCL_EP=1 pip install -e ".[nvep]"` (editable, from -`/host/flashinfer`). See `docker/install/build_flashinfer_ep_pytorch.sh` / -`docker/Dockerfile.flashinfer-ep-pytorch`. - -### 3.2 vLLM-from-source image (`vllm-flashinfer-ep.sqsh`) - -Clone vLLM at `feat/flashinfer-ep-all2all` into `$RW/vllm`. Build script `build_vllm.sh` -(strip torch pins, install build-deps under `--no-build-isolation`, add a Rust toolchain since -vLLM bundles `rust/Cargo.toml`, pin torch to the base version): - -```bash -#!/bin/bash -set -eo pipefail -cd /host/vllm -python use_existing_torch.py -TORCH_VER=$(python -c 'import torch;print(torch.__version__.split("+")[0])') -echo "torch==$TORCH_VER" > /tmp/tc.txt -command -v cargo >/dev/null 2>&1 || \ - curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal -export PATH="$HOME/.cargo/bin:$PATH" -PIP_CONSTRAINT=/tmp/tc.txt pip install --no-cache-dir -r requirements/build/cuda.txt -MAX_JOBS=32 VLLM_USE_PRECOMPILED=0 \ - PIP_CONSTRAINT=/tmp/tc.txt pip install --no-cache-dir --no-build-isolation -e . -v -# vLLM's deps pull `flashinfer-python` from PyPI and shadow the branch editable (its moe_ep -# lacks EpLayout/FleetAlgoKnobAllocator) — restore the branch editable. -pip install --no-cache-dir --no-build-isolation --no-deps -e /host/flashinfer -python -c 'import vllm, flashinfer; from flashinfer.moe_ep import EpLayout, FleetAlgoKnobAllocator; \ - print("vllm", vllm.__version__, "| flashinfer", flashinfer.__file__)' -``` -```bash -srun --account=coreai_libraries_cudnn --partition=batch -N1 --ntasks-per-node=1 --time=03:00:00 \ - --container-image=$RW/flashinfer-ep-pt2605.sqsh --container-writable \ - --container-save=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ - bash /host/build_vllm.sh -``` -(`--container-writable` is required so the vLLM install is captured by `--container-save`; the -base image here is a local `.sqsh` file, so no `nvcr.io#` registry prefix.) Equivalent one-shot -on a Docker host: `docker/Dockerfile.vllm-flashinfer-ep`. - -### 3.3 DeepEP image (`vllm-fi-ep-deepep.sqsh`) — for the comparison - -Layer DeepEP + NVSHMEM on the vLLM image via vLLM's own installer. **Two fixes are required -on CUDA 13.2:** (a) `UV_BREAK_SYSTEM_PACKAGES=1` (the installer's `uv pip install --system` -trips PEP-668), and (b) **`TORCH_CUDA_ARCH_LIST=10.0a`** — otherwise DeepEP compiles for -`sm_75` and `ptxas` fails (`Feature 'elect'/'mbarrier'/'cp.async.bulk' requires .target -sm_90 or higher`) because its kernels are Hopper/Blackwell-only. - -```bash -#!/bin/bash # build_deepep.sh -export UV_BREAK_SYSTEM_PACKAGES=1 PIP_BREAK_SYSTEM_PACKAGES=1 UV_SYSTEM_PYTHON=1 -set -eo pipefail -export TORCH_CUDA_ARCH_LIST="10.0a" # B200; sm_90+ features -command -v uv >/dev/null 2>&1 || pip install -q uv # installer uses `uv pip install --system` -export PATH="$HOME/.local/bin:/usr/local/bin:$PATH" # ensure uv is on PATH -TORCH_VER=$(python -c 'import torch;print(torch.__version__.split("+")[0])') -echo "torch==$TORCH_VER" > /tmp/tc.txt -PIP_CONSTRAINT=/tmp/tc.txt UV_CONSTRAINT=/tmp/tc.txt \ - bash /host/vllm/tools/ep_kernels/install_python_libraries.sh --workspace /host/ep_kernels_workspace -# DeepEP asserts torch nccl == the nvidia-nccl wheel (2.30.7); the NGC image also ships a system -# libnccl 2.30.4 that torch loads by default. Force the wheel + bake into /etc/profile.d so -# `bash -lc` runtime runs load 2.30.7 too. -NCCL_LIB=$(python -c 'import nvidia.nccl,os;print(os.path.join(list(nvidia.nccl.__path__)[0],"lib"))') -echo "export LD_LIBRARY_PATH=\"$NCCL_LIB:\${LD_LIBRARY_PATH:-}\"" > /etc/profile.d/zz_nccl_wheel.sh -export LD_LIBRARY_PATH="$NCCL_LIB:${LD_LIBRARY_PATH:-}" -# `deep_ep OK` = DeepEP's assert (loaded libnccl == wheel) passed. Don't judge by -# torch.cuda.nccl.version() — that's torch's BUILD-time NCCL (cosmetic), not the loaded .so. -python -c 'import ctypes, torch; torch.cuda.init(); import deep_ep; \ - m=sorted({l.split()[-1] for l in open("/proc/self/maps") if "libnccl.so" in l}); \ - lib=ctypes.CDLL(m[0]); v=ctypes.c_int(); lib.ncclGetVersion(ctypes.byref(v)); \ - print("deep_ep OK; loaded", m, "ncclGetVersion", v.value)' -``` -```bash -srun --account=coreai_libraries_cudnn --partition=batch -N1 --ntasks-per-node=1 --time=02:00:00 \ - --container-image=$RW/vllm-flashinfer-ep.sqsh --container-writable \ - --container-save=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ - bash /host/build_deepep.sh -``` -Installs NVSHMEM 3.3.24 + `DeepEP@d4f41e4e93` (with the installer's CUDA-13 cccl patch). -The DeepEP runs in §4.4/§4.5 use this image with `BACKEND=deepep_low_latency` / -`deepep_high_throughput`. - -> **Getting the FlashInfer branch code into the vLLM image at runtime.** `moe_ep` is pure -> Python, editable-installed from `/host/flashinfer`. To run branch code without rebuilding, -> mount a checkout of `feat/vllm-moe-ep-api` and prepend it to `PYTHONPATH`. In this -> validation a **git worktree** `$RW/fi-vllmep` was used (to avoid touching a dirty primary -> checkout); it needs `3rdparty/{cutlass,spdlog,cccl}` present (init or symlink to a populated -> checkout), `flashinfer/_build_meta.py`, and the `flashinfer/data/*` symlinks. Simplest for a -> fresh repro: check the branch out directly in `$RW/flashinfer` before building 3.1 so the -> editable install already points at branch code, and skip `PYTHONPATH`. - ---- - -## 4. Reproduce each result - -Common env inside every run: `export PYTHONPATH=/host/fi-vllmep` (only if using the worktree), -`HF_HOME=/host/hf_cache`, `FLASHINFER_WORKSPACE_BASE=/host/fi_cache`, -`FLASHINFER_CUBIN_DIR=/host/fi_cubins` (persist JIT + cubins across runs), `NCCL_GIN_TYPE=3`. - -### 4.1 GAP unit tests (host-only, ~1 min) -```bash -srun ... --container-image=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host bash -lc ' - export PYTHONPATH=/host/fi-vllmep; python -m pip install -q pytest - cd /host/fi-vllmep && python -m pytest \ - tests/moe_ep/nccl_ep/test_gaps_mock.py tests/moe_ep/nccl_ep/test_fleet_mock.py -q' -``` - -### 4.2 EP dispatch+combine correctness @ world=8 (`--validate`) -Per-rank runner `benchmarks/run_ep_matrix_one_pt.sh` (points `cd /host/flashinfer`; for the -worktree use a copy that `cd`s to `/host/fi-vllmep` + sets `PYTHONPATH`). Launch **8 tasks/node** -with a `file://` rendezvous: -```bash -srun ... --ntasks-per-node=8 --container-image=$RW/flashinfer-ep-pt2605.sqsh \ - --container-mounts=$RW:/host bash -lc \ - 'EP_SYNC=/host/sync_ht NCCL_GIN_TYPE=3 bash /host//benchmarks/run_ep_matrix_one_pt.sh \ - --algorithm ht --layout fl --tokens 4096 --hidden 7168 --top-k 8 --experts 256 \ - --warmup 5 --iters 10 --validate' -# LL: --algorithm ll --layout em --tokens 128 (same COMMON args) -``` -Expect: `[validate] ht/fl world=8 per_rank=4096 dispatch+combine OK` (and `ll/em ... OK`). - -> The pytest `tests/moe_ep/test_moe_ep_ht_correctness.py` (launched via `torchrun`) hangs on a -> default-PG collective on this image — use the comm-matrix `file://` `--validate` path above. - -### 4.3 vLLM e2e smoke (both backends) -Run a **file-based** program (vLLM uses `spawn`; a heredoc/`stdin` program fails the workers): -```python -# _vllm_smoke_prog.py -import os -from vllm import LLM, SamplingParams -def main(): - llm = LLM(model=os.environ["MODEL"], tensor_parallel_size=8, - enable_expert_parallel=True, all2all_backend=os.environ["BACKEND"], - trust_remote_code=True, enforce_eager=True, max_model_len=2048) - out = llm.generate(["San Francisco is a"], SamplingParams(max_tokens=32, temperature=0.0)) - print(repr(out[0].outputs[0].text)) -if __name__ == "__main__": - main() -``` -```bash -srun ... --container-image=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ - --export=ALL,MODEL=allenai/OLMoE-1B-7B-0924,BACKEND=flashinfer_ep_low_latency bash -lc ' - export PYTHONPATH=/host/fi-vllmep HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 \ - FLASHINFER_WORKSPACE_BASE=/host/fi_cache - python -u /host/fi-vllmep/benchmarks/_vllm_smoke_prog.py' -# repeat with BACKEND=flashinfer_ep_high_throughput -``` - -### 4.4 GSM8K 5-shot (accuracy gate) -```bash -srun ... --container-image=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ - --export=ALL,BACKEND=flashinfer_ep_low_latency bash -lc ' - export PYTHONPATH=/host/fi-vllmep HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 \ - FLASHINFER_WORKSPACE_BASE=/host/fi_cache - python -m pip install -q lm_eval - lm_eval --model vllm --tasks gsm8k --num_fewshot 5 --batch_size auto \ - --model_args pretrained=Qwen/Qwen3-30B-A3B,tensor_parallel_size=8,enable_expert_parallel=True,all2all_backend=$BACKEND,trust_remote_code=True,max_model_len=4096,enforce_eager=True' -# repeat with BACKEND=flashinfer_ep_high_throughput -``` -> ⚠ **GSM8K is an accuracy gate only — it does not exercise the all2all transport.** lm_eval's own -> `data_parallel_size` launches independent replica engines (each `dp_size=1` ⇒ monolithic path), -> so there is no unified EP group to dispatch through; keep `tensor_parallel_size=8`. The transport -> is validated separately by the §4.2 `--validate` round-trip and the nsys dispatch/combine capture -> (runbook §3e). - -### 4.5 Throughput sweep -**Use `--dataset-name random` with `--random-input-len`/`--random-output-len`** — otherwise -`vllm bench throughput` falls back to the `sonnet` dataset (~1024/128) and silently ignores -`--input-len`/`--output-len`. -```bash -# ISL/OSL ∈ {128/128, 2048/128, 128/2048}; backend ∈ {flashinfer_ep_low_latency, ...high_throughput} -srun ... --container-image=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ - --export=ALL,BACKEND=flashinfer_ep_low_latency,ISL=2048,OSL=128 bash -lc ' - export PYTHONPATH=/host/fi-vllmep HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 \ - FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins - cd /tmp; torchrun --nproc_per_node=8 /host/dprun/driver.py --model Qwen/Qwen3-30B-A3B \ - --dataset-name random --random-input-len $ISL --random-output-len $OSL --num-prompts 1000 \ - --data-parallel-size 8 --distributed-executor-backend external_launcher \ - --enable-expert-parallel --all2all-backend $BACKEND \ - --trust-remote-code --max-model-len 4096 --enforce-eager' -``` -> ⚠ **DP-EP via `torchrun` + `external_launcher`** (needs `$RW/dprun/driver.py`, runbook §3.0) — -> plain `vllm bench throughput --data-parallel-size 8` errors offline. With `--tensor-parallel-size -> 8` you get the monolithic path (§1 caveat). Each of the 8 ranks prints its own `Throughput:`; -> the deployment total ≈ their sum. Confirm the log says `Using FlashInferEP…/DeepEP…`, not -> `…Monolithic`. - -### 4.6 DeepEP comparison (§1.2/§1.3) -Identical to §4.4/§4.5 but use the DeepEP image and DeepEP backend names: -`--container-image=$RW/vllm-fi-ep-deepep.sqsh` and -`BACKEND=deepep_low_latency` / `deepep_high_throughput`. - -### 4.7 Multi-node (2-node / 16-GPU) -`benchmarks/_perf_2node.sh` stands up a Ray cluster across the two nodes' containers (rank-0 -`ray start --head`, rank-1 `ray start --address=` via a shared `/host/ray_head_ip.$JOBID` -file), waits for 16 GPUs, then runs `vllm bench throughput --data-parallel-size 16 ---enable-expert-parallel --distributed-executor-backend ray` with `NCCL_MNNVL_ENABLE=1` for the -cross-node EP fabric. (Data-parallel EP, not TP-16 — see §1 caveat / runbook §3.0.) -```bash -srun --account=coreai_libraries_cudnn --partition=batch -N2 --ntasks-per-node=1 \ - --container-image=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ - --export=ALL,BACKEND=flashinfer_ep_low_latency,ISL=128,OSL=128 \ - bash /host/fi-vllmep/benchmarks/_perf_2node.sh -``` -Warm `/host/fi_cubins` first (from the single-node runs) — a cold cubin cache makes the 16-way -init stall for tens of minutes while every worker downloads FlashInfer cubins. - -### 4.8 Memory footprint (§1.3b) -Same as §4.5 but `--num-prompts 8` (init does the memory profiling regardless) and grep the -KV-cache line from the **full** stream (don't `tail`-truncate): -```bash -cd /tmp; torchrun --nproc_per_node=8 /host/dprun/driver.py --model Qwen/Qwen3-30B-A3B \ - --dataset-name random --random-input-len 128 --random-output-len 128 --num-prompts 8 \ - --data-parallel-size 8 --distributed-executor-backend external_launcher \ - --enable-expert-parallel --all2all-backend $BACKEND \ - --gpu-memory-utilization 0.9 --trust-remote-code --max-model-len 4096 --enforce-eager 2>&1 \ - | grep -iE "Available KV cache|GPU KV cache size|Maximum concurrency|Using .*PrepareAndFinalize" -``` - ---- - -## 5. Gotchas (learned during this validation) - -- **⚠ TP-only EP silently disables the all2all backend (biggest gotcha).** vLLM only takes the - modular EP dispatch/combine path when `dp_size > 1` - (`fused_moe/config.py::use_all2all_kernels = dp_size > 1 and use_ep`). Running - `--tensor-parallel-size 8 --enable-expert-parallel` (dp_size=1) picks - `MoEPrepareAndFinalizeNoDPEPMonolithic` — experts run locally, reconciled by TP all-reduce, and - `--all2all-backend` is a **no-op**. Use `--data-parallel-size 8 --enable-expert-parallel` - instead. **Always confirm** the log prints `Using FlashInferEPLL/HT…` or `DeepEPLL/HT…` - `PrepareAndFinalize`, never `…Monolithic`. (This is why the first §1.2/§1.3 sweep showed - FI-EP ≈ DeepEP and identical nsys kernels — the transport was never on the GPU.) -- **No `--gres`** on this cluster (whole-node); `--container-name` doesn't persist across jobs. -- **`NCCL_GIN_TYPE=3`** for the EP GIN transport; multi-node also needs `NCCL_MNNVL_ENABLE=1`. -- **vLLM `spawn`** re-imports the main module → run a real `.py` file, never a heredoc/stdin. -- **`pytest` / `lm_eval`** are not in the image → `pip install` them in the job. -- Don't recursive-glob `/usr/**` inside the container (pathologically slow → looks like a hang). -- Persist JIT to `/host` (`FLASHINFER_WORKSPACE_BASE`) and cubins - (`FLASHINFER_CUBIN_DIR`) so the ~267-unit CUTLASS MoE compile happens once. -- First vLLM forward JIT-compiles the FlashInfer CUTLASS MoE expert kernel — needs the `cccl` - submodule present in the checkout used at runtime. -- **DeepEP build on CUDA 13.2:** the vLLM installer calls `uv pip install --system`, so `uv` - must be present — `pip install uv` first (the NGC/vLLM image has `pip`, not `uv`). Also set - `UV_BREAK_SYSTEM_PACKAGES=1` (PEP-668) **and** `TORCH_CUDA_ARCH_LIST=10.0a` — without the arch - it builds for `sm_75` and `ptxas` rejects the Hopper/Blackwell-only features (`elect`, - `mbarrier`, `cp.async.bulk`). -- **vLLM shadows the branch flashinfer:** `pip install -e .` (vLLM) pulls `flashinfer-python` - from PyPI and uninstalls the branch editable → `ImportError: cannot import name 'EpLayout'` - (and the `flashinfer_ep_*` backend breaks: no `FleetAlgoKnobAllocator`). Fix: re-run - `pip install --no-build-isolation --no-deps -e /host/flashinfer` after the vLLM install (or - set `PYTHONPATH=/host/flashinfer` on every run). Verify `flashinfer.__file__` → - `/host/flashinfer/...`. -- **DeepEP NCCL-version assert:** DeepEP requires torch's loaded NCCL to equal the nvidia-nccl - wheel (2.30.7), but the NGC image also has a system `libnccl.so.2.30.4` that torch loads by - default → `AssertionError: Invalid NCCL versions: ...2.30.4 (loaded) v.s. ...wheel...`. Fix: - prepend the wheel's `nvidia/nccl/lib` to `LD_LIBRARY_PATH` (2.30.x is ABI-compatible) and bake - it into `/etc/profile.d` so `bash -lc` runs inherit it. diff --git a/docs/design_docs/vllm_moe_ep_runbook.md b/docs/design_docs/vllm_moe_ep_runbook.md deleted file mode 100644 index 3d3eee06e02..00000000000 --- a/docs/design_docs/vllm_moe_ep_runbook.md +++ /dev/null @@ -1,566 +0,0 @@ -# Runbook — NCCL-EP (FlashInfer) vs DeepEP on Pre-Nyx (pyxis/enroot, no Docker) - -Linear, copy-paste steps to build the container images **with pyxis (no Docker daemon)** and run -the **NCCL-EP vs DeepEP** comparison in vLLM, on **1 node (8 GPU)** and **2 nodes (16 GPU)**. - -- **NCCL-EP** = the `nccl.ep`-backed vLLM all2all backends `flashinfer_ep_low_latency` / - `flashinfer_ep_high_throughput` (FlashInfer branch **`feat/nvep-default`** ≥ `fa09bc46` + - vLLM branch **`feat/flashinfer-ep-all2all`** ≥ `ab1415e` — these commits carry the HT - token-cap clamp, the HT recv-trim, the batched-DP cap membership and the fleet host-path - caches; older refs reproduce the pre-optimization 2–6× gap). The standalone raw-`nccl.ep` - backend is **not** in upstream vLLM (GitLab-fork only), so it's out of scope here. -- **DeepEP** = `deepep_low_latency` / `deepep_high_throughput`. -- Model: `Qwen/Qwen3-30B-A3B` (128 experts, bf16). Base: `nvcr.io/nvidia/pytorch:26.05-py3` - (CUDA 13.2 — required; older stacks abort cross-node HT at `nccl_ep.cc:2884`). - -Everything runs from a shared-FS work dir mounted `/host` inside the container. - -```bash -# --- run once per shell --- -RW=/lustre/fsw/coreai_libraries_cudnn/agopal/agopal-moe-ep-verif # a FRESH shared-FS work dir -ACCT=coreai_libraries_cudnn ; PART=batch -mkdir -p $RW/logs -``` - -Starting **from scratch** in a fresh `$RW`? Run §1 (clone) → §2 (build all three images, -nothing is reused) → §3 (single-node runs) → §4 (multi-node). Everything below is written to be -copy-pasted top-to-bottom into a shell that has already `export`ed `RW`/`ACCT`/`PART` above. - ---- - -## 0. Why pyxis, not Docker - -Pre-Nyx login/compute nodes have **no Docker daemon**. Container images are **enroot squashfs -(`.sqsh`) files**, built by running the install steps *inside* a container under `srun` and -snapshotting it with `--container-save`: - -- `--container-image="nvcr.io#nvidia/"` — enroot registry syntax (note the **`#`**, not - `/`). A local image is just a path: `--container-image=$RW/foo.sqsh`. -- `--container-writable` — **required** so `apt`/`pip`/build outputs are captured by the save. -- `--container-save=$RW/out.sqsh` — writes the resulting image. -- `--container-mounts=$RW:/host` — the shared FS shows up at `/host` in the container. -- Whole-node allocation only — **do not pass `--gres`** (rejected on this cluster). -- A `--container-name` does **not** persist across separate `srun` jobs; always pass - `--container-image=<...>.sqsh`. - ---- - -## 1. Prerequisites — create the fresh dir and clone the repos - -```bash -# (RW/ACCT/PART already exported above — e.g. RW=.../agopal-moe-ep-verif) -mkdir -p $RW/logs -# Pinned refs — these carry ALL the DP-EP fixes/optimizations the reference numbers were -# measured with (HT clamp+trim, batched-DP cap membership, fleet host-path caches): -git clone -b feat/nvep-default https://github.com/Anerudhan/flashinfer.git $RW/flashinfer -git clone -b feat/flashinfer-ep-all2all https://github.com/Anerudhan/vllm.git $RW/vllm -git -C $RW/flashinfer submodule update --init --recursive # cutlass, cccl, spdlog, nccl -# sanity: the perf-critical commits must be present -git -C $RW/flashinfer merge-base --is-ancestor fa09bc46 HEAD && echo "flashinfer ref OK" -git -C $RW/vllm merge-base --is-ancestor ab1415e HEAD && echo "vllm ref OK" -``` - ---- - -## 2. Build the three images with pyxis - -> **Rebuilds vs `git pull`:** the images install flashinfer and vLLM as *editable* installs -> pointing at `/host/flashinfer` / `/host/vllm` — i.e. at **your `$RW` clones, resolved at -> runtime**. Python-only changes (all the perf fixes in the pinned refs are Python) take effect -> by just updating the clones (§1); **no image rebuild needed**. Rebuild only for dependency/ -> native changes (e.g. a new nccl4py/nvidia-nccl pin). - -### 2a. FlashInfer-EP base — `flashinfer-ep-pt2605.sqsh` -Runs `docker/install/build_flashinfer_ep_pytorch.sh` (pins nvidia-nccl-cu13 2.30.7 / nccl4py -0.3.1 / cuda-core 1.0.1 / cuda-bindings 13.2.0, then `BUILD_NCCL_EP=1 pip install -e .[nvep]`). - -```bash -srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=03:00:00 \ - --container-image="nvcr.io#nvidia/pytorch:26.05-py3" --container-writable \ - --container-save=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host \ - bash -lc 'cd /host/flashinfer && bash docker/install/build_flashinfer_ep_pytorch.sh' -# sanity: -srun -A $ACCT -p $PART -N1 --container-image=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host \ - bash -lc "python -c \"from flashinfer.moe_ep import available_backends; print(available_backends())\"" -# expect: ['nccl_ep'] -``` - -### 2b. vLLM from source — `vllm-flashinfer-ep.sqsh` -Create `$RW/build_vllm.sh`: -```bash -cat > $RW/build_vllm.sh <<'EOS' -#!/bin/bash -set -eo pipefail -cd /host/vllm -python use_existing_torch.py # strip torch==2.11 pin -> use NGC torch -TORCH_VER=$(python -c 'import torch;print(torch.__version__.split("+")[0])'); echo "torch==$TORCH_VER" > /tmp/tc.txt -command -v cargo >/dev/null 2>&1 || curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y --profile minimal -export PATH="$HOME/.cargo/bin:$PATH" # vLLM bundles a Rust crate -PIP_CONSTRAINT=/tmp/tc.txt pip install --no-cache-dir -r requirements/build/cuda.txt -MAX_JOBS=32 VLLM_USE_PRECOMPILED=0 PIP_CONSTRAINT=/tmp/tc.txt \ - pip install --no-cache-dir --no-build-isolation -e . -v -# vLLM's deps pull `flashinfer-python` from PyPI, which UNINSTALLS/shadows our branch editable -# (the PyPI moe_ep lacks EpLayout / FleetAlgoKnobAllocator). Restore the branch editable so -# `import flashinfer` resolves to /host/flashinfer at runtime. -pip install --no-cache-dir --no-build-isolation --no-deps -e /host/flashinfer -python -c 'import vllm, flashinfer; from flashinfer.moe_ep import EpLayout, FleetAlgoKnobAllocator; \ - print("vllm", vllm.__version__, "| flashinfer", flashinfer.__file__)' -EOS -``` -Build: -```bash -srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=03:00:00 \ - --container-image=$RW/flashinfer-ep-pt2605.sqsh --container-writable \ - --container-save=$RW/vllm-flashinfer-ep.sqsh --container-mounts=$RW:/host \ - bash /host/build_vllm.sh -``` - -### 2c. DeepEP — `vllm-fi-ep-deepep.sqsh` -Four things to get right on CUDA 13.2: (1) `uv` must be installed (the vLLM installer runs -`uv pip install --system`); (2) `UV_BREAK_SYSTEM_PACKAGES=1` (PEP-668); (3) -**`TORCH_CUDA_ARCH_LIST=10.0a`** (else DeepEP builds `sm_75` and `ptxas` rejects its -`elect`/`mbarrier`/`cp.async.bulk` kernels); and (4) **torch must load the 2.30.7 wheel -libnccl** — DeepEP asserts *torch nccl == the nvidia-nccl wheel*, but the NGC image also has a -system libnccl 2.30.4 that torch loads by default, so the script forces the wheel onto -`LD_LIBRARY_PATH` and bakes it into `/etc/profile.d`. (Symptom if skipped: -`AssertionError: Invalid NCCL versions: ...2.30.4 (loaded) v.s. ...nvidia/nccl/lib/libnccl.so.2 -(expected)`.) Create `$RW/build_deepep.sh`: -```bash -cat > $RW/build_deepep.sh <<'EOS' -#!/bin/bash -export UV_BREAK_SYSTEM_PACKAGES=1 PIP_BREAK_SYSTEM_PACKAGES=1 UV_SYSTEM_PYTHON=1 -set -eo pipefail -export TORCH_CUDA_ARCH_LIST="10.0a" # B200 / sm_100 -command -v uv >/dev/null 2>&1 || pip install -q uv # installer calls `uv pip install --system` -export PATH="$HOME/.local/bin:/usr/local/bin:$PATH" # ensure `uv` is on PATH -TORCH_VER=$(python -c 'import torch;print(torch.__version__.split("+")[0])'); echo "torch==$TORCH_VER" > /tmp/tc.txt -PIP_CONSTRAINT=/tmp/tc.txt UV_CONSTRAINT=/tmp/tc.txt \ - bash /host/vllm/tools/ep_kernels/install_python_libraries.sh --workspace /host/ep_kernels_workspace -# guard: ensure the branch flashinfer editable is still the active install (nothing shadowed it) -pip install --no-cache-dir --no-build-isolation --no-deps -e /host/flashinfer -# DeepEP asserts torch's loaded NCCL == the nvidia-nccl wheel (2.30.7). The NGC image also has a -# system libnccl 2.30.4 that torch loads by default -> mismatch. Force the wheel's libnccl first, -# and bake it into /etc/profile.d so every `bash -lc` runtime run picks it up. -NCCL_LIB=$(python -c 'import nvidia.nccl,os;print(os.path.join(list(nvidia.nccl.__path__)[0],"lib"))') -echo "export LD_LIBRARY_PATH=\"$NCCL_LIB:\${LD_LIBRARY_PATH:-}\"" > /etc/profile.d/zz_nccl_wheel.sh -export LD_LIBRARY_PATH="$NCCL_LIB:${LD_LIBRARY_PATH:-}" -# `import deep_ep` runs DeepEP's NCCL assert (loaded libnccl must == the wheel); if it prints OK -# the correct lib is loaded. NOTE: torch.cuda.nccl.version() reports torch's BUILD-time NCCL -# (cosmetic) — to see the actually-loaded runtime lib use ncclGetVersion on /proc/self/maps. -python -c 'import ctypes, torch; torch.cuda.init(); import deep_ep; \ - m=sorted({l.split()[-1] for l in open("/proc/self/maps") if "libnccl.so" in l}); \ - lib=ctypes.CDLL(m[0]); v=ctypes.c_int(); lib.ncclGetVersion(ctypes.byref(v)); \ - print("deep_ep OK; loaded libnccl", m, "runtime ncclGetVersion", v.value)' -EOS -``` -Build: -```bash -srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=02:00:00 \ - --container-image=$RW/vllm-flashinfer-ep.sqsh --container-writable \ - --container-save=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ - bash /host/build_deepep.sh -``` - -> The DeepEP image contains vLLM + FlashInfer-EP + DeepEP, so you can run **all four backends -> from `vllm-fi-ep-deepep.sqsh`**. (The FI-EP image lacks DeepEP.) - -### Warm the caches once (avoids a 30-min first-run JIT/cubin storm, esp. multi-node) -The very first vLLM forward JIT-compiles the CUTLASS MoE kernels (~267 units) and downloads -cubins. Persist them to `/host` so every later run (and every worker) reuses them: -``` -export FLASHINFER_WORKSPACE_BASE=/host/fi_cache # JIT cache -export FLASHINFER_CUBIN_DIR=/host/fi_cubins # cubin cache -``` -(These are already in the run commands below. Run one single-node throughput first to populate.) - ---- - -## 3. Single node (8 GPU) — the comparison - -Common per-run env (put at the top of each `bash -lc '...'`): -``` -export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 \ - FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins -``` - -### ⚠ 3.0 CRITICAL — you MUST use **data-parallel EP** or the all-to-all backend is a no-op - -The `--all2all-backend` flag ONLY takes effect when vLLM selects the **modular EP** dispatch/combine -path. That selection is gated in `vllm/model_executor/layers/fused_moe/config.py`: - -```python -@property -def use_all2all_kernels(self): - return self.dp_size > 1 and self.use_ep # ← dp_size MUST be > 1 -``` - -and every backend predicate (`use_flashinfer_ep_ll_kernels`, `use_deepep_ll_kernels`, …) is -`use_all2all_kernels and all2all_backend == ""`. With **pure `--tensor-parallel-size 8` -(dp_size = 1)** `use_all2all_kernels` is `False`, so `maybe_make_prepare_finalize()` -(`all2all_utils.py:148`) returns the **monolithic** `MoEPrepareAndFinalizeNoDPEPMonolithic` — the -experts run locally per rank and are reconciled with the ordinary **TP all-reduce**. The dispatch/ -combine transport (nccl.ep for FlashInfer-EP, NVSHMEM for DeepEP) is **never launched**, and the -two backends produce byte-identical communication kernels. This is exactly what our first nsys -capture showed (both logged `Using MoEPrepareAndFinalizeNoDPEPMonolithic`). - -**Correct config for a single 8-GPU node** (EP=8 across the 8 ranks, all2all engaged): -`--data-parallel-size 8 --enable-expert-parallel` (with TP=1, so DP×TP = 8 GPUs). Qwen3-30B-A3B in -bf16 fits per-GPU on B200/GB200 (non-expert weights replicated per DP rank; the 128 experts shard -16-per-rank). When this engages, the log shows the expert backend flip to a *batched* one -(`Using BATCHED_TRITON …` or `FlashInfer CUTLASS`) and `Using FlashInferEPLLPrepareAndFinalize`. - -**⚠ Offline `vllm bench throughput` cannot take `--data-parallel-size` directly** — it errors -`Data parallel is only supported with external launcher mode with synchronous engine in offline -benchmark` (`benchmarks/throughput.py:914`). You must launch it under **`torchrun` (one process per -DP rank) with `--distributed-executor-backend external_launcher`**. Create a tiny driver once -(the `sys.path` scrub is required — torchrun prepends the script dir, and a sibling `vllm/` repo -dir would otherwise shadow the installed `vllm` package → `ModuleNotFoundError: -vllm.benchmarks.throughput`): -```bash -mkdir -p $RW/dprun -cat > $RW/dprun/driver.py <<'PY' -import sys -sys.path = [p for p in sys.path if p not in ("", "/host", "/host/dprun")] -from vllm.utils.argparse_utils import FlexibleArgumentParser -from vllm.benchmarks.throughput import add_cli_args, main -p = FlexibleArgumentParser(); add_cli_args(p); main(p.parse_args()) -PY -``` -Then every offline throughput/nsys run below uses `cd /tmp && torchrun --nproc_per_node=8 -/host/dprun/driver.py --data-parallel-size 8 --distributed-executor-backend -external_launcher …` (each rank prints its own `Throughput:`; the deployment total ≈ sum of the 8). -`lm_eval` (GSM8K, §3b) does its own dataset-sharding DP and does **not** build a unified EP group, -so it stays an *accuracy* check on the monolithic path — the transport is validated by §3a -`--validate` and the §3e nsys dispatch/combine capture, not by GSM8K. To exercise the transport -end-to-end with GSM8K, use the **server** path instead (`vllm serve --data-parallel-size 8 ---enable-expert-parallel --all2all-backend ` + `lm_eval --model local-completions ---model_args base_url=http://127.0.0.1:8000/v1/completions,...`; needs `pip install lm-eval[api]`). - -**HT backends (`*_high_throughput`) require `--max-num-batched-tokens 8192`.** nccl_ep HT hard-caps -`max_dispatch_tokens_per_rank` at `MAX_SUPPORTED_TOKENS_PER_RANK=8192`; the flashinfer fleet clamps -to it, but a single forward with more than 8192 tokens per rank raises at dispatch — so cap the -scheduler to match. (DeepEP-LL, conversely, *rejects* `--max-num-batched-tokens 8192`, so only pass -it for the HT runs.) - -**LL backends: do NOT pass `--max-num-batched-tokens` at all.** `flashinfer_ep_low_latency` is in -vLLM's `use_batched_dp_moe` set (with `deepep_low_latency`/`nixl_ep`), so when the flag is unset the -scheduler auto-caps to the 256-token batched-DP budget the BatchedExperts format needs — the padded -`[local_experts, max_tokens×world, N]` workspaces (and the LL transport slot buffers) are sized from -it, and an explicit large value silently makes every fill/activation/GEMM pad 32× (this was the -original 2–5× perf gap vs DeepEP; see results doc §1.1f). - -**Always verify the path** after every run — the log must NOT say `Monolithic`: -```bash -grep -h "Using .*PrepareAndFinalize" $RW/logs/ # expect one of: -# Using FlashInferEPLLPrepareAndFinalize (flashinfer_ep_low_latency) -# Using FlashInferEPHTPrepareAndFinalize (flashinfer_ep_high_throughput) -# Using DeepEPLLPrepareAndFinalize (deepep_low_latency) -# Using DeepEPHTPrepareAndFinalize (deepep_high_throughput) -# ✗ BUG (transport NOT exercised): Using MoEPrepareAndFinalizeNoDPEPMonolithic -``` -(vLLM logs this via the oracle at `oracle/unquantized.py:332` — `logger.info_once("Using %s", …)`.) - -### 3a. (optional) Correctness — EP dispatch/combine `--validate` @ world=8 -```bash -srun -A $ACCT -p $PART -N1 --ntasks-per-node=8 --time=00:25:00 \ - --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ - bash -lc 'EP_SYNC=/host/sync_ht NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ - bash /host/flashinfer/benchmarks/run_ep_matrix_one_pt.sh \ - --algorithm ht --layout fl --tokens 4096 --hidden 7168 --top-k 8 --experts 256 --validate' -# LL: --algorithm ll --layout em --tokens 128 --validate (expect "... dispatch+combine OK") -``` - -### 3b. GSM8K accuracy (5-shot) — loop the 4 backends -```bash -for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do - srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=01:00:00 \ - --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ - bash -lc "export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ - FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins; \ - python -m pip install -q lm_eval; \ - lm_eval --model vllm --tasks gsm8k --num_fewshot 5 --batch_size auto \ - --model_args pretrained=Qwen/Qwen3-30B-A3B,tensor_parallel_size=8,enable_expert_parallel=True,all2all_backend=$B,trust_remote_code=True,max_model_len=4096,enforce_eager=True" \ - > $RW/logs/gsm8k_${B}.log 2>&1 & -done; wait -grep -H "flexible-extract\|strict-match" $RW/logs/gsm8k_*.log -``` -> **GSM8K is an accuracy gate only — it does NOT exercise the all2all transport.** lm_eval's own -> `data_parallel_size` spins up independent replica engines (each `dp_size=1` internally ⇒ -> monolithic path), so there is no single cross-rank EP group to dispatch through. Keep -> `tensor_parallel_size=8` here; it confirms end-to-end model accuracy is correct with the backend -> selected. The **transport** is validated by §3a (`--validate`), §3b′ below, and the §3e nsys capture. - -### 3b′. GSM8K THROUGH the transport (real DP-EP server) — the one that exercises dispatch/combine -Serve a genuine DP-EP engine, then point lm_eval's OpenAI-compatible client at it. Unlike §3b this -runs GSM8K over the actual all2all transport. One srun starts the server, waits for `/health`, runs -the eval, and tears down. HT backends need `--max-num-batched-tokens 8192` (see §3.0). -```bash -for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput; do - case $B in *high_throughput*) CAP="--max-num-batched-tokens 8192";; *) CAP="";; esac - srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=01:00:00 \ - --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ - bash -lc "export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ - FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins; \ - python -m pip install -q lm_eval tenacity; \ - vllm serve Qwen/Qwen3-30B-A3B --port 8000 --data-parallel-size 8 --enable-expert-parallel \ - --all2all-backend $B $CAP --trust-remote-code --max-model-len 4096 --enforce-eager \ - > /host/logs/serve_${B}.log 2>&1 & SP=\$!; \ - for i in \$(seq 1 120); do curl -sf http://127.0.0.1:8000/health && break; \ - kill -0 \$SP || { echo SERVER_DIED; tail -30 /host/logs/serve_${B}.log; exit 1; }; sleep 10; done; \ - lm_eval --model local-completions --tasks gsm8k --num_fewshot 5 --batch_size 1 \ - --model_args model=Qwen/Qwen3-30B-A3B,base_url=http://127.0.0.1:8000/v1/completions,num_concurrent=128,tokenized_requests=False; \ - kill \$SP" \ - > $RW/logs/gsm8k_dpep_${B}.log 2>&1 - grep -H "flexible-extract\|strict-match\|Using .*PrepareAndFinalize" $RW/logs/gsm8k_dpep_${B}.log -done -# Expect (transport-exercised, at the pinned refs): LL 0.856/0.898, HT 0.857/0.898 -# (flex/strict; ±0.01 run-to-run). Log says FlashInferEPLL/HTPrepareAndFinalize (NOT -# Monolithic). Needs lm-eval[api] (tenacity). -``` - -### 3c. Throughput — 4 backends × 3 ISL/OSL shapes -**Must pass `--dataset-name random`** or vLLM silently uses the `sonnet` default (~1024/128) and -ignores `--input-len`/`--output-len`. -Launched under `torchrun` + `external_launcher` (see §3.0) so DP-EP engages. Each of the 8 ranks -prints its own `Throughput:`; sum them (or ×8 the mean) for the deployment total. -```bash -for SHAPE in "128 128" "2048 128" "128 2048"; do set -- $SHAPE; ISL=$1; OSL=$2 - for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do - # HT backends need the 8192 cap (nccl_ep HT limit); LL backends must NOT get it - # (deepep_low_latency rejects --max-num-batched-tokens 8192). See §3.0. - case $B in *high_throughput*) CAP="--max-num-batched-tokens 8192";; *) CAP="";; esac - srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=01:15:00 \ - --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ - bash -lc "export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ - FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins; cd /tmp; \ - torchrun --nproc_per_node=8 /host/dprun/driver.py --model Qwen/Qwen3-30B-A3B \ - --dataset-name random --random-input-len $ISL --random-output-len $OSL --num-prompts 256 \ - --data-parallel-size 8 --distributed-executor-backend external_launcher \ - --enable-expert-parallel --all2all-backend $B $CAP \ - --trust-remote-code --max-model-len 4096 --enforce-eager" \ - > $RW/logs/tp_${B}_${ISL}x${OSL}.log 2>&1 & - done; wait # (or drop `wait` to run shapes concurrently if you have the nodes) -done -# Deployment total = SUM of the 8 per-rank Throughput lines; this prints it per run: -for f in $RW/logs/tp_*.log; do - tot=$(grep -a "Throughput:" "$f" | grep -oE "[0-9.]+ total tokens/s" | grep -oE "^[0-9.]+" \ - | awk '{s+=$1} END{printf "%.0f", s}') - pf=$(grep -a "Using .*PrepareAndFinalize" "$f" | head -1 | grep -oE "Using \w+") - echo "$(basename $f) sum_total_tok/s=$tot [$pf]" # [..] must be FlashInferEP…/DeepEP…, NOT Monolithic -done -``` -> `--num-prompts 256` matches the reference matrix below exactly; larger NP (e.g. 1000) runs -> longer/steadier but shifts absolute numbers — keep it fixed when comparing backends. -> Requires the `$RW/dprun/driver.py` from §3.0. Plain `vllm bench throughput --data-parallel-size 8` -> does **not** work offline (it errors and tells you to use external launcher / serving). -> These are **eager-mode DP-EP** numbers (transport on the critical path); they are far below the -> CUDA-graph monolithic numbers — see the measured values + interpretation in -> `vllm_moe_ep_results_prenyx.md` §1.1e. - -### 3d. (optional) Memory footprint -Same as 3c but `--num-prompts 8`, and grep the KV-cache line from the **full** stream (add -`--max-num-batched-tokens 8192` for HT backends, see §3.0): -```bash -srun ... bash -lc "... cd /tmp; torchrun --nproc_per_node=8 /host/dprun/driver.py \ - --model Qwen/Qwen3-30B-A3B --dataset-name random \ - --random-input-len 128 --random-output-len 128 --num-prompts 8 --data-parallel-size 8 \ - --distributed-executor-backend external_launcher \ - --enable-expert-parallel --all2all-backend $B $CAP --gpu-memory-utilization 0.9 \ - --trust-remote-code --max-model-len 4096 --enforce-eager 2>&1 \ - | grep -iE 'Available KV cache|GPU KV cache size|Maximum concurrency|Using .*PrepareAndFinalize'" -``` - -### 3e. Capture the launched kernels with nsys, and list the all-to-all kernels - -Profile a **short** run under Nsight Systems, then dump the GPU kernel summary and filter for -the EP dispatch/combine (all-to-all) kernels. Use `--enforce-eager` (already set) so kernels are -launched individually (not hidden inside CUDA graphs), and a tiny `--num-prompts` so the report -is small. `nsys profile` follows the vLLM worker child processes, so all ranks' GPU kernels -land in one `.nsys-rep`. - -> **Use the DP-EP `torchrun` form** (see §3.0). With `--tensor-parallel-size 8` the run takes the -> monolithic path and the `.nsys-rep` will contain **no** dispatch/combine kernels — which defeats -> the purpose. `nsys` wraps `torchrun`; `--trace-fork-before-exec=true` makes it follow the 8 -> external-launcher rank processes so all ranks' kernels land in one `.nsys-rep`. - -```bash -for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do - case $B in *high_throughput*) CAP="--max-num-batched-tokens 8192";; *) CAP="";; esac # see §3.0 - srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=00:45:00 \ - --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ - bash -lc "export HF_HOME=/host/hf_cache NCCL_GIN_TYPE=3 FLASHINFER_DISABLE_VERSION_CHECK=1 \ - FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins; cd /tmp; \ - nsys profile -t cuda,nvtx,nccl --force-overwrite true --sample=none --cpuctxsw=none \ - --trace-fork-before-exec=true -o /host/logs/nsys_dpep_${B} \ - torchrun --nproc_per_node=8 /host/dprun/driver.py --model Qwen/Qwen3-30B-A3B \ - --dataset-name random --random-input-len 128 --random-output-len 128 --num-prompts 64 \ - --data-parallel-size 8 --distributed-executor-backend external_launcher \ - --enable-expert-parallel --all2all-backend $B $CAP \ - --trust-remote-code --max-model-len 4096 --enforce-eager" \ - > $RW/logs/nsys_dpep_${B}.log 2>&1 -done -# GATE: confirm every run actually took the modular EP path before trusting the kernel dump. -grep -H "Using .*PrepareAndFinalize" $RW/logs/nsys_dpep_*.log -# want: FlashInferEPLL/HT... or DeepEPLL/HT...PrepareAndFinalize ✗ reject: ...Monolithic -``` - -Dump the per-kernel GPU-time summary and **filter to the all-to-all / EP kernels**. `nsys` is -only inside the container (not on the login node), so run the parsing under `srun` too — it -needs no GPU, just reads the `.nsys-rep`. The `kern_*.txt` land in `$RW/logs/` (via `/host`), so -you can grep them afterward on the login node. -```bash -srun -A $ACCT -p $PART -N1 --ntasks-per-node=1 --time=00:15:00 \ - --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ - bash -lc 'for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do - echo "############ $B — GPU kernel summary ############" - nsys stats --report cuda_gpu_kern_sum --format table /host/logs/nsys_dpep_${B}.nsys-rep \ - | tee /host/logs/kern_dpep_${B}.txt | head -40 - echo "---- all-to-all / EP dispatch-combine kernels only ----" - grep -iE "nccl.?ep|gin|gdaki|hybridep|deep_?ep|nvshmem|intranode|internode|dispatch|combine|moe.?ep" \ - /host/logs/kern_dpep_${B}.txt | tee /host/logs/kern_a2a_${B}.txt - [ -s /host/logs/kern_a2a_${B}.txt ] || echo " (NONE — check §3.0: run took the monolithic path?)" - done' -# If nsys is not on PATH in the image: NSYS=$(ls /opt/nvidia/nsight-systems/*/bin/nsys | head -1); use "$NSYS" stats ... -``` - -**What to expect** — the exact symbols observed on this stack (Qwen3-30B-A3B, 8×DP-EP, 64 prompts, -eager): -- **FlashInfer-EP LL** (`nccl.ep`): `nccl_ep::internode_ll::dispatch` + `nccl_ep::internode_ll::combine`. -- **FlashInfer-EP HT** (`nccl.ep`, JIT): `nccl_ep_jit_ht_dispatch_kernel` + `nccl_ep_jit_ht_combine_kernel` - + `nccl_ep_jit_ht_scan_kernel` (FLAT metadata) + `nccl_ep::hybridep::{dense_to_sparse_prob, - convert_topk_to_routing_map,sparse_to_dense_prob}` helpers. -- **DeepEP LL** (`deep_ep`+NVSHMEM): `deep_ep::legacy::internode_ll::dispatch` + - `…::internode_ll::combine` + one-time `nvshmemi_init_array_kernel`. -- **DeepEP HT** (`deep_ep` intranode): `deep_ep::legacy::intranode::{notify_dispatch,dispatch, - combine,cached_notify_combine}` + `…::layout::get_dispatch_layout`. -- Both also show the shared **expert GEMM** and attention/norm kernels — those are *not* - all-to-all; the grep above narrows to transport. - -> **Reading the times (important):** the LL dispatch/combine kernels **busy-wait on the network**, -> so `Total`/`Avg`/`Max` are spin-dominated (multi-second `Max`, huge `StdDev`) — use the **median** -> per-launch. And **launch counts differ by design**: DeepEP LL issues ~2× the launches of -> FlashInfer-EP LL because DeepEP low-latency splits each dispatch/combine into a **send kernel + -> a deferred receive "hook"** (`low_latency_dispatch(return_recv_hook=True)` then `hook()`, for -> compute/comm overlap), whereas nccl.ep issues one fused kernel per call (`handle.dispatch()` + -> a `handle.complete()` stream-sync, not a 2nd launch). So compare DeepEP's (send+recv) sum vs -> FlashInfer's single launch, not launch-for-launch. HT is 1 launch each → directly comparable. -> Measured medians and the full table are in `vllm_moe_ep_results_prenyx.md` §1.1b; for a -> spin-free per-op latency use the standalone comm benchmark below. -- **If `kern_a2a_.txt` is empty for a backend, the transport did not run** — the summary will - instead be dominated by `multimem_all_reduce_kernel` / `vllm::cross_device_reduce_*` (TP - all-reduce). That means the run fell back to the monolithic path (§3.0); fix the DP flag and - re-capture. **This is the exact failure our first capture hit** (both backends identical, only - TP all-reduce, no dispatch/combine). - -Notes: -- **`nsys` lives only inside the container**, not on the Pre-Nyx login node — run *both* - `nsys profile` and `nsys stats` under `srun --container-image=...`. `nsys stats` needs no GPU - (it only reads the `.nsys-rep`), so it's a cheap short job. -- **Cleaner isolation (NCCL-EP only):** to see *just* the dispatch/combine kernels with no model - noise, profile the standalone comm benchmark instead of vLLM: - ```bash - srun ... --ntasks-per-node=8 bash -lc 'EP_SYNC=/host/sync_ns NCCL_GIN_TYPE=3 \ - nsys profile -t cuda,nvtx,nccl -o /host/logs/nsys_epcomm_r${SLURM_PROCID} \ - bash /host/flashinfer/benchmarks/run_ep_matrix_one_pt.sh \ - --algorithm ht --layout fl --tokens 4096 --hidden 7168 --top-k 8 --experts 256 --iters 20' - # then: nsys stats --report cuda_gpu_kern_sum /host/logs/nsys_epcomm_r0.nsys-rep - ``` -- If the `.nsys-rep` has no GPU kernels, nsys didn't follow the workers — re-run adding - `--trace-fork-before-exec=true`, or use the per-rank standalone form above (`-o ..._r${SLURM_PROCID}`). -- Copy `$RW/logs/*.nsys-rep` locally to open the timeline in the Nsight Systems GUI if you want - to see the dispatch→expert-GEMM→combine sequence visually. - ---- - -## 4. Multi-node (2 nodes / 16 GPU) - -Multi-node needs a **Ray cluster across both nodes' containers** + `--data-parallel-size 16 ---enable-expert-parallel --distributed-executor-backend ray` + `NCCL_MNNVL_ENABLE=1` for the -cross-node EP fabric. **Use data-parallel EP, not `--tensor-parallel-size 16`** — same reason as -§3.0: TP-only ⇒ `dp_size=1` ⇒ monolithic path ⇒ `all2all_backend` ignored (and the cross-node -all-to-all you're trying to measure never runs). -Create `$RW/run_2node.sh` (rank-0 starts the Ray head + runs the bench; rank-1 joins and blocks): - -```bash -cat > $RW/run_2node.sh <<'EOS' -#!/bin/bash -set -o pipefail -export HF_HOME=/host/hf_cache FLASHINFER_WORKSPACE_BASE=/host/fi_cache FLASHINFER_CUBIN_DIR=/host/fi_cubins -export NCCL_GIN_TYPE=3 NCCL_MNNVL_ENABLE=1 -# force the 2.30.7 wheel libnccl (DeepEP needs torch nccl == wheel; runs via bash, not -lc) -NCCL_LIB=$(python -c 'import nvidia.nccl,os;print(os.path.join(list(nvidia.nccl.__path__)[0],"lib"))') -export LD_LIBRARY_PATH="$NCCL_LIB:${LD_LIBRARY_PATH:-}" -B=${BACKEND:-flashinfer_ep_low_latency}; ISL=${ISL:-128}; OSL=${OSL:-128} -python -m pip install -q ray 2>/dev/null || true -HEADF=/host/ray_head.$SLURM_JOB_ID; TMP=/host/raylog_${SLURM_JOB_ID}_${SLURM_NODEID} -if [ "${SLURM_NODEID:-0}" = "0" ]; then - hostname -I | awk '{print $1}' > $HEADF - ray start --head --port=6379 --num-gpus=8 --disable-usage-stats --temp-dir=$TMP - for i in $(seq 1 60); do - n=$(python -c 'import ray;ray.init(address="auto");print(int(ray.cluster_resources().get("GPU",0)))' 2>/dev/null|tail -1) - [ "$n" = "16" ] && break; sleep 5; done; echo "cluster GPUs=$n" - vllm bench throughput --model Qwen/Qwen3-30B-A3B --dataset-name random \ - --random-input-len $ISL --random-output-len $OSL --num-prompts 1000 \ - --data-parallel-size 16 --enable-expert-parallel --all2all-backend $B \ - --distributed-executor-backend ray --trust-remote-code --max-model-len 4096 --enforce-eager - ray stop; rm -f $HEADF -else - for i in $(seq 1 60); do [ -f $HEADF ] && break; sleep 3; done; sleep 8 - ray start --address=$(cat $HEADF):6379 --num-gpus=8 --disable-usage-stats --temp-dir=$TMP --block -fi -EOS -``` -Run (one task per node; warm `/host/fi_cubins` from a single-node run first): -```bash -for B in flashinfer_ep_low_latency flashinfer_ep_high_throughput deepep_low_latency deepep_high_throughput; do - srun -A $ACCT -p $PART -N2 --ntasks-per-node=1 --time=01:00:00 \ - --container-image=$RW/vllm-fi-ep-deepep.sqsh --container-mounts=$RW:/host \ - --export=ALL,BACKEND=$B,ISL=128,OSL=128 \ - bash /host/run_2node.sh > $RW/logs/tp2n_${B}.log 2>&1 - grep -H "Throughput:" $RW/logs/tp2n_${B}.log -done -``` - -> ⚠ **Known issue (environmental, not the EP code):** in our runs the Ray 16-GPU cluster forms -> and vLLM launches, but engine-core init **stalls** — and a plain `--tensor-parallel-size -> 16` run **without** EP stalls identically. So it's the cluster's cross-node vLLM/NCCL bring-up, -> not FlashInfer-EP/DeepEP. If you hit it, debug the cross-node fabric first, e.g. export before -> the bench: `NCCL_DEBUG=INFO`, and set `NCCL_SOCKET_IFNAME` / `NCCL_IB_HCA` to the node's IB -> interfaces (`ibdev2netdev` / `ibv_devices`); verify a plain 2-node NCCL all-reduce works. - -> ⚠ **Offline-DP caveat also applies here.** `vllm bench throughput --data-parallel-size 16` with -> the Ray backend hits the same offline-DP guard as §3.0 (it needs `external_launcher`, not `ray`). -> For a working 2-node DP-EP throughput run once the fabric is up, either launch the §3.0 driver -> under `torchrun --nnodes=2 --nproc_per_node=8 --distributed-executor-backend external_launcher`, -> or use the **server** path (`vllm serve --data-parallel-size 16 --enable-expert-parallel` + -> `vllm bench serve`). The Ray script above is kept as the cluster-bring-up reference. - ---- - -## Expected reference numbers (8-GPU, Qwen3-30B-A3B; full detail in `vllm_moe_ep_results_prenyx.md`) - -**DP-EP, transport-exercised** (the numbers to trust for a backend comparison; §1.1c–f). -Measured at the pinned refs (flashinfer `fa09bc46`, vLLM `ab1415e`), §3c command verbatim -(NP=256, eager, HT capped at 8192, LL uncapped): - -| Backend | 128/128 | 2048/128 | 128/2048 | -|---|---|---|---| -| `flashinfer_ep_low_latency` | **9,088** | **23,106** | **5,825** | -| `deepep_low_latency` | 10,116 | 24,013 | 6,595 | -| `flashinfer_ep_high_throughput` | **6,797** | **45,224** | 3,795 | -| `deepep_high_throughput` | 5,736 | 35,623 | 4,539 | - -(total tok/s, sum of 8 ranks; expect ±3–5% run-to-run.) **FI-HT is ahead of DeepEP-HT by -19–27%** on 128/128 and 2048/128; FI-LL within 4–12% of DeepEP-LL; decode-heavy within -12–16% for both modes. -- **GSM8K over a real DP-EP server** (§3b′, flex/strict): FI-EP **LL 0.856/0.898**, - **HT 0.857/0.898** — both ≥0.80, on par with the ~0.88 reference. Transport genuinely exercised. -- ⚠ These numbers **require the pinned refs** (§1) — the HT token clamp + recv-trim, the - batched-DP cap membership and the fleet host-path caches. On older refs FI-EP lands 2–6× - behind DeepEP and HT SIGABRTs at group-create (root causes + fixes: results doc §1.1c/§1.1f). -- Also required: HT runs get `--max-num-batched-tokens 8192`; LL runs get **no** such flag. - -**Monolithic (`--tensor-parallel-size 8`, dp_size=1) — NOT a transport comparison** (§1.2, retained -for reference only): GSM8K strict ~0.89 all four; throughput 32.5k · 141k · 18.5k (FI) ≈ DeepEP -(both ran the identical TP-all-reduce path, so the ~1–2% closeness is an artifact); memory identical -150.45 GiB KV cache. The all2all backend had **no effect** here. From 767f5df084d8caf32a89bb8c88e1cf81c58be9ea Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Mon, 6 Jul 2026 13:03:52 -0700 Subject: [PATCH 11/13] docs(moe_ep): inline the content of the removed internal validation docs integration.md referenced the internal Pre-Nyx validation notes (vllm_moe_ep_results_prenyx.md / runbook) in 12 places; those files were dropped from the PR in the previous commit. Remove the dangling links and fold the load-bearing facts inline: the three perf root-cause fixes, the offline-DP torchrun/external_launcher requirement, the canonical Dockerfile build spec pointer, and the nccl_ep dispatch/combine kernel names to look for in an nsys capture. Co-Authored-By: Claude Opus 4.8 --- docs/design_docs/vllm_moe_ep_integration.md | 41 +++++++++++---------- 1 file changed, 22 insertions(+), 19 deletions(-) diff --git a/docs/design_docs/vllm_moe_ep_integration.md b/docs/design_docs/vllm_moe_ep_integration.md index 028d09cfe23..fe964ecbc0a 100644 --- a/docs/design_docs/vllm_moe_ep_integration.md +++ b/docs/design_docs/vllm_moe_ep_integration.md @@ -34,23 +34,26 @@ FlashInfer run from the branch. All checks below **pass**: Both backends clear the GSM8K ≥ 0.80 gate (reference ~0.88). Correctness above (GAP tests + `--validate` transport round-trip) directly exercises the dispatch/combine path. -**Transport-exercised results (DP-EP, the numbers that matter — results doc §1.1d–f):** +**Transport-exercised results (DP-EP — the numbers that matter):** GSM8K through a real DP-EP server: **LL 0.856/0.898, HT 0.857/0.898** (flex/strict). DP-EP eager throughput vs DeepEP after the perf iteration (Qwen3-30B-A3B, 8×GPU, total tok/s, 128/128 · 2048/128 · 128/2048): **FI-LL 9,088/23,106/5,825 (0.90/0.96/0.88× of DeepEP-LL)**; **FI-HT 6,797/45,224/3,795 (1.19/1.27/0.84× of DeepEP-HT — ahead on 2 of 3 shapes)**. -The initial DP-EP pass was 2–6× behind; the closure came from three root-cause fixes -(batched-DP scheduler cap membership, HT recv-count trim, fleet-level host-path caches) — -full iteration log in results doc §1.1f. +The initial DP-EP pass was 2–6× behind; the closure came from three root-cause fixes: +`flashinfer_ep_low_latency` added to vLLM's `use_batched_dp_moe` (batched-DP 256-token scheduler +cap, matching `deepep_low_latency`), the HT recv-count compute trim (§2 GAP 3), and fleet-level +host-path caches in `nccl_ep/handle.py`. > 🛑 **The throughput/GSM8K/memory numbers below are historical and do NOT compare the two > transports.** They were run with `--tensor-parallel-size 8` (`dp_size=1`), so vLLM took the > `MoEPrepareAndFinalizeNoDPEPMonolithic` path — experts computed locally, reconciled by TP > all-reduce — and **`--all2all-backend` was a no-op** (confirmed by nsys: identical kernels, only > TP all-reduce, no dispatch/combine, for both FI-EP and DeepEP). The all2all transport is only -> selected when `dp_size > 1` (`fused_moe/config.py::use_all2all_kernels`). The tables above / -> in results doc §1.1e use `--data-parallel-size 8 --enable-expert-parallel` (verify the log says -> `Using FlashInferEPLL/HT…PrepareAndFinalize`, not `…Monolithic`). See runbook §3.0. +> selected when `dp_size > 1` (`fused_moe/config.py::use_all2all_kernels`). The tables above use +> `--data-parallel-size 8 --enable-expert-parallel` (verify the log says +> `Using FlashInferEPLL/HT…PrepareAndFinalize`, not `…Monolithic`). Note offline +> `vllm bench throughput` rejects `--data-parallel-size` directly — launch it under +> `torchrun --nproc_per_node=8` with `--distributed-executor-backend external_launcher`. **Throughput vs DeepEP** *(provisional — monolithic path, transport not exercised)* (`vllm bench throughput --dataset-name random`, Qwen3-30B-A3B, 8-GPU, 1000 prompts; total tok/s): @@ -64,12 +67,10 @@ full iteration log in results doc §1.1f. All four numbers land within ~1–2% — but that is because all four ran the *same* monolithic TP-all-reduce path, not because the transports are equivalent. GSM8K accuracy likewise within noise (monolithic path — end-to-end accuracy, not transport). Memory identical across all four -(150.45 GiB / 6.57M-token KV cache at `--gpu-memory-utilization 0.9`; monolithic path). See -[`vllm_moe_ep_results_prenyx.md`](vllm_moe_ep_results_prenyx.md) for the full method, -per-backend req/s, GSM8K-vs-DeepEP table, multi-node (2-node/16-GPU), and reproduction. +(150.45 GiB / 6.57M-token KV cache at `--gpu-memory-utilization 0.9`; monolithic path). **Not measured:** raw NCCL-EP (N/A upstream), TTFT/TPOT via `bench serve`. **2-node/16-GPU:** Ray+TP=16 plumbing comes up but cross-node engine init stalls — reproduced with plain TP=16 -(no EP), so it's a cluster cross-node NCCL/fabric issue, not the EP integration (see results doc §1.4). +(no EP), so it's a cluster cross-node NCCL/fabric issue, not the EP integration. --- @@ -163,7 +164,8 @@ srun -A coreai_libraries_cudnn -p batch -N1 --time=03:00:00 \ # 2. vLLM (from source) and 3. DeepEP images are layered on the base the same way # (srun --container-image=.sqsh --container-save=.sqsh ...). -# Full recipes: see vllm_moe_ep_results_prenyx.md §3.2 (vLLM) and §3.3 (DeepEP). +# The canonical build spec is docker/Dockerfile.vllm-flashinfer-ep (directly usable +# on a Docker host; on pyxis clusters run its steps inside srun --container-save). ``` Notes: whole-node allocations only (**no `--gres`** on this cluster). `--container-writable` is @@ -207,7 +209,7 @@ GAP 3 HT `recv_total_counter` binding + `DispatchOutput` surfacing — all again ### 5.2 FlashInfer 8-GPU EP round-trip (single node, 8 GPU) Validate dispatch+combine correctness at world=8 via the **comm-matrix `--validate`** path (`srun --ntasks-per-node=8`, `file://` rendezvous, `NCCL_GIN_TYPE=3`; whole-node — **no -`--gres`** on this cluster). See `vllm_moe_ep_results_prenyx.md` §4.2 for the exact runner: +`--gres`** on this cluster). The exact runner: ```bash srun --ntasks-per-node=8 --container-image=$RW/flashinfer-ep-pt2605.sqsh --container-mounts=$RW:/host \ bash -lc 'EP_SYNC=/host/sync_ht NCCL_GIN_TYPE=3 bash /host//benchmarks/run_ep_matrix_one_pt.sh \ @@ -215,8 +217,7 @@ srun --ntasks-per-node=8 --container-image=$RW/flashinfer-ep-pt2605.sqsh --conta # LL: --algorithm ll --layout em --tokens 128 --validate ``` > The pytest `tests/moe_ep/test_moe_ep_ht_correctness.py` launched via `torchrun` **hangs** on a -> default-PG collective on this image — use the comm-matrix `--validate` path above instead -> (results doc §1.4). +> default-PG collective on this image — use the comm-matrix `--validate` path above instead. ### 5.3 vLLM smoke (single node, 8 GPU) — both backends ```bash @@ -242,9 +243,11 @@ lm_eval --model vllm \ > **Accuracy gate only — does not exercise the transport.** lm_eval's `data_parallel_size` spawns > independent replica engines (each `dp_size=1` ⇒ monolithic path), so it can't drive a unified EP > group; keep `tensor_parallel_size=8`. The dispatch/combine transport is validated by §5.2 -> (`--validate`) and the nsys capture (runbook §3e). To exercise the transport end-to-end in vLLM, -> use the **server** path (§5.3, `vllm serve --data-parallel-size 8 --enable-expert-parallel`, -> which *does* build a real DP-EP deployment) or offline `torchrun … external_launcher` (runbook §3.0). +> (`--validate`) and by an nsys kernel capture of a DP-EP run (look for the +> `nccl_ep::internode_ll` / `nccl_ep_jit_ht_*` dispatch+combine kernels). To exercise the +> transport end-to-end in vLLM, use the **server** path (§5.3, `vllm serve +> --data-parallel-size 8 --enable-expert-parallel`, which *does* build a real DP-EP deployment) +> or offline via `torchrun --nproc_per_node=8 … --distributed-executor-backend external_launcher`. ### 5.5 Multi-node (2 nodes, 16 GPU) Repeat 5.2–5.4 across 2 nodes. FlashInfer tests: `srun --nodes=2 --ntasks-per-node=1 @@ -290,7 +293,7 @@ torchrun --nproc_per_node=8 benchmarks/bench_moe_ep.py \ Fixed load: **Qwen3-30B-A3B BF16, ISL/OSL 128/128, `max_concurrency=32`, `NUM_PROMPTS=1000`**. Launch each `vllm serve` with **`--data-parallel-size 8 --enable-expert-parallel`** (NOT TP-only) so the all2all transport is actually on the path — -otherwise every cell collapses to the identical monolithic path (§0 caveat / runbook §3.0). For +otherwise every cell collapses to the identical monolithic path (§0 caveat). For each cell, start `vllm serve` with the backend, then: ```bash vllm bench serve \ From a68ab635dad133a748f2d69b51b331e19e4d1aa4 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Mon, 6 Jul 2026 22:59:54 -0700 Subject: [PATCH 12/13] ci: stop nvshmem4py-cu12 install from evicting CUDA torch; pin torch for job pips unit_test_gb300 [cu130] fails with "Torch not compiled with CUDA enabled": 1. task_run_unit_tests.sh unconditionally runs `pip install nvshmem4py-cu12`. The package is already in the image, but re-resolving its deps applies its cuda-python<=12.9 pin, downgrading the cu130 container's cuda-python 13.0 / cuda-bindings 13.0.3 to 12.9 (pip warns: torch 2.11.0+cu130 requires cuda-bindings<14,>=13.0.3). 2. The subsequent `pip install -r requirements.txt` then can't satisfy both the installed torch's cuda-bindings>=13.0.3 requirement and cuda-python 12.9's cuda-bindings~=12.9.0 chain, so the resolver replaces torch, backtracking 2.12.1 -> 2.12.0 -> 2.11.0 -> 2.10.0 -- and the plain-PyPI aarch64 torch wheel is CPU-only. This was previously masked by the cuda-python>=13.0 base-dep floor, which accidentally re-upgraded cuda-python during the requirements install; relaxing that floor to >=12.0 (needed to stop cu12 environments from being bulldozed) exposed the latent bug. Same disease as the earlier nvidia-nccl-cu13 floor that evicted torch on aarch64; same cure: - task_run_unit_tests.sh / task_test_single_node_comm_kernels.sh: install nvshmem4py-cu12 only when `import nvshmem.core` fails, and with --no-deps (the image ships the right-flavor cuda-python and nvidia-nvshmem). - test_utils.sh / setup_test_env.sh (idempotent, whichever is sourced first): export PIP_CONSTRAINT pinning the preinstalled torch== for every job-time pip install, so any future dep whose constraints would evict torch fails loudly at install time instead of silently degrading to a CPU wheel. Torch is not in [build-system].requires, so isolated build envs are unaffected. AI-assisted (Claude Code): root-caused from the GitLab job log (job 356145796). Co-Authored-By: Claude Opus 4.8 --- scripts/setup_test_env.sh | 19 ++++++++++++++++++ scripts/task_run_unit_tests.sh | 8 +++++++- scripts/task_test_single_node_comm_kernels.sh | 8 +++++++- scripts/test_utils.sh | 20 +++++++++++++++++++ 4 files changed, 53 insertions(+), 2 deletions(-) diff --git a/scripts/setup_test_env.sh b/scripts/setup_test_env.sh index 5cd61330f10..7697aa08a7e 100755 --- a/scripts/setup_test_env.sh +++ b/scripts/setup_test_env.sh @@ -9,6 +9,25 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" +# Pin the preinstalled CUDA torch for every job-time pip install (same guard as +# test_utils.sh; idempotent — whichever is sourced first wins). Prevents a dep's +# transitive constraints from making pip re-resolve torch and silently evict the +# CUDA build (on aarch64 pip backtracks to the CPU-only PyPI wheel -> "Torch not +# compiled with CUDA enabled"); with the constraint such a resolution fails +# loudly at install time. Torch is not in [build-system].requires, so isolated +# build envs are unaffected. +if [ -z "${PIP_CONSTRAINT:-}" ]; then + _torch_pin=$(python -c "import torch; print('torch=='+torch.__version__)" 2>/dev/null || true) + if [ -n "${_torch_pin}" ]; then + _constraint_file=$(mktemp /tmp/ci-torch-constraint.XXXXXX.txt) + echo "${_torch_pin}" > "${_constraint_file}" + export PIP_CONSTRAINT="${_constraint_file}" + echo "Pinning for all pip installs in this job: ${_torch_pin}" + unset _constraint_file + fi + unset _torch_pin +fi + # Source the environment override file if it exists if [ -f "${REPO_ROOT}/ci/setup_python.env" ]; then source "${REPO_ROOT}/ci/setup_python.env" diff --git a/scripts/task_run_unit_tests.sh b/scripts/task_run_unit_tests.sh index 751be2e28d7..876dcc31fb3 100755 --- a/scripts/task_run_unit_tests.sh +++ b/scripts/task_run_unit_tests.sh @@ -11,8 +11,14 @@ SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" # shellcheck disable=SC1091 # File exists, checked separately source "${SCRIPT_DIR}/test_utils.sh" +# nvshmem4py-cu12 pins cuda-python<=12.9; letting pip resolve its deps on a +# cu13 container downgrades cuda-python/cuda-bindings and makes the next +# requirements resolution evict CUDA torch (aarch64 backtracks to the CPU-only +# wheel -> "Torch not compiled with CUDA enabled"). Install only if missing, +# and --no-deps: the image already ships the right-flavor cuda-python and +# nvidia-nvshmem libraries. # TODO: Remove once CI container ships with nvshmem4py pre-installed. -pip install nvshmem4py-cu12 +python -c "import nvshmem.core" 2>/dev/null || pip install --no-deps nvshmem4py-cu12 # Find and filter test files based on pytest.ini exclusions find_test_files() { diff --git a/scripts/task_test_single_node_comm_kernels.sh b/scripts/task_test_single_node_comm_kernels.sh index 72bbb6bb6cf..071a4e45cb4 100644 --- a/scripts/task_test_single_node_comm_kernels.sh +++ b/scripts/task_test_single_node_comm_kernels.sh @@ -17,8 +17,14 @@ echo "" pip install -e . -v +# nvshmem4py-cu12 pins cuda-python<=12.9; letting pip resolve its deps on a +# cu13 container downgrades cuda-python/cuda-bindings and makes the next +# requirements resolution evict CUDA torch (aarch64 backtracks to the CPU-only +# wheel -> "Torch not compiled with CUDA enabled"). Install only if missing, +# and --no-deps: the image already ships the right-flavor cuda-python and +# nvidia-nvshmem libraries. # TODO: Remove once CI container ships with nvshmem4py pre-installed. -pip install nvshmem4py-cu12 +python -c "import nvshmem.core" 2>/dev/null || pip install --no-deps nvshmem4py-cu12 # vllm ar pytest -s tests/comm/test_vllm_custom_allreduce.py diff --git a/scripts/test_utils.sh b/scripts/test_utils.sh index e49d6723b3d..9889309674f 100755 --- a/scripts/test_utils.sh +++ b/scripts/test_utils.sh @@ -23,6 +23,26 @@ if [ -z "${MAX_JOBS:-}" ]; then fi export MAX_JOBS +# Pin the preinstalled CUDA torch for every job-time pip install. Twice now a +# runtime dep's transitive constraint has made pip re-resolve torch and evict +# the CUDA build (the nvidia-nccl-cu13 floor, then nvshmem4py-cu12's +# cuda-python<=12.9 pin downgrading cuda-bindings on cu13 images) — on aarch64 +# pip backtracks to the CPU-only PyPI wheel and tests fail later with "Torch +# not compiled with CUDA enabled". A constraints file makes any resolution that +# would replace torch fail loudly at install time instead. Build-isolation +# envs are unaffected (torch is not in [build-system].requires). +if [ -z "${PIP_CONSTRAINT:-}" ]; then + _torch_pin=$(python -c "import torch; print('torch=='+torch.__version__)" 2>/dev/null || true) + if [ -n "${_torch_pin}" ]; then + _constraint_file=$(mktemp /tmp/ci-torch-constraint.XXXXXX.txt) + echo "${_torch_pin}" > "${_constraint_file}" + export PIP_CONSTRAINT="${_constraint_file}" + echo "Pinning for all pip installs in this job: ${_torch_pin}" + unset _constraint_file + fi + unset _torch_pin +fi + # CUDA_VISIBLE_DEVICES: Not set by default - let detect_gpus() auto-detect via nvidia-smi : "${SAMPLE_RATE:=5}" # Run every Nth test in sanity mode (5 = ~20% coverage) : "${PARALLEL_TESTS:=false}" # Disable parallel test execution by default From ae8d8f65503ac2d10e9d31c65941e02cce518750 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Mon, 6 Jul 2026 23:13:02 -0700 Subject: [PATCH 13/13] ci: strip the +cuXXX local tag from the torch pip constraint The PIP_CONSTRAINT guard pinned torch== including the +cuNNN local tag. PEP-517 build environments inherit PIP_CONSTRAINT, and flashinfer-jit-cache's build-system.requires includes torch -- its isolated build env then tries to resolve torch==2.11.0+cu130 from PyPI, where local-version wheels do not exist: AOT Build Import (arm64, cu130) fails with ResolutionImpossible ("no matching distributions available: torch"). Pin torch== instead: PEP 440 lets the installed 2.X.Y+cuNNN satisfy ==2.X.Y, so the main-env protection is unchanged (the poisoned cuda-bindings scenario still fails loudly -- the plain PyPI wheel of the same version carries the same cuda-bindings requirement), while build envs resolve the plain wheel from PyPI. Co-Authored-By: Claude Opus 4.8 --- scripts/setup_test_env.sh | 9 ++++++--- scripts/test_utils.sh | 9 ++++++--- 2 files changed, 12 insertions(+), 6 deletions(-) diff --git a/scripts/setup_test_env.sh b/scripts/setup_test_env.sh index 7697aa08a7e..6500e0af58c 100755 --- a/scripts/setup_test_env.sh +++ b/scripts/setup_test_env.sh @@ -14,10 +14,13 @@ REPO_ROOT="$(cd "${SCRIPT_DIR}/.." && pwd)" # transitive constraints from making pip re-resolve torch and silently evict the # CUDA build (on aarch64 pip backtracks to the CPU-only PyPI wheel -> "Torch not # compiled with CUDA enabled"); with the constraint such a resolution fails -# loudly at install time. Torch is not in [build-system].requires, so isolated -# build envs are unaffected. +# loudly at install time. The +cuXXX local tag is stripped: PEP 440 lets the +# installed 2.X.Y+cuNNN satisfy ==2.X.Y, but PEP-517 build envs (flashinfer- +# jit-cache's build-system.requires includes torch) inherit PIP_CONSTRAINT and +# must be able to resolve the pin from PyPI, where local-version wheels don't +# exist. if [ -z "${PIP_CONSTRAINT:-}" ]; then - _torch_pin=$(python -c "import torch; print('torch=='+torch.__version__)" 2>/dev/null || true) + _torch_pin=$(python -c "import torch; print('torch=='+torch.__version__.split('+')[0])" 2>/dev/null || true) if [ -n "${_torch_pin}" ]; then _constraint_file=$(mktemp /tmp/ci-torch-constraint.XXXXXX.txt) echo "${_torch_pin}" > "${_constraint_file}" diff --git a/scripts/test_utils.sh b/scripts/test_utils.sh index 9889309674f..908d25484a1 100755 --- a/scripts/test_utils.sh +++ b/scripts/test_utils.sh @@ -29,10 +29,13 @@ export MAX_JOBS # cuda-python<=12.9 pin downgrading cuda-bindings on cu13 images) — on aarch64 # pip backtracks to the CPU-only PyPI wheel and tests fail later with "Torch # not compiled with CUDA enabled". A constraints file makes any resolution that -# would replace torch fail loudly at install time instead. Build-isolation -# envs are unaffected (torch is not in [build-system].requires). +# would replace torch fail loudly at install time instead. The +cuXXX local +# tag is stripped: PEP 440 lets the installed 2.X.Y+cuNNN satisfy ==2.X.Y, but +# PEP-517 build envs (flashinfer-jit-cache's build-system.requires includes +# torch) inherit PIP_CONSTRAINT and must be able to resolve the pin from PyPI, +# where local-version wheels don't exist. if [ -z "${PIP_CONSTRAINT:-}" ]; then - _torch_pin=$(python -c "import torch; print('torch=='+torch.__version__)" 2>/dev/null || true) + _torch_pin=$(python -c "import torch; print('torch=='+torch.__version__.split('+')[0])" 2>/dev/null || true) if [ -n "${_torch_pin}" ]; then _constraint_file=$(mktemp /tmp/ci-torch-constraint.XXXXXX.txt) echo "${_torch_pin}" > "${_constraint_file}"