Ep api design - Build Infra dependencies

[Anerudhan]

📌 Description

This PR allows the user to build the NCCL-EP and NIXL-EP as part Flashinfer install

🔍 Related Issues

🚀 Pull Request Checklist

Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.

✅ Pre-commit Checks

• I have installed pre-commit by running pip install pre-commit (or used your preferred method).
• I have installed the hooks with pre-commit install.
• I have run the hooks manually with pre-commit run --all-files and fixed any reported issues.

If you are unsure about how to set up pre-commit, see the pre-commit documentation.

🧪 Tests

• NA

Reviewer Notes

Added NCCL-EP and NIXL-EP as git sub-modules (and optional pip installs).

The following commits will have tests and FI logic.

  BUILD_NCCL_EP=1 pip install -e ".[nvep]"                                  # NCCL-EP only fails if build fails
  BUILD_NIXL_EP=1 pip install -e ".[nvep]"                                    # NIXL-EP only fails if build fails
  BUILD_NCCL_EP=1 BUILD_NIXL_EP=1 pip install -e ".[nvep]"  # both, explicit fails if either build fails
  BUILD_NVEP=1 pip install -e ".[nvep]"                                         # both, build what you can

If using build-isolation:

  uv pip install  \
      torch setuptools packaging \
      'apache-tvm-ffi>=0.1.6,<0.2,!=0.1.8,!=0.1.8.post0' \
      cython pybind11 meson ninja patchelf

  BUILD_NVEP=1 uv pip install --python .venv/bin/python \
      --no-build-isolation -e ".[nvep]"

Summary by CodeRabbit

• New Features

  • Added MOE-EP transport support (NCCL‑EP, NIXL‑EP) with import-time availability checks, lazy-loading stubs, and helpers to list/require backends.
  • Improved GPU architecture coverage for Hopper and Blackwell.

• Chores

  • Optional runtime dependency group for nvep and packaged backend shared libraries.
  • Build tooling can optionally build/vendor native EP components and stage them for installs.
  • New Dockerfile for CUDA13-based validation.
  • Updated ignore rules and added/updated third‑party submodules.

[coderabbitai]

Note

Reviews paused

It looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the reviews.auto_review.auto_pause_after_reviewed_commits setting.

Use the following commands to manage reviews:

• @coderabbitai resume to resume automatic reviews.
• @coderabbitai review to trigger a single review.

Use the checkboxes below for quick actions:

• ▶️ Resume reviews
• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Adds submodule pins and Meson patches, build orchestration to optionally build/stage NIXL-EP and NCCL-EP during wheel/editable preparation, runtime import-time probing/loading helpers that raise informative errors, packaging metadata to include staged libs, a Dockerfile for building nvep-enabled images, and .dockerignore/.gitignore updates.

Changes

MoE Expert-Parallel Transport Backends

Layer / File(s)	Summary
Submodule registration and NIXL patches .gitmodules, 3rdparty/nccl, 3rdparty/nixl, 3rdparty_patches/nixl/*	Adds nixl and nccl submodules with pinned commits and two nixl patches: one to add Blackwell gencodes (sm_100, sm_103) and one to support EP-only builds linking to wheel-provided libnixl.
Build backend orchestration and staging build_backend.py	Adds BUILD_NIXL_EP/BUILD_NCCL_EP/BUILD_NVEP parsing, CUDA detection, idempotent patch application, NIXL (Meson/Ninja) and NCCL (contrib) build flows, staging of shared objects and vendored Python wrapper sources, RPATH rewriting, runtime wheel install orchestration, and integrates these into wheel/editable preparation while skipping sdist.
Docker build recipe and context ignore docker/Dockerfile.flashinfer-nvep, .dockerignore	New Dockerfile that prepares a CUDA 13 image with UCX/DOCA/GDRCopy, pre-installs runtime wheels, runs editable .[nvep] installs wired to build flags and smoke-tests imports. .dockerignore reduces build context while preserving Meson wrap files.
Packaging metadata and VCS ignores .gitignore, pyproject.toml	Adds nvep optional dependency (cuda-python>=13.0) and includes _libs packaging for flashinfer.moe_ep.nixl_ep/nccl_ep; updates .gitignore to ignore in-tree moe_ep build artifacts and vendored outputs.
Runtime detection and lazy loaders flashinfer/moe_ep/__init__.py, flashinfer/moe_ep/nccl_ep/__init__.py, flashinfer/moe_ep/nixl_ep/__init__.py	Implements MoEEpNotBuiltError, have_nccl_ep(), have_nixl_ep(), available_backends(), _require_built(), and backend-specific helpers to locate/preload base libs and lazily load staged EP extensions with informative errors and warnings.

Estimated code review effort

🎯 4 (Complex) | ⏱️ ~45 minutes

Suggested labels

run-ci, op: moe

Suggested reviewers

• sricketts
• yongwww
• yzh119
• cyx-6
• kahyunnam
• bkryu
• samuellees
• jimmyzho

Poem

🐰 I hopped through patches, builds, and wheels,
I nudged libs into order with tiny heels,
If staged or missing I softly chide,
Load, warn, and guide — a rabbit's pride,
Now NCCL and NIXL stride.

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (1 warning, 1 inconclusive)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 75.76% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.
Title check	❓ Inconclusive	Title 'Ep api design - Build Infra dependencies' is vague and does not clearly convey the main change (adding NCCL-EP and NIXL-EP build infrastructure). It reads like a generic placeholder rather than a specific description of what the PR does.	Revise title to clearly state the main objective, e.g., 'Add build infrastructure for NCCL-EP and NIXL-EP backends' or 'Enable optional NCCL-EP and NIXL-EP builds via environment flags'.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Description check	✅ Passed	Description covers the main feature (enabling NCCL-EP/NIXL-EP builds), includes related context and reviewer notes, and checklist items are mostly completed. However, it lacks linked issues and has minimal testing details (marked 'NA').
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request adds infrastructure to build and integrate the moe_ep transport backends, NCCL-EP and NIXL-EP, into FlashInfer. It introduces git submodules, build logic in build_backend.py for compilation and patching, and a dedicated Dockerfile. Reviewers suggested improving build portability by checking for various library directory names, enhancing error handling for RPATH modifications, and strengthening dependency probes for git and nvcc. Further feedback recommended safeguarding submodule updates against non-git environments and ensuring import-time warnings account for all backend-specific build flags.

[coderabbitai]

Actionable comments posted: 4

🧹 Nitpick comments (3)

docker/Dockerfile.flashinfer-nvep (3)

109-109: ⚖️ Poor tradeoff

Acknowledge supply-chain risk of pipe-to-shell installation.

Piping curl to sh without verification is the officially recommended uv installation method, but it creates a supply-chain attack surface if the install script is compromised.

For enhanced security in production environments, consider:

1. Downloading the script, verifying its checksum, then executing it.
2. Using a pinned release URL instead of the latest version redirect.
3. Installing uv via a package manager if available for your base image.

For a reference Dockerfile intended for development use, the current approach is acceptable.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docker/Dockerfile.flashinfer-nvep` at line 109, The RUN line that pipes curl
to sh (RUN curl -LsSf https://astral.sh/uv/install.sh | sh) presents a
supply-chain risk; replace it by downloading a pinned release script or archive
first, verify its checksum/signature, and then execute it (or install uv from a
distro package if available). Specifically, change the current RUN step that
invokes https://astral.sh/uv/install.sh to a two-step process: fetch a
fixed-version URL or release artifact, validate integrity (checksum or
signature), and only then run the installer, ensuring the Dockerfile no longer
uses an unverified curl | sh pattern.

---

1-174: ⚡ Quick win

Consider adding a non-root USER for production deployments.

The static analysis warning about running as root (DS-0002) is valid. While acceptable for a development/reference image, production deployments should use a non-root user to limit the blast radius of potential container escapes or compromises.

🔒 Adding a non-root user

Add before the final CMD:

# Create non-root user with GPU access
RUN groupadd -r flashinfer -g 1000 && \
    useradd -r -u 1000 -g flashinfer -G video flashinfer && \
    chown -R flashinfer:flashinfer ${FLASHINFER_SRC} ${VENV}

USER flashinfer

Note: GPU access requires the video group or appropriate device permissions. Test with your specific --gpus runtime configuration.

If this image is only for development/testing, document that in the header comments and the current approach is acceptable.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docker/Dockerfile.flashinfer-nvep` around lines 1 - 174, The image runs as
root which triggers DS-0002; before the final CMD add creation of a non-root
user/group (e.g. group flashinfer gid 1000 and user flashinfer uid 1000 in group
flashinfer, supplementary group video for GPU access), chown the workspace and
venv (${FLASHINFER_SRC} and ${VENV}) to that user, and switch to it with USER
flashinfer so the container runs unprivileged; update any build steps that
require root (apt, installs, UCX build) to remain before this switch and keep
CMD ["bash"] unchanged.

---

47-49: ⚡ Quick win

Consider adding checksum verification and documenting URL stability expectations.

The hardcoded DOCA download URL could break if Mellanox reorganizes their downloads or removes old versions. While --tries=3 --waitretry=5 handles transient failures, it won't help if the URL becomes permanently unavailable.

💡 Suggested improvements

1. Add a comment documenting the expected URL stability and update policy:

 # 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.
+# NOTE: This URL may become unavailable as Mellanox updates their download
+# repository. If builds fail, check https://www.mellanox.com/downloads/DOCA
+# for the current download location.
 ARG DOCA_VERSION=3.2.0-125000-25.10

2. Optionally add SHA256 checksum verification:

 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 \
+    && echo "EXPECTED_SHA256_HERE /tmp/doca-host.deb" | sha256sum -c - \
     && dpkg -i /tmp/doca-host.deb \

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docker/Dockerfile.flashinfer-nvep` around lines 47 - 49, The Dockerfile's RUN
that wget's the DOCA package (the command downloading
https://.../doca-host_${DOCA_VERSION}-ubuntu2404_amd64.deb into
/tmp/doca-host.deb) should include a short comment about URL stability/update
policy and perform checksum verification to avoid silently using a
tampered/removed file; update the RUN to first download or embed a known SHA256
for the exact DOCA_VERSION and then verify the downloaded /tmp/doca-host.deb
with sha256sum (failing the build on mismatch) and remove the downloaded files
on failure/success, or alternatively download a vendor-signed checksum and
verify it before dpkg install—ensure you reference DOCA_VERSION and
/tmp/doca-host.deb in the verification steps and make the build error out if
checksums do not match.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@build_backend.py`:
- Around line 289-292: The subprocess.run call invoking [sys.executable, "-m",
"pip", "install", "--no-deps", *wheels] currently uses check=False and does not
inspect the result, which can silently ignore failed wheel installs; modify this
to detect failures by either using check=True or capturing the CompletedProcess
(result = subprocess.run(..., check=False)) and then checking result.returncode,
and on non-zero return code raise a clear exception or exit with a descriptive
error that includes the failing wheels (the wheels variable) and the return
code/message so runtime wheel install failures are surfaced rather than
swallowed.
- Around line 230-249: _nixl_buildable() currently misses a check for the CUDA
compiler which causes Meson to fail later; add a PATH check for "nvcc" similar
to the check in _nccl_buildable(), returning (False, "nvcc not on PATH (install
CUDA toolkit / add nvcc to PATH)") when shutil.which("nvcc") is falsy so the
function fails fast with a clear actionable message before running
meson/ninja/pkg-config checks.
- Around line 137-139: The staging loop only globs (build /
"examples/device/ep") for "nixl_ep_cpp*.so" but the comment notes Meson may
place the extension in build/ as well; update the logic in build_backend.py to
search both locations (e.g., both build.glob("nixl_ep_cpp*.so") and (build /
"examples/device/ep").glob("nixl_ep_cpp*.so") or iterate over a list of
candidate directories) and copy any found files using the existing
shutil.copy(cand, dst / cand.name) call so the nixl_ep_cpp*.so is reliably
staged from either location. Ensure the loop still handles multiple matches and
does not error if a directory yields no results.

In `@flashinfer/moe_ep/__init__.py`:
- Around line 96-104: The current import-time warning only checks BUILD_NVEP and
misses other per-backend build flags; update the check to consider BUILD_NCCL_EP
and BUILD_NIXL_EP as well (i.e., warn if any of these env vars is "1" and
available_backends() is empty). Modify the conditional around the warning in
__init__.py to test os.environ.get for BUILD_NVEP, BUILD_NCCL_EP, or
BUILD_NIXL_EP (instead of only BUILD_NVEP), and keep the warning text
referencing _pkg_dir and available_backends() unchanged so users still see the
same diagnostic when no backend libraries are found.

---

Nitpick comments:
In `@docker/Dockerfile.flashinfer-nvep`:
- Line 109: The RUN line that pipes curl to sh (RUN curl -LsSf
https://astral.sh/uv/install.sh | sh) presents a supply-chain risk; replace it
by downloading a pinned release script or archive first, verify its
checksum/signature, and then execute it (or install uv from a distro package if
available). Specifically, change the current RUN step that invokes
https://astral.sh/uv/install.sh to a two-step process: fetch a fixed-version URL
or release artifact, validate integrity (checksum or signature), and only then
run the installer, ensuring the Dockerfile no longer uses an unverified curl |
sh pattern.
- Around line 1-174: The image runs as root which triggers DS-0002; before the
final CMD add creation of a non-root user/group (e.g. group flashinfer gid 1000
and user flashinfer uid 1000 in group flashinfer, supplementary group video for
GPU access), chown the workspace and venv (${FLASHINFER_SRC} and ${VENV}) to
that user, and switch to it with USER flashinfer so the container runs
unprivileged; update any build steps that require root (apt, installs, UCX
build) to remain before this switch and keep CMD ["bash"] unchanged.
- Around line 47-49: The Dockerfile's RUN that wget's the DOCA package (the
command downloading https://.../doca-host_${DOCA_VERSION}-ubuntu2404_amd64.deb
into /tmp/doca-host.deb) should include a short comment about URL
stability/update policy and perform checksum verification to avoid silently
using a tampered/removed file; update the RUN to first download or embed a known
SHA256 for the exact DOCA_VERSION and then verify the downloaded
/tmp/doca-host.deb with sha256sum (failing the build on mismatch) and remove the
downloaded files on failure/success, or alternatively download a vendor-signed
checksum and verify it before dpkg install—ensure you reference DOCA_VERSION and
/tmp/doca-host.deb in the verification steps and make the build error out if
checksums do not match.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: e7db5da0-7b55-4bb6-a3e1-370917e76bd1

📥 Commits

Reviewing files that changed from the base of the PR and between 103fcf8 and c28bc150ca53c28f3351c4fba9045bf4d534dada.

📒 Files selected for processing (12)

• .dockerignore
• .gitignore
• .gitmodules
• 3rdparty/nccl
• 3rdparty/nixl
• 3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch
• build_backend.py
• docker/Dockerfile.flashinfer-nvep
• flashinfer/moe_ep/__init__.py
• flashinfer/moe_ep/nccl_ep/__init__.py
• flashinfer/moe_ep/nixl_ep/__init__.py
• pyproject.toml

[coderabbitai]

Actionable comments posted: 3

♻️ Duplicate comments (3)

flashinfer/moe_ep/__init__.py (1)

109-117: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Import-time warning should honor BUILD_NCCL_EP and BUILD_NIXL_EP too.

The warning only fires for BUILD_NVEP=1. Users who explicitly opted in via BUILD_NCCL_EP=1 or BUILD_NIXL_EP=1 (the per-backend flags introduced in build_backend.py) lose this diagnostic when the build silently produced no plugins.

Suggested fix

-if os.environ.get("BUILD_NVEP") == "1" and not available_backends():
+_requested_build = any(
+    os.environ.get(k) == "1"
+    for k in ("BUILD_NVEP", "BUILD_NCCL_EP", "BUILD_NIXL_EP")
+)
+if _requested_build and not available_backends():
     import warnings
 
     warnings.warn(
-        "BUILD_NVEP=1 was set, but no moe_ep backend libraries were found "
+        "A BUILD_*_EP=1 flag 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,
     )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/__init__.py` around lines 109 - 117, The import-time
warning currently only checks BUILD_NVEP; change the condition in
flashinfer/moe_ep/__init__.py so the warning fires if any per-backend build flag
is set (BUILD_NVEP, BUILD_NCCL_EP, BUILD_NIXL_EP) and available_backends()
returns false. Update the if expression that wraps warnings.warn to compute a
boolean like any(os.environ.get(flag) == "1" for flag in
("BUILD_NVEP","BUILD_NCCL_EP","BUILD_NIXL_EP")) and use that instead of only
BUILD_NVEP, keeping the same warning message and referencing _pkg_dir and
available_backends() as before.

build_backend.py (2)

140-143: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Staging only covers one of the two documented output locations.

The comment at Line 140 states the torch extension may land in build/ or build/examples/device/ep/, but the glob only searches the latter. If Meson places nixl_ep_cpp*.so under build/ (as the comment claims is possible), the build will succeed but the plugin will not be staged into flashinfer/moe_ep/nixl_ep/_libs/, and the runtime probe in flashinfer/moe_ep/__init__.py::_probe_nixl_ep will report NIXL-EP as unavailable.

Suggested fix

-    # 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}")
+    # The torch extension lands either in build/ or build/examples/device/ep/
+    seen: set[str] = set()
+    for root in (build, build / "examples/device/ep"):
+        for cand in root.glob("nixl_ep_cpp*.so"):
+            if cand.name in seen:
+                continue
+            shutil.copy(cand, dst / cand.name)
+            seen.add(cand.name)
+            print(f"[BUILD_NVEP] staged: {cand.name}")

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@build_backend.py` around lines 140 - 143, The staging loop only searches
(build / "examples/device/ep") for "nixl_ep_cpp*.so" but the comment says the
extension may also be under build/; update the staging logic in build_backend.py
(the loop that iterates over build /
"examples/device/ep".glob("nixl_ep_cpp*.so") and does shutil.copy to dst and
prints "[BUILD_NVEP] staged:") to also search build.glob("nixl_ep_cpp*.so") (or
iterate over both candidate directories), copy any matches into dst / cand.name,
and avoid duplicate copies so the runtime probe in
flashinfer/moe_ep/__init__.py::_probe_nixl_ep can find the plugin.

---

359-378: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Add nvcc availability check to _nixl_buildable().

NIXL EP compiles CUDA sources (the patch enables Blackwell gencode arches in meson.build), so nvcc is a hard build-time dep. Currently, all checks pass while nvcc is missing, and the failure surfaces later as an opaque Meson/Ninja error instead of a clear preflight message. _nccl_buildable() already does this check at Lines 391–395; mirror it here.

Suggested fix

 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)"
+    if not shutil.which("nvcc"):
+        return False, (
+            "nvcc not on PATH (install CUDA toolkit and put "
+            "/usr/local/cuda/bin on $PATH)"
+        )
     pkgconfig = shutil.which("pkg-config")

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@build_backend.py` around lines 359 - 378, The _nixl_buildable() probe is
missing an nvcc availability check so CUDA sources will fail later; add a check
like in _nccl_buildable(): call shutil.which("nvcc") inside _nixl_buildable()
and if missing return False with a clear message (e.g. "nvcc not on PATH
(install CUDA toolkit / nvcc)") so preflight fails early; place this check
alongside the other tool checks (meson, ninja, pkg-config) before invoking
pkg-config for ucx/libibverbs.

🧹 Nitpick comments (1)

docker/Dockerfile.flashinfer-nvep (1)

183-185: ⚡ Quick win

Smoke probes don't fail the build when no backends were produced.

available_backends() returns [] if both EP plugins are missing, but print('moe_ep backends:', []) still exits 0 — so a silently failed best-effort build (e.g., NIXL skipped on missing dep) would produce an image whose flashinfer-nvep tag promises both backends but ships neither. Consider asserting in the probe.

Suggested fix

-RUN python -c "from flashinfer.moe_ep import available_backends; print('moe_ep backends:', available_backends())"
+RUN python -c "from flashinfer.moe_ep import available_backends; \
+    b = available_backends(); print('moe_ep backends:', b); \
+    assert b, 'no moe_ep backends built — check the build log'"

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docker/Dockerfile.flashinfer-nvep` around lines 183 - 185, The smoke probe
currently prints available_backends() but doesn't fail the build if it returns
an empty list; change the RUN that invokes python -c "from flashinfer.moe_ep
import available_backends; print(...)" to check that available_backends() is
non-empty and exit non-zero (raise SystemExit or call sys.exit(1)) when it is
empty so the Docker build fails; keep the existing nccl_ep probe (python -c
"import nccl_ep; from nccl.core.communicator import Communicator; ...") as-is
but ensure the mojo EP check uses available_backends() to assert presence of at
least one backend before succeeding.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@docker/Dockerfile.flashinfer-nvep`:
- Around line 101-106: The comment block stating "DOCA gpunetio — NOT installed
here" is stale and contradicts the earlier apt install of doca-sdk-gpunetio,
libdoca-sdk-gpunetio-dev and libdoca-sdk-verbs-dev (installed from the DOCA host
deb); update or remove that comment to reflect the current state: either delete
the block or rephrase it to acknowledge that doca-sdk-gpunetio and related
packages are installed (referencing the apt install of doca-sdk-gpunetio,
libdoca-sdk-gpunetio-dev, libdoca-sdk-verbs-dev and the DOCA host deb
installation), so maintainers are not misled.

In `@flashinfer/moe_ep/nccl_ep/__init__.py`:
- Around line 71-101: The code currently calls ctypes.CDLL directly in two
places (when nccl_so is found and when loading libnccl_ep.so in
_load_libnccl_ep) which can raise OSError and leak raw errors; wrap both
ctypes.CDLL calls in try/except OSError as e and raise MoEEpNotBuiltError with
an actionable message (include the offending path, e.g. str(nccl_so) or the
_libs_dir / "libnccl_ep.so" value, and suggested install/build steps) using
"raise ... from e" to preserve the original error; update
_preload_libnccl/_load_libnccl_ep call sites accordingly so the module
consistently surfaces MoEEpNotBuiltError instead of raw OSError.

In `@flashinfer/moe_ep/nixl_ep/__init__.py`:
- Around line 74-111: The preload path may miss required base libs and load
failures aren't translated to MoEEpNotBuiltError; update _preload_libnixl to
validate every name in _NIXL_BASE_LIBS (either by trying ctypes.CDLL(libname,
RTLD_GLOBAL) when nixl_lib_dir is None, or by checking nixl_lib_dir / libname
exists and loading it), and if any lib is missing or a ctypes.CDLL raise OSError
capture that exception and raise MoEEpNotBuiltError with a clear message
(include the original exception). Also wrap the final ctypes.CDLL(...) in
_load_nixl_ep_cpp in a try/except OSError and re-raise MoEEpNotBuiltError
(including the underlying error) so extension load failures are reported as
MoEEpNotBuiltError; refer to functions _preload_libnixl, _load_nixl_ep_cpp,
symbol _NIXL_BASE_LIBS, variable nixl_lib_dir and _libs_dir when making the
changes.

---

Duplicate comments:
In `@build_backend.py`:
- Around line 140-143: The staging loop only searches (build /
"examples/device/ep") for "nixl_ep_cpp*.so" but the comment says the extension
may also be under build/; update the staging logic in build_backend.py (the loop
that iterates over build / "examples/device/ep".glob("nixl_ep_cpp*.so") and does
shutil.copy to dst and prints "[BUILD_NVEP] staged:") to also search
build.glob("nixl_ep_cpp*.so") (or iterate over both candidate directories), copy
any matches into dst / cand.name, and avoid duplicate copies so the runtime
probe in flashinfer/moe_ep/__init__.py::_probe_nixl_ep can find the plugin.
- Around line 359-378: The _nixl_buildable() probe is missing an nvcc
availability check so CUDA sources will fail later; add a check like in
_nccl_buildable(): call shutil.which("nvcc") inside _nixl_buildable() and if
missing return False with a clear message (e.g. "nvcc not on PATH (install CUDA
toolkit / nvcc)") so preflight fails early; place this check alongside the other
tool checks (meson, ninja, pkg-config) before invoking pkg-config for
ucx/libibverbs.

In `@flashinfer/moe_ep/__init__.py`:
- Around line 109-117: The import-time warning currently only checks BUILD_NVEP;
change the condition in flashinfer/moe_ep/__init__.py so the warning fires if
any per-backend build flag is set (BUILD_NVEP, BUILD_NCCL_EP, BUILD_NIXL_EP) and
available_backends() returns false. Update the if expression that wraps
warnings.warn to compute a boolean like any(os.environ.get(flag) == "1" for flag
in ("BUILD_NVEP","BUILD_NCCL_EP","BUILD_NIXL_EP")) and use that instead of only
BUILD_NVEP, keeping the same warning message and referencing _pkg_dir and
available_backends() as before.

---

Nitpick comments:
In `@docker/Dockerfile.flashinfer-nvep`:
- Around line 183-185: The smoke probe currently prints available_backends() but
doesn't fail the build if it returns an empty list; change the RUN that invokes
python -c "from flashinfer.moe_ep import available_backends; print(...)" to
check that available_backends() is non-empty and exit non-zero (raise SystemExit
or call sys.exit(1)) when it is empty so the Docker build fails; keep the
existing nccl_ep probe (python -c "import nccl_ep; from nccl.core.communicator
import Communicator; ...") as-is but ensure the mojo EP check uses
available_backends() to assert presence of at least one backend before
succeeding.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 6d86a23d-4939-4f38-8b8a-ca8ab5eab7e1

📥 Commits

Reviewing files that changed from the base of the PR and between c28bc150ca53c28f3351c4fba9045bf4d534dada and 8cf421fbcb95fb039c60c49e8d6117848b91c9de.

📒 Files selected for processing (5)

• build_backend.py
• docker/Dockerfile.flashinfer-nvep
• flashinfer/moe_ep/__init__.py
• flashinfer/moe_ep/nccl_ep/__init__.py
• flashinfer/moe_ep/nixl_ep/__init__.py

[coderabbitai]

Actionable comments posted: 2

🧹 Nitpick comments (2)

flashinfer/moe_ep/nixl_ep/__init__.py (1)

65-68: 💤 Low value

Narrow the exception clause to expected types.

Catching bare Exception masks unexpected errors. The expected failures here are AttributeError (non-package module without __path__), TypeError (non-subscriptable), or IndexError (empty __path__).

Suggested fix

         try:
             pkg_root = Path(mod.__path__[0])
-        except Exception:
+        except (AttributeError, TypeError, IndexError):
             continue

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/nixl_ep/__init__.py` around lines 65 - 68, Replace the
broad "except Exception" in the block around "pkg_root = Path(mod.__path__[0])"
with a narrow tuple of the expected exception types to avoid masking other
errors; catch (AttributeError, TypeError, IndexError) instead so only the cases
where mod has no __path__, is not subscriptable, or __path__ is empty are
swallowed and all other exceptions propagate.

docker/Dockerfile.flashinfer-nvep (1)

72-75: ⚡ Quick win

Pin UCX to a commit instead of the moving v1.21.x branch.

Using a branch head makes this image non-reproducible and can break future builds without any repo change. A tested commit SHA is safer for build infra.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docker/Dockerfile.flashinfer-nvep` around lines 72 - 75, The Dockerfile
currently uses ARG UCX_VERSION=v1.21.x and clones with --branch ${UCX_VERSION},
which pins to a moving branch; change this to a fixed commit SHA by replacing or
adding an ARG like UCX_COMMIT and using that SHA in the git clone/checkout step
(or clone default and run git checkout <UCX_COMMIT>) so the RUN git clone ...
--branch ${UCX_VERSION} line uses a stable commit instead of the v1.21.x branch;
update any references from UCX_VERSION to the new UCX_COMMIT symbol and document
the tested SHA in the ARG default.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@build_backend.py`:
- Around line 212-213: The current logic skips running Meson setup when
build.exists() is true, which reuses a stale build dir across different
configurations (BUILD_NIXL_EP_HERMETIC, -Dnixl_ep_only, -Dnixl_wheel_lib_dir);
change it so that if the build directory exists you either remove it or invoke
Meson with reconfigure semantics instead of skipping: detect the existing build
dir (build.exists()) and call subprocess.run(setup_args + ['--reconfigure'],
check=True) or remove the dir and call subprocess.run(setup_args, check=True) to
ensure Meson is reconfigured for the current options.

In `@docker/Dockerfile.flashinfer-nvep`:
- Around line 46-49: The wget URL hard-codes "DOCA_v3.2.0" so overriding ARG
DOCA_VERSION will break; update the Dockerfile so the path is derived from
DOCA_VERSION instead of a fixed string: either add a new ARG (for example
DOCA_SERIES) and set it to the corresponding "DOCA_vX.Y.Z" series and use
${DOCA_SERIES} in the wget URL, or compute the series from DOCA_VERSION inside
the RUN (e.g. use shell expansion ${DOCA_VERSION%%-*} to build
"DOCA_v${DOCA_VERSION%%-*}") so the wget line that currently contains
"https://.../DOCA_v3.2.0/host/doca-host_${DOCA_VERSION}-ubuntu2404_amd64.deb"
becomes dynamically constructed from DOCA_VERSION/DOCA_SERIES.

---

Nitpick comments:
In `@docker/Dockerfile.flashinfer-nvep`:
- Around line 72-75: The Dockerfile currently uses ARG UCX_VERSION=v1.21.x and
clones with --branch ${UCX_VERSION}, which pins to a moving branch; change this
to a fixed commit SHA by replacing or adding an ARG like UCX_COMMIT and using
that SHA in the git clone/checkout step (or clone default and run git checkout
<UCX_COMMIT>) so the RUN git clone ... --branch ${UCX_VERSION} line uses a
stable commit instead of the v1.21.x branch; update any references from
UCX_VERSION to the new UCX_COMMIT symbol and document the tested SHA in the ARG
default.

In `@flashinfer/moe_ep/nixl_ep/__init__.py`:
- Around line 65-68: Replace the broad "except Exception" in the block around
"pkg_root = Path(mod.__path__[0])" with a narrow tuple of the expected exception
types to avoid masking other errors; catch (AttributeError, TypeError,
IndexError) instead so only the cases where mod has no __path__, is not
subscriptable, or __path__ is empty are swallowed and all other exceptions
propagate.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 3294f496-4426-491d-a600-89e3a00712a3

📥 Commits

Reviewing files that changed from the base of the PR and between 8cf421fbcb95fb039c60c49e8d6117848b91c9de and 180f4a4973d1813954199f723b5535b1fd10a747.

📒 Files selected for processing (5)

• 3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch
• 3rdparty_patches/nixl/0002-ep-only-build.patch
• build_backend.py
• docker/Dockerfile.flashinfer-nvep
• flashinfer/moe_ep/nixl_ep/__init__.py

🚧 Files skipped from review as they are similar to previous changes (1)

• 3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch

[coderabbitai]

Actionable comments posted: 3

♻️ Duplicate comments (3)

flashinfer/moe_ep/nixl_ep/__init__.py (1)

91-112: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Preload all required NIXL base libs and wrap load failures consistently.

Current logic can proceed with a partial base-lib set and can leak raw OSError (including at Line 128), which undermines predictable MoEEpNotBuiltError diagnostics.

Suggested fix

 def _preload_libnixl() -> None:
@@
     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)
+    candidates = (
+        _NIXL_BASE_LIBS
+        if nixl_lib_dir is None
+        else tuple(str(nixl_lib_dir / libname) for libname in _NIXL_BASE_LIBS)
+    )
+    try:
+        for candidate in candidates:
+            ctypes.CDLL(candidate, mode=ctypes.RTLD_GLOBAL)
+    except OSError as e:
+        raise MoEEpNotBuiltError(
+            "Could not locate or load 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 the NIXL runtime libs."
+        ) from e
@@
 def _load_nixl_ep_cpp() -> ctypes.CDLL:
@@
     _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"Failed to load staged NIXL-EP extension at {so_files[0]}. "
+            "Rebuild with BUILD_NIXL_EP=1 or BUILD_NVEP=1."
+        ) from e

Also applies to: 127-128

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/nixl_ep/__init__.py` around lines 91 - 112, In
_preload_libnixl(), ensure every library in _NIXL_BASE_LIBS is loaded and any
ctypes.CDLL OSError is caught and re-raised as MoEEpNotBuiltError with a clear
message; if _find_nixl_lib_dir() returns None keep the current attempt to load
"libnixl.so" but wrap its OSError into MoEEpNotBuiltError, and in the loop over
_NIXL_BASE_LIBS catch OSError for each libpath (or missing file) and raise
MoEEpNotBuiltError rather than letting raw OSError propagate so callers of
_preload_libnixl get consistent diagnostics.

flashinfer/moe_ep/nccl_ep/__init__.py (1)

73-74: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Normalize NCCL loader failures to MoEEpNotBuiltError.

Line 73 and Line 101 can raise raw OSError, which bypasses the actionable error path this module is trying to provide.

Suggested fix

 def _preload_libnccl() -> None:
@@
     nccl_so = _find_libnccl()
     if nccl_so is not None:
-        ctypes.CDLL(str(nccl_so), mode=ctypes.RTLD_GLOBAL)
-        return
+        try:
+            ctypes.CDLL(str(nccl_so), mode=ctypes.RTLD_GLOBAL)
+            return
+        except OSError as e:
+            raise MoEEpNotBuiltError(
+                f"Failed to load NCCL runtime at {nccl_so}. "
+                "Ensure the nvidia-nccl-cu13 runtime is compatible with this environment, "
+                "or expose a compatible libnccl.so.2 via LD_LIBRARY_PATH."
+            ) from e
@@
 def _load_libnccl_ep() -> ctypes.CDLL:
@@
     _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"Failed to load staged NCCL-EP plugin at {so}. "
+            "Rebuild with BUILD_NCCL_EP=1 or BUILD_NVEP=1."
+        ) from e

Also applies to: 100-101

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/nccl_ep/__init__.py` around lines 73 - 74, Wrap the
ctypes.CDLL(...) calls that load nccl_so in a try/except that catches OSError
and raises MoEEpNotBuiltError instead (preserving the original error message as
context), replacing the raw return-on-success behavior; specifically update the
blocks that call ctypes.CDLL(str(nccl_so), mode=ctypes.RTLD_GLOBAL) so any
OSError is normalized to MoEEpNotBuiltError with the original exception chained
or included in the message (reference symbols: nccl_so, ctypes.CDLL, and
MoEEpNotBuiltError).

flashinfer/moe_ep/__init__.py (1)

109-109: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Include granular EP build flags in the import-time warning gate.

The warning currently only triggers for BUILD_NVEP=1; users setting BUILD_NCCL_EP=1 or BUILD_NIXL_EP=1 won’t get this diagnostic when no backend is staged.

Suggested fix

-if os.environ.get("BUILD_NVEP") == "1" and not available_backends():
+_requested_build = any(
+    os.environ.get(k) == "1" for k in ("BUILD_NVEP", "BUILD_NCCL_EP", "BUILD_NIXL_EP")
+)
+if _requested_build and not available_backends():

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@flashinfer/moe_ep/__init__.py` at line 109, The import-time warning gate
currently only checks BUILD_NVEP; update the condition that reads if
os.environ.get("BUILD_NVEP") == "1" and not available_backends(): to trigger
when any EP build flag is set by checking BUILD_NVEP, BUILD_NCCL_EP, and
BUILD_NIXL_EP (e.g. replace the single check with any(os.environ.get(v) == "1"
for v in ("BUILD_NVEP","BUILD_NCCL_EP","BUILD_NIXL_EP")) and not
available_backends()). Ensure available_backends() is still used to suppress the
warning when a backend is staged.

🧹 Nitpick comments (1)

docker/Dockerfile.flashinfer-nvep (1)

193-195: ⚡ Quick win

Make the smoke probe fail when NIXL-EP is missing.

The first probe only prints available_backends(), so this image still builds successfully if NIXL quietly disappears. Since this Dockerfile is meant to validate the full nvep path, assert the expected backend set instead of logging it.

Suggested check

-RUN python -c "from flashinfer.moe_ep import available_backends; print('moe_ep backends:', available_backends())"
+RUN python - <<'PY'
+from flashinfer.moe_ep import available_backends
+backends = set(available_backends())
+print("moe_ep backends:", sorted(backends))
+assert {"nccl_ep", "nixl_ep"} <= backends, backends
+PY

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docker/Dockerfile.flashinfer-nvep` around lines 193 - 195, The smoke probe
should fail when NIXL-EP is missing: replace the current RUN that just prints
available_backends() with a Python check that raises a nonzero exit when the
expected backend set is not present (e.g. assert 'nixl_ep' in
available_backends() or assert set(available_backends()) ==
{...expected_backends...}); update the RUN invoking available_backends() (from
flashinfer.moe_ep) so it asserts the presence/identity of NIXL-EP and exits
nonzero on mismatch, optionally keeping a print for success.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In `@build_backend.py`:
- Around line 223-236: The builder silently continues when no plugin .so files
are found; modify the staging loops (e.g., the loop iterating (build /
"examples/device/ep").glob("nixl_ep_cpp*.so") that copies into dst = _moe_ep_pkg
/ "nixl_ep" / "_libs") to count how many artifacts were actually copied and, if
the count is zero, raise an exception (RuntimeError or similar) with a clear
message before any code sets built_nixl or built_nccl to True; apply the same
pattern to the other builder loops referenced in the comment (the ranges around
lines 397-408 and 643-669) so each builder fails loudly when no expected
artifacts were staged.
- Around line 512-578: The runtime-wheel installation in
_install_nvep_runtime_wheels is currently executed during both wheel builds and
editable installs, which causes base libs to be installed into an isolated
PEP517 build env and not bundled with distributed wheels; change the call sites
instead of installing during wheel builds: remove or gate calls to
_install_nvep_runtime_wheels from the wheel-build path (the invocation in
_prepare_for_wheel) and ensure it is only invoked from the editable-install path
(the invocation in _prepare_for_editable), or alternatively add a boolean
parameter to _install_nvep_runtime_wheels (e.g. editable_only=True) and have
callers pass the appropriate flag so that the function no-ops during wheel
creation; update error/messages accordingly and keep the function signature
_install_nvep_runtime_wheels(built_nixl: bool, built_nccl: bool, editable_only:
bool=False) if you choose the flag approach.

In `@docker/Dockerfile.flashinfer-nvep`:
- Around line 66-74: The Dockerfile currently pins UCX with ARG
UCX_VERSION=v1.21.x which tracks a moving branch; replace that with the
immutable commit SHA by setting ARG
UCX_VERSION=c982cef7cdfc80008c838741ff4cc38a4ce54c89 (or otherwise use that SHA
as the ref in the git clone command) so the RUN git clone --depth=1 --branch
${UCX_VERSION} https://github.com/openucx/ucx.git /tmp/ucx will always fetch the
exact commit; update the ARG value and/or the clone invocation to reference the
provided SHA instead of the branch name.

---

Duplicate comments:
In `@flashinfer/moe_ep/__init__.py`:
- Line 109: The import-time warning gate currently only checks BUILD_NVEP;
update the condition that reads if os.environ.get("BUILD_NVEP") == "1" and not
available_backends(): to trigger when any EP build flag is set by checking
BUILD_NVEP, BUILD_NCCL_EP, and BUILD_NIXL_EP (e.g. replace the single check with
any(os.environ.get(v) == "1" for v in
("BUILD_NVEP","BUILD_NCCL_EP","BUILD_NIXL_EP")) and not available_backends()).
Ensure available_backends() is still used to suppress the warning when a backend
is staged.

In `@flashinfer/moe_ep/nccl_ep/__init__.py`:
- Around line 73-74: Wrap the ctypes.CDLL(...) calls that load nccl_so in a
try/except that catches OSError and raises MoEEpNotBuiltError instead
(preserving the original error message as context), replacing the raw
return-on-success behavior; specifically update the blocks that call
ctypes.CDLL(str(nccl_so), mode=ctypes.RTLD_GLOBAL) so any OSError is normalized
to MoEEpNotBuiltError with the original exception chained or included in the
message (reference symbols: nccl_so, ctypes.CDLL, and MoEEpNotBuiltError).

In `@flashinfer/moe_ep/nixl_ep/__init__.py`:
- Around line 91-112: In _preload_libnixl(), ensure every library in
_NIXL_BASE_LIBS is loaded and any ctypes.CDLL OSError is caught and re-raised as
MoEEpNotBuiltError with a clear message; if _find_nixl_lib_dir() returns None
keep the current attempt to load "libnixl.so" but wrap its OSError into
MoEEpNotBuiltError, and in the loop over _NIXL_BASE_LIBS catch OSError for each
libpath (or missing file) and raise MoEEpNotBuiltError rather than letting raw
OSError propagate so callers of _preload_libnixl get consistent diagnostics.

---

Nitpick comments:
In `@docker/Dockerfile.flashinfer-nvep`:
- Around line 193-195: The smoke probe should fail when NIXL-EP is missing:
replace the current RUN that just prints available_backends() with a Python
check that raises a nonzero exit when the expected backend set is not present
(e.g. assert 'nixl_ep' in available_backends() or assert
set(available_backends()) == {...expected_backends...}); update the RUN invoking
available_backends() (from flashinfer.moe_ep) so it asserts the
presence/identity of NIXL-EP and exits nonzero on mismatch, optionally keeping a
print for success.

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 0fdcc28f-85d6-472d-8c95-2f62a8366b31

📥 Commits

Reviewing files that changed from the base of the PR and between 180f4a4973d1813954199f723b5535b1fd10a747 and 12f7a12.

📒 Files selected for processing (13)

• .dockerignore
• .gitignore
• .gitmodules
• 3rdparty/nccl
• 3rdparty/nixl
• 3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch
• 3rdparty_patches/nixl/0002-ep-only-build.patch
• build_backend.py
• docker/Dockerfile.flashinfer-nvep
• flashinfer/moe_ep/__init__.py
• flashinfer/moe_ep/nccl_ep/__init__.py
• flashinfer/moe_ep/nixl_ep/__init__.py
• pyproject.toml

✅ Files skipped from review due to trivial changes (4)

• 3rdparty/nixl
• .gitmodules
• .gitignore
• .dockerignore

🚧 Files skipped from review as they are similar to previous changes (4)

• 3rdparty/nccl
• pyproject.toml
• 3rdparty_patches/nixl/0001-meson-add-blackwell-arches.patch
• 3rdparty_patches/nixl/0002-ep-only-build.patch

[coderabbitai]

♻️ Duplicate comments (4)

docker/Dockerfile.flashinfer-nvep (1)

46-49: ⚠️ Potential issue | 🟡 Minor | ⚡ Quick win

Keep the DOCA download path in sync with DOCA_VERSION.

DOCA_VERSION is exposed as an ARG, but the URL hardcodes DOCA_v3.2.0. Overriding DOCA_VERSION to a different series (e.g., 3.3.0-*) will cause a 404.

Suggested fix

 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 \
+        https://www.mellanox.com/downloads/DOCA/DOCA_v${DOCA_VERSION%%-*}/host/doca-host_${DOCA_VERSION}-ubuntu2404_amd64.deb \
         -O /tmp/doca-host.deb \

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@docker/Dockerfile.flashinfer-nvep` around lines 46 - 49, The Dockerfile
hardcodes "DOCA_v3.2.0" in the wget URL which won't match other DOCA_VERSION
values; update the RUN/wget invocation to derive the path from the DOCA_VERSION
ARG (e.g., interpolate DOCA_VERSION into the URL or build a DOCA_URL ARG) so the
segment "DOCA_v<version>" is generated from DOCA_VERSION rather than hardcoded,
ensuring the download path and the DOCA_VERSION variable stay in sync.

build_backend.py (3)

245-248: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Stage nixl_ep_cpp*.so from both documented output locations.

The comment at Line 245 indicates the torch extension can land in either build/ or build/examples/device/ep/, but Line 246 only globs the latter. This can leave NIXL built but undetected.

Suggested fix

-    # 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}")
+    # The torch extension lands either in build/ or build/examples/device/ep/
+    staged_count = 0
+    for root in (build, build / "examples/device/ep"):
+        for cand in root.glob("nixl_ep_cpp*.so"):
+            shutil.copy(cand, dst / cand.name)
+            staged_count += 1
+            print(f"[BUILD_NVEP] staged: {cand.name}")
+    if staged_count == 0:
+        raise RuntimeError(
+            "NIXL-EP build completed but no nixl_ep_cpp*.so was produced. "
+            "Check the Meson build output for errors."
+        )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@build_backend.py` around lines 245 - 248, The staging loop only searches
(build / "examples/device/ep") for "nixl_ep_cpp*.so" but the comment says the
extension may land in build/ as well; update the logic in build_backend.py (the
loop that uses (build / "examples/device/ep").glob(...) and copies to dst) to
search both locations — e.g., iterate over both build.glob("nixl_ep_cpp*.so")
and (build / "examples/device/ep").glob("nixl_ep_cpp*.so"), copy any matches to
dst, and avoid double-copying the same filename if found in both places.

---

541-607: ⚠️ Potential issue | 🟠 Major | 🏗️ Heavy lift

Wheel builds won't carry runtime base library dependencies.

This function runs during both _prepare_for_wheel() and _prepare_for_editable() (via _build_nvep_if_enabled()). For editable installs, the wheels are installed into the user's environment. For wheel builds, PEP 517 isolation means these install into a temporary build environment that doesn't transfer to the final installation.

Distributed wheels will ship the EP plugins (libnccl_ep.so, nixl_ep_cpp.so) but lack the base libraries (libnccl.so.2, libnixl.so) they depend on. The [nvep] optional dependency only declares cuda-python>=13.0.

Consider one of:

1. Add nvidia-nccl-cu13 and nixl-cu13 to [nvep] optional-dependencies in pyproject.toml
2. Skip _install_nvep_runtime_wheels() during wheel builds (keep for editable only)
3. Document that distributed wheels require manual installation of these runtime wheels

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@build_backend.py` around lines 541 - 607, The current
_install_nvep_runtime_wheels installs runtime base-library wheels even during
PEP517 wheel builds (via calls from _prepare_for_wheel), causing distributed
wheels to miss libnccl/libnixl; change the flow so runtime wheels are only
installed for editable installs: modify _build_nvep_if_enabled and its callers
(_prepare_for_editable and _prepare_for_wheel) to accept and propagate a flag
like install_runtime (default False), call _install_nvep_runtime_wheels only
when install_runtime is True (i.e., from _prepare_for_editable), and leave
wheel-build paths skipping that function; alternatively, if you prefer shipping
runtime deps, add nvidia-nccl-cu13 and nixl-cu13 to the [nvep]
optional-dependencies in pyproject.toml instead (pick one approach and apply
consistently).

---

416-421: ⚠️ Potential issue | 🟠 Major | ⚡ Quick win

Fail the build if libnccl_ep.so was not staged.

The staging loop at Lines 416-420 silently completes even if libnccl_ep.so doesn't exist in the build output. This can mark built_nccl = True upstream while producing a package with no usable EP plugin.

Suggested fix

+    staged_count = 0
     for soname in ("libnccl_ep.so",):
         sopath = build / "lib" / soname
         if sopath.exists():
             shutil.copy(sopath, dst / soname)
+            staged_count += 1
             print(f"[BUILD_NVEP] staged: {soname}")
+    if staged_count == 0:
+        raise RuntimeError(
+            "NCCL-EP build completed but libnccl_ep.so was not produced. "
+            "Check the contrib/nccl_ep Makefile output for errors."
+        )

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@build_backend.py` around lines 416 - 421, The staging loop over soname
("libnccl_ep.so",) currently does nothing when the file is missing, so
downstream code (e.g., built_nccl flag) can be set incorrectly; update the block
that iterates soname/sopath to explicitly fail the build when libnccl_ep.so is
not found by raising an exception or calling sys.exit(1) with a clear message
(e.g., "libnccl_ep.so not found, failing build"), or alternatively return/raise
from the function that contains this loop; ensure you reference the same
variables (soname, sopath, build, dst) and adjust any logic that sets built_nccl
so it only becomes True after a successful copy.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Duplicate comments:
In `@build_backend.py`:
- Around line 245-248: The staging loop only searches (build /
"examples/device/ep") for "nixl_ep_cpp*.so" but the comment says the extension
may land in build/ as well; update the logic in build_backend.py (the loop that
uses (build / "examples/device/ep").glob(...) and copies to dst) to search both
locations — e.g., iterate over both build.glob("nixl_ep_cpp*.so") and (build /
"examples/device/ep").glob("nixl_ep_cpp*.so"), copy any matches to dst, and
avoid double-copying the same filename if found in both places.
- Around line 541-607: The current _install_nvep_runtime_wheels installs runtime
base-library wheels even during PEP517 wheel builds (via calls from
_prepare_for_wheel), causing distributed wheels to miss libnccl/libnixl; change
the flow so runtime wheels are only installed for editable installs: modify
_build_nvep_if_enabled and its callers (_prepare_for_editable and
_prepare_for_wheel) to accept and propagate a flag like install_runtime (default
False), call _install_nvep_runtime_wheels only when install_runtime is True
(i.e., from _prepare_for_editable), and leave wheel-build paths skipping that
function; alternatively, if you prefer shipping runtime deps, add
nvidia-nccl-cu13 and nixl-cu13 to the [nvep] optional-dependencies in
pyproject.toml instead (pick one approach and apply consistently).
- Around line 416-421: The staging loop over soname ("libnccl_ep.so",) currently
does nothing when the file is missing, so downstream code (e.g., built_nccl
flag) can be set incorrectly; update the block that iterates soname/sopath to
explicitly fail the build when libnccl_ep.so is not found by raising an
exception or calling sys.exit(1) with a clear message (e.g., "libnccl_ep.so not
found, failing build"), or alternatively return/raise from the function that
contains this loop; ensure you reference the same variables (soname, sopath,
build, dst) and adjust any logic that sets built_nccl so it only becomes True
after a successful copy.

In `@docker/Dockerfile.flashinfer-nvep`:
- Around line 46-49: The Dockerfile hardcodes "DOCA_v3.2.0" in the wget URL
which won't match other DOCA_VERSION values; update the RUN/wget invocation to
derive the path from the DOCA_VERSION ARG (e.g., interpolate DOCA_VERSION into
the URL or build a DOCA_URL ARG) so the segment "DOCA_v<version>" is generated
from DOCA_VERSION rather than hardcoded, ensuring the download path and the
DOCA_VERSION variable stay in sync.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: f0e4ad11-9441-4966-9577-d3c9e89fab9a

📥 Commits

Reviewing files that changed from the base of the PR and between 12f7a12 and ae54260.

📒 Files selected for processing (5)

• build_backend.py
• docker/Dockerfile.flashinfer-nvep
• flashinfer/moe_ep/__init__.py
• flashinfer/moe_ep/nccl_ep/__init__.py
• flashinfer/moe_ep/nixl_ep/__init__.py

[Anerudhan]

/bot run

[flashinfer-bot]

@Anerudhan is not authorized to trigger this CI job. cc: @yzh119, @sricketts, @yongwww

[aleozlx]

/bot run

[flashinfer-bot]

GitLab MR !683 has been created, and the CI pipeline #51615095 is currently running. I'll report back once the pipeline job completes.

[dierksen]

I can't really evaluate the intent/implementation of these dependencies, but the infra side LGTM.

[aleozlx]

wanna talk to Ane at 4:30 before merging

[aleozlx]

So these aren't submodules but copies ? Oh this is docker ignore file... misunderstood

[aleozlx]

What's this file for ? Patching at runtime ?

[aleozlx]

Does torch exposed cuda access not satisfy it?

[aleozlx]

Im fuzzy what the package-data is used for...

[aleozlx]

Turned out I didn't submit my comments from mobile they were pending ... but reviewed in a meeting everything good