Ep api design -- Adding the actual code and tests - #3453
Merged
Merged
Conversation
Anerudhan
requested review from
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jimmyzho,
kahyunnam,
nv-yunzheq,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
May 29, 2026 06:16
Contributor
|
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:
📝 WalkthroughWalkthroughImplements a pluggable MoE Expert‑Parallel subsystem: typed configs and algo‑knobs, validators, abstract Fleet/Handle contracts, NCCL‑EP and NIXL‑EP backend implementations, MoEEpLayer/tensors, unit/integration tests, and build/test scripts. ChangesMoE-EP Subsystem Implementation
Sequence Diagram — MoEEpLayer forward (high-level): sequenceDiagram
participant Caller
participant Layer as MoEEpLayer
participant Fleet as Fleet(backend)
participant Handle as Handle(per-call)
Caller->>Layer: forward(MoEEpTensors)
Layer->>Layer: _ensure_fleet()
Layer->>Fleet: create_fleet(...) / create_handle(...)
Layer->>Handle: dispatch(tokens)
Handle->>Layer: DispatchOutput(expert_tensors, num_tokens)
Layer->>Layer: _inner_compute_identity(expert_tensors)
Layer->>Handle: combine(expert_outputs, weights)
Layer->>Handle: complete()
Handle->>Caller: CombineOutput(x=combined_tensor)
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ 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 |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request introduces a backend-agnostic Expert-Parallel (EP) layer (MoEEpLayer) supporting both NCCL-EP and NIXL-EP backends, along with associated configurations, validators, typed algorithm knobs, and tests. The review feedback highlights several critical bugs and robustness improvements: 1) fixing bfloat16 tensor reconstruction in NDTensor by mapping it to a supported 2-byte integer type and casting the view, 2) caching the loaded NCCL library reference to avoid repeated lookups and prevent interpreter shutdown crashes, 3) ensuring num_experts is divisible by the topology capacity in NixlEpFleet and importing the missing MoEEpConfigError, and 4) expanding the use_fp8 check to cover other FP8 variants like FP8E5M2 and NVFP8.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 7
🧹 Nitpick comments (2)
scripts/build_in_container.sh (1)
127-127: ⚡ Quick winRemove the no-op loop.
This loop only runs
:and never usessp; the actual subproject path is (re)assigned toNIXL_SPon the next line. It's a leftover/no-op (and the source of the SC2034 warning). Safe to delete.♻️ Proposed cleanup
-for sp in /tmp /var/tmp "${REPO_ROOT}/3rdparty/nixl/subprojects"; do : ; done NIXL_SP="${REPO_ROOT}/3rdparty/nixl/subprojects"🤖 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 `@scripts/build_in_container.sh` at line 127, Delete the no-op for-loop "for sp in /tmp /var/tmp \"${REPO_ROOT}/3rdparty/nixl/subprojects\"; do : ; done" because it only executes ':' and never uses the loop variable; instead keep the subsequent assignment that sets NIXL_SP (the real subproject path), removing the redundant loop to eliminate the SC2034 warning and dead code.scripts/task_test_moe_ep_smoke.sh (1)
16-41: 💤 Low valueOptional: validate
BACKENDearly.An unexpected
BACKENDvalue (e.g. a typo) silently skips both smoke blocks (Lines 22-32) and then forwards--backend=<typo>to pytest at Line 40, surfacing as a confusing pytest error rather than a clear message. A small guard up front would fail fast.🤖 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 `@scripts/task_test_moe_ep_smoke.sh` around lines 16 - 41, The script does not validate the BACKEND variable, so typos silently skip the smoke_nccl_ep/smoke_nixl_ep branches and later pass an invalid --backend to pytest; add an early guard after BACKEND is set that checks BACKEND is one of: "nccl_ep", "nixl_ep", or "both" (referencing the BACKEND variable and the for BE loop and the smoke_nccl_ep / smoke_nixl_ep blocks), and if not, echo a clear error and exit non‑zero to fail fast; implement the check before the smoke blocks so invalid values never reach torchrun or pytest.
🤖 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 `@flashinfer/moe_ep/__init__.py`:
- Around line 63-96: Add the missing re-exports validate_arch_for_backend and
validate_fleet_params to the module's public API by including their names in the
__all__ list in __init__.py so they are exported and no longer flagged as
unused; update the existing __all__ array (which already contains symbols like
create_fleet and have_nccl_ep) to also contain "validate_arch_for_backend" and
"validate_fleet_params".
In `@flashinfer/moe_ep/_validators.py`:
- Around line 43-46: The capability check is using
torch.cuda.get_device_capability(0) which hardcodes GPU 0; update both places
where this appears (in validate_arch_for_backend and the nixl_ep UE8M0 check) to
call torch.cuda.get_device_capability(torch.cuda.current_device()) (or
equivalent using torch.cuda.current_device()) so the check uses the rank’s
current CUDA device before raising MoEEpArchError.
In `@flashinfer/moe_ep/nccl_ep/fleet.py`:
- Around line 150-162: update_topology refreshes self._fleet_knobs via
_index_knobs(algo_knobs) but then reuses the stale self._cfg when calling
lib.ncclEpCreateGroup, so config-affecting knobs are ignored; fix by
reconstructing the cfg from the updated knobs before calling ncclEpCreateGroup
(i.e., after setting self._fleet_knobs call the same cfg-construction routine
used initially to produce self._cfg so fields like FleetAlgoKnobRdmaBufferSize /
FleetAlgoKnobNumQpsPerRank / FleetAlgoKnobNumChannelsPerRank reflect the new
knobs), or alternatively add explicit documentation on update_topology (or
method name that triggers this diff) stating those knobs are immutable and will
not affect self._cfg; ensure the change touches the same block that sets
self._fleet_knobs, self._cfg, and the call to lib.ncclEpCreateGroup so the group
is created with the rebuilt cfg.
In `@flashinfer/moe_ep/nccl_ep/handle.py`:
- Around line 186-194: The host read of recv_count_t.sum().item() must be
synchronized with the CUDA stream stored in self._stream; add an explicit wait
on self._stream before calling self._recv_count_t.sum().item() (e.g., record a
CUDA event on self._stream and synchronize it on the host or call
current_stream().wait_stream(...) with the External/torch Stream wrapping
self._stream) so num_tokens is not read before ncclEpDispatch/ncclEpComplete
populate _recv_count_t; update the code in the block after lib.ncclEpComplete
(where num_tokens is computed) to perform this wait, or alternatively document
that callers must ensure the current stream equals NcclEpHandle._stream.
In `@flashinfer/moe_ep/nixl_ep/fleet.py`:
- Around line 92-110: The code computes cap from FleetAlgoKnobTopologyCapacity
and then sets num_experts_per_rank = params.num_experts // cap before calling
self._buffer.update_memory_buffers(...) and self._buffer.connect_ranks(...),
which can silently drop experts if cap does not divide num_experts; add a
validation step (before using cap) that checks params.num_experts % cap == 0 and
cap <= params.num_experts (or equivalently num_experts_per_rank > 0) and raise a
clear error if the check fails, or alternatively handle the remainder explicitly
(e.g., distribute extra experts or fail fast); update references around
FleetAlgoKnobTopologyCapacity, num_experts_per_rank, update_memory_buffers, and
connect_ranks to ensure the chosen behavior is enforced.
In `@tests/moe_ep/nixl_ep/test_fleet_mock.py`:
- Around line 198-208: The test is incorrect to rely on divisibility checks;
instead ensure the fleet's topology capacity accommodates added ranks—update the
test to construct the NixlEpFleet with a topology capacity at least 6 (e.g.,
pass a FleetAlgoKnobTopologyCapacity or equivalent knob to the Bootstrap/Fleet
creation) so that NixlEpFleet.__init__ sets cap >= 6 and the buffer sizing
matches the later call to update_topology which triggers connect_ranks([4, 5]);
leave validate_fleet_params, update_topology, connect_ranks and disconnect_ranks
behavior unchanged.
In `@tests/moe_ep/test_layer_single_gpu.py`:
- Around line 49-53: The fixture stores a stub in _BACKEND_REGISTRY under key
"nccl_ep" but only restores it when saved is not None, so when saved is None the
stub remains; change the teardown to check saved and if saved is None remove the
"nccl_ep" key from _BACKEND_REGISTRY (otherwise restore saved) so _StubFleet
isn't leaked — update the block that uses saved/_BACKEND_REGISTRY/"nccl_ep" to
delete the key when saved is None.
---
Nitpick comments:
In `@scripts/build_in_container.sh`:
- Line 127: Delete the no-op for-loop "for sp in /tmp /var/tmp
\"${REPO_ROOT}/3rdparty/nixl/subprojects\"; do : ; done" because it only
executes ':' and never uses the loop variable; instead keep the subsequent
assignment that sets NIXL_SP (the real subproject path), removing the redundant
loop to eliminate the SC2034 warning and dead code.
In `@scripts/task_test_moe_ep_smoke.sh`:
- Around line 16-41: The script does not validate the BACKEND variable, so typos
silently skip the smoke_nccl_ep/smoke_nixl_ep branches and later pass an invalid
--backend to pytest; add an early guard after BACKEND is set that checks BACKEND
is one of: "nccl_ep", "nixl_ep", or "both" (referencing the BACKEND variable and
the for BE loop and the smoke_nccl_ep / smoke_nixl_ep blocks), and if not, echo
a clear error and exit non‑zero to fail fast; implement the check before the
smoke blocks so invalid values never reach torchrun or pytest.
🪄 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: bd0f178e-fb47-41aa-a26d-b22ac964a846
📥 Commits
Reviewing files that changed from the base of the PR and between c5a2b06 and 1ed4f627d945f6192798386c0a5069f036514c24.
📒 Files selected for processing (32)
flashinfer/moe_ep/__init__.pyflashinfer/moe_ep/_validators.pyflashinfer/moe_ep/algo_knobs.pyflashinfer/moe_ep/config.pyflashinfer/moe_ep/fleet.pyflashinfer/moe_ep/handle.pyflashinfer/moe_ep/layer.pyflashinfer/moe_ep/nccl_ep/fleet.pyflashinfer/moe_ep/nccl_ep/handle.pyflashinfer/moe_ep/nccl_ep/ndtensor.pyflashinfer/moe_ep/nixl_ep/fleet.pyflashinfer/moe_ep/nixl_ep/handle.pyflashinfer/moe_ep/split_backends/__init__.pyflashinfer/moe_ep/split_backends/nccl_ep_comm.pyflashinfer/moe_ep/split_backends/nixl_ep_comm.pyflashinfer/moe_ep/tensors.pyflashinfer/moe_ep/tests/__init__.pyflashinfer/moe_ep/tests/smoke_nccl_ep.pyflashinfer/moe_ep/tests/smoke_nixl_ep.pyscripts/build_in_container.shscripts/task_test_moe_ep_smoke.shtests/conftest.pytests/moe_ep/__init__.pytests/moe_ep/nccl_ep/__init__.pytests/moe_ep/nccl_ep/test_fleet_mock.pytests/moe_ep/nccl_ep/test_ndtensor.pytests/moe_ep/nixl_ep/__init__.pytests/moe_ep/nixl_ep/test_fleet_mock.pytests/moe_ep/test_config.pytests/moe_ep/test_constraints.pytests/moe_ep/test_layer_single_gpu.pytests/moe_ep/test_moe_ep_layer_multirank.py
3 tasks
aleozlx
added a commit
to aleozlx/flashinfer
that referenced
this pull request
May 31, 2026
…ashinfer-ai#3453) Cite the coordinating PR flashinfer-ai#3453 and name the integration seam: moe_ep's MoEEpLayer.forward (dispatch -> inner_compute -> combine) is where this unified fused_moe path becomes the per-rank expert compute (inner_compute is identity today). Detailed coordination review kept out of the doc (in var/log) to avoid polluting it. AI-assisted: drafted with Claude Code. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aleozlx
reviewed
Jun 1, 2026
aleozlx
reviewed
Jun 1, 2026
aleozlx
reviewed
Jun 1, 2026
aleozlx
reviewed
Jun 1, 2026
aleozlx
reviewed
Jun 1, 2026
aleozlx
reviewed
Jun 1, 2026
aleozlx
reviewed
Jun 1, 2026
Member
h100 test above flagged an issue (click to expand)The fix is to make these mocked NCCL-EP tests run without requiring a built native Root causeAll 4 failures come from
These tests already mock:
But they do not mock the package-level build probe, so the backend is rejected before the mocks are used. Best solutionPatch the tests to override The cleanest approach is a fixture that patches Suggested code changeUpdate import ctypes
from unittest import mock
import pytest
@pytest.fixture
def fake_nccl_ep_built():
import flashinfer.moe_ep as moe_ep
with mock.patch.object(moe_ep, "_probe_nccl_ep", return_value=True):
yieldThen add this fixture to the affected tests: def test_fleet_init_populates_group_config(
fake_nccl_ep_module, fake_nccl_ep_built, patched_lib
):
...def test_handle_create_passes_use_fp8(
fake_nccl_ep_module, fake_nccl_ep_built, patched_lib
):
...def test_dispatch_round_scales_from_ue8m0(
fake_nccl_ep_module, fake_nccl_ep_built, patched_lib
):
...def test_complete_called_internally_after_dispatch(
fake_nccl_ep_module, fake_nccl_ep_built, patched_lib
):
...Even better: patch
|
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
scripts/task_test_moe_ep_smoke.sh (1)
7-7:⚠️ Potential issue | 🟡 Minor | ⚡ Quick win
NPROCdefault of 8 conflicts with thegpu_4multi-rank marker.The multi-rank stage selects
-m "nvep and gpu_4", which targets a 4-GPU tray, yet it launches--nproc_per_node="${NPROC}"with the defaultNPROC=8. On a single 4-GPU node this oversubscribes devices (8 ranks → 4 GPUs), and the smoke entrypoints pin one device per rank viaLOCAL_RANK, so ranks will collide ontorch.cuda.set_device.The header comment is also self-contradictory: "GPUs per node (default 8 ... = 2 nodes × 4)" mixes a per-node count with a 2-node total —
--nproc_per_nodeis per-node, so 8 means 16 ranks across 2 nodes. Please clarify the intended per-node value (likely 4 for thegpu_4stage) or document the multi-node launch assumptions.Also applies to: 39-41
🤖 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 `@scripts/task_test_moe_ep_smoke.sh` at line 7, The NPROC default of 8 oversubscribes a 4-GPU tray used by the multi-rank stage (-m "nvep and gpu_4") and conflicts with per-node semantics; update the script to set NPROC=4 (per-node GPUs) or make the multi-rank selector match NPROC, and fix the header comment to state "NPROC = GPUs per node (default 4 for single 4‑GPU node)" so --nproc_per_node="${NPROC}" is correct; ensure callers that rely on LOCAL_RANK/torch.cuda.set_device will only get one rank per GPU by aligning NPROC with the gpu_4 marker or documenting multi-node expectations (2 nodes => total ranks = 2 * NPROC).tests/moe_ep/nccl_ep/test_fleet_mock.py (1)
149-155:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winArch gate still trips on unsupported GPUs — add an sm_90+ skip to these nccl_ep mock tests.
NcclEpFleet.__init__callsvalidate_arch_for_backend("nccl_ep")immediately after_require_built;validate_arch_for_backendonly “skips” whentorch.cuda.is_available()is false, but raisesMoEEpArchErrorwhen CUDA is present yetget_device_capability(0) < (9, 0). The fixture bypasses the build probe only, so heterogeneous CUDA shards with older GPUs can hard-fail these tests.Proposed guard (apply to all four tests)
import torch + from flashinfer.utils import get_compute_capability if not torch.cuda.is_available(): pytest.skip("needs CUDA") + if get_compute_capability(torch.device("cuda:0")) < (9, 0): + pytest.skip("nccl_ep backend requires sm_90+")This applies to
test_fleet_init_populates_group_config,test_handle_create_passes_use_fp8,test_dispatch_round_scales_from_ue8m0, andtest_complete_called_internally_after_dispatch.🤖 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 `@tests/moe_ep/nccl_ep/test_fleet_mock.py` around lines 149 - 155, The tests need an additional GPU-arch guard: after the existing torch.cuda.is_available() check in each of the four tests (test_fleet_init_populates_group_config, test_handle_create_passes_use_fp8, test_dispatch_round_scales_from_ue8m0, test_complete_called_internally_after_dispatch) call torch.cuda.get_device_capability(0) and if the returned capability is less than (9, 0) call pytest.skip("requires sm_90+ GPU"); place this check immediately after the CUDA availability skip so the mock nccl_ep initialization doesn't raise MoEEpArchError on older hardware.
🧹 Nitpick comments (1)
tests/moe_ep/smoke_nixl_ep.py (1)
32-40: ⚖️ Poor tradeoffSibling-port derivation can collide on shared nodes.
nixl_port = master_port + 1assumes the next port is free. On a node packed with multiple jobs/process-groups this can clash with another process' rendezvous store and cause a flaky bind failure. For an on-cluster smoke harness this is usually fine, but consider deriving the port from a job-unique value (e.g., a hash ofMASTER_ADDR:MASTER_PORTmapped into an ephemeral range, or an env override likeNIXL_RENDEZVOUS_PORT) to reduce the collision surface.🤖 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 `@tests/moe_ep/smoke_nixl_ep.py` around lines 32 - 40, The current fixed sibling-port logic (nixl_port = master_port + 1) can collide; change nixl_port derivation to first check an environment override (e.g., NIXL_RENDEZVOUS_PORT) and if not present derive a deterministic, job-unique port by hashing the MASTER_ADDR:MASTER_PORT string into an ephemeral port range (e.g., map hash % (max_port - min_port) + min_port), then use that computed nixl_port when constructing the dist.TCPStore (referencing nixl_port, master_port, MASTER_ADDR/MASTER_PORT, NIXL_RENDEZVOUS_PORT, and dist.TCPStore/TCPStore constructor).
🤖 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 `@tests/moe_ep/nccl_ep/test_fleet_mock.py`:
- Around line 30-42: The mocked NcclEP tests currently only skip based on
torch.cuda.is_available() but still invoke validate_arch_for_backend("nccl_ep")
(via NcclEpFleet.__init__) which raises MoEEpArchError on unsupported GPU
compute capability; update the four tests that use the bypass_moe_ep_build_check
fixture to also check GPU arch before running by using
flashinfer.utils.get_compute_capability(torch.device("cuda")) or the helper
is_sm90a_supported(...), and skip the tests unless CUDA is available AND the
compute capability meets the required threshold (i.e., only proceed when both
torch.cuda.is_available() and the compute-capability check returns true).
---
Outside diff comments:
In `@scripts/task_test_moe_ep_smoke.sh`:
- Line 7: The NPROC default of 8 oversubscribes a 4-GPU tray used by the
multi-rank stage (-m "nvep and gpu_4") and conflicts with per-node semantics;
update the script to set NPROC=4 (per-node GPUs) or make the multi-rank selector
match NPROC, and fix the header comment to state "NPROC = GPUs per node (default
4 for single 4‑GPU node)" so --nproc_per_node="${NPROC}" is correct; ensure
callers that rely on LOCAL_RANK/torch.cuda.set_device will only get one rank per
GPU by aligning NPROC with the gpu_4 marker or documenting multi-node
expectations (2 nodes => total ranks = 2 * NPROC).
In `@tests/moe_ep/nccl_ep/test_fleet_mock.py`:
- Around line 149-155: The tests need an additional GPU-arch guard: after the
existing torch.cuda.is_available() check in each of the four tests
(test_fleet_init_populates_group_config, test_handle_create_passes_use_fp8,
test_dispatch_round_scales_from_ue8m0,
test_complete_called_internally_after_dispatch) call
torch.cuda.get_device_capability(0) and if the returned capability is less than
(9, 0) call pytest.skip("requires sm_90+ GPU"); place this check immediately
after the CUDA availability skip so the mock nccl_ep initialization doesn't
raise MoEEpArchError on older hardware.
---
Nitpick comments:
In `@tests/moe_ep/smoke_nixl_ep.py`:
- Around line 32-40: The current fixed sibling-port logic (nixl_port =
master_port + 1) can collide; change nixl_port derivation to first check an
environment override (e.g., NIXL_RENDEZVOUS_PORT) and if not present derive a
deterministic, job-unique port by hashing the MASTER_ADDR:MASTER_PORT string
into an ephemeral port range (e.g., map hash % (max_port - min_port) +
min_port), then use that computed nixl_port when constructing the dist.TCPStore
(referencing nixl_port, master_port, MASTER_ADDR/MASTER_PORT,
NIXL_RENDEZVOUS_PORT, and dist.TCPStore/TCPStore constructor).
🪄 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: c5c5ad68-29c6-4c8b-8e22-d6776a73dbd9
📥 Commits
Reviewing files that changed from the base of the PR and between 1ed4f627d945f6192798386c0a5069f036514c24 and 96a7cee61707af58c802b08bcba0b02cdfeb8e62.
📒 Files selected for processing (13)
flashinfer/moe_ep/fleet.pyflashinfer/moe_ep/layer.pyflashinfer/moe_ep/nccl_ep/fleet.pyflashinfer/moe_ep/nccl_ep/handle.pyflashinfer/moe_ep/nixl_ep/fleet.pyflashinfer/moe_ep/nixl_ep/handle.pyflashinfer/moe_ep/split_backends/nccl_ep_comm.pyflashinfer/moe_ep/split_backends/nixl_ep_comm.pyscripts/task_test_moe_ep_smoke.shtests/moe_ep/nccl_ep/test_fleet_mock.pytests/moe_ep/nixl_ep/test_fleet_mock.pytests/moe_ep/smoke_nccl_ep.pytests/moe_ep/smoke_nixl_ep.py
🚧 Files skipped from review as they are similar to previous changes (7)
- flashinfer/moe_ep/fleet.py
- flashinfer/moe_ep/layer.py
- tests/moe_ep/nixl_ep/test_fleet_mock.py
- flashinfer/moe_ep/nccl_ep/handle.py
- flashinfer/moe_ep/nixl_ep/handle.py
- flashinfer/moe_ep/nixl_ep/fleet.py
- flashinfer/moe_ep/nccl_ep/fleet.py
Introduce the public Fleet / Handle / create_fleet surface from the
EP API design (§5.4). Backends register at module import time in
_BACKEND_REGISTRY; create_fleet() looks them up.
create_fleet() accepts either a string backend name ("nccl_ep" /
"nixl_ep") or a config object exposing a `.backend_name` str attribute,
so MoEEpLayer can take an NcclEpConfig / NvepConfig instance directly
in a later step.
Verified: `from flashinfer.moe_ep import Fleet, Handle, create_fleet`
imports cleanly; `create_fleet(None, None, backend='nope')` raises
`KeyError("unknown backend 'nope'; available: [])` until the
backend-register steps land.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Adds the dataclasses + enums + AlgoKnob types the Fleet / Handle surface needs: * flashinfer/moe_ep/config.py — BootstrapConfig (per-Fleet inputs), FleetParams (durable sizing), HandleParams (per-iteration topk), Dispatch/Combine I/O envelopes, EpAlgorithm + QuantType enums. * flashinfer/moe_ep/tensors.py — MoEEpTensors (mutable bundle for MoEEpLayer.forward()). * flashinfer/moe_ep/algo_knobs.py — typed AlgoKnob hierarchy plus _index_knobs() helper used by both backends to resolve a knob by class with ``.get(KnobClass)``. All frozen dataclasses, validation in __post_init__. Backend-specific constraints (hidden_size SUPPORTED_HIDDEN_SIZES for nixl_ep, max_tokens_per_rank ≤ 1024, num_experts % world_size == 0) land in B9 as _validators.py. Tests: 17 cases pass in tests/moe_ep/test_config.py (FleetParams validation, BootstrapConfig rank bounds, _index_knobs with marker knobs / quantization frozen sets / later-wins semantics). Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
flashinfer/moe_ep/nccl_ep/ndtensor.py — thin wrapper around the four ncclEpTensor* C functions (Create / Destroy / GetData / GetSizes). Upstream nccl_ep.NCCLLibrary defines these in `exported_functions` but doesn't expose Python methods; we drive them via `lib._funcs[name]` (same path the upstream wrapper uses internally). Tests: 5 cases pass in tests/moe_ep/nccl_ep/test_ndtensor.py — mocks nccl_ep + NCCLLibrary, verifies call arg shapes without needing a real ncclEpGroup_t. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
NIXL's meson tree pulls taskflow / asio / tomlplusplus / prometheus-cpp
via wraps. If a prior aborted build extracted one partially, meson
refuses to setup with "Subproject exists but has no meson.build file".
Sweep `3rdparty/nixl/subprojects/*/` and remove any subdir missing
meson.build (keep packagecache + packagefiles, which are meson-managed
sidecars; the tarball stays cached so re-extract is fast).
Also wipe build_nvep/{nixl,nccl} so meson doesn't try to --reconfigure
against a stale tree referencing the deleted subproject.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Upstream nccl_ep.get_nccl_comm_from_group() raises RuntimeError if called without an nccl_lib parameter — it needs the loaded NCCLLibrary handle to call ncclGetUniqueId + ncclCommInitRank. Pass the cached get_nccl_lib() result. Surfaced on Lyris GB200 torchrun --nproc_per_node=4 -m flashinfer.moe_ep.tests.smoke_nccl_ep: RuntimeError: Cannot create NCCL communicator without NCCLLibrary instance. Pass nccl_lib parameter to get_nccl_comm_from_group. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Surfaced on Lyris GB200: ncclEpCreateGroup asserts in_config->version==1 and fires SIGABRT at nccl_ep.cc:943 otherwise. Plan referenced version=0x10000 from the design draft; actual C ABI just wants 1. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
LL mode shapes per 3rdparty/nccl/contrib/nccl_ep/ep_test.py:
* dispatch input : 2D [num_tokens, hidden] bf16
* dispatch output : 3D [num_local_experts, max_per_rank*world, hidden]
bf16 (was 2D before — surfaced as "Assertion
recv_x->ndim == 3 failed" SIGABRT)
* dispatch local : 1D [num_local_experts] int32, tag
RECV_EXPERT_COUNTER_DEVICE — required by LL kernel
* combine input : same 3D shape as dispatch output
* combine output : 2D [num_tokens, hidden] bf16
* combine local : 2D [num_tokens, top_k] fp32, tag TOPK_WEIGHTS
(cast inside combine() if user hands us bf16)
ncclEpComplete runs unconditionally after each dispatch + combine
(matches upstream ep_test). Handle.complete() is a no-op now; the
HandleAlgoKnobSplitOperation hook is preserved for a future
HT-mode pipelined commit.
NcclEpHandle pre-allocates recv_count int32 tensor as handle-level
keepalive. NcclEpFleet exposes .bootstrap so the Handle can read
world_size for the 3D output sizing.
Mock tests updated for new dtypes + complete-inside-dispatch contract.
Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
Lyris GB200 has 4 B200 GPUs per compute tray; the 5-hr salloc held this session is single-node (4 GPUs). Relax the marker so the same test cell covers both 4-rank and 8-rank scenarios — the EP code paths are identical at either world size for LL bf16. Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
aleozlx flagged the GPU unittest shard going red: the 4 mocked tests in
tests/moe_ep/nccl_ep/test_fleet_mock.py call create_fleet(...,
backend="nccl_ep"), which raises MoEEpNotBuiltError because
NcclEpFleet.__init__ runs _require_built("nccl_ep") and the generic GPU
shard has no staged libnccl_ep.so. The nccl_ep-module + get_nccl_lib mocks
apply too late.
Add a bypass_moe_ep_build_check fixture that patches _require_built. fleet.py
does `from .. import _require_built` (a bound local), so we patch
flashinfer.moe_ep.nccl_ep.fleet._require_built — patching the parent package
wouldn't intercept the bound name. Applied to all 4 tests.
(The nixl_ep mock test already patches fleet._require_built in its
patched_loader fixture, so it needs no change.)
Also expands the module docstring (addresses aleozlx's "what does the mock
test aim to do?"): host-only tests validating config-struct marshaling +
call sequencing, not numerics — e2e correctness is covered by the
on-cluster smoke + multirank tests.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Mirror of the nccl_ep docstring expansion (aleozlx: "what does the mock test aim to do?"). States that these are host-only call-sequencing / arg-marshaling tests with a faked Buffer, not numerics — e2e is covered by the on-cluster smoke + multirank tests. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aleozlx: committing frozen=True on NcclEpConfig / NvepConfig now would make dropping it later a breaking change (it removes the synthesized __hash__ / __setattr__). They're only passed to MoEEpLayer(..., backend=...) and read via getattr(.backend_name) — never hashed or used as dict/set keys — so frozen buys nothing. Use plain @DataClass. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aleozlx: "do we want to put the smoke tests in tests/ folder as well?" The repo convention is top-level tests/ organized by subsystem. Move the two torchrun entry points out of the package: flashinfer/moe_ep/tests/smoke_nccl_ep.py -> tests/moe_ep/smoke_nccl_ep.py flashinfer/moe_ep/tests/smoke_nixl_ep.py -> tests/moe_ep/smoke_nixl_ep.py and drop the now-empty flashinfer/moe_ep/tests/ package. They're `if __name__ == "__main__"` scripts, so they run by path: torchrun --nproc_per_node=N tests/moe_ep/smoke_nccl_ep.py (pytest won't collect them — no test_ prefix.) Update scripts/task_test_moe_ep_smoke.sh to invoke by path, and align its multirank marker to gpu_4 (matches the test + the GB200 4-GPU tray). Fix the usage docstrings in both moved files. No `flashinfer.moe_ep.tests` refs remain. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aleozlx asked (4 threads on nixl_ep fleet/handle) whether these methods "are going to be @flashinfer_api". Answer: yes — wire it across the full user-facing surface symmetrically on both backends + the layer + the factory, so EP calls get the same repro-tracing/logging as the rest of FlashInfer. Decorated: * layer.py MoEEpLayer.forward * fleet.py create_fleet (factory) * nccl_ep/fleet.py NcclEpFleet.__init__, create_handle, update_topology, destroy * nccl_ep/handle.py NcclEpHandle.__init__, dispatch, combine, complete * nixl_ep/fleet.py NixlEpFleet.__init__, create_handle, update_topology, destroy * nixl_ep/handle.py NixlEpHandle.__init__, dispatch, combine, complete @flashinfer_api is zero-overhead at the default FLASHINFER_LOGLEVEL=0, so decorating ctors + None-returning methods is free. Verified at FLASHINFER_LOGLEVEL=3: forward / create_fleet / dispatch / combine / complete log without crashing; tests/moe_ep host suite stays green (35 passed, 1 skipped). Backend modules are one package level below flashinfer/, so they import the decorator as `from ...api_logging` (3 dots); layer.py + fleet.py use 2 dots. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Regression from moving the smoke scripts into tests/moe_ep/ (commit
9d47bfb2): `torchrun tests/moe_ep/smoke_nccl_ep.py` puts tests/moe_ep/ on
sys.path[0], and that dir contains the nccl_ep/ + nixl_ep/ *test*
subpackages, which shadow the installed nccl_ep / nixl_ep ctypes modules
the EP backends import. Surfaced on Lyris GB200:
ImportError: cannot import name 'get_nccl_comm_from_group' from 'nccl_ep'
(.../tests/moe_ep/nccl_ep/__init__.py)
AttributeError: module 'nixl_ep' has no attribute 'Buffer'
Fix: strip the script's own dir from sys.path at module top in both smoke
scripts, before any flashinfer / backend import runs. flashinfer itself is
in site-packages so it's unaffected.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Anerudhan
force-pushed
the
ep-api-design-partb
branch
from
June 1, 2026 18:51
c632f47 to
9096e7c
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 3
🧹 Nitpick comments (2)
flashinfer/moe_ep/nccl_ep/ndtensor.py (2)
263-270: ⚖️ Poor tradeoffReturned tensor view has dangling pointer risk if
NDTensoris collected.The
viewtensor borrows storage fromself, but nothing preventsselffrom being garbage collected whileviewis still in use. The comment documents this, but the caller has no way to keep theNDTensoralive once they only hold the view.Consider capturing
selfin the proxy or returning a tuple(view, self)so the caller must explicitly manage lifetime.🤖 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/ndtensor.py` around lines 263 - 270, The returned CUDA tensor view borrows storage from the NDTensor and can outlive it, causing a dangling pointer; modify the _CudaArrayProxy/return so the NDTensor is kept alive: either have _CudaArrayProxy capture and hold a strong reference to the NDTensor (e.g., store self_ndtensor on the proxy and reference it from the __cuda_array_interface__ provider) or change the API to return (view, ndtensor) so callers must keep the NDTensor alive; update the code paths that construct _CudaArrayProxy and the return site (symbols: _CudaArrayProxy, __cuda_array_interface__, view, NDTensor) accordingly.
272-280: 💤 Low valueSilent exception swallowing in
__del__is intentional but should log.The broad
except Exception: passis flagged by static analysis (S110, BLE001). While this pattern is valid for__del__during interpreter shutdown, consider logging at debug level for easier troubleshooting in non-shutdown scenarios.🤖 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/ndtensor.py` around lines 272 - 280, The destructor __del__ currently swallows all exceptions silently; instead catch Exception as e and log the exception at debug level so non-shutdown failures are visible: in the __del__ method (guarded by self._owns and self._handle), call get_nccl_lib() and invoke lib._funcs["ncclEpTensorDestroy"](self._group, self._handle) inside a try/except Exception as e and use a module logger (e.g., logging.getLogger(__name__)) to logger.debug a clear message that includes self._group, self._handle and the exception information before continuing to pass, so interpreter-shutdown cases still won’t raise but other errors are recorded.
🤖 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 `@flashinfer/moe_ep/nccl_ep/ndtensor.py`:
- Around line 185-199: The allocate function constructs a sizes list and calls
ncclEpTensorCreate without validating that len(shape) <= 5, which can overflow
like in from_torch; add the same validation used in from_torch to allocate:
raise an error (or assert) if shape has more than 5 dimensions, compute sizes =
list(shape) + [0] * (5 - len(shape)) only after validation, and pass
sizes[0]..sizes[4] to lib._funcs["ncclEpTensorCreate"] (keep tag, dtype mapping
via _torch_dtype_to_nccl, and ctypes.byref(handle) unchanged) so allocate
mirrors from_torch’s bounds checking and prevents >5-dim inputs.
- Around line 146-160: The call to ncclEpTensorCreate in ndtensor.py silently
drops dimensions when tensor.ndim > 5 because only five size args are passed;
add a validation in the code that constructs sizes (around the
ncclEpTensorCreate invocation and the sizes = list(tensor.shape) + [0] * (5 -
len(tensor.shape)) line) to check tensor.dim() (or len(tensor.shape)) and raise
a clear exception if ndim > 5 (or alternatively handle multi-dimensional
flattening per API contract), so tensors with >5 dims are rejected with an
explicit error instead of truncating their shape.
In `@scripts/build_in_container.sh`:
- Line 127: Remove the no-op for loop that iterates "for sp in /tmp /var/tmp
\"${REPO_ROOT}/3rdparty/nixl/subprojects\"; do :; done" since it does nothing
and leaves the unused variable sp; either delete this line entirely or, if the
original intent was to ensure those directories exist, replace it with an
explicit operation (e.g., mkdir -p on the listed paths) and update any logic
that expected sp to be used; locate the loop by the variable name sp and the
exact for-loop snippet in the script and remove or replace it accordingly.
---
Nitpick comments:
In `@flashinfer/moe_ep/nccl_ep/ndtensor.py`:
- Around line 263-270: The returned CUDA tensor view borrows storage from the
NDTensor and can outlive it, causing a dangling pointer; modify the
_CudaArrayProxy/return so the NDTensor is kept alive: either have
_CudaArrayProxy capture and hold a strong reference to the NDTensor (e.g., store
self_ndtensor on the proxy and reference it from the __cuda_array_interface__
provider) or change the API to return (view, ndtensor) so callers must keep the
NDTensor alive; update the code paths that construct _CudaArrayProxy and the
return site (symbols: _CudaArrayProxy, __cuda_array_interface__, view, NDTensor)
accordingly.
- Around line 272-280: The destructor __del__ currently swallows all exceptions
silently; instead catch Exception as e and log the exception at debug level so
non-shutdown failures are visible: in the __del__ method (guarded by self._owns
and self._handle), call get_nccl_lib() and invoke
lib._funcs["ncclEpTensorDestroy"](self._group, self._handle) inside a try/except
Exception as e and use a module logger (e.g., logging.getLogger(__name__)) to
logger.debug a clear message that includes self._group, self._handle and the
exception information before continuing to pass, so interpreter-shutdown cases
still won’t raise but other errors are recorded.
🪄 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: 9df13f94-efbd-4ab5-bcb7-b16341261543
📥 Commits
Reviewing files that changed from the base of the PR and between c632f47420ee561249f26e8948aeb83c730570ae and 9096e7c.
📒 Files selected for processing (31)
flashinfer/moe_ep/__init__.pyflashinfer/moe_ep/_validators.pyflashinfer/moe_ep/algo_knobs.pyflashinfer/moe_ep/config.pyflashinfer/moe_ep/fleet.pyflashinfer/moe_ep/handle.pyflashinfer/moe_ep/layer.pyflashinfer/moe_ep/nccl_ep/fleet.pyflashinfer/moe_ep/nccl_ep/handle.pyflashinfer/moe_ep/nccl_ep/ndtensor.pyflashinfer/moe_ep/nixl_ep/fleet.pyflashinfer/moe_ep/nixl_ep/handle.pyflashinfer/moe_ep/split_backends/__init__.pyflashinfer/moe_ep/split_backends/nccl_ep_comm.pyflashinfer/moe_ep/split_backends/nixl_ep_comm.pyflashinfer/moe_ep/tensors.pyscripts/build_in_container.shscripts/task_test_moe_ep_smoke.shtests/conftest.pytests/moe_ep/__init__.pytests/moe_ep/nccl_ep/__init__.pytests/moe_ep/nccl_ep/test_fleet_mock.pytests/moe_ep/nccl_ep/test_ndtensor.pytests/moe_ep/nixl_ep/__init__.pytests/moe_ep/nixl_ep/test_fleet_mock.pytests/moe_ep/smoke_nccl_ep.pytests/moe_ep/smoke_nixl_ep.pytests/moe_ep/test_config.pytests/moe_ep/test_constraints.pytests/moe_ep/test_layer_single_gpu.pytests/moe_ep/test_moe_ep_layer_multirank.py
🚧 Files skipped from review as they are similar to previous changes (24)
- flashinfer/moe_ep/tensors.py
- flashinfer/moe_ep/split_backends/nccl_ep_comm.py
- tests/moe_ep/test_config.py
- flashinfer/moe_ep/_validators.py
- flashinfer/moe_ep/handle.py
- flashinfer/moe_ep/split_backends/init.py
- flashinfer/moe_ep/nccl_ep/fleet.py
- flashinfer/moe_ep/split_backends/nixl_ep_comm.py
- flashinfer/moe_ep/algo_knobs.py
- tests/moe_ep/test_constraints.py
- flashinfer/moe_ep/init.py
- tests/moe_ep/smoke_nccl_ep.py
- flashinfer/moe_ep/nixl_ep/handle.py
- tests/moe_ep/test_moe_ep_layer_multirank.py
- flashinfer/moe_ep/layer.py
- tests/conftest.py
- tests/moe_ep/nixl_ep/test_fleet_mock.py
- tests/moe_ep/test_layer_single_gpu.py
- tests/moe_ep/nccl_ep/test_ndtensor.py
- flashinfer/moe_ep/config.py
- flashinfer/moe_ep/nccl_ep/handle.py
- tests/moe_ep/nccl_ep/test_fleet_mock.py
- flashinfer/moe_ep/nixl_ep/fleet.py
- tests/moe_ep/smoke_nixl_ep.py
…#3453 ndtensor.py: * _torch_typestr: map bfloat16 to "<i2" (was "<V2", which torch.as_tensor rejects); as_torch() reinterprets the 2-byte int view back to bfloat16 via .view(torch.bfloat16) (no copy). * as_torch(): coerce data_p.value to 0 when the device pointer is NULL so the __cuda_array_interface__ tuple stays valid. * Cache get_nccl_lib() as self._lib in __init__; use it in as_torch + __del__ so destruction doesn't re-resolve the lib during interpreter shutdown. (also fixes ruff SIM105 in __del__ via contextlib.suppress) nccl_ep/fleet.py: * Cache self._lib in __init__; use it in update_topology + destroy. * use_fp8 now matches any FP8 quant type {FP8E4M3, FP8E5M2, NVFP8} instead of only FP8E4M3 (silent-false otherwise). nccl_ep/handle.py: * Cache self._lib in __init__; use it in dispatch/combine/__del__. * Synchronize self._stream (via torch.cuda.ExternalStream) before the host read `recv_count_t.sum().item()`. .item() only syncs the *current* stream, but the EP work is enqueued on self._stream (which a caller may set to a different HandleAlgoKnobUserStream), so num_tokens could be read before ncclEpDispatch populated _recv_count_t. nixl_ep/fleet.py: * Import MoEEpConfigError and raise it when num_experts isn't a positive multiple of the topology capacity (cap), which is used as num_ranks for update_memory_buffers — integer division would otherwise silently truncate experts when a FleetAlgoKnobTopologyCapacity knob sets cap != world_size. * use_fp8 matches any FP8 quant type (same as nccl_ep). Host suite (tests/moe_ep, mocked): 35 passed, 1 skipped. ruff clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Two coderabbit follow-ups on PR flashinfer-ai#3453: * nccl_ep/fleet.py (r3322532556): update_topology refreshed self._fleet_knobs but rebuilt the EP group from the stale self._cfg, so config-affecting knobs (FleetAlgoKnobRdmaBufferSize / NumQpsPerRank / NumChannelsPerRank) were silently dropped while use_fp8/use_ue8m0 (read live from _fleet_knobs) did take effect — a partially-updated fleet. Extract _build_group_config() from __init__ and call it again in update_topology after refreshing the knobs, so the rebuilt group reflects the new config. * tests/moe_ep/nccl_ep/test_fleet_mock.py (r3336108650): the 4 mocked tests skip only on `not torch.cuda.is_available()`, but NcclEpFleet.__init__ also calls validate_arch_for_backend("nccl_ep"), which raises MoEEpArchError on sm < 9.0 (the A10G gpu-tests-a10g shard is sm_86). Rather than skip on low arch — which would zero out coverage on the only GPU CI shard these mocked tests target — extend bypass_moe_ep_build_check to also patch validate_arch_for_backend. These tests mock the whole NCCL library and never launch a kernel, so arch is irrelevant to what they assert (config marshaling + call sequencing). Host suite: 35 passed, 1 skipped. ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
aleozlx
reviewed
Jun 2, 2026
aleozlx
left a comment
aleozlx
left a comment
Member
There was a problem hiding this comment.
review looks good
will approve after some bot run checks
Member
|
/bot run |
Collaborator
…review) A reviewer asked to drop the @flashinfer_api decorators added in this PR. Per the request, comment them out rather than delete — both the 18 decorator usages and the 6 `from ...api_logging import flashinfer_api` imports (left live, the imports would trip ruff's unused-import check) — across layer.py, fleet.py, and the nccl_ep / nixl_ep fleet + handle modules. Each commented line is tagged "disabled per PR flashinfer-ai#3453 review" so the intent + how to re-enable is obvious. No behavior change beyond removing the call logging; the decorated methods run exactly as before. Host suite: 35 passed, 1 skipped. ruff + mypy clean. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Collaborator
|
[FAILED] Pipeline #53349926: 11/20 passed |
aleozlx
approved these changes
Jun 2, 2026
aleozlx
added a commit
that referenced
this pull request
Jun 10, 2026
… NVFP4 autotune (#3093) ## Summary Draft PR introducing a unified MoE layer that autotunes across NVFP4 backends — CuteDSL and TRTLLM FP4 routed — on a per-shape basis. ## Update — MVP cut (2026-05-31) Autonomous MVP-completion pass on a B200 (SM100). The cross-backend autotune objective is met and validated end-to-end. Full narrative in `docs/design_docs/flashinfer_moe_api.md` (Decision Log + the new "MVP As-Built Reference"). - **All MVP follow-ups (CR1–CR11) done & validated**: `tests/moe/test_unified_moe_api.py` **9/9** (layer + per-backend accuracy vs bf16, autotuner visits both candidates, CUDA-graph replay), `tests/moe/test_moe_api.py` **97/97** (CPU config + fail-fast validation), and the `unified_nvfp4_moe` sweep (128→16384 tokens) with `--refcheck` passing for both backends. - **Blocker found & fixed**: both runner adapters had never executed against the post-`main`-merge `core.py` (a stale raw-`moe_op` API + a class-vs-instance `tuning_config` bug). `TrtllmFp4RoutedRunner` now delegates to the canonical `core.MoERunner` (newly exported from `get_trtllm_moe_sm100_module()`); the unified adapters only translate Packs ⇄ the inner runner's native tensor list. Future direction (out of scope): make the low-level TVM-FFI ops take structured config objects (design doc §5) to kill positional-arg drift. - **New since first draft**: `local_expert_offset` wired into TRTLLM packing (+EP-offset test); fail-fast NVFP4/Swiglu scope validation; per-token-bucket winner cache; `tune_max_num_tokens` threaded into runner tuning; first-class NVFP4 weight prep (`prepare.py` + `*.prepare_weights`) replacing duplicated test/bench prep; `--refcheck` for the unified benchmark. ## What this adds - **`MoEConfig` / `MoEActivationPack` / `MoEWeightPack`** — frozen config dataclasses and per-call / long-lived tensor containers (`flashinfer/fused_moe/api.py`). Single `QuantVariant` enum replaces the old 3-axis dtype × granularity × variant split. - **`CuteDslNvfp4Runner` / `TrtllmFp4RoutedRunner`** — `TunableRunner` adapters with `pack_inputs(act, weights)` translating packs into the backend's native tensor list (`flashinfer/fused_moe/runners.py`). ~~each with its own `tuning_config`~~ → each **delegates to a canonical inner runner** (`CuteDslFusedMoENvfp4Runner` / `core.MoERunner`) and builds its `tuning_config` **per-instance**. - **`MoELayer`** — stateful dispatcher that builds one runner per compatible backend config, then on first call ~~per shape~~ **per token-bucket** runs per-runner `choose_one` (within-backend tactic tuning) + cross-runner `bench_gpu_time` comparison (cross-backend selection). Caches the winner **keyed by tuning bucket** (`flashinfer/fused_moe/layer.py`). - **`unified_nvfp4_moe` benchmark routine** — wired into `benchmarks/flashinfer_benchmark.py` (not a standalone script). Emits one result row per candidate backend with the winner marked; supports `--refcheck` against a shared bf16 reference. - **`bench_unified_moe_today.sh`** — convenience wrapper at repo root invoking the infra routine across the eight shapes we care about (EP=1 sweep + EP=16 wide-EP regime). - **`tests/moe/test_unified_moe_api.py`** — accuracy tests (`MoELayer` vs bf16 ref; each backend vs same ref), plumbing tests (autotuner visits all candidates; CUDA graph capture + replay), **and a pre-routed EP-offset packing test**. - **`flashinfer/fused_moe/prepare.py`** *(new)* — first-class NVFP4 weight-prep helpers exposed as `TrtllmFp4Config.prepare_weights` / `CuteDslConfig.prepare_weights`; the test and benchmark no longer carry duplicated prep copies. ## Design notes - **Per-runner `choose_one`** — `AutoTuner.choose_one` assumes all runners share one `inputs` list during profiling. Our backends' native schemas differ (CuteDSL 12 tensors with unpacked topk + trailing output buffer; TRTLLM 8-field packed `MoEInputs`). Resolution: call `choose_one` once per runner (within-backend tactic selection) and use `bench_gpu_time` to compare the per-runner winners. ~~Each runner owns its own `tuning_config` as a class attribute.~~ → **`tuning_config` is built per-instance** (the TRTLLM runner via `MoERunner._make_tuning_config`, which also threads `ExecutionConfig.tune_max_num_tokens` into the bucket set). - **Shared-reference accuracy testing** — each backend is tested against the same bf16 reference, not against the other backend. Catches shared-mode failures (both wrong in the same way) that cross-backend agreement would miss. `--refcheck` brings the same check to the benchmark. - **Weight layout: `Shuffled_MajorK` only** — the only NVFP4-compatible TRTLLM layout today. Multi-layout autotune (opt-in additional variants via a future `layouts` field on `TrtllmFp4Config`) is a V2 extension — design accommodates it without core changes. ## Out of scope for this PR - FP8 / BF16 / MxInt4 backends - Monolithic routing (our path is pre-routed) - Unpacked topk for TRTLLM — #2425 (when that lands, `TrtllmFp4RoutedRunner` drops the inline `(id << 16) | bf16_bits` packing) - Multi-layout TRTLLM autotune across `Shuffled_BlockMajorK` / `NoShuffle_MajorK` - SM120/SM121 support for CuteDSL (kernel is SM100/SM103 only today; `CuteDslConfig.supported` tightened to match) - First-class **activation** prep + structured-config TVM-FFI boundary (§5) — post-MVP carryover ## Status ~~Draft — benchmark numbers and observation notes are being collected on a Blackwell box. Will de-draft once the winner-flip table is filled in and accuracy tests pass on hardware.~~ **MVP scope complete and validated on B200 (SM100).** Kept as a **draft PR on purpose** so CI skips while further edits are pushed autonomously — flip to "Ready for review" when you want CI to run. ## Test plan - [x] `pytest tests/moe/test_unified_moe_api.py -v` (SM100/SM103) — **9/9** - [ ] ~~`./bench_unified_moe_today.sh` — winner column matches expected per-shape~~ → not run this pass; equivalent coverage via the `unified_nvfp4_moe` routine sweep (128→16384) + `--refcheck` (both backends pass). The script's EP=16 shapes are still worth a dedicated run. - [x] Pre-commit hooks pass ## Related PRs - #3453 🤖 Generated with [Claude Code](https://claude.com/claude-code) <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Unified MoE API surfaced package-wide: immutable configs, MoELayer (cached cross-backend autotune winner), backend runners, NVFP4 weight-prep utilities, and re-exported unified API symbols. * **Benchmarking** * New benchmark sweep script for unified NVFP4 MoE and registration to record token-size sweeps to CSV. * **Documentation** * Added comprehensive FlashInfer Unified MoE API design doc with migration plan. * **Tests** * Added CPU and gated-GPU tests for API, validation, accuracy, and autotune/dispatch behavior; improved enum reprs for round-trip logging. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Co-authored-by: Yang Xu <yanxu@nvidia.com>
aleozlx
added
the
unified_api
This was referenced Jul 2, 2026
Merged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
📌 Description
Ep api design -- Adding the actual code and tests
🔍 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
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes
Summary by CodeRabbit
New Features
Bug Fixes / Validation
Tests / Chores