From 3e4849a2c316a66f7f87c0680b3707c31aaa0203 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Tue, 12 May 2026 23:32:22 -0700 Subject: [PATCH 01/10] moe_ep: build infra for in-tree NIXL-EP + NCCL-EP (Part A) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire NVIDIA NIXL (ai-dynamo/nixl) and NCCL (NVIDIA/nccl) as git submodules under 3rdparty/ and add a BUILD_NVEP=1 build-time switch that produces the EP transport libraries in-tree: - 3rdparty/nixl — pinned to v1.1.0 (05e4243f) - 3rdparty/nccl — pinned to master HEAD (1933fdd6) - 3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch Replaces NIXL's `-arch=sm_90` flag with multi-gencode covering sm_90 + sm_100 + sm_103 (H100/B200/B300) plus sm_90 PTX for forward-compat. Applied automatically to the submodule worktree. build_backend.py grows _build_nvep_if_enabled() that runs meson on 3rdparty/nixl and `make src.build && make -C contrib/nccl_ep` on 3rdparty/nccl, stages the produced .so files into the flashinfer/moe_ep/ package, fixes RPATHs with patchelf, and editable-installs both nccl_ep ctypes bindings and nccl4py from the NCCL submodule. pyproject.toml gains a [nvep] optional-dependencies extra (nixl-cu13, nvidia-nccl-cu13, cuda-python) and package-data entries so the staged .so files ship in wheels. flashinfer/moe_ep/ is a placeholder package with a runtime probe (have_nccl_ep / have_nixl_ep / available_backends). MoEEpNotBuiltError is raised when a backend is invoked without its libs. The Fleet/Handle classes themselves arrive in Part B. docker/Dockerfile.flashinfer-nvep is a reference image showing the full system-dep stack (rdma-core, libibverbs-dev, UCX 1.21 with experimental API, GDRCopy, openmpi) required to make BUILD_NVEP=1 succeed end-to-end. Validated: - BUILD_NVEP=0 install: unchanged behavior, vanilla flashinfer works. - BUILD_NVEP=1 install: pipeline runs through patch overlay, meson configure (CUDA 12.8 detected), meson subproject downloads (taskflow, abseil, asio, tomlplusplus, prometheus-cpp), into ninja compile. Compile fails on the dev box at , which is the expected system-package gap — full build needs a host with libmlx5-dev / libibverbs-dev / DOCA gpunetio (the Dockerfile above installs these). --- .gitignore | 5 + .gitmodules | 6 + 3rdparty/nccl | 1 + 3rdparty/nixl | 1 + .../0001-meson-add-blackwell-arches.patch | 26 +++ build_backend.py | 170 ++++++++++++++++++ docker/Dockerfile.flashinfer-nvep | 79 ++++++++ flashinfer/moe_ep/__init__.py | 104 +++++++++++ flashinfer/moe_ep/nccl_ep/__init__.py | 8 + flashinfer/moe_ep/nixl_ep/__init__.py | 8 + pyproject.toml | 15 ++ 11 files changed, 423 insertions(+) create mode 160000 3rdparty/nccl create mode 160000 3rdparty/nixl create mode 100644 3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch create mode 100644 docker/Dockerfile.flashinfer-nvep create mode 100644 flashinfer/moe_ep/__init__.py create mode 100644 flashinfer/moe_ep/nccl_ep/__init__.py create mode 100644 flashinfer/moe_ep/nixl_ep/__init__.py diff --git a/.gitignore b/.gitignore index 771a46b29d3..4ab7544987c 100644 --- a/.gitignore +++ b/.gitignore @@ -201,3 +201,8 @@ cython_debug/ .cursor/ docs/tutorials/generated/ docs/sg_execution_times.rst + +# moe_ep build artifacts (BUILD_NVEP=1 in-tree build of NIXL-EP + NCCL-EP) +build_nvep/ +flashinfer/moe_ep/*/_libs/ +flashinfer/moe_ep/*/_vendored/ diff --git a/.gitmodules b/.gitmodules index 6bdbddf144a..50d7a9bfd07 100644 --- a/.gitmodules +++ b/.gitmodules @@ -7,3 +7,9 @@ [submodule "3rdparty/cccl"] path = 3rdparty/cccl url = https://github.com/NVIDIA/cccl.git +[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 new file mode 160000 index 00000000000..1933fdd6360 --- /dev/null +++ b/3rdparty/nccl @@ -0,0 +1 @@ +Subproject commit 1933fdd6360a8bfccaa0166bd71bce363d32e5b6 diff --git a/3rdparty/nixl b/3rdparty/nixl new file mode 160000 index 00000000000..05e4243f5ed --- /dev/null +++ b/3rdparty/nixl @@ -0,0 +1 @@ +Subproject commit 05e4243f5ed305a245912361ace20ba354ec0808 diff --git a/3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch b/3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch new file mode 100644 index 00000000000..f6e7ef57cca --- /dev/null +++ b/3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch @@ -0,0 +1,26 @@ +From: FlashInfer build infra +Subject: [PATCH] examples/device/ep: emit SASS for Hopper + Blackwell + +The upstream meson rule pins `-arch=sm_90`, which overrides any global +-gencode flags and ships an sm_90-only `nixl_ep_cpp.so`. FlashInfer needs +a single .so that runs natively on H100 (sm_90), B200 (sm_100), and B300 +(sm_103), plus sm_90 PTX for forward-compat onto future arches. + +This patch replaces the single `-arch=sm_90` flag with an explicit +multi-gencode list. Targets the NIXL pin 05e4243f (tag v1.1.0). + +diff --git a/examples/device/ep/meson.build b/examples/device/ep/meson.build +index a9ba19f..065dd13 100644 +--- a/examples/device/ep/meson.build ++++ b/examples/device/ep/meson.build +@@ -92,7 +92,9 @@ nixl_ep_cuda_args = [ + '-DHAVE_CUDA', + '-DTORCH_EXTENSION_NAME=nixl_ep_cpp', + '--expt-relaxed-constexpr', # Allow calling constexpr __host__ functions from __device__ functions +- '-arch=sm_90', # Only compile for sm90 (overrides global -gencode flags) ++ '-gencode=arch=compute_90,code=[sm_90,compute_90]', # H100 SASS + PTX forward-compat ++ '-gencode=arch=compute_100,code=sm_100', # B200 native SASS ++ '-gencode=arch=compute_103,code=sm_103', # B300 native SASS + '--ptxas-options=--register-usage-level=10', # Allow more register usage (matches setup.py) + '-Xcompiler', '-Wno-deprecated-declarations', + '-Xcompiler', '-Wno-unused-variable', diff --git a/build_backend.py b/build_backend.py index 9bcac12470b..a6521b40d3b 100644 --- a/build_backend.py +++ b/build_backend.py @@ -16,6 +16,8 @@ import os import shutil +import subprocess +import sys from pathlib import Path from setuptools import build_meta as orig @@ -24,6 +26,170 @@ _root = Path(__file__).parent.resolve() _data_dir = _root / "flashinfer" / "data" +# moe_ep build infra: gated by BUILD_NVEP=1 env var. +# When set, _build_nvep_if_enabled() runs meson on 3rdparty/nixl and make on +# 3rdparty/nccl, then stages the produced .so files under flashinfer/moe_ep/. +_BUILD_NVEP = os.environ.get("BUILD_NVEP", "0") == "1" +_nvep_build_root = _root / "build_nvep" +_moe_ep_pkg = _root / "flashinfer" / "moe_ep" + + +def _detect_cuda_major() -> int: + """Best-effort detection of the CUDA major version on the host.""" + try: + out = subprocess.check_output(["nvcc", "--version"]).decode() + for line in out.splitlines(): + if "release" in line: + # e.g. "Cuda compilation tools, release 13.0, V13.0.48" + token = line.split("release", 1)[1].split(",", 1)[0].strip() + return int(token.split(".")[0]) + except Exception: + pass + return 13 # default — pyproject's nvep extras pin cu13 packages + + +def _apply_patches(submodule_dir: Path, patches_dir: Path) -> None: + """Apply every *.patch in patches_dir to the submodule working tree. + + No-ops if patches_dir doesn't exist. Idempotent: skips patches that are + already applied by checking `git apply --reverse --check` first. + """ + if not patches_dir.is_dir(): + return + for patch in sorted(patches_dir.glob("*.patch")): + # Already applied? + already = subprocess.run( + ["git", "apply", "--reverse", "--check", str(patch)], + cwd=submodule_dir, capture_output=True, + ) + if already.returncode == 0: + print(f"[BUILD_NVEP] patch already applied, skipping: {patch.name}") + continue + # Check we *can* apply, then apply. + subprocess.run( + ["git", "apply", "--check", str(patch)], + cwd=submodule_dir, check=True, + ) + subprocess.run( + ["git", "apply", str(patch)], + cwd=submodule_dir, check=True, + ) + print(f"[BUILD_NVEP] applied patch: {patch.name}") + + +def _build_nixl_ep() -> None: + src = _root / "3rdparty" / "nixl" + build = _nvep_build_root / "nixl" + prefix = _nvep_build_root / "nixl_install" + _apply_patches(src, _root / "3rdparty_patches" / "nixl") + + if not build.exists(): + subprocess.run([ + "meson", "setup", str(build), str(src), + "-Dbuild_nixl_ep=true", + "-Dbuild_examples=true", + f"-Dprefix={prefix}", + "--buildtype=release", + ], check=True) + subprocess.run(["ninja", "-C", str(build), "install"], check=True) + + dst = _moe_ep_pkg / "nixl_ep" / "_libs" + dst.mkdir(parents=True, exist_ok=True) + + nixl_lib_src = prefix / "lib" / "x86_64-linux-gnu" + if nixl_lib_src.exists(): + shutil.copytree(nixl_lib_src, dst / "nixl_lib", dirs_exist_ok=True) + + # The torch extension lands either in build/ or build/examples/device/ep/ + for cand in (build / "examples/device/ep").glob("nixl_ep_cpp*.so"): + shutil.copy(cand, dst / cand.name) + print(f"[BUILD_NVEP] staged: {cand.name}") + + # Vendor the python wrapper sources so we can import from + # flashinfer.moe_ep.nixl_ep._vendored (Step B5). + vendored_src = src / "examples/device/ep/nixl_ep" + if vendored_src.exists(): + shutil.copytree( + vendored_src, + _moe_ep_pkg / "nixl_ep" / "_vendored", + dirs_exist_ok=True, + ) + + +def _build_nccl_ep() -> None: + src = _root / "3rdparty" / "nccl" + build = _nvep_build_root / "nccl" + _apply_patches(src, _root / "3rdparty_patches" / "nccl") + + subprocess.run( + ["make", "src.build", f"BUILDDIR={build}", "-j"], + cwd=src, check=True, + ) + subprocess.run( + ["make", "-C", "contrib/nccl_ep", f"BUILDDIR={build}", "-j"], + cwd=src, check=True, + ) + + dst = _moe_ep_pkg / "nccl_ep" / "_libs" + dst.mkdir(parents=True, exist_ok=True) + for soname in ("libnccl.so.2", "libnccl_ep.so"): + sopath = build / "lib" / soname + if sopath.exists(): + shutil.copy(sopath, dst / soname) + print(f"[BUILD_NVEP] staged: {soname}") + + # Editable-install the ctypes wrapper from contrib/nccl_ep/python so + # `import nccl_ep` resolves on the user's env. + subprocess.run([ + sys.executable, "-m", "pip", "install", "-e", + str(src / "contrib/nccl_ep/python"), + ], check=True) + + # Editable-install nccl4py — gives Cython bindings + Communicator(ptr=...) + # which is how the moe_ep NCCL backend bridges a torch.distributed + # process group's raw ncclComm_t pointer into ncclEpCreateGroup. + cuda_extra = f"cu{_detect_cuda_major()}" + env = os.environ.copy() + env.setdefault("CUDA_HOME", "/usr/local/cuda") + subprocess.run([ + sys.executable, "-m", "pip", "install", "-e", + f"{src / 'bindings' / 'nccl4py'}[{cuda_extra}]", + ], env=env, check=True) + + +def _fix_rpaths() -> None: + """Rewrite RPATHs on staged .so files so they find siblings without LD_LIBRARY_PATH.""" + patchelf_ok = shutil.which("patchelf") is not None + if not patchelf_ok: + print("[BUILD_NVEP] patchelf not found; skipping RPATH fix-up") + return + rpath = "$ORIGIN:$ORIGIN/_libs:$ORIGIN/_libs/nixl_lib" + for so in _moe_ep_pkg.rglob("*.so*"): + # Skip symlinks + if so.is_symlink(): + continue + subprocess.run( + ["patchelf", "--set-rpath", rpath, str(so)], + check=False, + ) + + +def _build_nvep_if_enabled() -> None: + if not _BUILD_NVEP: + return + print("[BUILD_NVEP] BUILD_NVEP=1 — building NIXL-EP + NCCL-EP from submodules") + # Make sure submodules are present (sdist installs won't have them + # initialized automatically). + if not (_root / "3rdparty/nixl/meson.build").exists(): + subprocess.run( + ["git", "submodule", "update", "--init", "--recursive"], + cwd=_root, check=True, + ) + _build_nixl_ep() + _build_nccl_ep() + _fix_rpaths() + print("[BUILD_NVEP] done") + def _create_build_metadata(): """Create build metadata file with version information.""" @@ -112,6 +278,7 @@ def ln(source: str, target: str) -> None: def _prepare_for_wheel(): # For wheel, copy actual files instead of symlinks so they are included in the wheel + _build_nvep_if_enabled() if _data_dir.exists(): shutil.rmtree(_data_dir) _create_data_dir(use_symlinks=False) @@ -128,6 +295,7 @@ def _prepare_for_wheel(): def _prepare_for_editable(): # For editable install, use symlinks so changes are reflected immediately + _build_nvep_if_enabled() if _data_dir.exists(): shutil.rmtree(_data_dir) _create_data_dir(use_symlinks=True) @@ -135,6 +303,8 @@ def _prepare_for_editable(): def _prepare_for_sdist(): # For sdist, copy actual files instead of symlinks so they are included in the tarball + # NOTE: do NOT build moe_ep here — submodules + patches travel in the sdist + # itself and get built during the *install* of the sdist. if _data_dir.exists(): shutil.rmtree(_data_dir) _create_data_dir(use_symlinks=False) diff --git a/docker/Dockerfile.flashinfer-nvep b/docker/Dockerfile.flashinfer-nvep new file mode 100644 index 00000000000..0d7bbf20f05 --- /dev/null +++ b/docker/Dockerfile.flashinfer-nvep @@ -0,0 +1,79 @@ +# SPDX-License-Identifier: Apache-2.0 +# +# Reference Dockerfile for building FlashInfer with the moe_ep transport +# backends enabled (NIXL-EP + NCCL-EP from in-tree git submodules). +# +# Usage: +# cd /path/to/flashinfer +# docker build -f docker/Dockerfile.flashinfer-nvep -t flashinfer-nvep:dev . +# docker run --gpus all --rm -it flashinfer-nvep:dev bash +# +# Build args: +# BUILD_NVEP — 0 to skip the moe_ep submodule builds (still installs +# FlashInfer). Defaults to 1. +# 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 +FROM ${CUDA_IMAGE} + +ENV DEBIAN_FRONTEND=noninteractive + +# System dependencies for NIXL-EP + NCCL-EP. CUDA toolkit + driver come from +# the base image; everything else is installed via apt + source builds. +RUN apt-get update && apt-get install -y --no-install-recommends \ + build-essential cmake ninja-build patchelf pkg-config \ + git git-lfs ca-certificates curl wget \ + python3 python3-pip python3-venv python3-dev \ + pybind11-dev \ + # RDMA / InfiniBand userspace stack + rdma-core libibverbs-dev libibverbs1 libibumad3 \ + librdmacm-dev libnl-3-dev libnl-route-3-dev \ + ibverbs-providers infiniband-diags \ + # MPI for NIXL bootstrap + libopenmpi-dev openmpi-bin \ + && rm -rf /var/lib/apt/lists/* + +# UCX 1.21 with experimental API (NIXL transport) +ARG UCX_VERSION=v1.21.0 +ARG UCX_PREFIX=/opt/ucx +RUN git clone --depth=1 --branch ${UCX_VERSION} https://github.com/openucx/ucx.git /tmp/ucx \ + && cd /tmp/ucx \ + && ./autogen.sh \ + && ./configure --prefix=${UCX_PREFIX} --enable-experimental-api \ + --with-cuda=/usr/local/cuda \ + && make -j"$(nproc)" install \ + && rm -rf /tmp/ucx +ENV PKG_CONFIG_PATH=${UCX_PREFIX}/lib/pkgconfig:${PKG_CONFIG_PATH} +ENV LD_LIBRARY_PATH=${UCX_PREFIX}/lib:${LD_LIBRARY_PATH} + +# GDRCopy ≥ 2.5.1 — note: requires the kernel module to be loaded on the host. +# Only the userspace library is built here. +ARG GDRCOPY_VERSION=v2.5.1 +RUN git clone --depth=1 --branch ${GDRCOPY_VERSION} https://github.com/NVIDIA/gdrcopy.git /tmp/gdrcopy \ + && cd /tmp/gdrcopy \ + && make -j"$(nproc)" lib lib_install \ + && rm -rf /tmp/gdrcopy + +# DOCA gpunetio — NOT installed here: it requires a separate Mellanox repo +# setup and a license-accepted MOFED package. On NVIDIA-managed hosts the +# /opt/mellanox/doca tree is provided by the base image. If your base image +# lacks DOCA, add the appropriate `apt-get install doca-sdk-gpunetio +# libdoca-sdk-gpunetio-dev` step here, gated on having the Mellanox apt repo +# configured. + +# Pre-stage uv for fast pip installs. +RUN curl -LsSf https://astral.sh/uv/install.sh | sh +ENV PATH=/root/.local/bin:${PATH} + +# Build & install FlashInfer + moe_ep backends. +ARG BUILD_NVEP=1 +ARG FLASHINFER_SRC=/workspace/flashinfer +COPY . ${FLASHINFER_SRC} +WORKDIR ${FLASHINFER_SRC} +RUN git submodule update --init --recursive +RUN BUILD_NVEP=${BUILD_NVEP} uv pip install --system -e ".[nvep]" + +# Smoke probe. +RUN python3 -c "from flashinfer.moe_ep import available_backends; print('moe_ep backends:', available_backends())" + +CMD ["bash"] diff --git a/flashinfer/moe_ep/__init__.py b/flashinfer/moe_ep/__init__.py new file mode 100644 index 00000000000..c3bd7adc99c --- /dev/null +++ b/flashinfer/moe_ep/__init__.py @@ -0,0 +1,104 @@ +"""flashinfer.moe_ep — MoE Expert-Parallel dispatch/combine over NCCL-EP and NIXL-EP. + +This package is a thin Python wrapper over two transport backends: + +- ``flashinfer.moe_ep.nccl_ep`` — primary backend, wraps NVIDIA's ``nccl_ep`` + (built in-tree from ``3rdparty/nccl/contrib/nccl_ep``). +- ``flashinfer.moe_ep.nixl_ep`` — alternate backend, wraps ai-dynamo's + ``nixl_ep`` (built in-tree from ``3rdparty/nixl/examples/device/ep``). + +The shared libraries that back these wrappers (``libnccl_ep.so``, +``nixl_ep_cpp*.so``, etc.) are produced by the FlashInfer build only when +``BUILD_NVEP=1`` is set in the env at install time: + + BUILD_NVEP=1 pip install -e ".[nvep]" + +Without ``BUILD_NVEP=1`` the package imports succeed but calling +:func:`create_fleet` raises :class:`MoEEpNotBuiltError` with rebuild +instructions. This file lays down only the import-time probe and the +``Fleet`` / ``Handle`` factory plumbing; the actual abstract classes and +backend implementations land in Part B of the integration plan. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +__all__ = [ + "MoEEpNotBuiltError", + "have_nccl_ep", + "have_nixl_ep", + "available_backends", +] + + +_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." +) + + +class MoEEpNotBuiltError(RuntimeError): + """Raised when an EP backend is invoked but its native libs are missing.""" + + +def _probe_nccl_ep() -> bool: + libs = _pkg_dir / "nccl_ep" / "_libs" + return (libs / "libnccl_ep.so").exists() and (libs / "libnccl.so.2").exists() + + +def _probe_nixl_ep() -> bool: + libs = _pkg_dir / "nixl_ep" / "_libs" + if not libs.is_dir(): + return False + return any(libs.glob("nixl_ep_cpp*.so")) + + +def have_nccl_ep() -> bool: + """Return True if the NCCL-EP backend native libs are present.""" + return _probe_nccl_ep() + + +def have_nixl_ep() -> bool: + """Return True if the NIXL-EP backend native libs are present.""" + return _probe_nixl_ep() + + +def available_backends() -> list[str]: + """Names of EP backends with both native libs and python wrappers present.""" + out: list[str] = [] + if have_nccl_ep(): + out.append("nccl_ep") + if have_nixl_ep(): + out.append("nixl_ep") + return out + + +def _require_built(backend: str) -> None: + """Raise MoEEpNotBuiltError if `backend` is missing its native libs.""" + probe = {"nccl_ep": _probe_nccl_ep, "nixl_ep": _probe_nixl_ep}.get(backend) + if probe is None: + raise ValueError( + f"unknown moe_ep backend {backend!r}; expected one of nccl_ep, nixl_ep" + ) + if not probe(): + raise MoEEpNotBuiltError( + f"moe_ep backend {backend!r} is not built.\n\n{_REBUILD_HINT}" + ) + + +# Quiet diagnostic at import time when BUILD_NVEP was set but the libs are +# absent — most likely cause is a partial build. Helpful for first-time users. +if os.environ.get("BUILD_NVEP") == "1" and not available_backends(): + import warnings + + warnings.warn( + "BUILD_NVEP=1 was set, but no moe_ep backend libraries were found " + f"under {_pkg_dir}. Check the build log for meson/make failures.", + RuntimeWarning, + stacklevel=2, + ) diff --git a/flashinfer/moe_ep/nccl_ep/__init__.py b/flashinfer/moe_ep/nccl_ep/__init__.py new file mode 100644 index 00000000000..3bd53091d17 --- /dev/null +++ b/flashinfer/moe_ep/nccl_ep/__init__.py @@ -0,0 +1,8 @@ +"""NCCL-EP backend stub. Real implementation lands in Part B (B3-B4). + +Importing this module succeeds even when the native libs are absent; calling +into the (yet-to-be-implemented) Fleet/Handle factory functions will raise +:class:`flashinfer.moe_ep.MoEEpNotBuiltError`. +""" + +from __future__ import annotations diff --git a/flashinfer/moe_ep/nixl_ep/__init__.py b/flashinfer/moe_ep/nixl_ep/__init__.py new file mode 100644 index 00000000000..44e365040ad --- /dev/null +++ b/flashinfer/moe_ep/nixl_ep/__init__.py @@ -0,0 +1,8 @@ +"""NIXL-EP backend stub. Real implementation lands in Part B (B5). + +Importing this module succeeds even when the native libs are absent; calling +into the (yet-to-be-implemented) Fleet/Handle factory functions will raise +:class:`flashinfer.moe_ep.MoEEpNotBuiltError`. +""" + +from __future__ import annotations diff --git a/pyproject.toml b/pyproject.toml index 9ea503dbd53..36c774e1d6b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -26,6 +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 in-tree submodule build (NIXL meson + NCCL make), this extra +# pulls the matching pip-installable agents/runtimes. +nvep = [ + "nixl-cu13>=1.0.1", + "nvidia-nccl-cu13>=2.30.4", + "cuda-python>=13.0", +] [project.scripts] flashinfer = "flashinfer.__main__:cli" @@ -64,6 +73,12 @@ 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 +# build_backend._build_nvep_if_enabled). +"flashinfer.moe_ep.nixl_ep" = ["_libs/**"] +"flashinfer.moe_ep.nccl_ep" = ["_libs/*.so*"] [tool.mypy] files = ["flashinfer"] From 77cecc62ccc6e43cf9463efac9a19658b2acc7de Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Tue, 12 May 2026 23:32:22 -0700 Subject: [PATCH 02/10] moe_ep: build infra for in-tree NIXL-EP + NCCL-EP MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Wire NVIDIA NIXL (ai-dynamo/nixl) and NCCL (NVIDIA/nccl) as git submodules under 3rdparty/ and add a BUILD_NVEP=1 build-time switch that produces the EP transport libraries in-tree: - 3rdparty/nixl — pinned to v1.1.0 (05e4243f) - 3rdparty/nccl — pinned to master HEAD (1933fdd6) - 3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch Replaces NIXL's `-arch=sm_90` flag with multi-gencode covering sm_90 + sm_100 + sm_103 (H100/B200/B300) plus sm_90 PTX for forward-compat. Applied automatically to the submodule worktree. build_backend.py grows _build_nvep_if_enabled() that runs meson on 3rdparty/nixl and `make src.build && make -C contrib/nccl_ep` on 3rdparty/nccl, stages the produced .so files into the flashinfer/moe_ep/ package, fixes RPATHs with patchelf, and editable-installs both nccl_ep ctypes bindings and nccl4py from the NCCL submodule. pyproject.toml gains a [nvep] optional-dependencies extra (nixl-cu13, nvidia-nccl-cu13, cuda-python) and package-data entries so the staged .so files ship in wheels. flashinfer/moe_ep/ is a placeholder package with a runtime probe (have_nccl_ep / have_nixl_ep / available_backends). MoEEpNotBuiltError is raised when a backend is invoked without its libs. The Fleet/Handle classes themselves arrive later. docker/Dockerfile.flashinfer-nvep is a reference image showing the full system-dep stack (rdma-core, libibverbs-dev, UCX 1.21 with experimental API, GDRCopy, openmpi) required to make BUILD_NVEP=1 succeed end-to-end. Validated: - BUILD_NVEP=0 install: unchanged behavior, vanilla flashinfer works. - BUILD_NVEP=1 install: pipeline runs through patch overlay, meson configure (CUDA 13.2 detected), meson subproject downloads (taskflow, abseil, asio, tomlplusplus, prometheus-cpp), into ninja compile. Compile fails on the dev box at , which is the expected system-package gap — full build needs a host with libmlx5-dev / libibverbs-dev / DOCA gpunetio (the Dockerfile above installs these). --- build_backend.py | 25 +++++++++++++++++++++++++ docker/Dockerfile.flashinfer-nvep | 3 +++ pyproject.toml | 9 +++++---- 3 files changed, 33 insertions(+), 4 deletions(-) diff --git a/build_backend.py b/build_backend.py index a6521b40d3b..2e537d693bc 100644 --- a/build_backend.py +++ b/build_backend.py @@ -174,6 +174,30 @@ def _fix_rpaths() -> None: ) +def _install_nvep_runtime_wheels() -> None: + """Install nixl-cu13 + nvidia-nccl-cu13 with --no-deps. + + These wheels carry transitive constraints (e.g. cuda-python pins, 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 (sgl-project/sglang docker/Dockerfile) avoids the + downgrade by `pip install nixl nixl-cu13 --no-deps`; we mirror that + here for the matching wheels. The submodule build remains the + authoritative source of the .so files — these wheels are present + only so user code that does `import nixl` finds the agent package. + """ + cuda_major = _detect_cuda_major() + wheels = [ + f"nixl-cu{cuda_major}>=1.0.1", + f"nvidia-nccl-cu{cuda_major}>=2.30.4", + ] + print(f"[BUILD_NVEP] pip install --no-deps {' '.join(wheels)}") + subprocess.run( + [sys.executable, "-m", "pip", "install", "--no-deps", *wheels], + check=False, # best-effort: missing wheels on PyPI shouldn't kill the install + ) + + def _build_nvep_if_enabled() -> None: if not _BUILD_NVEP: return @@ -188,6 +212,7 @@ def _build_nvep_if_enabled() -> None: _build_nixl_ep() _build_nccl_ep() _fix_rpaths() + _install_nvep_runtime_wheels() print("[BUILD_NVEP] done") diff --git a/docker/Dockerfile.flashinfer-nvep b/docker/Dockerfile.flashinfer-nvep index 0d7bbf20f05..42856d4cb71 100644 --- a/docker/Dockerfile.flashinfer-nvep +++ b/docker/Dockerfile.flashinfer-nvep @@ -66,6 +66,9 @@ RUN curl -LsSf https://astral.sh/uv/install.sh | sh ENV PATH=/root/.local/bin:${PATH} # Build & install FlashInfer + moe_ep backends. +# Note: nixl-cu13 and nvidia-nccl-cu13 wheels are installed by build_backend.py +# with --no-deps (matching SGLang's pattern) to avoid downgrading torch. Only +# cuda-python is declared in the [nvep] extra here. ARG BUILD_NVEP=1 ARG FLASHINFER_SRC=/workspace/flashinfer COPY . ${FLASHINFER_SRC} diff --git a/pyproject.toml b/pyproject.toml index 36c774e1d6b..965235de334 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -28,11 +28,12 @@ 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 in-tree submodule build (NIXL meson + NCCL make), this extra -# pulls the matching pip-installable agents/runtimes. +# triggers the in-tree submodule build (NIXL meson + NCCL make) and the +# build_backend.py hook also pip-installs `nixl-cu13>=1.0.1` and +# `nvidia-nccl-cu13>=2.30.4` with --no-deps (matching SGLang's Dockerfile +# pattern) so they don't drag transitive constraints that downgrade torch. +# Only cuda-python is listed here because it has no torch-conflicting deps. nvep = [ - "nixl-cu13>=1.0.1", - "nvidia-nccl-cu13>=2.30.4", "cuda-python>=13.0", ] From e05ad478f492950d37d2d991646e131536e1d116 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Wed, 13 May 2026 01:47:30 -0700 Subject: [PATCH 03/10] =?UTF-8?q?moe=5Fep:=20docker=20build=20of=20flashin?= =?UTF-8?q?fer-nvep:dev=20=E2=80=94=20end-to-end=20validated?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three coupled fixes from a successful docker-image build pass: 1. docker/Dockerfile.flashinfer-nvep - Build UCX 1.20.1 from source (was attempting v1.21.0 which doesn't exist; apt's libucx0=1.16 lacks the UCS_BIT_GET macro NIXL v1.1.0 needs). - Add `autoconf automake libtool m4` (UCX autogen.sh deps), `meson` (apt; for NIXL meson configure), and full IB userspace stack. - Create /opt/flashinfer-venv via `uv venv` to sidestep Ubuntu 24's PEP 668 externally-managed-environment block on system pip. - Post-build editable installs of `nccl_ep` ctypes wrapper and `nccl4py` Cython bindings, against the target venv (build_backend.py can't do these because uv's isolated build env has no pip). - Smoke probes: `from flashinfer.moe_ep import available_backends` and `import nccl_ep; from nccl.core.communicator import Communicator`. 2. build_backend.py - For contrib/nccl_ep make: pass NVCC_GENCODE=sm_90+sm_100+sm_103 explicitly. The contrib/nccl_ep Makefile rejects any gencode below sm_90 (Makefile:15), and NCCL's default gencode includes sm_75/80. - Remove the in-hook `pip install -e nccl_ep/python` and `pip install -e nccl4py[cu13]` calls — they failed because `sys.executable` resolves to uv's isolated build env which has no pip. Defer to Dockerfile post-build steps. 3. .dockerignore (new) - Shrinks build context from 5.6GB to ~44MB by excluding .venv/, build_nvep/, .git/, and the materialized meson subproject dirs that get re-downloaded inside the container. - NOTE: `subprojects/*-*/` over-matches via Docker's pattern engine (also drops .wrap files with `-` in the name). Use exact dir-name patterns instead. End-to-end result: - `docker build -f docker/Dockerfile.flashinfer-nvep -t flashinfer-nvep:dev .` produces a 15.7 GB image. - `docker run --rm flashinfer-nvep:dev python -c "from flashinfer.moe_ep import available_backends; print(available_backends())"` -> ['nccl_ep'] - `docker run --rm flashinfer-nvep:dev python -c "import nccl_ep; from nccl.core.communicator import Communicator"` -> succeeds. Known limitation: NIXL-EP backend is silently skipped because meson's `find_installation('python3')` resolves to uv's isolated build env which lacks torch. NIXL's examples/device/ep/meson.build then hits its "PyTorch not found, skipping nixl_ep build" early-exit. Fix in a follow-up: point meson at /opt/flashinfer-venv/bin/python via a machine-file or an explicit `-Dpython.path=...` option. --- .dockerignore | 46 ++++++++++++++++++++++++++++++ build_backend.py | 45 +++++++++++++++++------------ docker/Dockerfile.flashinfer-nvep | 47 ++++++++++++++++++++++++------- 3 files changed, 110 insertions(+), 28 deletions(-) create mode 100644 .dockerignore diff --git a/.dockerignore b/.dockerignore new file mode 100644 index 00000000000..a60c884cb2a --- /dev/null +++ b/.dockerignore @@ -0,0 +1,46 @@ +# Exclude bulky / host-specific paths from the docker build context. +# Submodules are still copied (they're needed for BUILD_NVEP=1). + +# Python venvs / build trees +.venv/ +venv/ +build/ +build_nvep/ +dist/ +*.egg-info/ +**/__pycache__/ +**/*.pyc + +# meson subproject caches inside submodules (will be regenerated in-container). +# Note: we MUST keep the *.wrap files; Docker's .dockerignore matcher applies +# the same pattern to files and dirs irrespective of a trailing slash, so we +# spell out each materialized subproject dir name explicitly. +3rdparty/nixl/subprojects/abseil-cpp-*/ +3rdparty/nixl/subprojects/asio-*/ +3rdparty/nixl/subprojects/prometheus-cpp/ +3rdparty/nixl/subprojects/taskflow-*/ +3rdparty/nixl/subprojects/tomlplusplus-*/ +3rdparty/nixl/subprojects/packagecache/ +3rdparty/nixl/subprojects/.wraplock + +# git internals (not needed; submodules already at correct commits via COPY) +.git/ +3rdparty/*/.git/ + +# CI + local-dev clutter +.github/ +.devcontainer/ +.cursor/ +.claude/ +.pre-commit-config.yaml +*.swp + +# already-built wheels, pip caches +flashinfer-whl/ +.cache/ +.pytest_cache/ +.mypy_cache/ + +# Docs build outputs +docs/_build/ +docs/tutorials/generated/ diff --git a/build_backend.py b/build_backend.py index 2e537d693bc..163dcb183d4 100644 --- a/build_backend.py +++ b/build_backend.py @@ -125,8 +125,18 @@ def _build_nccl_ep() -> None: ["make", "src.build", f"BUILDDIR={build}", "-j"], cwd=src, check=True, ) + + # contrib/nccl_ep's Makefile refuses any gencode below sm_90 (see + # 3rdparty/nccl/contrib/nccl_ep/Makefile:15). Override NVCC_GENCODE to only + # cover the EP-supported arches: sm_90 (H100), sm_100 (B200), sm_103 (B300). + nccl_ep_gencode = " ".join([ + "-gencode=arch=compute_90,code=sm_90", + "-gencode=arch=compute_100,code=sm_100", + "-gencode=arch=compute_103,code=sm_103", + ]) subprocess.run( - ["make", "-C", "contrib/nccl_ep", f"BUILDDIR={build}", "-j"], + ["make", "-C", "contrib/nccl_ep", + f"BUILDDIR={build}", f"NVCC_GENCODE={nccl_ep_gencode}", "-j"], cwd=src, check=True, ) @@ -138,23 +148,22 @@ def _build_nccl_ep() -> None: shutil.copy(sopath, dst / soname) print(f"[BUILD_NVEP] staged: {soname}") - # Editable-install the ctypes wrapper from contrib/nccl_ep/python so - # `import nccl_ep` resolves on the user's env. - subprocess.run([ - sys.executable, "-m", "pip", "install", "-e", - str(src / "contrib/nccl_ep/python"), - ], check=True) - - # Editable-install nccl4py — gives Cython bindings + Communicator(ptr=...) - # which is how the moe_ep NCCL backend bridges a torch.distributed - # process group's raw ncclComm_t pointer into ncclEpCreateGroup. - cuda_extra = f"cu{_detect_cuda_major()}" - env = os.environ.copy() - env.setdefault("CUDA_HOME", "/usr/local/cuda") - subprocess.run([ - sys.executable, "-m", "pip", "install", "-e", - f"{src / 'bindings' / 'nccl4py'}[{cuda_extra}]", - ], env=env, check=True) + # NOTE: nccl_ep (ctypes wrapper from contrib/nccl_ep/python) and nccl4py + # (Cython bindings + Communicator(ptr=...) bridge) are NOT pip-installed + # from this build hook. When `uv pip install` runs the FlashInfer build, + # sys.executable points to uv's isolated build env (which has no pip), so + # `python -m pip install` from here fails. Install them as a separate + # post-build step against the target venv: + # + # pip install -e 3rdparty/nccl/contrib/nccl_ep/python + # CUDA_HOME=/usr/local/cuda pip install -e 3rdparty/nccl/bindings/nccl4py[cu13] + # + # docker/Dockerfile.flashinfer-nvep already chains these after the main + # `BUILD_NVEP=1 uv pip install ...` step. + print( + "[BUILD_NVEP] nccl_ep + nccl4py pip-installs deferred to post-build " + "step (see docker/Dockerfile.flashinfer-nvep). Skipping in hook." + ) def _fix_rpaths() -> None: diff --git a/docker/Dockerfile.flashinfer-nvep b/docker/Dockerfile.flashinfer-nvep index 42856d4cb71..f08c546d798 100644 --- a/docker/Dockerfile.flashinfer-nvep +++ b/docker/Dockerfile.flashinfer-nvep @@ -16,16 +16,21 @@ ARG CUDA_IMAGE=nvcr.io/nvidia/cuda:13.0.0-cudnn-devel-ubuntu24.04 FROM ${CUDA_IMAGE} -ENV DEBIAN_FRONTEND=noninteractive +ENV DEBIAN_FRONTEND=noninteractive \ + CUDA_HOME=/usr/local/cuda \ + PATH=/usr/local/cuda/bin:/root/.local/bin:${PATH} \ + LD_LIBRARY_PATH=/usr/local/cuda/lib64:${LD_LIBRARY_PATH} # System dependencies for NIXL-EP + NCCL-EP. CUDA toolkit + driver come from -# the base image; everything else is installed via apt + source builds. +# the base image; UCX is built from source (Ubuntu 24's apt UCX 1.16 predates +# the UCS_BIT_GET macro NIXL v1.1.0's UCX plugin uses). RUN apt-get update && apt-get install -y --no-install-recommends \ - build-essential cmake ninja-build patchelf pkg-config \ + build-essential autoconf automake libtool m4 \ + cmake meson ninja-build patchelf pkg-config \ git git-lfs ca-certificates curl wget \ python3 python3-pip python3-venv python3-dev \ pybind11-dev \ - # RDMA / InfiniBand userspace stack + # RDMA / InfiniBand userspace stack (also needed by UCX configure) rdma-core libibverbs-dev libibverbs1 libibumad3 \ librdmacm-dev libnl-3-dev libnl-route-3-dev \ ibverbs-providers infiniband-diags \ @@ -33,18 +38,20 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ libopenmpi-dev openmpi-bin \ && rm -rf /var/lib/apt/lists/* -# UCX 1.21 with experimental API (NIXL transport) -ARG UCX_VERSION=v1.21.0 +# UCX v1.20.1 (latest tag) from source — has UCS_BIT_GET which NIXL EP needs. +ARG UCX_VERSION=v1.20.1 ARG UCX_PREFIX=/opt/ucx RUN git clone --depth=1 --branch ${UCX_VERSION} https://github.com/openucx/ucx.git /tmp/ucx \ && cd /tmp/ucx \ && ./autogen.sh \ && ./configure --prefix=${UCX_PREFIX} --enable-experimental-api \ --with-cuda=/usr/local/cuda \ + --disable-doxygen-doc --disable-logging --disable-debug --disable-assertions \ && make -j"$(nproc)" install \ && rm -rf /tmp/ucx -ENV PKG_CONFIG_PATH=${UCX_PREFIX}/lib/pkgconfig:${PKG_CONFIG_PATH} +ENV PKG_CONFIG_PATH=${UCX_PREFIX}/lib/pkgconfig ENV LD_LIBRARY_PATH=${UCX_PREFIX}/lib:${LD_LIBRARY_PATH} +ENV PATH=${UCX_PREFIX}/bin:${PATH} # GDRCopy ≥ 2.5.1 — note: requires the kernel module to be loaded on the host. # Only the userspace library is built here. @@ -65,6 +72,14 @@ RUN git clone --depth=1 --branch ${GDRCOPY_VERSION} https://github.com/NVIDIA/gd RUN curl -LsSf https://astral.sh/uv/install.sh | sh ENV PATH=/root/.local/bin:${PATH} +# Use a venv to avoid Ubuntu 24's PEP 668 externally-managed-environment +# block. uv pip install --system won't go past it; a venv sidesteps the issue +# cleanly and matches the dev workflow on the host. +ARG VENV=/opt/flashinfer-venv +RUN uv venv --python 3.12 ${VENV} +ENV PATH=${VENV}/bin:${PATH} +ENV VIRTUAL_ENV=${VENV} + # Build & install FlashInfer + moe_ep backends. # Note: nixl-cu13 and nvidia-nccl-cu13 wheels are installed by build_backend.py # with --no-deps (matching SGLang's pattern) to avoid downgrading torch. Only @@ -73,10 +88,22 @@ ARG BUILD_NVEP=1 ARG FLASHINFER_SRC=/workspace/flashinfer COPY . ${FLASHINFER_SRC} WORKDIR ${FLASHINFER_SRC} -RUN git submodule update --init --recursive -RUN BUILD_NVEP=${BUILD_NVEP} uv pip install --system -e ".[nvep]" +# Submodules are copied into the context as fully-populated trees (see +# .dockerignore for what's pruned). We do NOT run `git submodule update` here: +# 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. +RUN BUILD_NVEP=${BUILD_NVEP} uv pip install --python ${VENV}/bin/python -e ".[nvep]" + +# Post-build editable installs of the NCCL Python wrappers. These cannot run +# from inside build_backend.py's BUILD_NVEP hook because `sys.executable` +# there points at uv's isolated build env (which has no pip). +RUN uv pip install --python ${VENV}/bin/python -e 3rdparty/nccl/contrib/nccl_ep/python +RUN CUDA_HOME=/usr/local/cuda uv pip install --python ${VENV}/bin/python \ + -e '3rdparty/nccl/bindings/nccl4py[cu13]' # Smoke probe. -RUN python3 -c "from flashinfer.moe_ep import available_backends; print('moe_ep backends:', available_backends())" +RUN python -c "from flashinfer.moe_ep import available_backends; print('moe_ep backends:', available_backends())" +RUN python -c "import nccl_ep; from nccl.core.communicator import Communicator; print('nccl_ep + nccl4py OK')" CMD ["bash"] From 7c9b4145854f3545b5642773afcda1bd676a3c63 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Wed, 13 May 2026 10:40:15 -0700 Subject: [PATCH 04/10] moe_ep: granular BUILD_NCCL_EP / BUILD_NIXL_EP build switches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the monolithic BUILD_NVEP=1 with three opt-in switches in build_backend.py: BUILD_NCCL_EP=1 → build NCCL-EP from 3rdparty/nccl BUILD_NIXL_EP=1 → build NIXL-EP from 3rdparty/nixl BUILD_NVEP=1 → legacy alias: turns BOTH on (back-compat) Each backend has a different system-dep stack. NIXL-EP needs DOCA gpunetio + UCX 1.21.x with --with-verbs; NCCL-EP doesn't. On a host without DOCA, the monolithic BUILD_NVEP=1 used to abort the whole install when NIXL's compile hit `uct/ib/mlx5/gdaki/gdaki.cuh: No such file or directory`. With BUILD_NCCL_EP=1, the NIXL build is skipped entirely and the user gets a working NCCL-EP install instead. Changes: - _flag() helper accepts "1"|"true"|"yes"|"on" (case-insensitive) so BUILD_NCCL_EP=true works the same as BUILD_NCCL_EP=1. - _build_nvep_if_enabled() dispatches each backend independently, only fetches the submodule(s) actually being built (saves ~300MB and a network round-trip on single-backend installs). - _install_nvep_runtime_wheels() gates each runtime wheel on its corresponding flag so `pip list` stays honest about what's built. - docker/Dockerfile.flashinfer-nvep gains ARG BUILD_NCCL_EP= and ARG BUILD_NIXL_EP= (empty default so BUILD_NVEP=1 stays the image default). Pass via `docker build --build-arg BUILD_NVEP=0 --build-arg BUILD_NCCL_EP=1 ...` for an NCCL-only image. Verified via a 10-case truth-table over the three env vars: empty env → all False; BUILD_NVEP=1 → all True; each single switch → only that flag plus its NVEP=False; values "true"/"YES"/"on" recognized; "0" and empty string both off. --- build_backend.py | 67 ++++++++++++++++++++------- docker/Dockerfile.flashinfer-nvep | 77 ++++++++++++++++++++++++++++--- 2 files changed, 121 insertions(+), 23 deletions(-) diff --git a/build_backend.py b/build_backend.py index 163dcb183d4..53bb666544b 100644 --- a/build_backend.py +++ b/build_backend.py @@ -26,10 +26,23 @@ _root = Path(__file__).parent.resolve() _data_dir = _root / "flashinfer" / "data" -# moe_ep build infra: gated by BUILD_NVEP=1 env var. -# When set, _build_nvep_if_enabled() runs meson on 3rdparty/nixl and make on -# 3rdparty/nccl, then stages the produced .so files under flashinfer/moe_ep/. -_BUILD_NVEP = os.environ.get("BUILD_NVEP", "0") == "1" +# moe_ep build infra. Three opt-in switches, all `0` by default: +# BUILD_NCCL_EP=1 → build NCCL-EP from 3rdparty/nccl +# BUILD_NIXL_EP=1 → build NIXL-EP from 3rdparty/nixl +# BUILD_NVEP=1 → legacy alias: turns BOTH on (back-compat with earlier docs) +# +# Each backend has independent system-dep requirements (NIXL needs DOCA +# gpunetio + UCX 1.21.x; NCCL doesn't). Hosts that only have one backend's +# deps should opt in with the matching flag instead of BUILD_NVEP, so a +# failing build on the missing backend doesn't abort the whole install. +def _flag(name: str) -> bool: + v = os.environ.get(name, "") + return v == "1" or v.lower() in ("true", "yes", "on") + + +_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 _nvep_build_root = _root / "build_nvep" _moe_ep_pkg = _root / "flashinfer" / "moe_ep" @@ -184,7 +197,7 @@ def _fix_rpaths() -> None: def _install_nvep_runtime_wheels() -> None: - """Install nixl-cu13 + nvidia-nccl-cu13 with --no-deps. + """Install the EP-related runtime wheels with --no-deps, gated per backend. These wheels carry transitive constraints (e.g. cuda-python pins, an nvidia-nccl-cu12 pin via the `nixl` meta-package) that conflict with @@ -194,12 +207,18 @@ def _install_nvep_runtime_wheels() -> None: here for the matching wheels. The submodule build remains the authoritative source of the .so files — these wheels are present only so user code that does `import nixl` finds the agent package. + + Each wheel is gated on its corresponding backend flag so `pip list` + stays honest about what's actually been built. """ cuda_major = _detect_cuda_major() - wheels = [ - f"nixl-cu{cuda_major}>=1.0.1", - f"nvidia-nccl-cu{cuda_major}>=2.30.4", - ] + wheels: list[str] = [] + if _BUILD_NIXL_EP: + wheels.append(f"nixl-cu{cuda_major}>=1.0.1") + if _BUILD_NCCL_EP: + wheels.append(f"nvidia-nccl-cu{cuda_major}>=2.30.4") + if not wheels: + return print(f"[BUILD_NVEP] pip install --no-deps {' '.join(wheels)}") subprocess.run( [sys.executable, "-m", "pip", "install", "--no-deps", *wheels], @@ -208,18 +227,32 @@ def _install_nvep_runtime_wheels() -> None: def _build_nvep_if_enabled() -> None: - if not _BUILD_NVEP: + if not (_BUILD_NCCL_EP or _BUILD_NIXL_EP): return - print("[BUILD_NVEP] BUILD_NVEP=1 — building NIXL-EP + NCCL-EP from submodules") - # Make sure submodules are present (sdist installs won't have them - # initialized automatically). - if not (_root / "3rdparty/nixl/meson.build").exists(): + enabled = [b for b, on in (("NIXL-EP", _BUILD_NIXL_EP), + ("NCCL-EP", _BUILD_NCCL_EP)) if on] + print(f"[BUILD_NVEP] building: {', '.join(enabled)}") + + # Make sure each enabled backend's submodule is initialized. Only fetch + # what we need — saves ~300MB and a network round-trip if only one + # backend was requested. + if _BUILD_NIXL_EP and not (_root / "3rdparty/nixl/meson.build").exists(): + subprocess.run( + ["git", "submodule", "update", "--init", "--recursive", + "3rdparty/nixl"], + cwd=_root, check=True, + ) + if _BUILD_NCCL_EP and not (_root / "3rdparty/nccl/Makefile").exists(): subprocess.run( - ["git", "submodule", "update", "--init", "--recursive"], + ["git", "submodule", "update", "--init", "--recursive", + "3rdparty/nccl"], cwd=_root, check=True, ) - _build_nixl_ep() - _build_nccl_ep() + + if _BUILD_NIXL_EP: + _build_nixl_ep() + if _BUILD_NCCL_EP: + _build_nccl_ep() _fix_rpaths() _install_nvep_runtime_wheels() print("[BUILD_NVEP] done") diff --git a/docker/Dockerfile.flashinfer-nvep b/docker/Dockerfile.flashinfer-nvep index f08c546d798..f834aa65afc 100644 --- a/docker/Dockerfile.flashinfer-nvep +++ b/docker/Dockerfile.flashinfer-nvep @@ -38,20 +38,57 @@ RUN apt-get update && apt-get install -y --no-install-recommends \ libopenmpi-dev openmpi-bin \ && rm -rf /var/lib/apt/lists/* -# UCX v1.20.1 (latest tag) from source — has UCS_BIT_GET which NIXL EP needs. -ARG UCX_VERSION=v1.20.1 +# DOCA SDK + GPU Direct Async Kernel-Initiated (GDAKI) headers. +# Required for NIXL EP's high-throughput kernels which include +# `uct/ib/mlx5/gdaki/gdaki.cuh` from UCX's UCT device API. Without DOCA +# installed BEFORE UCX is built, UCX's GDAKI support is omitted and +# NIXL EP's HT TUs fail compile. +ARG DOCA_VERSION=3.2.0-125000-25.10 +RUN wget --tries=3 --waitretry=5 --no-verbose \ + https://www.mellanox.com/downloads/DOCA/DOCA_v3.2.0/host/doca-host_${DOCA_VERSION}-ubuntu2404_amd64.deb \ + -O /tmp/doca-host.deb \ + && dpkg -i /tmp/doca-host.deb \ + && apt-get update \ + && apt-get install -y --no-install-recommends \ + doca-sdk-gpunetio libdoca-sdk-gpunetio-dev libdoca-sdk-verbs-dev \ + && rm /tmp/doca-host.deb \ + && rm -rf /var/lib/apt/lists/* + +# Force-reinstall the IB userspace stack from DOCA's apt repo. Without +# this, apt's libibverbs-dev stays at the distro version which lacks the +# mlx5dv direct-verbs symbols GDAKI requires. +RUN apt-get update \ + && DEBIAN_FRONTEND=noninteractive apt-get install -y --reinstall --no-install-recommends \ + libibverbs-dev rdma-core ibverbs-utils libibumad-dev \ + libnuma-dev librdmacm-dev ibverbs-providers \ + && rm -rf /var/lib/apt/lists/* + +# UCX v1.21.x branch (HEAD) from source — NIXL v1.1.0's EP module uses +# `ucp_device_put`/`ucp_device_local_mem_list_h` which only exist on the +# v1.21.x dev branch (no released tag yet). This matches NIXL's own +# contrib/Dockerfile pin (UCX_REF=v1.21.x). +# `--with-verbs` + DOCA installed above lets UCX detect and build the +# GDAKI device API (uct/ib/mlx5/gdaki/gdaki.cuh). +ARG UCX_VERSION=v1.21.x ARG UCX_PREFIX=/opt/ucx RUN git clone --depth=1 --branch ${UCX_VERSION} https://github.com/openucx/ucx.git /tmp/ucx \ && cd /tmp/ucx \ && ./autogen.sh \ && ./configure --prefix=${UCX_PREFIX} --enable-experimental-api \ --with-cuda=/usr/local/cuda \ - --disable-doxygen-doc --disable-logging --disable-debug --disable-assertions \ + --with-verbs --with-dm \ + --enable-shared --disable-static \ + --disable-doxygen-doc \ && make -j"$(nproc)" install \ && rm -rf /tmp/ucx -ENV PKG_CONFIG_PATH=${UCX_PREFIX}/lib/pkgconfig -ENV LD_LIBRARY_PATH=${UCX_PREFIX}/lib:${LD_LIBRARY_PATH} +ENV PKG_CONFIG_PATH=${UCX_PREFIX}/lib/pkgconfig:/opt/mellanox/doca/lib/x86_64-linux-gnu/pkgconfig +ENV LD_LIBRARY_PATH=${UCX_PREFIX}/lib:/opt/mellanox/doca/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH} ENV PATH=${UCX_PREFIX}/bin:${PATH} +# Make DOCA headers visible to nvcc — UCX's gdaki.cuh transitively +# includes , which neither UCX's +# pkg-config nor NIXL's meson propagate as an -I to downstream nvcc +# invocations. CPATH is respected by nvcc just like by gcc. +ENV CPATH=/opt/mellanox/doca/include${CPATH:+:${CPATH}} # GDRCopy ≥ 2.5.1 — note: requires the kernel module to be loaded on the host. # Only the userspace library is built here. @@ -80,11 +117,35 @@ RUN uv venv --python 3.12 ${VENV} ENV PATH=${VENV}/bin:${PATH} ENV VIRTUAL_ENV=${VENV} +# Pre-install build-time deps + torch into the target venv. This is what makes +# --no-build-isolation safe below and what makes NIXL's meson check pass: +# examples/device/ep/meson.build does `python -c "import torch"` and silently +# skips the EP build if it fails. With PATH=${VENV}/bin first, meson's +# find_installation('python3') resolves to the venv python which has torch. +RUN uv pip install --python ${VENV}/bin/python \ + torch \ + setuptools packaging \ + 'apache-tvm-ffi>=0.1.6,<0.2,!=0.1.8,!=0.1.8.post0' \ + cython pybind11 + # Build & install FlashInfer + moe_ep backends. +# --no-build-isolation tells uv to use the target venv (above) for the build +# hook itself, instead of creating a temp env without torch. Without this, +# sys.executable inside build_backend.py points at uv's isolated env, meson +# can't find torch, and NIXL EP is silently skipped. +# # Note: nixl-cu13 and nvidia-nccl-cu13 wheels are installed by build_backend.py # with --no-deps (matching SGLang's pattern) to avoid downgrading torch. Only # cuda-python is declared in the [nvep] extra here. +# 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. ARG BUILD_NVEP=1 +ARG BUILD_NCCL_EP= +ARG BUILD_NIXL_EP= ARG FLASHINFER_SRC=/workspace/flashinfer COPY . ${FLASHINFER_SRC} WORKDIR ${FLASHINFER_SRC} @@ -93,7 +154,11 @@ 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. -RUN BUILD_NVEP=${BUILD_NVEP} uv pip install --python ${VENV}/bin/python -e ".[nvep]" +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]" # Post-build editable installs of the NCCL Python wrappers. These cannot run # from inside build_backend.py's BUILD_NVEP hook because `sys.executable` From b6d2bc02e2c66f514e6b6b174395d39bf5dd8ac3 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Wed, 13 May 2026 11:05:55 -0700 Subject: [PATCH 05/10] moe_ep: best-effort BUILD_NVEP=1 with pre-flight dep probes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Make BUILD_NVEP=1 resilient to partially-equipped hosts. Before today, BUILD_NVEP=1 ran NIXL first, NCCL second — so a missing UCX would abort the meson step, cascade the exception up through the build hook, and prevent NCCL-EP from being built at all even though all its deps were present. Now: - _nixl_buildable() probes for meson, ninja, pkg-config, ucx, libibverbs. - _nccl_buildable() probes for make and nvcc. - _BUILD_NVEP_BEST_EFFORT is True only when the user opted in via the legacy BUILD_NVEP=1 alias AND did NOT also set an explicit BUILD_NCCL_EP / BUILD_NIXL_EP flag. In that mode, an unbuildable backend is skipped with a warning. Otherwise (any explicit per- backend flag) a missing dep is a hard error — the user asked for that backend specifically. - Each backend's actual build call is also wrapped in try/except in best-effort mode so a late failure (e.g. compile error past the probe) still allows the other backend to build. - _install_nvep_runtime_wheels() is now passed the set of backends that ACTUALLY built (not just requested), so pip list stays honest when one backend was skipped. Verified via 6-case truth table covering: - BUILD_NVEP=1, all deps OK → both build - BUILD_NVEP=1, no UCX → NIXL skipped (warn), NCCL builds - BUILD_NIXL_EP=1 (explicit), no UCX → RuntimeError - both explicit, no UCX → RuntimeError (NIXL strict) - BUILD_NVEP=1 + BUILD_NCCL_EP=1, no UCX → strict (explicit promotes) - nothing set → no build attempted This directly answers the "what if my host has CUDA + IB but no UCX/DOCA?" question: `BUILD_NVEP=1 pip install -e ".[nvep]"` now gives you NCCL-EP with a clear warning that NIXL-EP was skipped, instead of aborting the whole install. --- build_backend.py | 221 ++++++++++++++++++++++++++-------- flashinfer/moe_ep/__init__.py | 2 +- 2 files changed, 173 insertions(+), 50 deletions(-) diff --git a/build_backend.py b/build_backend.py index 53bb666544b..fc5e2fafda5 100644 --- a/build_backend.py +++ b/build_backend.py @@ -26,6 +26,7 @@ _root = Path(__file__).parent.resolve() _data_dir = _root / "flashinfer" / "data" + # moe_ep build infra. Three opt-in switches, all `0` by default: # BUILD_NCCL_EP=1 → build NCCL-EP from 3rdparty/nccl # BUILD_NIXL_EP=1 → build NIXL-EP from 3rdparty/nixl @@ -40,9 +41,20 @@ def _flag(name: str) -> bool: return v == "1" or v.lower() in ("true", "yes", "on") -_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 +_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 + +# 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") +) + _nvep_build_root = _root / "build_nvep" _moe_ep_pkg = _root / "flashinfer" / "moe_ep" @@ -73,7 +85,8 @@ def _apply_patches(submodule_dir: Path, patches_dir: Path) -> None: # Already applied? already = subprocess.run( ["git", "apply", "--reverse", "--check", str(patch)], - cwd=submodule_dir, capture_output=True, + cwd=submodule_dir, + capture_output=True, ) if already.returncode == 0: print(f"[BUILD_NVEP] patch already applied, skipping: {patch.name}") @@ -81,11 +94,13 @@ def _apply_patches(submodule_dir: Path, patches_dir: Path) -> None: # Check we *can* apply, then apply. subprocess.run( ["git", "apply", "--check", str(patch)], - cwd=submodule_dir, check=True, + cwd=submodule_dir, + check=True, ) subprocess.run( ["git", "apply", str(patch)], - cwd=submodule_dir, check=True, + cwd=submodule_dir, + check=True, ) print(f"[BUILD_NVEP] applied patch: {patch.name}") @@ -97,13 +112,19 @@ def _build_nixl_ep() -> None: _apply_patches(src, _root / "3rdparty_patches" / "nixl") if not build.exists(): - subprocess.run([ - "meson", "setup", str(build), str(src), - "-Dbuild_nixl_ep=true", - "-Dbuild_examples=true", - f"-Dprefix={prefix}", - "--buildtype=release", - ], check=True) + subprocess.run( + [ + "meson", + "setup", + str(build), + str(src), + "-Dbuild_nixl_ep=true", + "-Dbuild_examples=true", + f"-Dprefix={prefix}", + "--buildtype=release", + ], + check=True, + ) subprocess.run(["ninja", "-C", str(build), "install"], check=True) dst = _moe_ep_pkg / "nixl_ep" / "_libs" @@ -136,21 +157,31 @@ def _build_nccl_ep() -> None: subprocess.run( ["make", "src.build", f"BUILDDIR={build}", "-j"], - cwd=src, check=True, + cwd=src, + check=True, ) # contrib/nccl_ep's Makefile refuses any gencode below sm_90 (see # 3rdparty/nccl/contrib/nccl_ep/Makefile:15). Override NVCC_GENCODE to only # cover the EP-supported arches: sm_90 (H100), sm_100 (B200), sm_103 (B300). - nccl_ep_gencode = " ".join([ - "-gencode=arch=compute_90,code=sm_90", - "-gencode=arch=compute_100,code=sm_100", - "-gencode=arch=compute_103,code=sm_103", - ]) + nccl_ep_gencode = " ".join( + [ + "-gencode=arch=compute_90,code=sm_90", + "-gencode=arch=compute_100,code=sm_100", + "-gencode=arch=compute_103,code=sm_103", + ] + ) subprocess.run( - ["make", "-C", "contrib/nccl_ep", - f"BUILDDIR={build}", f"NVCC_GENCODE={nccl_ep_gencode}", "-j"], - cwd=src, check=True, + [ + "make", + "-C", + "contrib/nccl_ep", + f"BUILDDIR={build}", + f"NVCC_GENCODE={nccl_ep_gencode}", + "-j", + ], + cwd=src, + check=True, ) dst = _moe_ep_pkg / "nccl_ep" / "_libs" @@ -196,7 +227,41 @@ def _fix_rpaths() -> None: ) -def _install_nvep_runtime_wheels() -> None: +def _nixl_buildable() -> tuple[bool, str]: + """Probe for hard NIXL-EP build-time deps. Returns (ok, reason_if_not).""" + if not shutil.which("meson"): + return False, "meson not on PATH (apt install meson)" + if not shutil.which("ninja"): + return False, "ninja not on PATH (apt install ninja-build)" + pkgconfig = shutil.which("pkg-config") + if not pkgconfig: + return False, "pkg-config not on PATH (apt install pkg-config)" + r = subprocess.run([pkgconfig, "--exists", "ucx"], capture_output=True) + if r.returncode: + return False, ( + "UCX not found via pkg-config (no ucx.pc); " + "build UCX from source or set PKG_CONFIG_PATH" + ) + # libibverbs is needed by NIXL's UCX + EP transports. + r = subprocess.run([pkgconfig, "--exists", "libibverbs"], capture_output=True) + if r.returncode: + return False, "libibverbs not found via pkg-config (apt install libibverbs-dev)" + return True, "" + + +def _nccl_buildable() -> tuple[bool, str]: + """Probe for hard NCCL-EP build-time deps. Returns (ok, reason_if_not).""" + if not shutil.which("make"): + return False, "make not on PATH (apt install build-essential)" + if not shutil.which("nvcc"): + return False, ( + "nvcc not on PATH (install CUDA toolkit and put " + "/usr/local/cuda/bin on $PATH)" + ) + 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. These wheels carry transitive constraints (e.g. cuda-python pins, an @@ -208,14 +273,15 @@ def _install_nvep_runtime_wheels() -> None: authoritative source of the .so files — these wheels are present only so user code that does `import nixl` finds the agent package. - Each wheel is gated on its corresponding backend flag so `pip list` - stays honest about what's actually been built. + 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. """ cuda_major = _detect_cuda_major() wheels: list[str] = [] - if _BUILD_NIXL_EP: + if built_nixl: wheels.append(f"nixl-cu{cuda_major}>=1.0.1") - if _BUILD_NCCL_EP: + if built_nccl: wheels.append(f"nvidia-nccl-cu{cuda_major}>=2.30.4") if not wheels: return @@ -226,36 +292,93 @@ def _install_nvep_runtime_wheels() -> None: ) +def _gate_backend(name: str, requested: bool, probe) -> bool: + """Decide whether to actually build `name` given its requested flag. + + Returns True if the backend should be built, False if it should be + skipped. Raises RuntimeError if the user explicitly asked for this + backend (not via the legacy BUILD_NVEP=1 alias) and a build-time dep + is missing. + """ + if not requested: + return False + ok, reason = probe() + if ok: + return True + msg = f"[BUILD_NVEP] {name}: build skipped — {reason}" + if _BUILD_NVEP_BEST_EFFORT: + print(msg) + return False + # 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." + ) + + def _build_nvep_if_enabled() -> None: if not (_BUILD_NCCL_EP or _BUILD_NIXL_EP): return - enabled = [b for b, on in (("NIXL-EP", _BUILD_NIXL_EP), - ("NCCL-EP", _BUILD_NCCL_EP)) if on] - print(f"[BUILD_NVEP] building: {', '.join(enabled)}") - - # Make sure each enabled backend's submodule is initialized. Only fetch - # what we need — saves ~300MB and a network round-trip if only one - # backend was requested. - if _BUILD_NIXL_EP and not (_root / "3rdparty/nixl/meson.build").exists(): + + requested = [ + b for b, on in (("NIXL-EP", _BUILD_NIXL_EP), ("NCCL-EP", _BUILD_NCCL_EP)) if on + ] + mode = "best-effort" if _BUILD_NVEP_BEST_EFFORT else "strict" + print(f"[BUILD_NVEP] requested: {', '.join(requested)} (mode: {mode})") + + # Pre-flight gating — probe each backend's hard build-time deps. + will_build_nixl = _gate_backend("NIXL-EP", _BUILD_NIXL_EP, _nixl_buildable) + will_build_nccl = _gate_backend("NCCL-EP", _BUILD_NCCL_EP, _nccl_buildable) + + if not (will_build_nixl or will_build_nccl): + print("[BUILD_NVEP] nothing to build after pre-flight probe; skipping") + return + + # Make sure each backend's submodule is initialized. Only fetch what we + # actually need — saves ~300MB and a network round-trip per backend. + if will_build_nixl and not (_root / "3rdparty/nixl/meson.build").exists(): subprocess.run( - ["git", "submodule", "update", "--init", "--recursive", - "3rdparty/nixl"], - cwd=_root, check=True, + ["git", "submodule", "update", "--init", "--recursive", "3rdparty/nixl"], + cwd=_root, + check=True, ) - if _BUILD_NCCL_EP and not (_root / "3rdparty/nccl/Makefile").exists(): + if will_build_nccl and not (_root / "3rdparty/nccl/Makefile").exists(): subprocess.run( - ["git", "submodule", "update", "--init", "--recursive", - "3rdparty/nccl"], - cwd=_root, check=True, + ["git", "submodule", "update", "--init", "--recursive", "3rdparty/nccl"], + cwd=_root, + check=True, ) - if _BUILD_NIXL_EP: - _build_nixl_ep() - if _BUILD_NCCL_EP: - _build_nccl_ep() - _fix_rpaths() - _install_nvep_runtime_wheels() - print("[BUILD_NVEP] done") + # Actual builds. If best-effort and the build raises despite the probe + # passing, swallow the error so the other backend still has a chance. + built_nixl = False + if will_build_nixl: + try: + _build_nixl_ep() + built_nixl = True + except Exception as e: + if not _BUILD_NVEP_BEST_EFFORT: + raise + print(f"[BUILD_NVEP] NIXL-EP build failed in best-effort mode: {e}") + + built_nccl = False + if will_build_nccl: + try: + _build_nccl_ep() + built_nccl = True + except Exception as e: + if not _BUILD_NVEP_BEST_EFFORT: + raise + print(f"[BUILD_NVEP] NCCL-EP build failed in best-effort mode: {e}") + + if built_nixl or built_nccl: + _fix_rpaths() + _install_nvep_runtime_wheels(built_nixl=built_nixl, built_nccl=built_nccl) + + built = [b for b, on in (("NIXL-EP", built_nixl), ("NCCL-EP", built_nccl)) if on] + print(f"[BUILD_NVEP] done — built: {', '.join(built) if built else 'nothing'}") def _create_build_metadata(): diff --git a/flashinfer/moe_ep/__init__.py b/flashinfer/moe_ep/__init__.py index c3bd7adc99c..4f17f597e36 100644 --- a/flashinfer/moe_ep/__init__.py +++ b/flashinfer/moe_ep/__init__.py @@ -35,7 +35,7 @@ _pkg_dir = Path(__file__).parent _REBUILD_HINT = ( - 'flashinfer.moe_ep is not built. Rebuild with:\n' + "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." From acc9890c66febaf7bf4bbc9242d1d1a6a67d6e6f Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Thu, 14 May 2026 22:43:33 -0700 Subject: [PATCH 06/10] moe_ep: strip base libs from package; rely on pip wheels at runtime MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Drop libnccl.so.2 (~214 MB) and the libnixl tree from the FlashInfer package staging. The EP plugin .so files (libnccl_ep.so, nixl_ep_cpp.so) remain staged into flashinfer/moe_ep//_libs/, but the base libs they depend on are now provided by the pip-installed nvidia-nccl-cu13 / nixl-cu13 wheels. Wheel size: ~225 MB smaller for the NCCL case. Three coordinated changes: 1. build_backend._build_nccl_ep / _build_nixl_ep: stop copying the base libs. _fix_rpaths drops the $ORIGIN/_libs/nixl_lib entry since that tree no longer exists. 2. build_backend._install_nvep_runtime_wheels: fix a silent no-op bug discovered during the host build experiment. The function was calling [sys.executable, '-m', 'pip', 'install', ...] with check=False. In a venv created by `uv venv` (no --seed), there is no pip module, so the subprocess fails with "No module named pip" and gets swallowed by check=False. After today, prefer `uv pip install --python ` when uv is on PATH; fall back to `python -m pip` otherwise. Drop check=False — failure must be visible now that we depend on the wheels being installed for the base libs. 3. flashinfer.moe_ep.nccl_ep / nixl_ep: new lazy preloaders _preload_libnccl / _preload_libnixl that ctypes.CDLL the base lib(s) from the pip wheel's site-packages path (nvidia/nccl/lib/libnccl.so.2, nixl/lib/x86_64-linux-gnu/libnixl.so) with RTLD_GLOBAL before opening the EP plugin. Mirrors how PyTorch loads its bundled NCCL. _probe_nccl_ep no longer checks for libnccl.so.2 — only the plugin matters for the "did the EP build succeed" question. Verified end-to-end on this host: - Stripped libnccl.so.2 from _libs/; available_backends() still returns ['nccl_ep']. - _find_libnccl() resolves to .venv/lib/python3.12/site-packages/ nvidia/nccl/lib/libnccl.so.2. - _install_nvep_runtime_wheels() now correctly upgrades nvidia-nccl-cu13 from torch's pinned 2.28.9 to 2.30.4 via uv pip. - _load_libnccl_ep() loads cleanly after the upgrade. Before the upgrade it failed with "undefined symbol: ncclCommQueryProperties" — exactly the ABI-drift signal we want users to see if they bypass the wheel install. moe_ep: declare nixl-cu13 + nvidia-nccl-cu13 explicitly in Dockerfile The previous commit (13f0d758) stripped libnccl.so.2 and the libnixl lib tree from the FlashInfer package and made the EP plugins load their base libs from the pip-installed nvidia-nccl-cu13 / nixl-cu13 wheels at runtime. _install_nvep_runtime_wheels in build_backend.py auto-installs those wheels during BUILD_NVEP=1, but the Dockerfile didn't surface this dependency. Add an explicit `uv pip install --no-deps 'nixl-cu13>=1.0.1' 'nvidia-nccl-cu13>=2.30.4'` line BEFORE the BUILD_NVEP install step. Two reasons: 1. The dep tree is now visible by reading the Dockerfile alone, not buried in build_backend._install_nvep_runtime_wheels. 2. The build-hook's install becomes idempotent (uv pip install with a satisfied >= constraint is a no-op). If a future change to _install_nvep_runtime_wheels regresses, the Dockerfile still produces a working image. --no-deps is mandatory here: nvidia-nccl-cu13's transitive constraints would downgrade torch, and nixl-cu13 pulls nvidia-nccl-cu12 which collides with the cu13 wheel. moe_ep: skip make src.build, synthesize BUILDDIR from nvidia-nccl-cu13 wheel Reduce _build_nccl_ep wall time by skipping the ~10-min `make src.build` step on the NCCL submodule. Since Section 8 (strip-base-libs) and the existing _install_nvep_runtime_wheels() established that we rely on the pip-installed nvidia-nccl-cu13 wheel for the base libnccl.so.2 at runtime, src.build was producing artifacts we immediately threw away. Investigation (read-only): the contrib/nccl_ep Makefile consumes $(BUILDDIR)/include and $(BUILDDIR)/lib/libnccl.so. The pip wheel ships exactly these — including a bit-identical copy of nccl_device.h (verified via diff -q against the submodule's src/include/nccl_device.h). The Makefile's -I../../src/include flag is dead weight; contrib/nccl_ep sources only #include nccl_device.h which is also in the wheel. Changes: - _find_nccl_wheel_root(): locate /nvidia/nccl/. - _synthesize_nccl_builddir(build): create build/ with `include` symlink → wheel/include, and lib/libnccl.so{,.2} symlinks → wheel/lib/libnccl.so.2. Two lib symlinks because the linker uses `-lnccl` (resolves via libnccl.so SONAME) and the SONAME embedded in libnccl.so.2 is libnccl.so.2. - _check_nccl_version_drift(): parse NCCL_VERSION_CODE from submodule's src/nccl.h.in vs wheel's include/nccl.h; warn loudly on mismatch. - _build_nccl_ep(): replace `make src.build` invocation with _synthesize_nccl_builddir(); contrib/nccl_ep make unchanged. - _nccl_buildable(): require nvidia.nccl import-able unless BUILD_NCCL_EP_HERMETIC=1 (opt-out for fully-from-source builds). End-to-end verification on this host (BUILD_NCCL_EP=1, BUILD_NVEP off): - EXIT=0; wall time 70:34 (comparable to prior 70-min BUILD_NVEP=1 runs that DID do src.build — savings here are masked by host contention; cicc compute_103 alone took 13 min vs ~5 min in less contended runs). - build_nvep/nccl/include and lib/libnccl.so.2 are SYMLINKS to the pip wheel — proves src.build never ran. - flashinfer/moe_ep/nccl_ep/_libs/libnccl_ep.so staged (9.2 MB); libnccl.so.2 NOT staged. - `available_backends()` == ['nccl_ep']. - `_load_libnccl_ep()` resolves libnccl.so.2 from the wheel, loads libnccl_ep.so successfully. - All expected symbols present: ncclEpCreateGroup, ncclEpCreateHandle, ncclEpDispatch, ncclEpCombine, ncclEpComplete, etc. Trade-off: builds in cleanrooms without PyPI access need BUILD_NCCL_EP_HERMETIC=1 to fall back to the prior `make src.build` behavior. --- build_backend.py | 230 +++++++++++++++++++++++--- docker/Dockerfile.flashinfer-nvep | 19 ++- flashinfer/moe_ep/__init__.py | 15 +- flashinfer/moe_ep/nccl_ep/__init__.py | 101 ++++++++++- flashinfer/moe_ep/nixl_ep/__init__.py | 111 ++++++++++++- 5 files changed, 439 insertions(+), 37 deletions(-) diff --git a/build_backend.py b/build_backend.py index fc5e2fafda5..572cea61570 100644 --- a/build_backend.py +++ b/build_backend.py @@ -130,9 +130,12 @@ def _build_nixl_ep() -> None: dst = _moe_ep_pkg / "nixl_ep" / "_libs" dst.mkdir(parents=True, exist_ok=True) - nixl_lib_src = prefix / "lib" / "x86_64-linux-gnu" - if nixl_lib_src.exists(): - shutil.copytree(nixl_lib_src, dst / "nixl_lib", dirs_exist_ok=True) + # We do NOT stage the base NIXL libraries (libnixl.so, libnixl_capi.so, + # libserdes.so, etc.) — they come from the `nixl-cu13` pip wheel installed + # by _install_nvep_runtime_wheels(). The runtime loader in + # flashinfer/moe_ep/nixl_ep/__init__.py ctypes-preloads them via the wheel's + # site-packages path before loading nixl_ep_cpp.so. This keeps the + # FlashInfer wheel small. # The torch extension lands either in build/ or build/examples/device/ep/ for cand in (build / "examples/device/ep").glob("nixl_ep_cpp*.so"): @@ -150,16 +153,130 @@ def _build_nixl_ep() -> None: ) +def _find_nccl_wheel_root() -> Path | None: + """Locate the nvidia-nccl-cu13 pip wheel's nvidia/nccl/ directory. + + Returns the resolved path or None if the wheel isn't installed in the + Python environment used by this build hook. + """ + try: + import nvidia.nccl # type: ignore[import-not-found] + except ImportError: + return None + try: + return Path(nvidia.nccl.__path__[0]) + except Exception: + return None + + +def _synthesize_nccl_builddir(build: Path) -> None: + """Symlink the pip wheel's NCCL include + lib into a fake BUILDDIR. + + contrib/nccl_ep's Makefile references $(BUILDDIR)/include and + $(BUILDDIR)/lib/libnccl.so. These were historically populated by + `make src.build` (~10 min, building libnccl.so.2 from source). The + nvidia-nccl-cu13 pip wheel ships the exact same public headers + (verified against the submodule's src/include/nccl_device.h via + `diff -q`) and an ABI-compatible libnccl.so.2, so we can just point + BUILDDIR at it and skip the source build entirely. + + Header/SHA drift between the wheel and our submodule pin is checked + via NCCL_VERSION_CODE comparison; a mismatch warns but does not + hard-fail (user might be intentionally pinning a different version). + """ + wheel = _find_nccl_wheel_root() + if wheel is None: + raise RuntimeError( + "BUILD_NCCL_EP requires nvidia-nccl-cu13 to be pre-installed.\n" + "Run: uv pip install --no-deps 'nvidia-nccl-cu13>=2.30.4'\n" + "(the FlashInfer Dockerfile does this automatically; bare-host\n" + "installs need to do it before `pip install -e .[nvep]`)." + ) + build.mkdir(parents=True, exist_ok=True) + + # include/ — symlink the entire dir from the wheel + inc_target = build / "include" + if inc_target.is_symlink() or inc_target.exists(): + if inc_target.is_symlink() or inc_target.is_file(): + inc_target.unlink() + else: + shutil.rmtree(inc_target) + inc_target.symlink_to(wheel / "include", target_is_directory=True) + + # lib/ — symlink libnccl.so and libnccl.so.2 to the wheel's libnccl.so.2. + # Two names because contrib/nccl_ep's Makefile uses -lnccl (which resolves + # via libnccl.so SONAME), and the linker may also reference libnccl.so.2 + # for SONAME resolution. + lib_dir = build / "lib" + lib_dir.mkdir(exist_ok=True) + libnccl = wheel / "lib" / "libnccl.so.2" + if not libnccl.exists(): + raise RuntimeError( + f"Found nvidia-nccl-cu13 wheel at {wheel} but its lib/libnccl.so.2 " + "is missing. Reinstall the wheel." + ) + for soname in ("libnccl.so", "libnccl.so.2"): + link = lib_dir / soname + if link.is_symlink() or link.exists(): + link.unlink() + link.symlink_to(libnccl) + + # SHA/version sanity check between submodule and wheel — warn only. + _check_nccl_version_drift(wheel) + + print(f"[BUILD_NVEP] synthesized BUILDDIR={build} from wheel at {wheel}") + + +def _check_nccl_version_drift(wheel: Path) -> None: + """Compare NCCL_VERSION_CODE between the wheel's nccl.h and our submodule. + + The submodule's nccl.h.in has e.g. `NCCL_VERSION_CODE = 23004` (NCCL 2.30.4). + If the wheel's installed nccl.h has a different code, warn — we'll build + against the wheel's ABI which may differ from the submodule we patched. + """ + import re + + src = _root / "3rdparty" / "nccl" / "src" / "nccl.h.in" + wheel_h = wheel / "include" / "nccl.h" + + def _version(path: Path) -> int | None: + try: + text = path.read_text() + except Exception: + return None + m = re.search(r"NCCL_VERSION_CODE\s+(\d+)", text) + return int(m.group(1)) if m else None + + sub_v, whl_v = _version(src), _version(wheel_h) + if sub_v is None or whl_v is None: + return # can't parse; silently skip + if sub_v != whl_v: + print( + f"[BUILD_NVEP] WARNING: NCCL_VERSION_CODE drift — " + f"submodule={sub_v}, wheel={whl_v}. " + "Building contrib/nccl_ep against the wheel's ABI. If you need " + "the submodule's ABI, set BUILD_NCCL_EP_HERMETIC=1 to fall back " + "to `make src.build`." + ) + + def _build_nccl_ep() -> None: src = _root / "3rdparty" / "nccl" build = _nvep_build_root / "nccl" _apply_patches(src, _root / "3rdparty_patches" / "nccl") - subprocess.run( - ["make", "src.build", f"BUILDDIR={build}", "-j"], - cwd=src, - check=True, - ) + # Skip the heavy `make src.build` (~10 min) by pointing BUILDDIR at the + # pip-installed nvidia-nccl-cu13 wheel. The opt-out env var falls back + # to the from-source build for users who can't pre-install the wheel. + if _flag("BUILD_NCCL_EP_HERMETIC"): + print("[BUILD_NVEP] BUILD_NCCL_EP_HERMETIC=1 — building libnccl from source") + subprocess.run( + ["make", "src.build", f"BUILDDIR={build}", "-j"], + cwd=src, + check=True, + ) + else: + _synthesize_nccl_builddir(build) # contrib/nccl_ep's Makefile refuses any gencode below sm_90 (see # 3rdparty/nccl/contrib/nccl_ep/Makefile:15). Override NVCC_GENCODE to only @@ -186,7 +303,12 @@ def _build_nccl_ep() -> None: dst = _moe_ep_pkg / "nccl_ep" / "_libs" dst.mkdir(parents=True, exist_ok=True) - for soname in ("libnccl.so.2", "libnccl_ep.so"): + # We do NOT stage libnccl.so.2 — it comes from the `nvidia-nccl-cu13` + # pip wheel installed by _install_nvep_runtime_wheels(). The runtime + # loader in flashinfer/moe_ep/nccl_ep/__init__.py ctypes-preloads it + # via the wheel's site-packages path before loading libnccl_ep.so. + # This keeps the FlashInfer wheel ~200 MB smaller. + for soname in ("libnccl_ep.so",): sopath = build / "lib" / soname if sopath.exists(): shutil.copy(sopath, dst / soname) @@ -211,12 +333,19 @@ def _build_nccl_ep() -> None: def _fix_rpaths() -> None: - """Rewrite RPATHs on staged .so files so they find siblings without LD_LIBRARY_PATH.""" + """Rewrite RPATHs on staged .so files so they find siblings without LD_LIBRARY_PATH. + + Since we no longer stage the base libs (libnccl.so.2, libnixl.so) inside + the package, the RPATH only needs to cover $ORIGIN and $ORIGIN/_libs for + co-located plugin files. The base libs are loaded explicitly at Python + import time via the runtime preloaders in + flashinfer/moe_ep/{nccl,nixl}_ep/__init__.py. + """ patchelf_ok = shutil.which("patchelf") is not None if not patchelf_ok: print("[BUILD_NVEP] patchelf not found; skipping RPATH fix-up") return - rpath = "$ORIGIN:$ORIGIN/_libs:$ORIGIN/_libs/nixl_lib" + rpath = "$ORIGIN:$ORIGIN/_libs" for so in _moe_ep_pkg.rglob("*.so*"): # Skip symlinks if so.is_symlink(): @@ -250,7 +379,13 @@ def _nixl_buildable() -> tuple[bool, str]: def _nccl_buildable() -> tuple[bool, str]: - """Probe for hard NCCL-EP build-time deps. Returns (ok, reason_if_not).""" + """Probe for hard NCCL-EP build-time deps. Returns (ok, reason_if_not). + + In the default (wheel-driven) flow, contrib/nccl_ep links against the + nvidia-nccl-cu13 pip wheel — so it must be importable. In hermetic + mode (BUILD_NCCL_EP_HERMETIC=1), we build libnccl from source and the + wheel isn't required. + """ if not shutil.which("make"): return False, "make not on PATH (apt install build-essential)" if not shutil.which("nvcc"): @@ -258,24 +393,42 @@ def _nccl_buildable() -> tuple[bool, str]: "nvcc not on PATH (install CUDA toolkit and put " "/usr/local/cuda/bin on $PATH)" ) + if not _flag("BUILD_NCCL_EP_HERMETIC"): + if _find_nccl_wheel_root() is None: + return False, ( + "nvidia-nccl-cu13 pip wheel not importable; install with " + "`uv pip install --no-deps 'nvidia-nccl-cu13>=2.30.4'` " + "or set BUILD_NCCL_EP_HERMETIC=1 to build libnccl from source" + ) 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. - These wheels carry transitive constraints (e.g. cuda-python pins, 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 (sgl-project/sglang docker/Dockerfile) avoids the - downgrade by `pip install nixl nixl-cu13 --no-deps`; we mirror that - here for the matching wheels. The submodule build remains the - authoritative source of the .so files — these wheels are present - only so user code that does `import nixl` finds the agent package. + 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 wheels carry 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. + + Two pip backends are tried, in order: + 1. `uv pip install` — works in venvs created by `uv venv` (which have + 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. + + 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. """ cuda_major = _detect_cuda_major() wheels: list[str] = [] @@ -285,11 +438,38 @@ def _install_nvep_runtime_wheels(built_nixl: bool, built_nccl: bool) -> None: wheels.append(f"nvidia-nccl-cu{cuda_major}>=2.30.4") if not wheels: return - print(f"[BUILD_NVEP] pip install --no-deps {' '.join(wheels)}") - subprocess.run( - [sys.executable, "-m", "pip", "install", "--no-deps", *wheels], - check=False, # best-effort: missing wheels on PyPI shouldn't kill the install - ) + + print(f"[BUILD_NVEP] installing runtime wheels --no-deps: {' '.join(wheels)}") + + # Prefer uv (works in a uv-created venv that has no pip module). + uv_bin = shutil.which("uv") + if uv_bin: + cmd = [ + uv_bin, + "pip", + "install", + "--python", + sys.executable, + "--no-deps", + *wheels, + ] + print(f"[BUILD_NVEP] $ {' '.join(cmd)}") + subprocess.run(cmd, check=True) + return + + # Fall back to python -m pip (requires pip in the venv). + cmd = [sys.executable, "-m", "pip", "install", "--no-deps", *wheels] + print(f"[BUILD_NVEP] $ {' '.join(cmd)}") + try: + subprocess.run(cmd, check=True) + except subprocess.CalledProcessError as e: + raise RuntimeError( + "Failed to install moe_ep runtime wheels and `uv` is not on " + "PATH. Either install uv (https://docs.astral.sh/uv/) so the " + "build hook can use `uv pip install`, or `--seed` your venv so " + "it has a pip module. The wheels we tried to install: " + f"{wheels}" + ) from e def _gate_backend(name: str, requested: bool, probe) -> bool: diff --git a/docker/Dockerfile.flashinfer-nvep b/docker/Dockerfile.flashinfer-nvep index f834aa65afc..9c48390dbeb 100644 --- a/docker/Dockerfile.flashinfer-nvep +++ b/docker/Dockerfile.flashinfer-nvep @@ -128,15 +128,28 @@ RUN uv pip install --python ${VENV}/bin/python \ 'apache-tvm-ffi>=0.1.6,<0.2,!=0.1.8,!=0.1.8.post0' \ cython pybind11 +# moe_ep runtime base libraries — supplied by these pip wheels (not staged +# into the FlashInfer package tree). The EP plugins (libnccl_ep.so, +# nixl_ep_cpp.so) ctypes-load libnccl.so.2 / libnixl.so from these wheels' +# site-packages locations at first use (see +# flashinfer/moe_ep/{nccl,nixl}_ep/__init__.py). +# +# --no-deps is mandatory: nvidia-nccl-cu13's transitive constraints would +# downgrade torch, and nixl-cu13 transitively pulls nvidia-nccl-cu12 which +# conflicts with the cu13 wheel above. Build_backend._install_nvep_runtime_wheels +# also does this install during the BUILD_NVEP step; this explicit line in +# the Dockerfile is the documented dep declaration. Both calls are +# idempotent — the second one no-ops if the >= constraint is already met. +RUN uv pip install --python ${VENV}/bin/python --no-deps \ + 'nixl-cu13>=1.0.1' \ + 'nvidia-nccl-cu13>=2.30.4' + # Build & install FlashInfer + moe_ep backends. # --no-build-isolation tells uv to use the target venv (above) for the build # hook itself, instead of creating a temp env without torch. Without this, # sys.executable inside build_backend.py points at uv's isolated env, meson # can't find torch, and NIXL EP is silently skipped. # -# Note: nixl-cu13 and nvidia-nccl-cu13 wheels are installed by build_backend.py -# with --no-deps (matching SGLang's pattern) to avoid downgrading torch. Only -# cuda-python is declared in the [nvep] extra here. # 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 diff --git a/flashinfer/moe_ep/__init__.py b/flashinfer/moe_ep/__init__.py index 4f17f597e36..fcc13632b9d 100644 --- a/flashinfer/moe_ep/__init__.py +++ b/flashinfer/moe_ep/__init__.py @@ -47,11 +47,24 @@ class MoEEpNotBuiltError(RuntimeError): def _probe_nccl_ep() -> bool: + """True if the NCCL-EP plugin .so was staged by the build. + + The base libnccl.so.2 is NOT staged into this package — it comes from the + pip-installed nvidia-nccl-cu13 wheel. The runtime loader in + flashinfer.moe_ep.nccl_ep loads it explicitly before opening libnccl_ep.so. + """ libs = _pkg_dir / "nccl_ep" / "_libs" - return (libs / "libnccl_ep.so").exists() and (libs / "libnccl.so.2").exists() + return (libs / "libnccl_ep.so").exists() def _probe_nixl_ep() -> bool: + """True if the NIXL-EP plugin .so was staged by the build. + + The base libnixl.so + plugins are NOT staged into this package — they + come from the pip-installed nixl-cu13 wheel. The runtime loader in + flashinfer.moe_ep.nixl_ep loads them explicitly before opening + nixl_ep_cpp.so. + """ libs = _pkg_dir / "nixl_ep" / "_libs" if not libs.is_dir(): return False diff --git a/flashinfer/moe_ep/nccl_ep/__init__.py b/flashinfer/moe_ep/nccl_ep/__init__.py index 3bd53091d17..f216aa6dcbf 100644 --- a/flashinfer/moe_ep/nccl_ep/__init__.py +++ b/flashinfer/moe_ep/nccl_ep/__init__.py @@ -1,8 +1,101 @@ -"""NCCL-EP backend stub. Real implementation lands in Part B (B3-B4). +"""NCCL-EP backend. -Importing this module succeeds even when the native libs are absent; calling -into the (yet-to-be-implemented) Fleet/Handle factory functions will raise -:class:`flashinfer.moe_ep.MoEEpNotBuiltError`. +Two pieces matter for import-time success: + +1. The base NCCL runtime library, ``libnccl.so.2`` — *not* shipped inside this + package. It's expected to come from the ``nvidia-nccl-cu13`` pip wheel, + installed automatically when the user runs ``BUILD_NVEP=1 pip install ...`` + (see ``build_backend._install_nvep_runtime_wheels``). + +2. The EP plugin, ``libnccl_ep.so`` — built in-tree from + ``3rdparty/nccl/contrib/nccl_ep`` and staged into ``_libs/`` here. + +The two ``_preload_*`` helpers below are intentionally module-level but +**not invoked at import time** — calling them eagerly would force libnccl +to load whenever Python touches this package (e.g. on ``from flashinfer +import *``), which we don't want. The Fleet/Handle wrapper code (Part B, +not yet landed) will call ``_load_libnccl_ep()`` lazily on first use. + +Until that lands, importing this module always succeeds. Attempting to +actually exercise the backend will raise ``MoEEpNotBuiltError`` (see +``flashinfer.moe_ep``) if the lib staging is incomplete. """ from __future__ import annotations + +import ctypes +import os +from pathlib import Path + +from .. import MoEEpNotBuiltError + +_pkg_dir = Path(__file__).resolve().parent +_libs_dir = _pkg_dir / "_libs" + + +def _find_libnccl() -> Path | None: + """Locate libnccl.so.2 via the pip-installed nvidia-nccl-cu13 wheel. + + Returns the resolved path or None if it can't be found via the wheel. + Falls back to letting the dynamic linker's default search find it via + LD_LIBRARY_PATH, ldconfig, etc. + """ + # The wheel installs the lib at /nvidia/nccl/lib/libnccl.so.2. + try: + import nvidia.nccl # type: ignore[import-not-found] + except ImportError: + return None + try: + # Newer wheels expose .lib as a subpackage; older ones place files + # directly under nvidia/nccl/. Probe both. + candidates = [ + Path(nvidia.nccl.__path__[0]) / "lib" / "libnccl.so.2", + Path(nvidia.nccl.__path__[0]) / "libnccl.so.2", + ] + except Exception: + return None + for c in candidates: + if c.exists(): + return c + return None + + +def _preload_libnccl() -> None: + """ctypes-load libnccl.so.2 with RTLD_GLOBAL before opening libnccl_ep.so. + + libnccl_ep.so links against ``libnccl.so.2`` by SONAME but doesn't have + the wheel's site-packages location in its RPATH. Loading it explicitly + here exports the symbols globally so the subsequent dlopen of + libnccl_ep.so resolves them. + """ + nccl_so = _find_libnccl() + if nccl_so is not None: + ctypes.CDLL(str(nccl_so), mode=ctypes.RTLD_GLOBAL) + return + # No wheel found; try the dynamic linker's default search. + try: + ctypes.CDLL("libnccl.so.2", mode=ctypes.RTLD_GLOBAL) + except OSError as e: + raise MoEEpNotBuiltError( + "Could not locate libnccl.so.2. Install it with one of:\n" + " uv pip install --no-deps 'nvidia-nccl-cu13>=2.30.4'\n" + " pip install --no-deps 'nvidia-nccl-cu13>=2.30.4'\n" + "or set LD_LIBRARY_PATH to a directory containing libnccl.so.2." + ) from e + + +def _load_libnccl_ep() -> ctypes.CDLL: + """Load the EP plugin .so, preloading its libnccl.so.2 dep first. + + Returns the opened CDLL handle. Caller is responsible for keeping a + reference (the dynamic linker won't unload while the handle is alive). + """ + so = _libs_dir / "libnccl_ep.so" + if not so.exists(): + raise MoEEpNotBuiltError( + f"libnccl_ep.so is not staged at {so}. Rebuild with:\n" + ' BUILD_NVEP=1 pip install -e ".[nvep]"\n' + "or BUILD_NCCL_EP=1 for an NCCL-EP-only build." + ) + _preload_libnccl() + return ctypes.CDLL(str(so), mode=ctypes.RTLD_GLOBAL) diff --git a/flashinfer/moe_ep/nixl_ep/__init__.py b/flashinfer/moe_ep/nixl_ep/__init__.py index 44e365040ad..6fc6b2f549d 100644 --- a/flashinfer/moe_ep/nixl_ep/__init__.py +++ b/flashinfer/moe_ep/nixl_ep/__init__.py @@ -1,8 +1,111 @@ -"""NIXL-EP backend stub. Real implementation lands in Part B (B5). +"""NIXL-EP backend. -Importing this module succeeds even when the native libs are absent; calling -into the (yet-to-be-implemented) Fleet/Handle factory functions will raise -:class:`flashinfer.moe_ep.MoEEpNotBuiltError`. +Two pieces matter for import-time success: + +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``). + +2. The EP torch extension, ``nixl_ep_cpp*.so`` — built in-tree from + ``3rdparty/nixl/examples/device/ep`` and staged into ``_libs/`` here. + +The two ``_preload_*`` helpers below are intentionally module-level but +**not invoked at import time** — calling them eagerly would force libnixl +to load whenever Python touches this package. The Fleet/Handle wrapper +code (Part B, not yet landed) will call ``_load_nixl_ep_cpp()`` lazily on +first use. + +Until that lands, importing this module always succeeds. Attempting to +actually exercise the backend will raise ``MoEEpNotBuiltError`` (see +``flashinfer.moe_ep``) if the lib staging is incomplete. """ from __future__ import annotations + +import ctypes +import os +from pathlib import Path + +from .. import MoEEpNotBuiltError + +_pkg_dir = Path(__file__).resolve().parent +_libs_dir = _pkg_dir / "_libs" + + +# Order matters: libnixl_common must be available to libnixl, libnixl_capi +# depends on libnixl, etc. We load with RTLD_GLOBAL so each preceding lib +# exports symbols visible to subsequent dlopens. +_NIXL_BASE_LIBS = ( + "libnixl_common.so", + "libserdes.so", + "libnixl_build.so", + "libnixl.so", + "libnixl_capi.so", +) + + +def _find_nixl_lib_dir() -> Path | None: + """Locate the NIXL base-lib directory via the pip-installed nixl-cu13 wheel. + + Returns the resolved path or None if it can't be found. + """ + # The wheel typically installs libs at /nixl/lib/x86_64-linux-gnu/. + try: + import nixl # type: ignore[import-not-found] + except ImportError: + return None + try: + nixl_root = Path(nixl.__path__[0]) + except Exception: + return None + candidates = [ + nixl_root / "lib" / "x86_64-linux-gnu", + nixl_root / "lib", + nixl_root, # last-resort: libs directly under nixl/ + ] + for c in candidates: + if c.is_dir() and (c / "libnixl.so").exists(): + return c + return None + + +def _preload_libnixl() -> None: + """ctypes-load the NIXL base libs with RTLD_GLOBAL before opening nixl_ep_cpp.so.""" + nixl_lib_dir = _find_nixl_lib_dir() + if nixl_lib_dir is None: + # Try the dynamic linker's default search for the minimum lib. + try: + ctypes.CDLL("libnixl.so", mode=ctypes.RTLD_GLOBAL) + return + except OSError as e: + raise MoEEpNotBuiltError( + "Could not locate the NIXL runtime libraries. Install with " + "one of:\n" + " uv pip install --no-deps 'nixl-cu13>=1.0.1'\n" + " pip install --no-deps 'nixl-cu13>=1.0.1'\n" + "or set LD_LIBRARY_PATH to a directory containing libnixl.so." + ) from e + + for libname in _NIXL_BASE_LIBS: + libpath = nixl_lib_dir / libname + if libpath.exists(): + ctypes.CDLL(str(libpath), mode=ctypes.RTLD_GLOBAL) + + +def _load_nixl_ep_cpp() -> ctypes.CDLL: + """Load the EP torch extension, preloading its NIXL base-lib deps first. + + Returns the opened CDLL handle. Caller is responsible for keeping a + reference. + """ + 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." + ) + _preload_libnixl() + return ctypes.CDLL(str(so_files[0]), mode=ctypes.RTLD_GLOBAL) From bab61e3153ddf2067182d4edcd10bdf539d0c999 Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Fri, 15 May 2026 01:48:16 -0700 Subject: [PATCH 07/10] moe_ep: NIXL-EP wheel-driven build; skip parent libnixl compile MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Mirror the Section 9 NCCL-EP change for NIXL: build only the nixl_ep_cpp.so torch extension and link it against the libnixl.so shipped in the nixl-cu13 pip wheel, instead of compiling the full parent NIXL meson tree (src/ + plugins) on every install. A new patch overlay 3rdparty_patches/nixl/0002-ep-only-build.patch adds two meson options to the v1.1.0 submodule pin: -Dnixl_ep_only=true skips subdir('src'), subdir('test'), install_headers; routes directly to examples/device/ep. -Dnixl_wheel_lib_dir=PATH feeds cc.find_library('nixl', dirs:[PATH]) so the EP example links against an external libnixl.so. build_backend._build_nixl_ep now probes for nixl-cu13 via the meson-python sidecar `/.nixl_cu13.mesonpy.libs/` and passes those options; the fallback flag BUILD_NIXL_EP_HERMETIC=1 restores the original full-tree build for hosts without the wheel. The runtime loader `_find_nixl_lib_dir` is fixed to look at the same `.nixl_cu13.mesonpy.libs/` path — the prior `import nixl` probe was silently broken against the meson-python-packaged wheel (which installs as `nixl_cu13`, not `nixl`) and only worked because the preload fell back to ldconfig's default search. A `_time_phase` context manager prints per-phase wall times so future build logs distinguish NIXL vs NCCL compile costs. uv currently swallows the build hook's stdout under `pip install -e`, so the prints don't appear in docker layer logs today — but they do appear under `python -m pip install -v` and inside the BUILD_NIXL_EP_HERMETIC path on the host. Per-backend timings for this commit (docker, derived from artifact mtimes in flashinfer-nvep:dev image sha 77b25256): NIXL-EP compile (nixl_ep_only) 102 s ~1m 42s NCCL-EP compile (wheel-linked) 1599 s ~26m 39s step #17 total (incl. uv I/O) 2305 s ~38m 25s full docker build ~2302 s ~38m 22s End-to-end smoke probe inside the image: available_backends() -> ['nccl_ep', 'nixl_ep'] _load_nixl_ep_cpp() -> loads OK _load_libnccl_ep() -> loads OK nccl_ep/_libs/ -> libnccl_ep.so (8.6 MB) only nixl_ep/_libs/ -> nixl_ep_cpp.cpython-*.so (10.4 MB) only Co-Authored-By: Claude Opus 4.7 (1M context) --- .../nixl/0002-ep-only-build.patch | 104 ++++++++++++ build_backend.py | 155 +++++++++++++++--- docker/Dockerfile.flashinfer-nvep | 10 ++ flashinfer/moe_ep/nixl_ep/__init__.py | 55 ++++--- 4 files changed, 286 insertions(+), 38 deletions(-) create mode 100644 3rdparty_patches/nixl/0002-ep-only-build.patch diff --git a/3rdparty_patches/nixl/0002-ep-only-build.patch b/3rdparty_patches/nixl/0002-ep-only-build.patch new file mode 100644 index 00000000000..c671c550b43 --- /dev/null +++ b/3rdparty_patches/nixl/0002-ep-only-build.patch @@ -0,0 +1,104 @@ +From: FlashInfer build infra +Subject: [PATCH] meson: add nixl_ep_only option to skip parent libnixl build + +When `nixl_ep_only=true`, the parent NIXL library (and tests, headers, +non-EP examples) is not built. The EP example links against an +externally-provided libnixl.so (typically the nixl-cu13 pip wheel), +mirroring the wheel-driven build path used by FlashInfer's NCCL-EP. + +Two new meson options: + - nixl_ep_only (bool, default false): skip subdir('src')/subdir('test') + and install_headers; build only examples/device/ep. + - nixl_wheel_lib_dir (string, default ''): when nixl_ep_only=true, the + directory containing libnixl.so (used via cc.find_library()). + +Targets the NIXL pin 05e4243f (tag v1.1.0). + +diff --git a/examples/device/ep/meson.build b/examples/device/ep/meson.build +index a9ba19f..a02d583 100644 +--- a/examples/device/ep/meson.build ++++ b/examples/device/ep/meson.build +@@ -52,8 +52,21 @@ if not pybind_dep.found() + subdir_done() + endif + +-nixl_dep = declare_dependency(link_with: nixl_lib, include_directories: nixl_inc_dirs) +-nixl_lib_dir = join_paths(meson.project_build_root(), 'src', 'core') ++if get_option('nixl_ep_only') ++ cc_ep = meson.get_compiler('cpp') ++ wheel_lib_dir = get_option('nixl_wheel_lib_dir') ++ if wheel_lib_dir == '' ++ error('nixl_ep_only=true requires -Dnixl_wheel_lib_dir=') ++ endif ++ libnixl_external = cc_ep.find_library('nixl', dirs: [wheel_lib_dir], required: true) ++ nixl_dep = declare_dependency(dependencies: libnixl_external, include_directories: nixl_inc_dirs) ++ nixl_lib_dir = wheel_lib_dir ++ nixl_ep_link_with = [] ++else ++ nixl_dep = declare_dependency(link_with: nixl_lib, include_directories: nixl_inc_dirs) ++ nixl_lib_dir = join_paths(meson.project_build_root(), 'src', 'core') ++ nixl_ep_link_with = [nixl_lib] ++endif + + ucx_build_deps = [] + if ucx_dep.found() +@@ -147,7 +160,7 @@ nixl_ep_ext = py.extension_module('nixl_ep_cpp', + include_directories: nixl_ep_inc_dirs, + cpp_args: nixl_ep_cpp_args, + cuda_args: nixl_ep_cuda_args, +- link_with: [nixl_lib], ++ link_with: nixl_ep_link_with, + build_rpath: nixl_ep_rpath, + install_rpath: nixl_ep_install_rpath, + override_options: nixl_ep_override_options, +diff --git a/meson.build b/meson.build +index 2812ce8..46519e2 100644 +--- a/meson.build ++++ b/meson.build +@@ -366,18 +366,27 @@ nixl_gpu_inc_dirs = include_directories('src/api/gpu/ucx') + plugins_inc_dirs = include_directories('src/plugins') + utils_inc_dirs = include_directories('src/utils') + +-subdir('src') +-if get_option('build_tests') and get_option('buildtype') != 'release' +- subdir('test') ++if not get_option('nixl_ep_only') ++ subdir('src') ++ if get_option('build_tests') and get_option('buildtype') != 'release' ++ subdir('test') ++ endif + endif + + # nixl_ep currently lives under examples/device/ep. Build that subtree when + # either full examples are requested or nixl_ep is explicitly requested. +-if get_option('build_examples') or get_option('build_nixl_ep') ++# In nixl_ep_only mode, skip the rest of examples/ (cpp/ etc. depend on ++# nixl_lib from subdir('src')) and route straight to examples/device/ep. ++if get_option('nixl_ep_only') ++ if not get_option('build_nixl_ep') ++ error('nixl_ep_only=true requires -Dbuild_nixl_ep=true') ++ endif ++ subdir('examples/device/ep') ++elif get_option('build_examples') or get_option('build_nixl_ep') + subdir('examples') + endif + +-if get_option('install_headers') ++if get_option('install_headers') and not get_option('nixl_ep_only') + install_headers('src/api/cpp/nixl.h', install_dir: prefix_inc) + install_headers('src/api/cpp/nixl_types.h', install_dir: prefix_inc) + install_headers('src/api/cpp/nixl_params.h', install_dir: prefix_inc) +diff --git a/meson_options.txt b/meson_options.txt +index 4db4845..5b3750a 100644 +--- a/meson_options.txt ++++ b/meson_options.txt +@@ -37,3 +37,9 @@ option('build_tests', type: 'boolean', value: true, description: 'Build all test + option('build_examples', type: 'boolean', value: true, description: 'Build all examples') + option('build_nixl_ep', type: 'boolean', value: false, description: 'Build nixl_ep example (requires sm_90)') + option('test_all_plugins', type: 'boolean', value: false, description: 'Testing all plugins in addition to the mocks..') ++ ++# FlashInfer wheel-driven build: skip the parent libnixl build entirely and ++# link the EP example against an externally-provided libnixl.so (typically ++# from the nixl-cu13 pip wheel). ++option('nixl_ep_only', type: 'boolean', value: false, description: 'Skip parent libnixl build; build only examples/device/ep linked against an external libnixl.so') ++option('nixl_wheel_lib_dir', type: 'string', value: '', description: 'Path to directory containing external libnixl.so (used when nixl_ep_only=true)') diff --git a/build_backend.py b/build_backend.py index 572cea61570..7ee4e063041 100644 --- a/build_backend.py +++ b/build_backend.py @@ -18,6 +18,8 @@ import shutil import subprocess import sys +import time +from contextlib import contextmanager from pathlib import Path from setuptools import build_meta as orig @@ -41,6 +43,23 @@ def _flag(name: str) -> bool: 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. + + Docker buffers the entire `RUN` layer's output, so without these markers + you can't tell from a finished build log how long each backend took. The + lines are flushed explicitly so they survive a SIGKILL on the parent. + """ + print(f"[BUILD_NVEP] {label}: start", flush=True) + t0 = time.monotonic() + try: + yield + finally: + dt = time.monotonic() - t0 + 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 @@ -105,27 +124,101 @@ def _apply_patches(submodule_dir: Path, patches_dir: Path) -> None: print(f"[BUILD_NVEP] applied patch: {patch.name}") +def _find_nixl_wheel_lib_dir() -> Path | None: + """Locate the nixl-cu* pip wheel's libnixl.so directory. + + Layout of the current nixl-cu13 wheel (meson-python packaging): + /nixl_cu13/ — importable python module + /.nixl_cu13.mesonpy.libs/ — libnixl.so + sibling libs + /nixl_cu13.libs/ — auditwheel deps + plugins + + We probe by importing `nixl_cu13` / `nixl_cu12` / `nixl` (legacy), then + look for `.{name}.mesonpy.libs/libnixl.so` next to it. + """ + for pkg_name in ("nixl_cu13", "nixl_cu12", "nixl"): + try: + mod = __import__(pkg_name) + except ImportError: + continue + try: + pkg_root = Path(mod.__path__[0]) + except Exception: + continue + site_packages = pkg_root.parent + # Prefer the meson-python sibling layout. + for candidate in ( + site_packages / f".{pkg_name}.mesonpy.libs", + pkg_root / "lib" / "x86_64-linux-gnu", + pkg_root / "lib", + pkg_root, + ): + if (candidate / "libnixl.so").exists(): + return candidate + # Last resort: a glob-walk of site-packages for any .nixl_*.mesonpy.libs/. + try: + sp = Path(__import__("site").getsitepackages()[0]) # type: ignore[no-untyped-call] + except Exception: + return None + for candidate in sp.glob(".nixl*.mesonpy.libs"): + if (candidate / "libnixl.so").exists(): + return candidate + return None + + def _build_nixl_ep() -> None: src = _root / "3rdparty" / "nixl" build = _nvep_build_root / "nixl" prefix = _nvep_build_root / "nixl_install" _apply_patches(src, _root / "3rdparty_patches" / "nixl") - if not build.exists(): - subprocess.run( - [ - "meson", - "setup", - str(build), - str(src), - "-Dbuild_nixl_ep=true", - "-Dbuild_examples=true", - f"-Dprefix={prefix}", - "--buildtype=release", - ], - check=True, + # Default path: skip the parent libnixl build and link the EP example + # against the libnixl.so shipped by the nixl-cu13 pip wheel — mirrors the + # contrib/nccl_ep wheel-driven path (Section 9 of the integration plan). + # The hermetic env var falls back to the full parent build for hosts + # without the wheel pre-installed. + hermetic = _flag("BUILD_NIXL_EP_HERMETIC") + setup_args = [ + "meson", + "setup", + str(build), + str(src), + "-Dbuild_nixl_ep=true", + f"-Dprefix={prefix}", + "--buildtype=release", + ] + if hermetic: + print("[BUILD_NVEP] BUILD_NIXL_EP_HERMETIC=1 — building full NIXL tree") + setup_args.append("-Dbuild_examples=true") + else: + 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" + "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 += [ + "-Dbuild_examples=false", + "-Dnixl_ep_only=true", + f"-Dnixl_wheel_lib_dir={wheel_lib_dir}", + ] + print( + f"[BUILD_NVEP] nixl_ep_only=true; linking against wheel libnixl.so " + f"at {wheel_lib_dir}. Set BUILD_NIXL_EP_HERMETIC=1 to opt out." ) - subprocess.run(["ninja", "-C", str(build), "install"], check=True) + + if not build.exists(): + subprocess.run(setup_args, check=True) + + # `install` only makes sense in hermetic mode (it populates `prefix/` with + # libnixl + headers + plugins). In ep_only mode there's nothing to install + # — we just compile and pluck nixl_ep_cpp.so out of the build tree. + ninja_cmd = ["ninja", "-C", str(build)] + if hermetic: + ninja_cmd.append("install") + subprocess.run(ninja_cmd, check=True) dst = _moe_ep_pkg / "nixl_ep" / "_libs" dst.mkdir(parents=True, exist_ok=True) @@ -357,7 +450,13 @@ def _fix_rpaths() -> None: def _nixl_buildable() -> tuple[bool, str]: - """Probe for hard NIXL-EP build-time deps. Returns (ok, reason_if_not).""" + """Probe for hard NIXL-EP build-time deps. Returns (ok, reason_if_not). + + In the default (wheel-driven) flow, the EP example links against the + nixl-cu13 pip wheel — so it must be importable. In hermetic mode + (BUILD_NIXL_EP_HERMETIC=1), we build libnixl from source and the wheel + isn't required. + """ if not shutil.which("meson"): return False, "meson not on PATH (apt install meson)" if not shutil.which("ninja"): @@ -375,6 +474,13 @@ def _nixl_buildable() -> tuple[bool, str]: r = subprocess.run([pkgconfig, "--exists", "libibverbs"], capture_output=True) if r.returncode: return False, "libibverbs not found via pkg-config (apt install libibverbs-dev)" + if not _flag("BUILD_NIXL_EP_HERMETIC"): + if _find_nixl_wheel_lib_dir() is None: + return False, ( + "nixl pip wheel not importable (or libnixl.so missing); install with " + "`uv pip install --no-deps 'nixl-cu13>=1.0.1'` " + "or set BUILD_NIXL_EP_HERMETIC=1 to build the full NIXL tree" + ) return True, "" @@ -533,10 +639,12 @@ def _build_nvep_if_enabled() -> None: # Actual builds. If best-effort and the build raises despite the probe # passing, swallow the error so the other backend still has a chance. + overall_t0 = time.monotonic() built_nixl = False if will_build_nixl: try: - _build_nixl_ep() + with _time_phase("_build_nixl_ep"): + _build_nixl_ep() built_nixl = True except Exception as e: if not _BUILD_NVEP_BEST_EFFORT: @@ -546,7 +654,8 @@ def _build_nvep_if_enabled() -> None: built_nccl = False if will_build_nccl: try: - _build_nccl_ep() + with _time_phase("_build_nccl_ep"): + _build_nccl_ep() built_nccl = True except Exception as e: if not _BUILD_NVEP_BEST_EFFORT: @@ -554,8 +663,16 @@ def _build_nvep_if_enabled() -> None: print(f"[BUILD_NVEP] NCCL-EP build failed in best-effort mode: {e}") if built_nixl or built_nccl: - _fix_rpaths() - _install_nvep_runtime_wheels(built_nixl=built_nixl, built_nccl=built_nccl) + with _time_phase("_fix_rpaths"): + _fix_rpaths() + with _time_phase("_install_nvep_runtime_wheels"): + _install_nvep_runtime_wheels(built_nixl=built_nixl, built_nccl=built_nccl) + + print( + f"[BUILD_NVEP] total build phase wall time: " + f"{time.monotonic() - overall_t0:.1f}s", + flush=True, + ) built = [b for b, on in (("NIXL-EP", built_nixl), ("NCCL-EP", built_nccl)) if on] print(f"[BUILD_NVEP] done — built: {', '.join(built) if built else 'nothing'}") diff --git a/docker/Dockerfile.flashinfer-nvep b/docker/Dockerfile.flashinfer-nvep index 9c48390dbeb..22d10568f70 100644 --- a/docker/Dockerfile.flashinfer-nvep +++ b/docker/Dockerfile.flashinfer-nvep @@ -156,6 +156,16 @@ RUN uv pip install --python ${VENV}/bin/python --no-deps \ # 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. +# +# 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). ARG BUILD_NVEP=1 ARG BUILD_NCCL_EP= ARG BUILD_NIXL_EP= diff --git a/flashinfer/moe_ep/nixl_ep/__init__.py b/flashinfer/moe_ep/nixl_ep/__init__.py index 6fc6b2f549d..fdc6844035b 100644 --- a/flashinfer/moe_ep/nixl_ep/__init__.py +++ b/flashinfer/moe_ep/nixl_ep/__init__.py @@ -47,27 +47,44 @@ def _find_nixl_lib_dir() -> Path | None: - """Locate the NIXL base-lib directory via the pip-installed nixl-cu13 wheel. + """Locate the NIXL base-lib directory via the pip-installed nixl-cu* wheel. - Returns the resolved path or None if it can't be found. + Layout of the meson-python-packaged wheel: + /nixl_cu13/ — importable python module + /.nixl_cu13.mesonpy.libs/ — libnixl.so + sibling libs + + We probe known package names (`nixl_cu13`, `nixl_cu12`, legacy `nixl`), + look for the meson-python `.{name}.mesonpy.libs/` sidecar, and fall back + to a glob over site-packages. """ - # The wheel typically installs libs at /nixl/lib/x86_64-linux-gnu/. - try: - import nixl # type: ignore[import-not-found] - except ImportError: - return None - try: - nixl_root = Path(nixl.__path__[0]) - except Exception: - return None - candidates = [ - nixl_root / "lib" / "x86_64-linux-gnu", - nixl_root / "lib", - nixl_root, # last-resort: libs directly under nixl/ - ] - for c in candidates: - if c.is_dir() and (c / "libnixl.so").exists(): - return c + for pkg_name in ("nixl_cu13", "nixl_cu12", "nixl"): + try: + mod = __import__(pkg_name) + except ImportError: + continue + try: + pkg_root = Path(mod.__path__[0]) + except Exception: + continue + site_packages = pkg_root.parent + for candidate in ( + site_packages / f".{pkg_name}.mesonpy.libs", + pkg_root / "lib" / "x86_64-linux-gnu", + pkg_root / "lib", + pkg_root, + ): + if candidate.is_dir() and (candidate / "libnixl.so").exists(): + return candidate + # Last resort: glob for any `.nixl_*.mesonpy.libs/` under site-packages. + import site as _site + + for sp_str in _site.getsitepackages() + [_site.getusersitepackages()]: + sp = Path(sp_str) + if not sp.is_dir(): + continue + for candidate in sp.glob(".nixl*.mesonpy.libs"): + if (candidate / "libnixl.so").exists(): + return candidate return None From 12f7a12c549f7b7c56f3c485b6f640e46dfd08fd Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Fri, 15 May 2026 15:34:24 -0700 Subject: [PATCH 08/10] moe_ep: nixl_ep: emit SASS for sm_100/sm_103 from device-link, not just sm_90 MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Patch 0001 only added the multi-arch -gencode flags to the EP module's nvcc cuda_args (compile-side). With NIXL's -rdc=true device-link configuration, the project-wide nvcc_flags_link in the parent meson.build still listed only -gencode=arch=compute_90,code=sm_90 -- so nvcc correctly produced .o files containing sm_90/sm_100/sm_103 cubins but nvlink then discarded the sm_100 and sm_103 sections during the final device link, leaving an sm_90-only nixl_ep_cpp.so. Verified via cuobjdump --list-elf nixl_ep_cpp.cpython-*.so on the prior commit: only sm_90.cubin was present. Extend 0001 to also patch the parent meson.build so the project-wide nvcc_flags and nvcc_flags_link include -gencode entries for sm_100 and sm_103 when build_nixl_ep=true. Gated behind the EP option so non-EP NIXL builds (and other consumers of the submodule) are unaffected. Post-fix verification (flashinfer-nvep:dev image after this commit): cuobjdump --list-elf nixl_ep_cpp.cpython-*.so -> sm_90, sm_100, sm_103 available_backends() -> ['nccl_ep', 'nixl_ep'] Wall time impact (docker step #17): 2304.8s -> 2319.7s (~+15s). The compile is fast because nixl_ep only has 4 .cu + 2 .cpp files, all small kernels; nvcc reuses preprocessing across gencode targets. Apprx build time ┌──────────────────────────────────────────────────────────────────┬───────────────────┐ │ phase │ wall time │ ├──────────────────────────────────────────────────────────────────┼───────────────────┤ │ NIXL-EP compile (nixl_ep_only) │ 102 s — 1m 42s │ ├──────────────────────────────────────────────────────────────────┼───────────────────┤ │ NCCL-EP compile (wheel-linked) │ 1599 s — 26m 39s │ ├──────────────────────────────────────────────────────────────────┼───────────────────┤ │ step #17 total (incl. uv resolve + downloads + editable install) │ 2305 s — 38m 25s │ ├──────────────────────────────────────────────────────────────────┼───────────────────┤ │ full docker build │ ~2302 s — 38m 22s │ └──────────────────────────────────────────────────────────────────┴───────────────────┘ Co-Authored-By: Claude Opus 4.7 (1M context) --- .../0001-meson-add-blackwell-arches.patch | 45 +++++++++++++++++-- 1 file changed, 42 insertions(+), 3 deletions(-) diff --git a/3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch b/3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch index f6e7ef57cca..c1346f930d6 100644 --- a/3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch +++ b/3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch @@ -1,13 +1,21 @@ From: FlashInfer build infra -Subject: [PATCH] examples/device/ep: emit SASS for Hopper + Blackwell +Subject: [PATCH] meson: emit nixl_ep SASS for Hopper + Blackwell The upstream meson rule pins `-arch=sm_90`, which overrides any global -gencode flags and ships an sm_90-only `nixl_ep_cpp.so`. FlashInfer needs a single .so that runs natively on H100 (sm_90), B200 (sm_100), and B300 (sm_103), plus sm_90 PTX for forward-compat onto future arches. -This patch replaces the single `-arch=sm_90` flag with an explicit -multi-gencode list. Targets the NIXL pin 05e4243f (tag v1.1.0). +This patch: + 1. Replaces the single `-arch=sm_90` flag in examples/device/ep with + an explicit multi-gencode list (compile-side). + 2. Extends the project-wide nvcc_flags / nvcc_flags_link in the + top-level meson.build so the `-rdc=true` device link step also + covers sm_100 / sm_103. Without (2), nvlink discards the sm_100 / + sm_103 cubins generated by (1) during the final device link and + the resulting nixl_ep_cpp.so ends up sm_90-only. + +Targets the NIXL pin 05e4243f (tag v1.1.0). diff --git a/examples/device/ep/meson.build b/examples/device/ep/meson.build index a9ba19f..065dd13 100644 @@ -24,3 +32,34 @@ index a9ba19f..065dd13 100644 '--ptxas-options=--register-usage-level=10', # Allow more register usage (matches setup.py) '-Xcompiler', '-Wno-deprecated-declarations', '-Xcompiler', '-Wno-unused-variable', +diff --git a/meson.build b/meson.build +index 2812ce8..b96da85 100644 +--- a/meson.build ++++ b/meson.build +@@ -202,6 +202,15 @@ if cuda_dep.found() + nvcc_flags += ['-gencode', 'arch=compute_80,code=sm_80'] + endif + nvcc_flags += ['-gencode', 'arch=compute_90,code=sm_90'] ++ if get_option('build_nixl_ep') ++ # FlashInfer: nixl_ep targets H100 + B200 + B300. The compile-side ++ # gencode for those arches is set in examples/device/ep/meson.build ++ # via 0001-meson-add-blackwell-arches.patch; we also need them on ++ # the project-wide device-link step so nvlink doesn't drop sm_100 / ++ # sm_103 SASS during the final -rdc=true device link. ++ nvcc_flags += ['-gencode', 'arch=compute_100,code=sm_100'] ++ nvcc_flags += ['-gencode', 'arch=compute_103,code=sm_103'] ++ endif + add_project_arguments(nvcc_flags, language: 'cuda') + + # Refer to https://mesonbuild.com/Cuda-module.html +@@ -214,6 +223,10 @@ if cuda_dep.found() + nvcc_flags_link += ['-gencode=arch=compute_80,code=sm_80'] + endif + nvcc_flags_link += ['-gencode=arch=compute_90,code=sm_90'] ++ if get_option('build_nixl_ep') ++ nvcc_flags_link += ['-gencode=arch=compute_100,code=sm_100'] ++ nvcc_flags_link += ['-gencode=arch=compute_103,code=sm_103'] ++ endif + add_project_link_arguments(nvcc_flags_link, language: 'cuda') + message('nvcc version: ' + nvcc.version()) + if nvcc.version().version_compare('>=12.8') From 7a03b504812fb64378aeb3a98605f0a937c86aeb Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sun, 17 May 2026 10:29:35 -0700 Subject: [PATCH 09/10] moe_ep: address PR #3315 bot review comments MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Real bugs flagged by coderabbitai + gemini-code-assist on PR #3315: * _nixl_buildable / _nccl_buildable now probe for nvcc and git up front, matching what the build actually runs (nvcc for CUDA kernels, git for `git apply` of the patch overlays). Previously a missing nvcc on PATH led to "best-effort" silent skips that produced an EP-less wheel. * `git submodule update --init` is now guarded on `.git` existing in the superproject. Without the guard, sdist installs (which travel with submodule trees expanded inline and no `.git` metadata) crash with "not a git repository" instead of either succeeding or surfacing a clear error. * `_build_nixl_ep` re-runs `meson setup --reconfigure` when the build dir already exists. Previously a stale config persisted across patch / option changes (e.g. flipping BUILD_NIXL_EP_HERMETIC, bumping the wheel path) and required a manual rm -rf to take effect. * `_fix_rpaths` no longer swallows patchelf failures silently. It still uses check=False (patchelf legitimately exits nonzero on files that already have the desired RPATH), but now prints the stderr/stdout on nonzero rc so real failures (e.g. missing .dynamic) surface in the build log. * Import-time warning in flashinfer.moe_ep now fires for any of BUILD_NVEP / BUILD_NCCL_EP / BUILD_NIXL_EP being set rather than only BUILD_NVEP. The previous heuristic missed the per-backend flags entirely. * Narrowed `except Exception` in _find_nixl_lib_dir to the actual expected error set (AttributeError / IndexError / TypeError) so an unrelated failure (e.g. permission error reading site-packages) doesn't get masked. * ctypes.CDLL failure paths in _preload_libnccl / _load_libnccl_ep / _preload_libnixl / _load_nixl_ep_cpp now translate raw OSError into MoEEpNotBuiltError with an actionable rebuild hint (BUILD_*_EP_HERMETIC=1 or --force-reinstall the matching wheel), so callers don't have to interpret cryptic dlopen messages. * Dockerfile: removed the stale "DOCA gpunetio — NOT installed here" comment block that contradicted the apt install step right above it. * Dockerfile smoke probe now asserts that both 'nccl_ep' and 'nixl_ep' appear in available_backends(); previously it just printed the list and a silent best-effort skip in the build hook produced an empty list without failing the image build. Pushed back on (not addressed in this commit): * DOCA URL sha256 checksum — Mellanox URL is version-pinned; risk is bounded by version, adding/maintaining a checksum is more friction than it's worth for an internal dev image. * UCX `v1.21.x` branch -> commit SHA pin — matches upstream NIXL's own contrib/Dockerfile pin; deviating would create drift. * `curl | sh` for the uv installer — standard install pattern per astral.sh/uv docs. * Run as root in container — NVIDIA's official CUDA base images do the same; switching to a non-root user breaks the GPU device permissions on most clusters that mount /dev/nvidia* with root ownership. * x86_64-linux-gnu hardcode in _find_nixl_wheel_lib_dir — the probe has a fallback chain that covers the alternate aarch64 layout via the bare `lib/` and meson-python `.{pkg}.mesonpy.libs/` candidates; no functional gap on aarch64. Co-Authored-By: Claude Opus 4.7 (1M context) --- build_backend.py | 48 ++++++++++++++++++++++++--- docker/Dockerfile.flashinfer-nvep | 20 ++++++----- flashinfer/moe_ep/__init__.py | 21 +++++++++--- flashinfer/moe_ep/nccl_ep/__init__.py | 21 ++++++++++-- flashinfer/moe_ep/nixl_ep/__init__.py | 39 ++++++++++++++++++++-- 5 files changed, 126 insertions(+), 23 deletions(-) diff --git a/build_backend.py b/build_backend.py index 7ee4e063041..6701b1e011c 100644 --- a/build_backend.py +++ b/build_backend.py @@ -209,8 +209,13 @@ def _build_nixl_ep() -> None: f"at {wheel_lib_dir}. Set BUILD_NIXL_EP_HERMETIC=1 to opt out." ) - if not build.exists(): - subprocess.run(setup_args, check=True) + # Re-run meson setup with --reconfigure when the build dir already + # exists, so patch / option changes (e.g. a different wheel path, + # flipping HERMETIC mode) take effect on subsequent installs without + # requiring the user to `rm -rf build_nvep/nixl` manually. + if build.exists(): + setup_args.append("--reconfigure") + subprocess.run(setup_args, check=True) # `install` only makes sense in hermetic mode (it populates `prefix/` with # libnixl + headers + plugins). In ep_only mode there's nothing to install @@ -443,10 +448,18 @@ def _fix_rpaths() -> None: # Skip symlinks if so.is_symlink(): continue - subprocess.run( + # Use check=False because patchelf legitimately exits nonzero on + # files that already have the desired RPATH or that aren't ELFs + # we care about. Surface anything else as a warning so a real + # failure (e.g. binary lacks .dynamic section) isn't silently lost. + r = subprocess.run( ["patchelf", "--set-rpath", rpath, str(so)], - check=False, + capture_output=True, + text=True, ) + if r.returncode != 0: + err = (r.stderr or r.stdout or "").strip() + print(f"[BUILD_NVEP] WARNING: patchelf failed on {so.name}: {err}") def _nixl_buildable() -> tuple[bool, str]: @@ -461,6 +474,13 @@ def _nixl_buildable() -> tuple[bool, str]: return False, "meson not on PATH (apt install meson)" if not shutil.which("ninja"): return False, "ninja not on PATH (apt install ninja-build)" + if not shutil.which("nvcc"): + return False, ( + "nvcc not on PATH (install CUDA toolkit and put " + "/usr/local/cuda/bin on $PATH); needed for nixl_ep CUDA kernels" + ) + if not shutil.which("git"): + return False, "git not on PATH; needed for `git apply` of patch overlays" pkgconfig = shutil.which("pkg-config") if not pkgconfig: return False, "pkg-config not on PATH (apt install pkg-config)" @@ -499,6 +519,8 @@ def _nccl_buildable() -> tuple[bool, str]: "nvcc not on PATH (install CUDA toolkit and put " "/usr/local/cuda/bin on $PATH)" ) + if not shutil.which("git"): + return False, "git not on PATH; needed for `git apply` of patch overlays" if not _flag("BUILD_NCCL_EP_HERMETIC"): if _find_nccl_wheel_root() is None: return False, ( @@ -624,13 +646,31 @@ def _build_nvep_if_enabled() -> None: # Make sure each backend's submodule is initialized. Only fetch what we # actually need — saves ~300MB and a network round-trip per backend. + # Guarded on `.git` because an sdist install has no git metadata: the + # submodule trees travel inside the sdist as plain directories and + # `git submodule update` would fail with "not a git repository". + in_git_repo = (_root / ".git").exists() if will_build_nixl and not (_root / "3rdparty/nixl/meson.build").exists(): + if not in_git_repo: + raise RuntimeError( + "3rdparty/nixl/meson.build is missing and this is not a git " + "checkout (likely an sdist install where the submodule wasn't " + "packaged). Either install from a git clone or fetch the " + "submodule tree manually into 3rdparty/nixl." + ) subprocess.run( ["git", "submodule", "update", "--init", "--recursive", "3rdparty/nixl"], cwd=_root, check=True, ) if will_build_nccl and not (_root / "3rdparty/nccl/Makefile").exists(): + if not in_git_repo: + raise RuntimeError( + "3rdparty/nccl/Makefile is missing and this is not a git " + "checkout (likely an sdist install where the submodule wasn't " + "packaged). Either install from a git clone or fetch the " + "submodule tree manually into 3rdparty/nccl." + ) subprocess.run( ["git", "submodule", "update", "--init", "--recursive", "3rdparty/nccl"], cwd=_root, diff --git a/docker/Dockerfile.flashinfer-nvep b/docker/Dockerfile.flashinfer-nvep index 22d10568f70..3ed49b2619b 100644 --- a/docker/Dockerfile.flashinfer-nvep +++ b/docker/Dockerfile.flashinfer-nvep @@ -98,13 +98,6 @@ RUN git clone --depth=1 --branch ${GDRCOPY_VERSION} https://github.com/NVIDIA/gd && make -j"$(nproc)" lib lib_install \ && rm -rf /tmp/gdrcopy -# DOCA gpunetio — NOT installed here: it requires a separate Mellanox repo -# setup and a license-accepted MOFED package. On NVIDIA-managed hosts the -# /opt/mellanox/doca tree is provided by the base image. If your base image -# lacks DOCA, add the appropriate `apt-get install doca-sdk-gpunetio -# libdoca-sdk-gpunetio-dev` step here, gated on having the Mellanox apt repo -# configured. - # Pre-stage uv for fast pip installs. RUN curl -LsSf https://astral.sh/uv/install.sh | sh ENV PATH=/root/.local/bin:${PATH} @@ -190,8 +183,17 @@ RUN uv pip install --python ${VENV}/bin/python -e 3rdparty/nccl/contrib/nccl_ep/ RUN CUDA_HOME=/usr/local/cuda uv pip install --python ${VENV}/bin/python \ -e '3rdparty/nccl/bindings/nccl4py[cu13]' -# Smoke probe. -RUN python -c "from flashinfer.moe_ep import available_backends; print('moe_ep backends:', available_backends())" +# Smoke probe. Assertions fail the build if the EP backends weren't +# actually produced — without these, a silent skip in BUILD_NVEP=1's +# 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. +RUN python -c "\ +from flashinfer.moe_ep import available_backends; \ +b = available_backends(); \ +print('moe_ep backends:', b); \ +assert 'nccl_ep' in b, 'nccl_ep backend missing'; \ +assert 'nixl_ep' in b, 'nixl_ep backend missing'" RUN python -c "import nccl_ep; from nccl.core.communicator import Communicator; print('nccl_ep + nccl4py OK')" CMD ["bash"] diff --git a/flashinfer/moe_ep/__init__.py b/flashinfer/moe_ep/__init__.py index fcc13632b9d..a42be420d5e 100644 --- a/flashinfer/moe_ep/__init__.py +++ b/flashinfer/moe_ep/__init__.py @@ -104,14 +104,25 @@ def _require_built(backend: str) -> None: ) -# Quiet diagnostic at import time when BUILD_NVEP was set but the libs are -# absent — most likely cause is a partial build. Helpful for first-time users. -if os.environ.get("BUILD_NVEP") == "1" and not available_backends(): +# 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. +_set_build_flags = [ + name + for name in ("BUILD_NVEP", "BUILD_NCCL_EP", "BUILD_NIXL_EP") + if os.environ.get(name, "").lower() in ("1", "true", "yes", "on") +] +if _set_build_flags and not available_backends(): import warnings warnings.warn( - "BUILD_NVEP=1 was set, but no moe_ep backend libraries were found " - f"under {_pkg_dir}. Check the build log for meson/make failures.", + f"{'/'.join(_set_build_flags)} was set, but no moe_ep backend " + f"libraries were found under {_pkg_dir}. Check the build log " + "for pre-flight probe misses (meson/make/nvcc/git on PATH, " + "ucx/libibverbs via pkg-config, nixl-cu13 / nvidia-nccl-cu13 " + "wheels importable) or meson/make compile failures.", RuntimeWarning, stacklevel=2, ) diff --git a/flashinfer/moe_ep/nccl_ep/__init__.py b/flashinfer/moe_ep/nccl_ep/__init__.py index f216aa6dcbf..d1fe1048225 100644 --- a/flashinfer/moe_ep/nccl_ep/__init__.py +++ b/flashinfer/moe_ep/nccl_ep/__init__.py @@ -70,7 +70,15 @@ def _preload_libnccl() -> None: """ nccl_so = _find_libnccl() if nccl_so is not None: - ctypes.CDLL(str(nccl_so), mode=ctypes.RTLD_GLOBAL) + try: + ctypes.CDLL(str(nccl_so), mode=ctypes.RTLD_GLOBAL) + except OSError as e: + raise MoEEpNotBuiltError( + f"dlopen({nccl_so}) failed: {e}. The nvidia-nccl-cu13 wheel " + "may be corrupted or built against an incompatible glibc/CUDA. " + "Reinstall with: " + "uv pip install --force-reinstall --no-deps 'nvidia-nccl-cu13>=2.30.4'" + ) from e return # No wheel found; try the dynamic linker's default search. try: @@ -98,4 +106,13 @@ def _load_libnccl_ep() -> ctypes.CDLL: "or BUILD_NCCL_EP=1 for an NCCL-EP-only build." ) _preload_libnccl() - return ctypes.CDLL(str(so), mode=ctypes.RTLD_GLOBAL) + try: + return ctypes.CDLL(str(so), mode=ctypes.RTLD_GLOBAL) + except OSError as e: + raise MoEEpNotBuiltError( + f"dlopen({so}) failed: {e}. Most likely the wheel's NCCL " + "version doesn't match the one FlashInfer was built against " + "(check NCCL_VERSION_CODE). Reinstall with " + "BUILD_NCCL_EP_HERMETIC=1 to build libnccl from the pinned " + "submodule instead." + ) from e diff --git a/flashinfer/moe_ep/nixl_ep/__init__.py b/flashinfer/moe_ep/nixl_ep/__init__.py index fdc6844035b..8ffc4e4d2b1 100644 --- a/flashinfer/moe_ep/nixl_ep/__init__.py +++ b/flashinfer/moe_ep/nixl_ep/__init__.py @@ -62,9 +62,14 @@ def _find_nixl_lib_dir() -> Path | None: mod = __import__(pkg_name) except ImportError: continue + # `mod.__path__` is a `_NamespacePath` for namespace packages or a + # plain list for regular ones; element access can raise IndexError + # if it's empty (unusual but possible for malformed installs), and + # `__path__` itself may be missing (AttributeError) on a module + # imported from a single .py file. try: pkg_root = Path(mod.__path__[0]) - except Exception: + except (AttributeError, IndexError, TypeError): continue site_packages = pkg_root.parent for candidate in ( @@ -105,10 +110,28 @@ def _preload_libnixl() -> None: "or set LD_LIBRARY_PATH to a directory containing libnixl.so." ) from e + # libnixl itself is required; the others are best-effort siblings that + # may or may not ship in the wheel depending on its version. + primary = nixl_lib_dir / "libnixl.so" + if not primary.exists(): + raise MoEEpNotBuiltError( + f"libnixl.so is missing from the NIXL wheel lib dir at " + f"{nixl_lib_dir}. Reinstall: " + "uv pip install --no-deps 'nixl-cu13>=1.0.1'" + ) for libname in _NIXL_BASE_LIBS: libpath = nixl_lib_dir / libname - if libpath.exists(): + if not libpath.exists(): + continue + try: ctypes.CDLL(str(libpath), mode=ctypes.RTLD_GLOBAL) + except OSError as e: + raise MoEEpNotBuiltError( + f"Failed to load NIXL base lib {libpath}: {e}. The wheel " + "may be corrupted or built against a different glibc/CUDA. " + "Reinstall with: " + "uv pip install --force-reinstall --no-deps 'nixl-cu13>=1.0.1'" + ) from e def _load_nixl_ep_cpp() -> ctypes.CDLL: @@ -125,4 +148,14 @@ def _load_nixl_ep_cpp() -> ctypes.CDLL: "or BUILD_NIXL_EP=1 for a NIXL-EP-only build." ) _preload_libnixl() - return ctypes.CDLL(str(so_files[0]), mode=ctypes.RTLD_GLOBAL) + try: + return ctypes.CDLL(str(so_files[0]), mode=ctypes.RTLD_GLOBAL) + except OSError as e: + raise MoEEpNotBuiltError( + f"dlopen({so_files[0]}) failed: {e}. Most likely the NIXL " + "base libs preloaded above don't export every symbol nixl_ep " + "needs — check that the nixl-cu13 wheel and the FlashInfer " + "build were built against compatible NIXL revisions. Rebuild " + "with BUILD_NIXL_EP_HERMETIC=1 to pin against the submodule's " + "headers + libnixl instead of the wheel." + ) from e From ae542608c4c65437e12d1276847da26458d729fe Mon Sep 17 00:00:00 2001 From: Anerudhan Gopal Date: Sun, 17 May 2026 10:56:42 -0700 Subject: [PATCH 10/10] moe_ep: detect Debian multiarch dir via platform.machine() Replace the hardcoded `x86_64-linux-gnu` fallback in both _find_nixl_wheel_lib_dir (build-time probe) and _find_nixl_lib_dir (runtime loader) with `-linux-gnu`, so the probe resolves to `aarch64-linux-gnu` on NVIDIA Grace / AWS Graviton hosts without falling through to the less-specific bare `lib/` candidate. No functional change on x86_64 (the probe already worked via the meson-python `.nixl_*.mesonpy.libs/` sidecar, with the multiarch candidate as a backup). Addresses one of the bot review comments on PR #3315 that flagged the x86_64 hardcode as Ubuntu-specific; this makes the fallback chain explicitly arch-aware while keeping the same behavior on x86_64 hosts. Co-Authored-By: Claude Opus 4.7 (1M context) --- build_backend.py | 9 ++++++++- flashinfer/moe_ep/nixl_ep/__init__.py | 9 ++++++++- 2 files changed, 16 insertions(+), 2 deletions(-) diff --git a/build_backend.py b/build_backend.py index 6701b1e011c..e9bbdc39124 100644 --- a/build_backend.py +++ b/build_backend.py @@ -134,7 +134,14 @@ def _find_nixl_wheel_lib_dir() -> Path | None: We probe by importing `nixl_cu13` / `nixl_cu12` / `nixl` (legacy), then look for `.{name}.mesonpy.libs/libnixl.so` next to it. + + `-linux-gnu` follows the Debian multiarch convention used by + auditwheel-tagged wheels — `x86_64-linux-gnu` on x86_64, and + `aarch64-linux-gnu` on ARM64 (e.g. NVIDIA Grace / AWS Graviton hosts). """ + import platform + + multiarch = f"{platform.machine()}-linux-gnu" for pkg_name in ("nixl_cu13", "nixl_cu12", "nixl"): try: mod = __import__(pkg_name) @@ -148,7 +155,7 @@ def _find_nixl_wheel_lib_dir() -> Path | None: # Prefer the meson-python sibling layout. for candidate in ( site_packages / f".{pkg_name}.mesonpy.libs", - pkg_root / "lib" / "x86_64-linux-gnu", + pkg_root / "lib" / multiarch, pkg_root / "lib", pkg_root, ): diff --git a/flashinfer/moe_ep/nixl_ep/__init__.py b/flashinfer/moe_ep/nixl_ep/__init__.py index 8ffc4e4d2b1..eba06db9754 100644 --- a/flashinfer/moe_ep/nixl_ep/__init__.py +++ b/flashinfer/moe_ep/nixl_ep/__init__.py @@ -56,7 +56,14 @@ def _find_nixl_lib_dir() -> Path | None: We probe known package names (`nixl_cu13`, `nixl_cu12`, legacy `nixl`), look for the meson-python `.{name}.mesonpy.libs/` sidecar, and fall back to a glob over site-packages. + + `-linux-gnu` follows the Debian multiarch convention — resolves + to `x86_64-linux-gnu` on x86_64 and `aarch64-linux-gnu` on ARM64 + (e.g. NVIDIA Grace / AWS Graviton hosts). """ + import platform + + multiarch = f"{platform.machine()}-linux-gnu" for pkg_name in ("nixl_cu13", "nixl_cu12", "nixl"): try: mod = __import__(pkg_name) @@ -74,7 +81,7 @@ def _find_nixl_lib_dir() -> Path | None: site_packages = pkg_root.parent for candidate in ( site_packages / f".{pkg_name}.mesonpy.libs", - pkg_root / "lib" / "x86_64-linux-gnu", + pkg_root / "lib" / multiarch, pkg_root / "lib", pkg_root, ):