Skip to content

feat: make moe_ep (EP) part of the default install; drop nccl submodule - #3821

Merged
Anerudhan merged 13 commits into
flashinfer-ai:mainfrom
Anerudhan:feat/nvep-default
Jul 8, 2026
Merged

Anerudhan merged 13 commits into
flashinfer-ai:mainfrom
Anerudhan:feat/nvep-default

Conversation

@Anerudhan

@Anerudhan Anerudhan commented Jul 3, 2026

Copy link
Copy Markdown
Collaborator

📌 Description

Makes the MoE Expert-Parallel (flashinfer.moe_ep) backends part of the default install: a plain pip install . now pulls the NCCL-EP runtime (nccl4py) and builds NIXL-EP, replacing the previous opt-in BUILD_NVEP=1 pip install -e ".[nvep]" flow. The dead 3rdparty/nccl submodule is removed.

Changes

  • Base deps: cuda-python>=13.0, nccl4py>=0.3.1, nvidia-nccl-cu13>=2.30.7 move from the [nvep] extra into requirements.txt. The [nvep] extra stays as an empty deprecated alias so existing commands keep working.
  • Build hook defaults (build_backend.py): NCCL-EP and NIXL-EP default ON with tri-state env flags:
    • unset → build best-effort (missing build deps like meson/UCX skip NIXL-EP with a warning instead of failing the install)
    • BUILD_NIXL_EP=1 → strict, missing deps abort the install (previous behavior)
    • BUILD_NVEP=0 / BUILD_NIXL_EP=0 / BUILD_NCCL_EP=0 → opt out
  • _ensure_nixl_wheel(): the hook pre-installs the nixl-cu13 wheel (--no-deps) that the default NIXL-EP build links against — no manual pre-install step.
  • Remove 3rdparty/nccl submodule: NCCL-EP has been provided by the released nccl4py wheel since nccl-ep-v0.1.0; nothing built from the submodule anymore.
  • CUDA-12 guard: validate_arch_for_backend() raises a clear error when torch is built for CUDA < 13 (EP wheels are cu13-only), instead of a cryptic dlopen failure.
  • Dockerfiles, install scripts, docs, test markers, and rebuild-hint messages updated for the new default.

🧪 Tests

Verified end-to-end in nvcr.io/nvidia/pytorch:26.05-py3 (CUDA 13.2, torch 2.12, B-series GPU):

  • Stock image: pip install --no-build-isolation . succeeds in ~2m; nccl4py 0.3.1 / nvidia-nccl-cu13 2.30.7 / nixl-cu13 1.3.0 installed automatically; NIXL-EP skipped gracefully (image ships UCX 1.20, NIXL v1.1.0 needs the UCX 1.21 device API); available_backends() == ['nccl_ep'], import nccl.ep OK.

  • Image + DOCA 3.2 gpunetio + UCX v1.21.x from source: pip install . builds and stages nixl_ep_cpp.so; available_backends() == ['nccl_ep', 'nixl_ep']. Smoke pass.

  • Flag-resolution matrix (8 env combinations of BUILD_NVEP/BUILD_NCCL_EP/BUILD_NIXL_EP) unit-tested against the hook.

  • Tests have been added / updated as needed

  • Documentation has been updated as needed

  • AI-assisted (Claude Code)

🤖 Generated with Claude Code

Summary by CodeRabbit

  • New Features
    • MoE EP backends are enabled by default, with GPU runtime dependencies included by default (the old nvep extra is now a deprecated alias).
    • Builds now better handle backend enablement via explicit per-backend options, including stricter behavior when BUILD_NIXL_EP=1 is set.
  • Bug Fixes
    • Improved compatibility validation for CUDA 13+ (including stronger NCCL version floors when applicable).
    • Added runtime safeguards for High Throughput token budgeting (with clearer error guidance).
  • Chores
    • Updated Docker/container build instructions, docs, and test messaging to reflect the new default EP installation flow.

@coderabbitai

coderabbitai Bot commented Jul 3, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

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

This PR updates MoE-EP packaging and build selection so NCCL-EP is wheel-provided, NIXL-EP builds by default with auto wheel installation, the 3rdparty/nccl submodule is removed, related install/docs are rewritten, and backend validation now checks CUDA 13 and NCCL version constraints.

Changes

MoE-EP build and packaging migration

Layer / File(s) Summary
Tri-state backend resolution and NIXL wheel install
build_backend.py, pyproject.toml, requirements.txt
Updates backend flag resolution, NIXL wheel pre-installation, build gating, and dependency metadata/comments for the new default-on NIXL-EP and wheel-based NCCL-EP model.
Container install flow and submodule wiring
.gitmodules, .dockerignore, .gitignore, docker/Dockerfile.flashinfer-ep-pytorch, docker/Dockerfile.flashinfer-nvep, docker/install/build_flashinfer_ep_pytorch.sh, scripts/build_in_container.sh, benchmarks/MoE_benchmarks.md, benchmarks/bench_moe_ep.py
Removes the NCCL submodule entry, rewrites build/install commands and comments to use editable installs and backend env flags, and updates benchmark/container references.
Backend docs and error messages
flashinfer/moe_ep/__init__.py, flashinfer/moe_ep/nccl_ep/__init__.py, flashinfer/moe_ep/nixl_ep/__init__.py, flashinfer/moe_ep/nixl_ep/fleet.py, tests/conftest.py, tests/moe_ep/smoke_nccl_ep.py
Rewords package docs and runtime error hints to describe default-enabled backends, nccl4py-based NCCL-EP, and revised rebuild/skip messaging.
CUDA-13 backend validation
flashinfer/moe_ep/_validators.py
Adds CUDA-major gating and NCCL floor checks for backend validation.
NCCL-EP fleet and handle limits
flashinfer/moe_ep/nccl_ep/fleet.py, flashinfer/moe_ep/nccl_ep/handle.py
Changes NCCL-EP import guidance, clamps HT token budgets, and adds a dispatch-time per-rank token limit check.
NIXL mock test gating
tests/moe_ep/nixl_ep/test_fleet_mock.py
Consolidates host/CUDA capability skips into a shared helper for the NIXL fleet mock tests.

Estimated code review effort: 4 (Complex) | ~45 minutes

Possibly related PRs

  • flashinfer-ai/flashinfer#3315: Both PRs directly modify build_backend.py and flashinfer/moe_ep/* to control NCCL-EP/NIXL-EP builds through environment flags and wheel/submodule handling.
  • flashinfer-ai/flashinfer#3453: Both PRs touch flashinfer/moe_ep/__init__.py and flashinfer/moe_ep/_validators.py, overlapping on backend documentation and validation logic.
  • flashinfer-ai/flashinfer#3686: Both PRs change build wiring around NCCL-EP wheel usage and the removal of in-tree 3rdparty/nccl build assumptions.

Suggested labels: op: moe

Suggested reviewers: yzh119, sricketts, yongwww, cyx-6, saltyminty, bkryu, yyihuang, kahyunnam, jimmyzho, nv-yunzheq, aleozlx, dhiraj113, jiahanc

🚥 Pre-merge checks | ✅ 5
✅ Passed checks (5 passed)
Check name Status Explanation
Title check ✅ Passed The title accurately summarizes the main change: default MoE EP install and removal of the NCCL submodule.
Description check ✅ Passed The description follows the template with a clear summary, change breakdown, and tests; optional related issues and some checklist items are not filled.
Docstring Coverage ✅ Passed No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
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.
✨ 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

Comment @coderabbitai help to get the list of available commands.

@gemini-code-assist gemini-code-assist Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Code Review

This pull request refactors the moe_ep build infrastructure to enable both NCCL-EP and NIXL-EP backends by default during standard installation, moving their runtime dependencies into the base requirements. Feedback highlights a potential issue with PEP 517 build isolation where pre-installed wheels may not persist in the target environment, and suggests defensive parsing of the CUDA version string to prevent runtime errors.

Important

The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.

Comment thread build_backend.py
Comment thread flashinfer/moe_ep/_validators.py Outdated

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🤖 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 94-100: The strict-mode check in _BUILD_NVEP_BEST_EFFORT is
incorrectly keyed off BUILD_NCCL_EP even though NCCL-EP only affects logging and
never goes through _gate_backend. Update the backend gating logic in
build_backend.py so only BUILD_NIXL_EP controls whether NIXL-EP build deps are
treated as hard errors, and keep BUILD_NCCL_EP out of that strictness decision.
Use the existing _tri_flag and _gate_backend flow to ensure BUILD_NCCL_EP=1
alone does not flip the NIXL-EP path into RuntimeError behavior.
🪄 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: 3cb94c7c-8bad-40aa-a0b1-fe556e2c89fa

📥 Commits

Reviewing files that changed from the base of the PR and between b5ac097 and 9d20bae68a7b0650bdbadc8aba07c4df8a57336b.

📒 Files selected for processing (21)
  • .dockerignore
  • .gitignore
  • .gitmodules
  • 3rdparty/nccl
  • benchmarks/MoE_benchmarks.md
  • benchmarks/bench_moe_ep.py
  • build_backend.py
  • docker/Dockerfile.flashinfer-ep-pytorch
  • docker/Dockerfile.flashinfer-nvep
  • docker/install/build_flashinfer_ep_pytorch.sh
  • flashinfer/moe_ep/__init__.py
  • flashinfer/moe_ep/_validators.py
  • flashinfer/moe_ep/nccl_ep/__init__.py
  • flashinfer/moe_ep/nccl_ep/fleet.py
  • flashinfer/moe_ep/nixl_ep/__init__.py
  • flashinfer/moe_ep/nixl_ep/fleet.py
  • pyproject.toml
  • requirements.txt
  • scripts/build_in_container.sh
  • tests/conftest.py
  • tests/moe_ep/smoke_nccl_ep.py
💤 Files with no reviewable changes (2)
  • 3rdparty/nccl
  • .gitmodules

Comment thread build_backend.py Outdated
@Anerudhan

Copy link
Copy Markdown
Collaborator Author

CI fix: aarch64 AOT job torch downgrade (Torch not compiled with CUDA enabled)

Root cause: the nvidia-nccl-cu13>=2.30.7 base dependency conflicted with torch's exact nvidia-nccl-cu13 pins (2.11.0 pins ==2.28.9, 2.12.1 pins ==2.29.7), so on the aarch64 AOT runner pip's resolver evicted the installed torch 2.11.0+cu126 and backtracked to torch-2.10.0 — whose aarch64 PyPI wheel is CPU-only (its CUDA deps are all platform_machine == "x86_64"-gated), producing AssertionError: Torch not compiled with CUDA enabled.

Fixed in 90e67d15:

  • nvidia-nccl-cu13 removed from base deps — torch's own cu13 pin supplies libnccl, so pip never has a reason to touch torch.
  • The build hook installs nvidia-nccl-cu13>=2.30.7 with --no-deps on source installs (same pattern as nixl-cu13; never enters the resolver; best-effort).
  • The >=2.30.7 floor is enforced at runtime where it actually matters: NCCL-EP Fleet construction on Blackwell (sm_100+), where group-create fails with older NCCL — with an actionable upgrade message. Version probed from the wheel metadata with a ncclGetVersion ctypes fallback; undeterminable versions never block.

Verified in nvcr.io/nvidia/pytorch:26.05-py3 on a sm_100 GPU: plain pip install . leaves torch untouched, nvidia-nccl-cu13 2.30.7 still lands via the hook, and the validator passes on real Blackwell.

🤖 Generated with Claude Code

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 1

🧹 Nitpick comments (2)
flashinfer/moe_ep/_validators.py (2)

46-74: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick win

Silent detection failure defeats the purpose of the check.

Both branches swallow all exceptions with except Exception: pass. If either probe has a real bug (wrong lib name, ctypes signature mismatch, importlib.metadata API change) rather than just "package/lib absent," _installed_nccl_version() silently returns None, and validate_arch_for_backend then skips the Blackwell NCCL-floor check entirely — reproducing exactly the cryptic nccl_ep.cc:1438 failure this code was added to prevent, with no diagnostic trail. Ruff also flags this (S110/BLE001).

Consider at least a logger.debug(...) on the exception before falling through, so unexpected failures are distinguishable from "no wheel installed" in support/debug logs.

🔍 Suggested diff
+import logging
+
+logger = logging.getLogger(__name__)
+
 def _installed_nccl_version() -> "tuple[int, int, int] | None":
     ...
     try:
         from importlib.metadata import version

         parts = version("nvidia-nccl-cu13").split(".")[:3]
         return tuple(int(p) for p in parts)  # type: ignore[return-value]
-    except Exception:
-        pass
+    except Exception:
+        logger.debug("nvidia-nccl-cu13 metadata probe failed", exc_info=True)
     try:
         import ctypes

         lib = ctypes.CDLL("libnccl.so.2")
         out = ctypes.c_int()
         if lib.ncclGetVersion(ctypes.byref(out)) == 0:
             code = out.value
             return (code // 10000, (code // 100) % 100, code % 100)
-    except Exception:
-        pass
+    except Exception:
+        logger.debug("ncclGetVersion ctypes probe failed", exc_info=True)
     return None
🤖 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/_validators.py` around lines 46 - 74, The silent exception
handling in _installed_nccl_version() hides real probe failures and causes
validate_arch_for_backend to skip the NCCL floor check without any clue why.
Update the exception handling in both the importlib.metadata and ctypes.CDLL
probes to log a debug message with the caught exception before falling through,
so unexpected issues are distinguishable from a genuinely missing NCCL
installation.

Source: Linters/SAST tools


77-123: 🎯 Functional Correctness | 🔵 Trivial | ⚡ Quick win

CUDA-13 gate and Blackwell NCCL-floor check look correct.

Logic order (CUDA major check → cuda.is_available() mock bypass → sm_90+ → Blackwell NCCL floor), the encoding math for ncclGetVersion, and the actionable error messages all line up with the documented rationale and the NcclEpFleet.__init__ call site.

One gap: test_fleet_mock.py's bypass_build_checks fixture mocks validate_arch_for_backend entirely, so this new CUDA-13/NCCL-floor branch logic doesn't appear to be exercised by a unit test in the provided context. Consider adding a focused test for _installed_nccl_version() and the two new MoEEpConfigError branches.

🤖 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/_validators.py` around lines 77 - 123, The new CUDA-13 and
Blackwell NCCL validation paths in validate_arch_for_backend are not covered
because bypass_build_checks in test_fleet_mock.py mocks the validator away. Add
focused unit tests around validate_arch_for_backend and _installed_nccl_version
to exercise the CUDA < 13 MoEEpConfigError path and the Blackwell nccl_ep
NCCL-floor MoEEpConfigError path, while keeping the existing mock/test bypass
behavior intact.
🤖 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 674-676: The NCCL floor upgrade in _ensure_nccl_floor() should not
run unconditionally during PEP 517 isolated builds, because it only affects the
temporary build environment. Add the same isolation check used for the NIXL-EP
path around the _ensure_nccl_floor() call in build_backend.py, and either skip
it under isolation or emit a clear warning with a manual-install hint so users
know how to install the required NCCL version themselves.

---

Nitpick comments:
In `@flashinfer/moe_ep/_validators.py`:
- Around line 46-74: The silent exception handling in _installed_nccl_version()
hides real probe failures and causes validate_arch_for_backend to skip the NCCL
floor check without any clue why. Update the exception handling in both the
importlib.metadata and ctypes.CDLL probes to log a debug message with the caught
exception before falling through, so unexpected issues are distinguishable from
a genuinely missing NCCL installation.
- Around line 77-123: The new CUDA-13 and Blackwell NCCL validation paths in
validate_arch_for_backend are not covered because bypass_build_checks in
test_fleet_mock.py mocks the validator away. Add focused unit tests around
validate_arch_for_backend and _installed_nccl_version to exercise the CUDA < 13
MoEEpConfigError path and the Blackwell nccl_ep NCCL-floor MoEEpConfigError
path, while keeping the existing mock/test bypass behavior intact.
🪄 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: 57906707-0037-4970-bd91-da5050b06a31

📥 Commits

Reviewing files that changed from the base of the PR and between ff3bfc3223a81f60c7170635b1c0eb8d7a9f75cf and 90e67d158be40cc7715af240c477b18b8bfaa758.

📒 Files selected for processing (4)
  • build_backend.py
  • flashinfer/moe_ep/_validators.py
  • pyproject.toml
  • requirements.txt
🚧 Files skipped from review as they are similar to previous changes (1)
  • pyproject.toml

Comment thread build_backend.py

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 Nitpick comments (1)
flashinfer/moe_ep/nccl_ep/fleet.py (1)

39-82: 🗄️ Data Integrity & Integration | 🔵 Trivial | 🏗️ Heavy lift

Hardcoded mirror of an external wheel constant risks drift.

_HT_MAX_SUPPORTED_TOKENS_PER_RANK = 8192 mirrors a build-time constant from the separately-versioned nccl4py/nccl_ep wheel. If a future wheel release changes MAX_SUPPORTED_TOKENS_PER_RANK, this code will silently diverge (over-clamping or under-clamping and re-exposing the abort this change is meant to prevent), with no version check tying the two together.

Consider probing the actual limit from the imported nccl.ep module (if exposed) at runtime, or at minimum asserting/logging the nccl4py version this constant was mirrored from so drift is detectable.

Please check if nccl.ep (from the nccl4py package) exposes MAX_SUPPORTED_TOKENS_PER_RANK or an equivalent introspectable constant/attribute that could replace this hardcoded mirror.

🤖 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/fleet.py` around lines 39 - 82, The hardcoded
`_HT_MAX_SUPPORTED_TOKENS_PER_RANK` mirror in `_clamp_ht_max_tokens` can drift
from the real `nccl_ep` wheel constant, so replace it with a runtime value from
the imported `nccl.ep` module if `MAX_SUPPORTED_TOKENS_PER_RANK` (or an
equivalent exported attribute) exists. Update the clamp logic and warning
message to use that introspected limit, and keep a safe fallback plus a
version/assertion check if the module does not expose it so drift is detectable.
🤖 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.

Nitpick comments:
In `@flashinfer/moe_ep/nccl_ep/fleet.py`:
- Around line 39-82: The hardcoded `_HT_MAX_SUPPORTED_TOKENS_PER_RANK` mirror in
`_clamp_ht_max_tokens` can drift from the real `nccl_ep` wheel constant, so
replace it with a runtime value from the imported `nccl.ep` module if
`MAX_SUPPORTED_TOKENS_PER_RANK` (or an equivalent exported attribute) exists.
Update the clamp logic and warning message to use that introspected limit, and
keep a safe fallback plus a version/assertion check if the module does not
expose it so drift is detectable.

ℹ️ Review info
⚙️ Run configuration

Configuration used: defaults

Review profile: CHILL

Plan: Pro

Run ID: 81147dab-8d66-4cc4-af9d-88da9de6654c

📥 Commits

Reviewing files that changed from the base of the PR and between 90e67d158be40cc7715af240c477b18b8bfaa758 and 60ff0fc1fc9a8a26856a65ed5d439ead1bf9d40f.

📒 Files selected for processing (2)
  • flashinfer/moe_ep/nccl_ep/fleet.py
  • flashinfer/moe_ep/nccl_ep/handle.py

@Anerudhan
Anerudhan requested a review from qiching as a code owner July 6, 2026 19:50
@Anerudhan

Copy link
Copy Markdown
Collaborator Author

/bot run tests/moe

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

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

@Anerudhan

Copy link
Copy Markdown
Collaborator Author

/bot run

@flashinfer-bot

Copy link
Copy Markdown
Collaborator

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

@Anerudhan
Anerudhan force-pushed the feat/nvep-default branch from a902182 to 2c351bf Compare July 7, 2026 17:38
@Anerudhan
Anerudhan enabled auto-merge (squash) July 7, 2026 17:42
Anerudhan and others added 13 commits July 7, 2026 11:55
Previously the EP transport backends required an opt-in install:
`BUILD_NVEP=1 pip install -e ".[nvep]"`. Now a plain `pip install .`
enables them by default:

- Move the [nvep] extra's runtime deps (cuda-python>=13.0, nccl4py>=0.3.1,
  nvidia-nccl-cu13>=2.30.7) into the base dependencies (requirements.txt).
  The [nvep] extra remains as an empty deprecated alias.
- build_backend.py: NCCL-EP and NIXL-EP default ON with tri-state env
  flags. Unset -> build best-effort (missing build deps skip the backend
  with a warning instead of failing the install). BUILD_NIXL_EP=1 ->
  strict (missing deps abort). BUILD_NVEP=0 / BUILD_NIXL_EP=0 /
  BUILD_NCCL_EP=0 -> opt out.
- New _ensure_nixl_wheel(): the build hook pre-installs the nixl-cu13
  wheel (--no-deps) that the default NIXL-EP build links against, so no
  manual pre-install step is needed.
- Remove the 3rdparty/nccl submodule: NCCL-EP has been provided by the
  released nccl4py wheel since nccl-ep-v0.1.0; nothing builds from the
  submodule anymore.
- validate_arch_for_backend(): raise a clear error when torch is built
  for CUDA < 13 (the EP wheels ship CUDA-13 binaries only).
- Update Dockerfiles, install scripts, docs, test markers, and rebuild
  hints for the new default.

Verified in nvcr.io/nvidia/pytorch:26.05-py3 (CUDA 13.2, torch 2.12):
- Stock image: `pip install .` succeeds; NIXL-EP is skipped gracefully
  (image UCX 1.20 lacks the 1.21 device API); backends = ['nccl_ep'].
- With DOCA 3.2 + UCX v1.21.x from source: backends =
  ['nccl_ep', 'nixl_ep'], nixl_ep_cpp.so staged. Smoke pass.

AI-assisted.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
- Key NIXL-EP strictness solely off BUILD_NIXL_EP: only NIXL-EP goes
  through _gate_backend (NCCL-EP has no build step), so an explicit
  BUILD_NCCL_EP=1 no longer forces NIXL-EP into strict mode (coderabbit).
- Detect PEP 517 isolated build envs (pip-build-env-* / uv builds-v0) and
  print a prominent warning that the NIXL-EP build needs
  --no-build-isolation, since hook-installed wheels don't persist into
  the target environment (gemini-code-assist).
- Parse torch.version.cuda defensively in validate_arch_for_backend so
  custom/nightly version strings can't crash the CUDA-13 check
  (gemini-code-assist).

AI-assisted.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
create_fleet now runs the CUDA-13 check in validate_arch_for_backend
(the EP runtime wheels ship CUDA-13 binaries only), so on CI runners
with a CUDA-12 torch these mocked tests failed in validation before
reaching the fake Buffer. Replace the bare torch.cuda.is_available()
skips with a _skip_unless_ep_capable() helper that also skips when
torch.version.cuda < 13.

Verified in nvcr.io/nvidia/pytorch:26.05-py3 (CUDA 13.2): all three
tests run and pass; helper matrix-tested against stubbed torch
versions (12.8 -> skip, 13.2/None/unparseable -> run, no CUDA -> skip).

AI-assisted.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
The nvidia-nccl-cu13>=2.30.7 base dependency broke the aarch64 AOT CI
job: torch's cu13 wheels pin nvidia-nccl-cu13 EXACTLY (2.11.0 ==2.28.9,
2.12.1 ==2.29.7), so the floor made pip's resolver evict the installed
CUDA torch and backtrack to the CPU-only aarch64 torch-2.10.0 wheel
("Torch not compiled with CUDA enabled").

torch's own cu13 pin supplies libnccl, so the base dep is only needed
for the B200 EP floor — which can't be expressed in metadata without
fighting torch's exact pins. Instead:

- requirements.txt: drop nvidia-nccl-cu13 (keep nccl4py, cuda-python)
  with a comment explaining why a floor there is a footgun.
- build_backend.py: new _ensure_nccl_floor() installs
  nvidia-nccl-cu13>=2.30.7 with --no-deps on source installs (mirrors
  the nixl-cu13 pattern; never enters the resolver). Best-effort:
  failures warn and defer to the runtime check.
- moe_ep/_validators.py: enforce the floor where it actually matters —
  NCCL-EP Fleet construction on Blackwell (sm_100+), where group-create
  fails with NCCL < 2.30.7. _installed_nccl_version() probes the
  nvidia-nccl-cu13 wheel metadata, falling back to ncclGetVersion via
  ctypes; undeterminable versions never block.

Verified in nvcr.io/nvidia/pytorch:26.05-py3 on a sm_100 GPU: plain
`pip install .` leaves torch untouched, the hook installs
nvidia-nccl-cu13 2.30.7 out-of-band, available_backends() has nccl_ep,
the validator passes on real Blackwell, and the moe_ep mock tests pass.
Floor logic matrix-tested (old NCCL on sm_100 raises with actionable
message; Hopper and unknown-version paths don't block).

AI-assisted.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
…KENS_PER_RANK

ncclEpCreateGroup asserts (SIGABRT, nccl_ep.cc:1253) when a HIGH_THROUGHPUT
group is created with max_dispatch_tokens_per_rank > MAX_SUPPORTED_TOKENS_PER_RANK
(build-time 8192 in the nccl4py wheel). vLLM sizes the HT fleet from
moe.max_num_tokens = scheduler max_num_batched_tokens (e.g. 16384), which tripped
the assert and aborted all EP ranks the moment the modular EP (DP-EP) path was
actually exercised — masked until now because the monolithic (TP-only) path never
builds the HT prepare/finalize. LL is unbounded and unaffected.

Clamp the HT fleet's max_tokens_per_rank to the cap (warn once) in NcclEpFleet so
group creation succeeds; clamp the stored FleetParams (not just GroupConfig) so the
handle's recv-buffer sizing agrees. Add a clear MoEEpConfigError guard in
_dispatch_ht for the case where a single forward genuinely dispatches more than the
cap per rank (previously the C++ abort / buffer overflow), pointing at
--max-num-batched-tokens.

AI-assisted (Claude Code): root-caused from an 8-rank per-rank stderr capture of
the SIGABRT under DP-EP on Pre-Nyx.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
… flow

vLLM creates a fresh NcclEpHandle every MoE layer x step (routing binds at
create_handle), so the existing NV_FI_EP_FAST_PATH per-handle caches never hit
and every forward paid ~149us of host time (EP_PROFILE_HOST, LL decode:
FFI descriptor builds 32.6+31.3us, handle setup 35.9us, allocs, plus ~45us of
C calls). At decode the GPU is host-paced, and nccl.ep's fused send+recv
dispatch kernel absorbs the resulting inter-rank lag as in-kernel spin (median
256us/launch, 33% of GPU time in the capped LL profile).

Anchor the caches on the long-lived Fleet instead:
- recv buffers (LL + HT) and the recv-count tensor (not re-zeroed across
  forwards; the dispatch metadata fully overwrites it -- the same contract the
  per-handle _FAST reuse relied on),
- static FFI descriptor tuples (DispatchOutputs/LayoutInfo/configs),
- a (data_ptr, dtype, shape)-keyed memo for per-call Tensor wraps, restricted
  to tensors <= 2 MiB: the nccl.ep Tensor wrapper keeps the torch tensor
  alive, so memoizing large prefill activations pinned GBs across allocator
  addresses and OOM'd at --gpu-memory-utilization 0.9 (small tensors are
  exactly the host-bound decode path this cache targets).
Also adds EP_PROFILE_HOST timers around handle create/destroy.

Host path 149 -> ~119us/layer/step; measured DP-EP throughput (Qwen3-30B-A3B,
8xB200, eager): LL 128/2048 5416 -> 5654 tok/s, HT 128/2048 3755 -> 3926,
HT 2048/128 44467 (stable). GSM8K through the DP-EP transport re-validated:
LL 0.8560/0.8976, HT 0.8567/0.8984 (flex/strict).

AI-assisted (Claude Code).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Runbook is now a faithful manual-reproduction script for the final numbers:
- pin clone refs (flashinfer feat/nvep-default >= fa09bc46, vLLM
  feat/flashinfer-ep-all2all >= ab1415e) with merge-base sanity checks;
  note that editable installs resolve /host clones at runtime (no image
  rebuild for the Python-only perf fixes)
- per-backend --max-num-batched-tokens handling everywhere (HT needs 8192,
  LL must leave it unset so the batched-DP 256 auto-cap engages)
- 3b' GSM8K-through-the-transport step with final expected scores
- 3c fixed to NP=256 (what the reference matrix used) + per-run sum helper
- reference-numbers section replaced with the final measured matrix
  (FI-HT ahead of DeepEP-HT 19-27% on 2 of 3 shapes; FI-LL within 4-12%)

Results doc: final same-day 4-backend matrix (1.1e) + the perf-iteration log
(1.1f: batched-DP cap membership, HT recv-trim, fleet host-path caches, the
ep.Tensor pinning OOM lesson, remaining decode levers). Integration doc: 0
headline updated to the transport-exercised results.

AI-assisted (Claude Code).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
CI pre-commit (full-repo mypy) fails with 'NcclEpFleet has no attribute
"_hot_cache"' -- the cross-handle host-path cache was injected dynamically
from NcclEpHandle.__init__ via getattr/setattr. Declare it in
NcclEpFleet.__init__ instead (where it belongs), clear it in
update_topology() (world-size change invalidates cached buffer shapes),
and have the handle reference it directly.

No behavior change on the hot path; mypy on both files now passes with the
pre-commit config.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The cuda-python>=13.0 base-dep floor bulldozes CUDA-12 environments. On the
H100 CI runner (cu12 container) pip reports

    nvshmem4py-cu12 requires cuda-python<=12.9, but you have cuda-python 13.3.1

and then re-resolves torch to the PyPI-default cu13 build, whose
cuda-toolkit[cudart] dependency drops libcudart.so.13 into an env that already
carries libcudart.so.12 (nvidia-cutlass-dsl-libs-cu12). cudnn-frontend's
cuda-pathfinder loader (pathfinder itself newly present via cuda-python 13)
then fails every graph.check_support() with
"RuntimeError: Multiple libcudart libraries found" -- the JIT Unittest (H100)
failures on PR flashinfer-ai#3821.

Base deps must not force a CUDA major (the same rule requirements.txt already
documents for nvidia-nccl-cu13). With >=12.0 a cu12 env keeps cuda-python 12.x
-- exactly the configuration main's CI already passes with -- while cu13 envs
resolve 13.x and get the full nccl.ep stack. moe_ep itself stays CUDA-13-only,
enforced with a clear error at runtime (moe_ep/_validators.py); availability
probing is find_spec-based so nothing imports nccl.ep on cu12.

scripts/build_in_container.sh keeps its cuda-python>=13.0 pin deliberately:
that is the single-CUDA cu13 dev-container flow (installs nixl-cu13 /
nvidia-nccl-cu13 explicitly).

AI-assisted (Claude Code): root-caused by diffing the installed-package sets
of the failing PR job vs main's passing H100 job.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
vllm_moe_ep_results_prenyx.md and vllm_moe_ep_runbook.md are internal
cluster-specific validation notes (Pre-Nyx runbook + measurement log); keep
them out of the upstream PR. Files remain in local working trees, untracked.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
integration.md referenced the internal Pre-Nyx validation notes
(vllm_moe_ep_results_prenyx.md / runbook) in 12 places; those files were
dropped from the PR in the previous commit. Remove the dangling links and
fold the load-bearing facts inline: the three perf root-cause fixes, the
offline-DP torchrun/external_launcher requirement, the canonical Dockerfile
build spec pointer, and the nccl_ep dispatch/combine kernel names to look
for in an nsys capture.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
…for job pips

unit_test_gb300 [cu130] fails with "Torch not compiled with CUDA enabled":

1. task_run_unit_tests.sh unconditionally runs `pip install nvshmem4py-cu12`.
   The package is already in the image, but re-resolving its deps applies its
   cuda-python<=12.9 pin, downgrading the cu130 container's cuda-python 13.0 /
   cuda-bindings 13.0.3 to 12.9 (pip warns: torch 2.11.0+cu130 requires
   cuda-bindings<14,>=13.0.3).
2. The subsequent `pip install -r requirements.txt` then can't satisfy both the
   installed torch's cuda-bindings>=13.0.3 requirement and cuda-python 12.9's
   cuda-bindings~=12.9.0 chain, so the resolver replaces torch, backtracking
   2.12.1 -> 2.12.0 -> 2.11.0 -> 2.10.0 -- and the plain-PyPI aarch64 torch
   wheel is CPU-only.

This was previously masked by the cuda-python>=13.0 base-dep floor, which
accidentally re-upgraded cuda-python during the requirements install; relaxing
that floor to >=12.0 (needed to stop cu12 environments from being bulldozed)
exposed the latent bug. Same disease as the earlier nvidia-nccl-cu13 floor
that evicted torch on aarch64; same cure:

- task_run_unit_tests.sh / task_test_single_node_comm_kernels.sh: install
  nvshmem4py-cu12 only when `import nvshmem.core` fails, and with --no-deps
  (the image ships the right-flavor cuda-python and nvidia-nvshmem).
- test_utils.sh / setup_test_env.sh (idempotent, whichever is sourced first):
  export PIP_CONSTRAINT pinning the preinstalled torch==<version> for every
  job-time pip install, so any future dep whose constraints would evict torch
  fails loudly at install time instead of silently degrading to a CPU wheel.
  Torch is not in [build-system].requires, so isolated build envs are
  unaffected.

AI-assisted (Claude Code): root-caused from the GitLab job log (job 356145796).

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
The PIP_CONSTRAINT guard pinned torch==<exact installed version> including
the +cuNNN local tag. PEP-517 build environments inherit PIP_CONSTRAINT, and
flashinfer-jit-cache's build-system.requires includes torch -- its isolated
build env then tries to resolve torch==2.11.0+cu130 from PyPI, where
local-version wheels do not exist: AOT Build Import (arm64, cu130) fails
with ResolutionImpossible ("no matching distributions available: torch").

Pin torch==<public version> instead: PEP 440 lets the installed
2.X.Y+cuNNN satisfy ==2.X.Y, so the main-env protection is unchanged (the
poisoned cuda-bindings scenario still fails loudly -- the plain PyPI wheel of
the same version carries the same cuda-bindings requirement), while build
envs resolve the plain wheel from PyPI.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
@Anerudhan
Anerudhan force-pushed the feat/nvep-default branch from 2c351bf to ae8d8f6 Compare July 7, 2026 18:55
@Anerudhan
Anerudhan requested a review from IwakuraRein as a code owner July 7, 2026 18:55
@Anerudhan
Anerudhan disabled auto-merge July 7, 2026 18:58
@Anerudhan
Anerudhan enabled auto-merge (squash) July 7, 2026 18:59
@Anerudhan
Anerudhan disabled auto-merge July 7, 2026 19:34
@Anerudhan
Anerudhan enabled auto-merge (squash) July 7, 2026 19:34
@Anerudhan
Anerudhan merged commit af979cc into flashinfer-ai:main Jul 8, 2026
46 of 50 checks passed
@mhoqueanik

Copy link
Copy Markdown
Collaborator

@Anerudhan

A suggested fix to address the nvep default cu13 env mutation issue.

fix-nvep-default-cu13-env-mutation.patch

Claude generated description

Created a fix for the environment-mutation bug in the EP default-install build hook (_ensure_nccl_floor() and friends): on this PR's head, every PEP-517 hook invocation — including plain metadata preparation (pip wheel, pip download) — installs nvidia-nccl-cu13>=2.30.7 and nixl-cu13 into the target environment, silently breaking torch's exact nvidia-nccl-cu13==2.29.7 pin and causing an install/repair thrash cycle where each pip install undoes the previous one (on aarch64, pip can backtrack torch to the CPU-only wheel). The uv pip install --python sys.executable path also escapes PEP-517 build isolation, and _detect_cuda_major() defaulting to 13 when nvcc is absent injects cu13 wheels into cu12 environments.

The new fix does four things:

  1. Deletes _ensure_nccl_floor() — a build hook must never mutate the target env's resolved dependency set. The B200 >=2.30.7 floor is already enforced at runtime by moe_ep/_validators.py with an actionable error, which is the only place it matters.
  2. Makes _detect_cuda_major() derive from torch.version.cuda (nvcc as fallback) and return None instead of guessing 13; wheel installs are skipped when the major is unknown, so cu12 envs never get cu13 packages. _ensure_nixl_wheel() also refuses to run from inside an isolated build env instead of escaping it.
  3. Makes the metadata hooks side-effect free — get_requires_for_build_* / prepare_metadata_for_build_* no longer trigger the EP build or any --no-deps installs; those only run from build_wheel / build_editable.
  4. Restores the full +cuXXX tag in the CI torch pin (the stripped torch==2.X.Y pin lets pip silently swap to the PyPI-default flavor), and fixes the build-env ResolutionImpossible that motivated stripping it by exporting the matching download.pytorch.org/whl/ index as
    PIP_EXTRA_INDEX_URL.

Affected files: build_backend.py, scripts/setup_test_env.sh, scripts/test_utils.sh.

Validated in the pt26.05 cu13 container: on the unpatched head, metadata prep alone mutated the env (ran the NCCL floor install, added nixl-cu13==1.3.2); with the fix, metadata prep leaves pip freeze byte-identical, a full pip install --no-build-isolation . leaves torch/nvidia-nccl-cu13 untouched, and the moe_ep unit subset passes (test_config.py + test_constraints.py: 21 passed / 1 skipped; test_layer_single_gpu.py: 2 passed).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

5 participants