flashinfer-python 0.6.11.post3__py3-none-any.whl → 0.6.12__py3-none-any.whl

build_backend.py

@@ -16,6 +16,10 @@ limitations under the License.
 
 import os
 import shutil
+import subprocess
+import sys
+import time
+from contextlib import contextmanager
 from pathlib import Path
 
 from setuptools import build_meta as orig
@@ -25,6 +29,702 @@ _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
+#   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")
+
+
+@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
+
+# 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"
+
+
+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 _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):
+        <site-packages>/nixl_cu13/                  — importable python module
+        <site-packages>/.nixl_cu13.mesonpy.libs/    — libnixl.so + sibling libs
+        <site-packages>/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.
+
+    `<machine>-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)
+        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" / multiarch,
+            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")
+
+    # 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."
+        )
+
+    # 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
+    # — 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)
+
+    # 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"):
+        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 _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")
+
+    # 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
+    # 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}",
+            f"NVCC_GENCODE={nccl_ep_gencode}",
+            "-j",
+        ],
+        cwd=src,
+        check=True,
+    )
+
+    dst = _moe_ep_pkg / "nccl_ep" / "_libs"
+    dst.mkdir(parents=True, exist_ok=True)
+    # 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)
+            print(f"[BUILD_NVEP] staged: {soname}")
+
+    # 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:
+    """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"
+    for so in _moe_ep_pkg.rglob("*.so*"):
+        # Skip symlinks
+        if so.is_symlink():
+            continue
+        # 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)],
+            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]:
+    """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"):
+        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)"
+    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)"
+    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, ""
+
+
+def _nccl_buildable() -> tuple[bool, str]:
+    """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"):
+        return False, (
+            "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, (
+                "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 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] = []
+    if built_nixl:
+        wheels.append(f"nixl-cu{cuda_major}>=1.0.1")
+    if built_nccl:
+        wheels.append(f"nvidia-nccl-cu{cuda_major}>=2.30.4")
+    if not wheels:
+        return
+
+    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:
+    """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
+
+    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.
+    # 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,
+            check=True,
+        )
+
+    # 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:
+            with _time_phase("_build_nixl_ep"):
+                _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:
+            with _time_phase("_build_nccl_ep"):
+                _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:
+        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'}")
+
+
 def _create_build_metadata():
     """Create build metadata file with version information."""
     version_file = _root / "version.txt"
@@ -112,6 +812,7 @@ def _create_data_dir(use_symlinks=True):
 
 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 +829,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 +837,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)