feat: make moe_ep (EP) part of the default install; drop nccl submodule - #3821
Conversation
|
Note Reviews pausedIt 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 Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughThis PR updates MoE-EP packaging and build selection so NCCL-EP is wheel-provided, NIXL-EP builds by default with auto wheel installation, the ChangesMoE-EP build and packaging migration
Estimated code review effort: 4 (Complex) | ~45 minutes Possibly related PRs
Suggested labels: Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
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.gitmodules3rdparty/ncclbenchmarks/MoE_benchmarks.mdbenchmarks/bench_moe_ep.pybuild_backend.pydocker/Dockerfile.flashinfer-ep-pytorchdocker/Dockerfile.flashinfer-nvepdocker/install/build_flashinfer_ep_pytorch.shflashinfer/moe_ep/__init__.pyflashinfer/moe_ep/_validators.pyflashinfer/moe_ep/nccl_ep/__init__.pyflashinfer/moe_ep/nccl_ep/fleet.pyflashinfer/moe_ep/nixl_ep/__init__.pyflashinfer/moe_ep/nixl_ep/fleet.pypyproject.tomlrequirements.txtscripts/build_in_container.shtests/conftest.pytests/moe_ep/smoke_nccl_ep.py
💤 Files with no reviewable changes (2)
- 3rdparty/nccl
- .gitmodules
CI fix: aarch64 AOT job torch downgrade (
|
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
flashinfer/moe_ep/_validators.py (2)
46-74: 🩺 Stability & Availability | 🔵 Trivial | ⚡ Quick winSilent 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.metadataAPI change) rather than just "package/lib absent,"_installed_nccl_version()silently returnsNone, andvalidate_arch_for_backendthen skips the Blackwell NCCL-floor check entirely — reproducing exactly the crypticnccl_ep.cc:1438failure 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 winCUDA-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 forncclGetVersion, and the actionable error messages all line up with the documented rationale and theNcclEpFleet.__init__call site.One gap:
test_fleet_mock.py'sbypass_build_checksfixture mocksvalidate_arch_for_backendentirely, 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 newMoEEpConfigErrorbranches.🤖 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.pyflashinfer/moe_ep/_validators.pypyproject.tomlrequirements.txt
🚧 Files skipped from review as they are similar to previous changes (1)
- pyproject.toml
There was a problem hiding this comment.
🧹 Nitpick comments (1)
flashinfer/moe_ep/nccl_ep/fleet.py (1)
39-82: 🗄️ Data Integrity & Integration | 🔵 Trivial | 🏗️ Heavy liftHardcoded mirror of an external wheel constant risks drift.
_HT_MAX_SUPPORTED_TOKENS_PER_RANK = 8192mirrors a build-time constant from the separately-versionednccl4py/nccl_epwheel. If a future wheel release changesMAX_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.epmodule (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 thenccl4pypackage) exposesMAX_SUPPORTED_TOKENS_PER_RANKor 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.pyflashinfer/moe_ep/nccl_ep/handle.py
|
/bot run tests/moe |
|
@Anerudhan is not authorized to trigger this CI job. cc: @yzh119, @sricketts, @yongwww |
|
/bot run |
|
@Anerudhan is not authorized to trigger this CI job. cc: @yzh119, @sricketts, @yongwww |
a902182 to
2c351bf
Compare
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>
2c351bf to
ae8d8f6
Compare
|
A suggested fix to address the nvep default cu13 env mutation issue. fix-nvep-default-cu13-env-mutation.patch Claude generated descriptionCreated 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:
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). |
📌 Description
Makes the MoE Expert-Parallel (
flashinfer.moe_ep) backends part of the default install: a plainpip install .now pulls the NCCL-EP runtime (nccl4py) and builds NIXL-EP, replacing the previous opt-inBUILD_NVEP=1 pip install -e ".[nvep]"flow. The dead3rdparty/ncclsubmodule is removed.Changes
cuda-python>=13.0,nccl4py>=0.3.1,nvidia-nccl-cu13>=2.30.7move from the[nvep]extra intorequirements.txt. The[nvep]extra stays as an empty deprecated alias so existing commands keep working.build_backend.py): NCCL-EP and NIXL-EP default ON with tri-state env flags: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 thenixl-cu13wheel (--no-deps) that the default NIXL-EP build links against — no manual pre-install step.3rdparty/ncclsubmodule: NCCL-EP has been provided by the releasednccl4pywheel since nccl-ep-v0.1.0; nothing built from the submodule anymore.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.🧪 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.0installed 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.epOK.Image + DOCA 3.2 gpunetio + UCX v1.21.x from source:
pip install .builds and stagesnixl_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
nvepextra is now a deprecated alias).BUILD_NIXL_EP=1is set.