build_backend.py

"""
Copyright (c) 2023 by FlashInfer team.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

  http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
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
from build_utils import get_git_version

_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"
    if version_file.exists():
        with open(version_file, "r") as f:
            version = f.read().strip()
    else:
        version = "0.0.0+unknown"

    # Add dev suffix if specified
    dev_suffix = os.environ.get("FLASHINFER_DEV_RELEASE_SUFFIX", "")
    if dev_suffix:
        version = f"{version}.dev{dev_suffix}"

    # Get git version
    git_version = get_git_version(cwd=_root)

    # Append local version suffix if available
    local_version = os.environ.get("FLASHINFER_LOCAL_VERSION")
    if local_version:
        # Use + to create a local version identifier that will appear in wheel name
        version = f"{version}+{local_version}"

    # Create build metadata in the source tree
    package_dir = Path(__file__).parent / "flashinfer"
    build_meta_file = package_dir / "_build_meta.py"

    # Check if we're in a git repository
    git_dir = Path(__file__).parent / ".git"
    in_git_repo = git_dir.exists()

    # If file exists and not in git repo (installing from sdist), keep existing file
    if build_meta_file.exists() and not in_git_repo:
        print("Build metadata file already exists (not in git repo), keeping it")
        return version

    # In git repo (editable) or file doesn't exist, create/update it
    with open(build_meta_file, "w") as f:
        f.write('"""Build metadata for flashinfer package."""\n')
        f.write(f'__version__ = "{version}"\n')
        f.write(f'__git_version__ = "{git_version}"\n')

    print(f"Created build metadata file with version {version}")
    return version


# Create build metadata as soon as this module is imported
_create_build_metadata()


def write_if_different(path: Path, content: str) -> None:
    if path.exists() and path.read_text() == content:
        return
    path.parent.mkdir(parents=True, exist_ok=True)
    path.write_text(content)


def _create_data_dir(use_symlinks=True):
    _data_dir.mkdir(parents=True, exist_ok=True)

    def ln(source: str, target: str) -> None:
        src = _root / source
        dst = _data_dir / target
        if dst.exists():
            if dst.is_symlink():
                dst.unlink()
            elif dst.is_dir():
                shutil.rmtree(dst)
            else:
                dst.unlink()

        if use_symlinks:
            dst.symlink_to(src, target_is_directory=True)
        else:
            # For wheel/sdist, copy actual files instead of symlinks
            if src.exists():
                shutil.copytree(src, dst, symlinks=False, dirs_exist_ok=True)

    ln("3rdparty/cutlass", "cutlass")
    ln("3rdparty/spdlog", "spdlog")
    ln("3rdparty/cccl", "cccl")
    ln("csrc", "csrc")
    ln("include", "include")


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)

    # Copy license files from licenses/ to root to avoid nested path in wheel
    licenses_dir = _root / "licenses"
    if licenses_dir.exists():
        for license_file in licenses_dir.glob("*.txt"):
            shutil.copy2(
                license_file,
                _root / f"LICENSE.{license_file.stem.removeprefix('LICENSE.')}.txt",
            )


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)


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)


def get_requires_for_build_wheel(config_settings=None):
    _prepare_for_wheel()
    return []


def get_requires_for_build_sdist(config_settings=None):
    _prepare_for_sdist()
    return []


def get_requires_for_build_editable(config_settings=None):
    _prepare_for_editable()
    return []


def prepare_metadata_for_build_wheel(metadata_directory, config_settings=None):
    _prepare_for_wheel()
    return orig.prepare_metadata_for_build_wheel(metadata_directory, config_settings)


def prepare_metadata_for_build_editable(metadata_directory, config_settings=None):
    _prepare_for_editable()
    return orig.prepare_metadata_for_build_editable(metadata_directory, config_settings)


def build_editable(wheel_directory, config_settings=None, metadata_directory=None):
    _prepare_for_editable()
    return orig.build_editable(wheel_directory, config_settings, metadata_directory)


def build_sdist(sdist_directory, config_settings=None):
    _prepare_for_sdist()
    return orig.build_sdist(sdist_directory, config_settings)


def build_wheel(wheel_directory, config_settings=None, metadata_directory=None):
    _prepare_for_wheel()
    return orig.build_wheel(wheel_directory, config_settings, metadata_directory)