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..3dafd349955 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,24 +74,51 @@ 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 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: @@ -201,10 +241,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 +346,75 @@ 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 _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). @@ -349,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. @@ -367,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 @@ -524,9 +630,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,16 +648,38 @@ 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. + 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 + # 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." ) + # 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 + # 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) @@ -598,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/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/docs/design_docs/vllm_moe_ep_integration.md b/docs/design_docs/vllm_moe_ep_integration.md index 03e82788066..fe964ecbc0a 100644 --- a/docs/design_docs/vllm_moe_ep_integration.md +++ b/docs/design_docs/vllm_moe_ep_integration.md @@ -31,10 +31,32 @@ 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):** +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: +`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 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): | ISL/OSL | FI-EP LL | FI-EP HT | DeepEP LL | DeepEP HT | |---|---|---|---|---| @@ -42,15 +64,13 @@ 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 -[`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. +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). **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. --- @@ -144,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 @@ -188,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 \ @@ -196,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 @@ -216,10 +236,18 @@ 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 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 @@ -263,7 +291,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). For +each cell, start `vllm serve` with the backend, then: ```bash vllm bench serve \ --model Qwen/Qwen3-30B-A3B \ 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..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,10 +43,59 @@ 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 ``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. + # 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 + 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 " + 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) @@ -46,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/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..2ab975af382 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.""" @@ -46,9 +89,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 @@ -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( @@ -121,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 @@ -241,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 750aa9190c2..7e284511475 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, @@ -129,8 +130,20 @@ 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 (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 @@ -162,15 +175,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 @@ -216,6 +236,7 @@ def __init__( config=None, stream=self._stream, ) + _t = _hp("hinit.create_handle_c", _t) # ----------------------------------------------------------------- knobs @@ -223,6 +244,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 @@ -243,25 +296,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) @@ -269,11 +327,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( @@ -414,6 +478,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( @@ -435,8 +514,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 @@ -444,15 +529,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), @@ -461,13 +546,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) @@ -528,15 +611,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) @@ -568,31 +649,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) @@ -621,9 +700,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() 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..ffb137ca373 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,25 +26,15 @@ 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) 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] flashinfer = "flashinfer.__main__:cli" @@ -94,9 +84,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..985081ec2b8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,7 +1,30 @@ apache-tvm-ffi>=0.1.6,!=0.1.8,!=0.1.8.post0,<0.2 click +# 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). +# +# 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 ninja numpy nvidia-cudnn-frontend>=1.13.0 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/scripts/setup_test_env.sh b/scripts/setup_test_env.sh index 5cd61330f10..6500e0af58c 100755 --- a/scripts/setup_test_env.sh +++ b/scripts/setup_test_env.sh @@ -9,6 +9,28 @@ 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. 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__.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}" + 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..908d25484a1 100755 --- a/scripts/test_utils.sh +++ b/scripts/test_utils.sh @@ -23,6 +23,29 @@ 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. 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__.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}" + 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 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/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, 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