MoE-EP: wire unified MoE compute into NCCL-EP / NIXL-EP expert parallel (LL + HT) - #3686
Merged
Anerudhan merged 12 commits intoJun 30, 2026
Merged
Conversation
Anerudhan
requested review from
IwakuraRein,
aleozlx,
bkryu,
cyx-6,
dhiraj113,
jiahanc,
jimmyzho,
kahyunnam,
nv-yunzheq,
qsang-nv,
saltyminty,
samuellees,
sricketts,
yongwww,
yyihuang and
yzh119
as code owners
June 22, 2026 06:15
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:
📝 WalkthroughWalkthroughUpdates NCCL-EP to the ChangesNCCL-EP Migration and MoE EP Compute Integration
Benchmarking Harnesses and Design Documentation
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 3 | ❌ 2❌ Failed checks (2 warnings)
✅ Passed checks (3 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 |
![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 refactors the MoE Expert-Parallel (MoE-EP) path by transitioning the NCCL-EP backend to the released nccl4py wheel, adding comprehensive performance benchmarks, introducing a PyTorch-based Dockerfile to fix cross-node HT crashes, and implementing the RANK_MAJOR and FLAT layouts with a layout bridge to the unified compute API. The review feedback highlights critical improvement opportunities: validating shape, dtype, and device for cached buffers (_ll_recv_buf and _ht_recv_bufs) to prevent correctness bugs under dynamic batching; defensively validating that received expert indices are within the local expert range in the compute bridge; and explicitly specifying the device for _recv_count_t allocation to avoid device mismatch 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.
![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: 10
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
docker/Dockerfile.flashinfer-ep-pytorch (1)
44-108:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRun the final container as a non-root user.
There is no
USERdirective, so runtime stays root. This increases blast radius if the process or mounted paths are compromised.Proposed fix
@@ RUN python -c "\ from flashinfer.moe_ep import available_backends; \ b = available_backends(); print('moe_ep backends:', b); \ assert 'nccl_ep' in b, 'nccl_ep backend missing'" + +# Drop privileges for runtime. +RUN useradd --create-home --uid 10001 --shell /bin/bash appuser \ + && chown -R appuser:appuser /workspace +USER appuser CMD ["bash"]🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@docker/Dockerfile.flashinfer-ep-pytorch` around lines 44 - 108, The Dockerfile currently runs as root (no USER directive before the CMD instruction), which is a security risk. Add a RUN instruction before the CMD line to create a non-root user account, then add a USER directive to switch to that user. This ensures the container runs with reduced privileges, limiting the impact if the process or mounted paths are compromised. Make sure the new user has necessary read/write permissions for any required directories in the container.Source: Linters/SAST tools
🧹 Nitpick comments (10)
docs/design_docs/MoE_EP_impl.md (1)
28-28: 🧹 Nitpick | 🔵 Trivial | 💤 Low valueSpecify language for fenced code block.
The diagram at line 28 uses a fenced code block without a language identifier, which causes a markdown linter warning. Add a language tag (e.g.,
bash,text) to comply with markdown best practices.-``` +```text forward(t: MoEEpTensors)🤖 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 `@docs/design_docs/MoE_EP_impl.md` at line 28, The fenced code block containing forward(t: MoEEpTensors) at line 28 is missing a language identifier after the opening triple backticks, which causes markdown linter warnings. Add a language tag such as `text` immediately after the opening fence (change ``` to ```text) to comply with markdown standards.Source: Linters/SAST tools
benchmarks/run_ep_matrix_one.sh (1)
49-50: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winAdd error handling for
cdcommand.If the
cdfails (e.g., mount not ready), the script continues and runs Python from the wrong directory, potentially with confusing errors.🔧 Proposed fix
-cd /host/flashinfer +cd /host/flashinfer || { echo "ERROR: cannot cd to /host/flashinfer" >&2; exit 1; } exec python benchmarks/bench_ep_matrix.py "$@"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@benchmarks/run_ep_matrix_one.sh` around lines 49 - 50, The cd command at the beginning of the script can fail silently if the mount is not ready, causing the subsequent python command to execute from the wrong directory. Add error handling to the cd command by appending || exit 1 after cd /host/flashinfer to ensure the script exits immediately if the directory change fails, preventing execution of the python benchmarks from an incorrect location.benchmarks/run_ep_matrix_one_pt.sh (1)
36-37: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winAdd error handling for
cdcommand.Same issue as
run_ep_matrix_one.sh— if the mount isn't ready, the script runs from the wrong directory.🔧 Proposed fix
-cd /host/flashinfer +cd /host/flashinfer || { echo "ERROR: cannot cd to /host/flashinfer" >&2; exit 1; } exec python benchmarks/bench_ep_matrix.py "$@"🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@benchmarks/run_ep_matrix_one_pt.sh` around lines 36 - 37, The cd command to /host/flashinfer lacks error handling, which means if the directory change fails due to an unmounted path, the script will continue executing the subsequent exec python benchmarks/bench_ep_matrix.py command from the wrong directory. Add error handling immediately after the cd command to check if the directory change was successful, and exit with an error message if it fails.benchmarks/run_httest_torchrun.sh (2)
49-51: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winAdd error handling for
cdcommand.Same issue as the other launcher scripts.
🔧 Proposed fix
-cd /host/flashinfer +cd /host/flashinfer || { echo "ERROR: cannot cd to /host/flashinfer" >&2; exit 1; } exec torchrun --nproc_per_node=8 --master_addr=127.0.0.1 --master_port=29555 \ tests/moe_ep/test_moe_ep_ht_correctness.py🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@benchmarks/run_httest_torchrun.sh` around lines 49 - 51, The cd command to /host/flashinfer lacks error handling, which means the subsequent torchrun command will execute even if the directory change fails. Add error handling after the cd command by either using the || operator to exit on failure (cd /host/flashinfer || exit 1) or by wrapping it in a conditional check that exits if the cd fails. This ensures that the script terminates immediately if the directory change is unsuccessful, preventing the torchrun command from executing in an unexpected working directory.
1-48: 🧹 Nitpick | 🔵 Trivial | ⚖️ Poor tradeoffSignificant code duplication with
run_ep_matrix_one.sh.Lines 1-48 are nearly identical to
run_ep_matrix_one.sh. Consider extracting the shared NCCL/JIT setup into a sourced helper script to reduce maintenance burden.🤖 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 `@benchmarks/run_httest_torchrun.sh` around lines 1 - 48, The NCCL and JIT setup logic (including NCCLLIB detection, LD_LIBRARY_PATH configuration, WHEEL_EP/NCCL_INC/JIT_INC initialization, SLURM_LOCALID synchronization, and environment variable exports for NCCL_EP_JIT_* and torch.distributed) is duplicated between run_httest_torchrun.sh and run_ep_matrix_one.sh. Extract this shared setup block into a separate helper shell script, then source that helper script from both run_httest_torchrun.sh and run_ep_matrix_one.sh at the appropriate location to eliminate the duplication and reduce maintenance burden.tests/moe_ep/test_compute_bridge.py (2)
123-127: 🧹 Nitpick | 🔵 Trivial | 💤 Low valueRedundant
pytest.importorskip("torch").
torchis already imported at module level (line 22 viapytest.importorskip). This second skip is unreachable—if torch wasn't available, the module-level import would have already skipped.♻️ Remove redundant import skip
def test_build_activation_pack_rejects_2d(): - pytest.importorskip("torch") bad = torch.zeros(8, 128, dtype=torch.bfloat16) with pytest.raises(ValueError): build_activation_pack(bad, is_nvfp4=False)🤖 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/test_compute_bridge.py` around lines 123 - 127, The test function test_build_activation_pack_rejects_2d contains a redundant pytest.importorskip("torch") call since torch is already imported at the module level with the same skip mechanism. Remove the importorskip call from inside the test function as it is unreachable and unnecessary—if torch was unavailable, the module-level import would have already prevented the test from running.
130-133: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winUse
flashinfer.utils.is_sm100a_supported()for architecture check.Per coding guidelines, tests should use
flashinfer.utilsfunctions to skip tests based on CUDA architecture rather than raw capability checks.♻️ Use utility function for SM100+ check
+from flashinfer.utils import is_sm100a_supported + `@pytest.mark.skipif`( - not torch.cuda.is_available() or torch.cuda.get_device_capability()[0] < 10, + not torch.cuda.is_available() or not is_sm100a_supported(), reason="NVFP4 fp4_quantize needs SM100+", ) def test_nvfp4_pack_quantizes():🤖 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/test_compute_bridge.py` around lines 130 - 133, Replace the raw CUDA capability check in the pytest.mark.skipif decorator with the utility function from flashinfer.utils. Import is_sm100a_supported from flashinfer.utils at the top of the test file, then modify the skipif condition to use not flashinfer.utils.is_sm100a_supported() instead of the raw torch.cuda.get_device_capability() check. This ensures the test uses the standard architecture validation function as per coding guidelines.Source: Coding guidelines
tests/moe_ep/test_moe_ep_compute_correctness.py (1)
327-327: 🧹 Nitpick | 🔵 Trivial | 💤 Low valueUnused variable from unpacking.
Static analysis flagged
ep_vs_torchas unused. It's captured for the print statement but only referenced inside theif rank == 0block in_run_one_layout. Consider prefixing with underscore per the Ruff suggestion.♻️ Prefix unused variable with underscore
- rank, ep_vs_kernel, kernel_vs_torch, ep_vs_torch = _run_one_layout(layout) + rank, ep_vs_kernel, kernel_vs_torch, _ep_vs_torch = _run_one_layout(layout)🤖 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/test_moe_ep_compute_correctness.py` at line 327, The variable ep_vs_torch being unpacked from the _run_one_layout() call is flagged as unused by static analysis. Prefix this unused variable with an underscore (change ep_vs_torch to _ep_vs_torch) in the unpacking statement to follow the Python convention for intentionally unused variables and satisfy Ruff linting rules.flashinfer/moe_ep/config.py (1)
168-171: 🧹 Nitpick | 🔵 Trivial | ⚡ Quick winMemoize deferred
num_tokensresolution inget_num_tokens().
get_num_tokens()re-invokes the thunk on every call, so repeated reads can re-trigger stream sync/readback despite the “paid at most once” contract in the docstring/comments.Suggested patch
def get_num_tokens(self) -> int: """Resolve ``num_tokens`` to an int, evaluating a deferred thunk if present.""" nt = self.num_tokens - return nt() if callable(nt) else nt + if callable(nt): + nt = nt() + self.num_tokens = nt + return nt🤖 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/config.py` around lines 168 - 171, The get_num_tokens() method re-invokes the callable thunk on every call, causing repeated stream sync/readback operations that violate the "paid at most once" contract. Memoize the resolved value by storing it in an instance variable after the first resolution, then check for and return this cached value on subsequent calls instead of re-invoking the thunk. This ensures the thunk is only called once and the result is reused for all future get_num_tokens() invocations.flashinfer/moe_ep/nccl_ep/handle.py (1)
397-398: 🧹 Nitpick | 🔵 Trivial | 💤 Low valueConsider using
self._world_sizedirectly for clarity.The derived
world = num_experts // num_local_expertsequalsworld_sizefor valid inputs (divisibility is validated upstream). Usingself._world_sizedirectly would be clearer and avoid the indirect calculation.♻️ Suggested simplification
- world = self._fleet.params.num_experts // self._num_local_experts + world = self._world_size🤖 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/handle.py` around lines 397 - 398, The code calculates `world` by dividing `self._fleet.params.num_experts` by `self._num_local_experts`, but this value should equal `self._world_size` for valid inputs. Replace the derived calculation of `world` with `self._world_size` directly to improve code clarity and avoid the indirect computation. Update the line where `world` is used in the `num_recv` calculation to use `self._world_size` instead.
🤖 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 543-552: The NCCL-EP build status reporting at lines 547 and 609
relies only on the `_BUILD_NCCL_EP` flag without verifying that `nccl.ep` is
actually importable at runtime. Add an importability check for `nccl.ep` (using
a try/except pattern similar to the existing `_BUILD_NVEP_BEST_EFFORT` pattern
used for NIXL-EP error handling), and store the result in a variable. Then gate
both the print statement around line 547 and the reporting at line 609 to only
execute when both `_BUILD_NCCL_EP` is true AND `nccl.ep` is confirmed to be
importable, preventing false reporting of NCCL-EP availability when it will
later raise `MoEEpNotBuiltError`.
In `@flashinfer/fused_moe/layer.py`:
- Around line 64-67: The _RUNNER_QUANTS dictionary is currently only used for
validation but does not prevent instantiation of incompatible runners when mixed
candidates are provided. When constructing runners in the loop around lines
89-97, filter the runner instantiation to only create runners whose
configuration type supports the currently selected quant variant by checking if
the variant exists in the tuple of supported variants from _RUNNER_QUANTS. This
ensures that runners like CuteDslNvfp4Runner are not created when the selected
variant is QuantVariant.BF16, preventing mismatches between the runner's
expected quantization format and the actual weights or activations.
In `@flashinfer/fused_moe/runners.py`:
- Around line 342-351: The docstring for the TrtllmBf16RoutedRunner class
incorrectly documents the packing format of pre-routed ids as ((expert_id -
local_offset) << 16), but the actual pack_inputs() implementation now packs
GLOBAL expert_id << 16 and passes local_expert_offset separately. Update the
docstring in the TrtllmBf16RoutedRunner class to accurately reflect this change,
replacing the description of ((expert_id - local_offset) << 16) with the correct
format showing GLOBAL expert_id being shifted left by 16 bits and noting that
local_expert_offset is passed separately.
In `@flashinfer/moe_ep/_compute_bridge.py`:
- Around line 146-171: The is_local condition at line 165 only checks if idx >=
0, which allows out-of-range indices to incorrectly be treated as local experts.
Modify the is_local condition to also validate that idx < num_local_experts to
properly bound the rank-major local IDs to the valid range [0,
num_local_experts). Additionally, enhance the shape validation for recv_topk_idx
and recv_topk_weights to check the complete [M, top_k] shape rather than just
the first dimension, ensuring they have matching dimensions before using them.
In `@flashinfer/moe_ep/_validators.py`:
- Around line 112-136: The validator function needs to add a check to reject
compute configs where do_finalize is False, since the EP bridge expects
finalized tensors and MoEEpLayer consumes inputs assuming they are finalized.
Add a validation check after the existing expert configuration validations that
verifies compute_config.execution.do_finalize is True, and raise a
MoEEpConfigError with a descriptive message if it is False, following the same
error pattern as the other validation checks in this function.
In `@flashinfer/moe_ep/layer.py`:
- Around line 7-13: The docstring describing the Inner compute behavior
incorrectly states "no routing, no weighted finalize" without clarifying that
the RANK_MAJOR and HT routing layout branches actually DO apply routing weights
and perform top_k finalized local reductions during compute. Update the
docstring in the Inner compute section (around lines 7-13) to explicitly clarify
that while EXPERT_MAJOR layout uses top_k=1 without routing weight application,
the RANK_MAJOR and HT layout branches do apply the received routing weights and
perform proper top_k finalization. Apply the same clarification to the related
documentation section at lines 143-150 that mentions similar compute behavior
descriptions.
- Around line 217-250: The handle.complete() call can be skipped if an exception
occurs during dispatch, compute, or combine operations, leaving the EP handle
lifecycle incomplete. Wrap the dispatch/compute/combine operations in both the
non-timing path (handle.dispatch, self._inner_compute, and handle.combine calls)
and the profiling path (CUDA event recording and the same operations) with a
try/finally block. Move the handle.complete() call to the finally block in both
code paths so it always executes regardless of whether an exception occurs
during the dispatch, compute, or combine stages.
In `@flashinfer/moe_ep/nccl_ep/fleet.py`:
- Around line 29-30: The import statement for the flashinfer_api decorator is
commented out, leaving the public methods in the NcclEpFleet class
uninstrumented. Uncomment the import statement "from ...api_logging import
flashinfer_api" at the top of the file, then apply the `@flashinfer_api` decorator
to all public methods in the NcclEpFleet class to ensure they are properly
instrumented according to project coding guidelines. Verify that all public API
surface methods (those not prefixed with underscore) have the decorator applied.
- Around line 54-66: The _resolve_comm function silently ignores the bootstrap
parameter's nccl_comm attribute and always resolves from the default torch
process group, which can cause incorrect communicator selection and hangs. Add a
fail-fast guard at the beginning of _resolve_comm that checks if
bootstrap.nccl_comm is set and raises an error with a clear message indicating
that explicit communicator support is not yet implemented, ensuring the function
fails immediately rather than silently using incorrect topology.
In `@flashinfer/moe_ep/split_backends/__init__.py`:
- Around line 6-12: Sort the `__all__` list in the __init__.py file in
alphabetical order to satisfy the Ruff RUF022 rule. The list currently contains
NcclEpConfig, NCCLEPConfig, NvepConfig, NIXLEPConfig, and NixlEpConfig in an
unsorted order. Reorder these strings alphabetically so that the linter check
passes and prevents CI failures.
---
Outside diff comments:
In `@docker/Dockerfile.flashinfer-ep-pytorch`:
- Around line 44-108: The Dockerfile currently runs as root (no USER directive
before the CMD instruction), which is a security risk. Add a RUN instruction
before the CMD line to create a non-root user account, then add a USER directive
to switch to that user. This ensures the container runs with reduced privileges,
limiting the impact if the process or mounted paths are compromised. Make sure
the new user has necessary read/write permissions for any required directories
in the container.
---
Nitpick comments:
In `@benchmarks/run_ep_matrix_one_pt.sh`:
- Around line 36-37: The cd command to /host/flashinfer lacks error handling,
which means if the directory change fails due to an unmounted path, the script
will continue executing the subsequent exec python benchmarks/bench_ep_matrix.py
command from the wrong directory. Add error handling immediately after the cd
command to check if the directory change was successful, and exit with an error
message if it fails.
In `@benchmarks/run_ep_matrix_one.sh`:
- Around line 49-50: The cd command at the beginning of the script can fail
silently if the mount is not ready, causing the subsequent python command to
execute from the wrong directory. Add error handling to the cd command by
appending || exit 1 after cd /host/flashinfer to ensure the script exits
immediately if the directory change fails, preventing execution of the python
benchmarks from an incorrect location.
In `@benchmarks/run_httest_torchrun.sh`:
- Around line 49-51: The cd command to /host/flashinfer lacks error handling,
which means the subsequent torchrun command will execute even if the directory
change fails. Add error handling after the cd command by either using the ||
operator to exit on failure (cd /host/flashinfer || exit 1) or by wrapping it in
a conditional check that exits if the cd fails. This ensures that the script
terminates immediately if the directory change is unsuccessful, preventing the
torchrun command from executing in an unexpected working directory.
- Around line 1-48: The NCCL and JIT setup logic (including NCCLLIB detection,
LD_LIBRARY_PATH configuration, WHEEL_EP/NCCL_INC/JIT_INC initialization,
SLURM_LOCALID synchronization, and environment variable exports for
NCCL_EP_JIT_* and torch.distributed) is duplicated between
run_httest_torchrun.sh and run_ep_matrix_one.sh. Extract this shared setup block
into a separate helper shell script, then source that helper script from both
run_httest_torchrun.sh and run_ep_matrix_one.sh at the appropriate location to
eliminate the duplication and reduce maintenance burden.
In `@docs/design_docs/MoE_EP_impl.md`:
- Line 28: The fenced code block containing forward(t: MoEEpTensors) at line 28
is missing a language identifier after the opening triple backticks, which
causes markdown linter warnings. Add a language tag such as `text` immediately
after the opening fence (change ``` to ```text) to comply with markdown
standards.
In `@flashinfer/moe_ep/config.py`:
- Around line 168-171: The get_num_tokens() method re-invokes the callable thunk
on every call, causing repeated stream sync/readback operations that violate the
"paid at most once" contract. Memoize the resolved value by storing it in an
instance variable after the first resolution, then check for and return this
cached value on subsequent calls instead of re-invoking the thunk. This ensures
the thunk is only called once and the result is reused for all future
get_num_tokens() invocations.
In `@flashinfer/moe_ep/nccl_ep/handle.py`:
- Around line 397-398: The code calculates `world` by dividing
`self._fleet.params.num_experts` by `self._num_local_experts`, but this value
should equal `self._world_size` for valid inputs. Replace the derived
calculation of `world` with `self._world_size` directly to improve code clarity
and avoid the indirect computation. Update the line where `world` is used in the
`num_recv` calculation to use `self._world_size` instead.
In `@tests/moe_ep/test_compute_bridge.py`:
- Around line 123-127: The test function test_build_activation_pack_rejects_2d
contains a redundant pytest.importorskip("torch") call since torch is already
imported at the module level with the same skip mechanism. Remove the
importorskip call from inside the test function as it is unreachable and
unnecessary—if torch was unavailable, the module-level import would have already
prevented the test from running.
- Around line 130-133: Replace the raw CUDA capability check in the
pytest.mark.skipif decorator with the utility function from flashinfer.utils.
Import is_sm100a_supported from flashinfer.utils at the top of the test file,
then modify the skipif condition to use not
flashinfer.utils.is_sm100a_supported() instead of the raw
torch.cuda.get_device_capability() check. This ensures the test uses the
standard architecture validation function as per coding guidelines.
In `@tests/moe_ep/test_moe_ep_compute_correctness.py`:
- Line 327: The variable ep_vs_torch being unpacked from the _run_one_layout()
call is flagged as unused by static analysis. Prefix this unused variable with
an underscore (change ep_vs_torch to _ep_vs_torch) in the unpacking statement to
follow the Python convention for intentionally unused variables and satisfy Ruff
linting rules.
🪄 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: 01778636-c0ef-4bac-8e39-e031267c0053
📥 Commits
Reviewing files that changed from the base of the PR and between 9c5ed7c and 4bcb5fc11ff075f346796214a0073573b9c366e8.
📒 Files selected for processing (34)
3rdparty/ncclbenchmarks/MoE_benchmarks.mdbenchmarks/bench_ep_matrix.pybenchmarks/bench_moe_ep.pybenchmarks/run_ep_matrix.shbenchmarks/run_ep_matrix_one.shbenchmarks/run_ep_matrix_one_pt.shbenchmarks/run_httest_torchrun.shbuild_backend.pydocker/Dockerfile.flashinfer-ep-pytorchdocker/Dockerfile.flashinfer-nvepdocker/install/build_flashinfer_ep_pytorch.shdocs/design_docs/MoE_EP_impl.mdflashinfer/fused_moe/layer.pyflashinfer/fused_moe/runners.pyflashinfer/moe_ep/__init__.pyflashinfer/moe_ep/_compute_bridge.pyflashinfer/moe_ep/_validators.pyflashinfer/moe_ep/config.pyflashinfer/moe_ep/layer.pyflashinfer/moe_ep/nccl_ep/__init__.pyflashinfer/moe_ep/nccl_ep/fleet.pyflashinfer/moe_ep/nccl_ep/handle.pyflashinfer/moe_ep/nccl_ep/ndtensor.pyflashinfer/moe_ep/split_backends/__init__.pyflashinfer/moe_ep/split_backends/nccl_ep_comm.pyflashinfer/moe_ep/split_backends/nixl_ep_comm.pypyproject.tomlscripts/build_in_container.shtests/moe_ep/nccl_ep/test_fleet_mock.pytests/moe_ep/nccl_ep/test_ndtensor.pytests/moe_ep/test_compute_bridge.pytests/moe_ep/test_moe_ep_compute_correctness.pytests/moe_ep/test_moe_ep_ht_correctness.py
💤 Files with no reviewable changes (2)
- flashinfer/moe_ep/nccl_ep/ndtensor.py
- tests/moe_ep/nccl_ep/test_ndtensor.py
aleozlx
reviewed
Jun 26, 2026
aleozlx
reviewed
Jun 26, 2026
aleozlx
reviewed
Jun 26, 2026
aleozlx
reviewed
Jun 26, 2026
aleozlx
reviewed
Jun 26, 2026
Anerudhan
force-pushed
the
feat/moe_ep/enable_compute
branch
from
June 29, 2026 18:32
9da578c to
a21e1ba
Compare
Contributor
There was a problem hiding this comment.
Actionable comments posted: 11
🤖 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 `@benchmarks/bench_ep_matrix.py`:
- Around line 126-128: Update the benchmark layout selection in
bench_ep_matrix.py so HT runs map to EpLayout.FLAT when args.layout is "fl"
instead of falling through to EpLayout.EXPERT_MAJOR. Keep the existing
EpAlgorithm and Rank-major handling intact, but make the layout choice in the
args.layout branch explicit so the benchmark exercises the FLAT path for
high-throughput runs.
- Around line 217-223: The cached combine params in do_combine() are reusing a
stale input tensor when _reuse_params is enabled, because CombineInputParams is
only initialized once with the first cin. Update the existing cached params in
_comb_params["p"] so its x field is refreshed with the current cin on every call
before handle.combine(p), while still preserving the reused object. This should
be fixed in the do_combine function and the CombineInputParams usage.
- Line 4: The docstring in the matrix benchmark text uses a Unicode
multiplication sign that triggers Ruff’s RUF002 lint error. Update the wording
in the benchmark docstring for the ep matrix description to use ASCII text
instead of the `×` character, keeping the same meaning while ensuring lint-gated
runs pass.
In `@benchmarks/bench_moe_ep.py`:
- Line 42: The docstring in the benchmark module uses a Unicode multiplication
symbol, which Ruff flags as non-ASCII text. Update the `ep_bench` send-side
convention description to use an ASCII replacement for the `×` character so the
docstring stays lint-clean. Make the change in the docstring text itself and
keep the wording otherwise unchanged.
- Around line 279-284: The HT benchmark path is still forcing
EpLayout.EXPERT_MAJOR in the ep_layout selection logic, which conflicts with the
advertised ht_flat configuration. Update the ep_layout handling in the benchmark
setup around ep_layout and ep_algorithm so that HIGH_THROUGHPUT uses
EpLayout.FLAT instead of EXPERT_MAJOR, while preserving the existing rank_major
behavior for non-HT cases.
In `@benchmarks/run_ep_matrix_one_pt.sh`:
- Around line 15-27: The wheel path and nvcc resolution in
run_ep_matrix_one_pt.sh can fail silently, leaving empty values in
LD_LIBRARY_PATH/CPATH and causing later JIT or loading errors. Update the script
to validate the nccl.ep and nvidia.nccl python lookups and the nvcc command in
the setup block before exporting paths, and exit immediately with a clear
message if any lookup fails. Use the existing WHEEL_EP, NCCL_LIB, NCCL_INC, and
NCCL_EP_JIT_NVCC variables to keep the checks close to the affected setup.
In `@benchmarks/run_ep_matrix_one.sh`:
- Around line 9-40: The JIT include setup in run_ep_matrix_one.sh can mark the
shared include tree ready even when WHEEL_EP or NCCL_INC is missing or empty, so
other ranks continue with a broken JIT include path. Update the setup around
WHEEL_EP, NCCL_INC, and JIT_INC to validate both source directories exist and
contain the expected headers before creating symlinks and touching .ready. If
discovery or linking fails, exit immediately on localid 0 so waiting ranks do
not proceed with an invalid include tree.
In `@benchmarks/run_httest_torchrun.sh`:
- Around line 9-40: The HT JIT bootstrap in the shell script can publish .ready
even when WHEEL_EP or NCCL_INC is empty because the symlink commands hide
failures. Add explicit validation in the setup block before touching
$JIT_INC/.ready, and fail the script early if either include source is missing
or the combined include tree was not created correctly. Use the WHEEL_EP,
NCCL_INC, JIT_INC, and the localid 0 readiness path to locate the fix.
- Around line 2-6: The wrapper still double-launches workers because the script
invoked by srun is calling torchrun with --nproc_per_node=8, which creates 8
processes per SLURM task and causes port contention. Fix run_httest_torchrun.sh
so it is either a single torchrun-per-node launcher or a pure per-rank
entrypoint; in either case, remove the inner torchrun behavior and adjust the
launch flow around the existing SLURM/NCCL setup and benchmark argument
passthrough so only one local process group is started per node.
In `@docker/Dockerfile.flashinfer-ep-pytorch`:
- Around line 44-107: The Dockerfile.flashinfer-ep-pytorch image stays root all
the way through the final runtime stage, so add a non-root runtime user before
the final CMD and switch to it after the install/probe steps. Update the image
setup around the existing RUN blocks and CMD so the new user owns or can
read/write only the workspace and any required cache/output paths used by the
FlashInfer install and benchmark runs.
In `@docker/install/build_flashinfer_ep_pytorch.sh`:
- Around line 34-39: The install commands in build_flashinfer_ep_pytorch.sh are
using an empty environment assignment that triggers Shellcheck SC1007; update
the two pip invocations in the FlashInfer setup block to spell the empty
PIP_CONSTRAINT value explicitly. Keep the behavior the same by using the
explicit empty-string form before the pip command, and ensure both
package-install lines are updated consistently.
🪄 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: 45874c72-40c1-4a0f-962f-b01683db6901
📥 Commits
Reviewing files that changed from the base of the PR and between 9da578c9e99b992c48473f42befe98ae9f72ed1c and a21e1bae67484fe8f1707a7a952c29ce8bb8095f.
📒 Files selected for processing (35)
3rdparty/ncclbenchmarks/MoE_benchmarks.mdbenchmarks/bench_ep_matrix.pybenchmarks/bench_moe_ep.pybenchmarks/run_ep_matrix.shbenchmarks/run_ep_matrix_one.shbenchmarks/run_ep_matrix_one_pt.shbenchmarks/run_httest_torchrun.shbuild_backend.pydocker/Dockerfile.flashinfer-ep-pytorchdocker/Dockerfile.flashinfer-nvepdocker/install/build_flashinfer_ep_pytorch.shdocs/design_docs/MoE_EP_impl.mdflashinfer/fused_moe/layer.pyflashinfer/fused_moe/runners.pyflashinfer/moe_ep/__init__.pyflashinfer/moe_ep/_compute_bridge.pyflashinfer/moe_ep/_validators.pyflashinfer/moe_ep/config.pyflashinfer/moe_ep/layer.pyflashinfer/moe_ep/nccl_ep/__init__.pyflashinfer/moe_ep/nccl_ep/fleet.pyflashinfer/moe_ep/nccl_ep/handle.pyflashinfer/moe_ep/nccl_ep/ndtensor.pyflashinfer/moe_ep/split_backends/__init__.pyflashinfer/moe_ep/split_backends/nccl_ep_comm.pyflashinfer/moe_ep/split_backends/nixl_ep_comm.pypyproject.tomlscripts/build_in_container.shtests/moe/test_unified_moe.pytests/moe_ep/nccl_ep/test_fleet_mock.pytests/moe_ep/nccl_ep/test_ndtensor.pytests/moe_ep/test_compute_bridge.pytests/moe_ep/test_moe_ep_compute_correctness.pytests/moe_ep/test_moe_ep_ht_correctness.py
💤 Files with no reviewable changes (7)
- tests/moe/test_unified_moe.py
- tests/moe_ep/nccl_ep/test_ndtensor.py
- flashinfer/moe_ep/nccl_ep/ndtensor.py
- tests/moe_ep/test_moe_ep_compute_correctness.py
- tests/moe_ep/nccl_ep/test_fleet_mock.py
- tests/moe_ep/test_moe_ep_ht_correctness.py
- tests/moe_ep/test_compute_bridge.py
✅ Files skipped from review due to trivial changes (5)
- 3rdparty/nccl
- flashinfer/moe_ep/split_backends/nixl_ep_comm.py
- docs/design_docs/MoE_EP_impl.md
- benchmarks/MoE_benchmarks.md
- flashinfer/moe_ep/nccl_ep/init.py
🚧 Files skipped from review as they are similar to previous changes (14)
- flashinfer/moe_ep/split_backends/nccl_ep_comm.py
- flashinfer/fused_moe/layer.py
- flashinfer/fused_moe/runners.py
- docker/Dockerfile.flashinfer-nvep
- flashinfer/moe_ep/_validators.py
- pyproject.toml
- flashinfer/moe_ep/init.py
- scripts/build_in_container.sh
- flashinfer/moe_ep/config.py
- flashinfer/moe_ep/_compute_bridge.py
- flashinfer/moe_ep/nccl_ep/fleet.py
- build_backend.py
- flashinfer/moe_ep/layer.py
- flashinfer/moe_ep/nccl_ep/handle.py
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 (1)
benchmarks/run_httest_torchrun.sh (1)
12-23: 🩺 Stability & Availability | 🟠 Major | ⚡ Quick winFail fast when NCCL discovery misses the pinned version. If no
libnccl.so.2*here satisfies the version probe,NCCLLIBstays empty andLD_LIBRARY_PATHstarts with an empty component, sotorchruncan pick up an unintended NCCL from later paths. Build the path conditionally and exit early on discovery failure.🤖 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 `@benchmarks/run_httest_torchrun.sh` around lines 12 - 23, The NCCL discovery in the benchmark launch script can leave NCCLLIB empty when no matching libnccl.so.2* passes the version check, which then produces an unsafe LD_LIBRARY_PATH. Update the shell logic around the NCCLLIB probe to explicitly fail fast if no pinned NCCL is found, and only export LD_LIBRARY_PATH after a valid NCCLLIB directory has been discovered. Keep the fix localized to the NCCLLIB lookup and LD_LIBRARY_PATH export block so torchrun cannot fall back to an unintended NCCL.
🤖 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 `@benchmarks/run_httest_torchrun.sh`:
- Line 52: The NCCL JIT setup currently exports NCCL_EP_JIT_NVCC directly from
command -v nvcc, which leaves the variable empty and defers failure when nvcc is
missing. Update the script around the NCCL_EP_JIT_NVCC export to first resolve
nvcc into a local check, and if it is not found on PATH, exit immediately with a
clear message before continuing. Keep the fix in the same shell section where
NCCL_EP_JIT_NVCC is set so the behavior fails fast instead of propagating an
invalid environment value.
---
Outside diff comments:
In `@benchmarks/run_httest_torchrun.sh`:
- Around line 12-23: The NCCL discovery in the benchmark launch script can leave
NCCLLIB empty when no matching libnccl.so.2* passes the version check, which
then produces an unsafe LD_LIBRARY_PATH. Update the shell logic around the
NCCLLIB probe to explicitly fail fast if no pinned NCCL is found, and only
export LD_LIBRARY_PATH after a valid NCCLLIB directory has been discovered. Keep
the fix localized to the NCCLLIB lookup and LD_LIBRARY_PATH export block so
torchrun cannot fall back to an unintended NCCL.
🪄 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: 3b4c17ef-1639-41b4-a7d6-9158acc19d2a
📥 Commits
Reviewing files that changed from the base of the PR and between dc541fab4e74d61d5ecd331ed11ea4ff5c490e1f and 64c04b62a7c33a5c14e6b83045509bcb68151d22.
📒 Files selected for processing (13)
benchmarks/bench_ep_matrix.pybenchmarks/bench_moe_ep.pybenchmarks/run_ep_matrix_one.shbenchmarks/run_ep_matrix_one_pt.shbenchmarks/run_httest_torchrun.shdocker/install/build_flashinfer_ep_pytorch.shflashinfer/fused_moe/layer.pyflashinfer/fused_moe/runners.pyflashinfer/moe_ep/_compute_bridge.pyflashinfer/moe_ep/_validators.pyflashinfer/moe_ep/layer.pyflashinfer/moe_ep/split_backends/__init__.pytests/moe_ep/test_moe_ep_ht_correctness.py
🚧 Files skipped from review as they are similar to previous changes (11)
- flashinfer/moe_ep/split_backends/init.py
- benchmarks/run_ep_matrix_one_pt.sh
- flashinfer/fused_moe/layer.py
- flashinfer/fused_moe/runners.py
- tests/moe_ep/test_moe_ep_ht_correctness.py
- benchmarks/run_ep_matrix_one.sh
- benchmarks/bench_ep_matrix.py
- flashinfer/moe_ep/_compute_bridge.py
- docker/install/build_flashinfer_ep_pytorch.sh
- benchmarks/bench_moe_ep.py
- flashinfer/moe_ep/layer.py
Anerudhan
force-pushed
the
feat/moe_ep/enable_compute
branch
from
June 30, 2026 04:06
5c63cac to
bfdc7cc
Compare
aleozlx
added
the
unified_api
…on fast path Squashes the cross-node High-Throughput enablement with the FlashInfer-vs-ep_bench host-call analysis and Python-side optimization work. Cross-node HT: - Size the HT recv buffer to the max_recv budget (max_tokens_per_rank * world), not num_local_experts — fixes the nccl_ep.cc:3269 'invalid argument' at combine for num_experts > world^2 geometries. - Rebase the EP container on the PyTorch base image (Dockerfile.flashinfer-ep-pytorch), which provides the IB-GDAKI/GPUDirect stack cross-node HT actually needs (was nccl_ep.cc:2884 on the CUDA-13.0 image). - Add a pure-nccl.ep cross-node repro (tests/moe_ep/repro_ht_crossnode.py) and a continue-on-error matrix driver. Benchmark + analysis: - bench_ep_matrix.py: EP_TIMING=pipeline and EP_REUSE_PARAMS timing modes; nsys rank-0 wrapper, cross-node driver, and trace gap analysis (fi_gap_analysis.py). - nsys profiling confirms FlashInfer and ep_bench launch byte-identical dispatch/combine kernels with identical GPU time, single-node and to 64 GPU; the larger FI "kernel-only" number is the per-call host launch path, not kernel work. Burn-down: LL is Python-bound (recv-count .sum().item() readback + FFI-object churn), HT is library-bound (blocking nccl.ep dispatch). Python fast path (EP_FAST_PATH=1, env-gated; EP_PROFILE_HOST adds the per-step host-wall burn-down): - Lazy num_tokens via DispatchOutput.get_num_tokens() (defers the recv-count readback + stream sync; ints still work everywhere). - Cache the per-call FFI wrapper objects over stable tensors; rebuild only the input-token wrap. - Cache the LL recv buffer. Result: LL em dispatch host-wall 101 -> 17 us (-84%), bench measured 112 -> 67 us (-40%, approaching the 45 us pure kernel); HT unchanged (library-bound). Round-trip --validate passes for LL and HT. See benchmarks/HOST_CALL_BURNDOWN.md and benchmarks/KERNEL_MATCH_FI_vs_epbench.md. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…fast-path flag
Documentation + cleanup for the MoE-EP host-call/benchmark work:
- Consolidate the four MoE-EP markdown docs into two:
* benchmarks/MoE_benchmarks.md — container/Dockerfile, how to run the
benchmarks, and the 28-case comm results vs ep_bench (kernel + event-measured,
single + multi-node) plus the host-call burn-down / fast path and full-compute
results.
* docs/design_docs/MoE_EP_impl.md — the MoEEpLayer API + call stack and a
walkthrough of tests/moe_ep/test_moe_ep_compute_correctness.py.
(Removes HOST_CALL_BURNDOWN.md, KERNEL_MATCH_FI_vs_epbench.md,
MoE_EP_validation_results.md, MoE_EP_verif.md.)
- Drop optional profiling/diagnostic helpers that are not benchmark dependencies:
fi_gap_analysis.py, run_ep_matrix_one_pt_nsys.sh, run_fi_crossnode_nsys.sh,
the manual repro tests/moe_ep/repro_ht_crossnode.py, and its dead launcher
run_repro_one.sh (doc reference trimmed).
- Rename the opt-in host-call fast-path env var EP_FAST_PATH -> NV_FI_EP_FAST_PATH
(flashinfer/moe_ep/nccl_ep/handle.py). Validated multi-node (--validate
dispatch+combine OK): LL EXPERT_MAJOR 8/16/32 GPU, HT FLAT 8/16/32/64 GPU.
Still off by default (opt-in); see MoE_benchmarks.md §3.2.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Neither failure is Hopper-specific; both assert pre-PR behavior this branch intentionally changed. Validated on B200 (arch-independent logic): - tests/moe/test_unified_moe.py::...::test_non_nvfp4_quant_rejected[BF16]: this PR added BF16 to the MoELayer MVP scope (TrtllmBf16Config, the EP grouped-GEMM path), so BF16 is no longer rejected. Drop BF16 from the rejected-variants set (FP8/MxFp8/MXFP4/MxInt4 still rejected). - tests/moe_ep/test_compute_bridge.py::test_rank_major_pack_faithful_routing_and_masking: the test fed GLOBAL expert ids, but build_activation_pack_rank_major was corrected to the real library ABI (RANK_MAJOR returns LOCAL ids, -1 = non-local, converts local->global). Update the test to the local-index convention. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…back DispatchOutput.num_tokens was functionally dead: the only reader (MoEEpLayer._inner_compute_identity) ignored it, and the real compute path (build_activation_pack / _rank_major) derives everything from expert_tensors.shape + the recv_topk_idx/weights routing. Remove the field (and the get_num_tokens() lazy thunk added earlier). Producers drop the per-dispatch device->host recv-count readback that only existed to compute it (recv_count.sum().item() + ExternalStream.synchronize() in the LL / RANK_MAJOR paths; the constant in HT; nixl's sum().item()). The recv-count buffer is still allocated and passed to the library via LayoutInfo (the library writes counts there) — we just stop reading it. Net: a small unconditional host-side win and the NV_FI_EP_FAST_PATH no longer needs a num_tokens branch (it now gates only the FFI-wrapper + recv-buffer caching). NOTE: num_tokens is a pre-existing field on upstream's public DispatchOutput, so this is a breaking API change — flagged for maintainer review. Verified on Pre-Nyx B200: bench_ep_matrix.py --validate dispatch+combine OK for LL EXPERT_MAJOR and HT at 8 and 16 GPU (single + 2-node); test_moe_ep_compute_correctness EP-vs-kernel rel-err = 0.0045 for both LL layouts (8 GPU), matching the prior baseline. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Functional: - fused_moe/layer.py: build a runner only when the backend supports the configured quant variant (skip e.g. CuteDslNvfp4Runner under BF16 with mixed candidates). - moe_ep/_compute_bridge.py: bound rank-major/HT local ids to [0, num_local_experts) (treat out-of-range as non-local, weight 0) and validate idx/weights share [M,top_k]. - moe_ep/_validators.py: reject compute_config with execution.do_finalize=False (the EP bridge consumes a finalized output; RANK_MAJOR/HT need the weighted pre-reduce). - moe_ep/layer.py: run dispatch/compute/combine under try/finally so handle.complete() always drains the handle lifecycle if compute/combine raises (both timing modes). Docs/lint: - moe_ep/layer.py: correct docstrings — RANK_MAJOR and HT apply the received routing weights at real top_k (only EXPERT_MAJOR is the top_k=1 pre-routed pack). - fused_moe/runners.py: docstring now reflects GLOBAL-id packing. - split_backends/__init__.py: sort __all__ (RUF022). - bench_moe_ep.py: ASCII 'x' instead of Unicode multiplication sign (RUF002). - build_flashinfer_ep_pytorch.sh: PIP_CONSTRAINT="" instead of bare = (SC1007). - test_moe_ep_ht_correctness.py: spell out HT (high-throughput) in the docstring. - bench_*.py: clarify that EpLayout has no FLAT member; HT uses the library FLAT internally and FleetParams.layout is an inert placeholder for HT (not changed to a nonexistent EpLayout.FLAT). Benchmark scripts (stability): - run_ep_matrix_one_pt.sh / run_ep_matrix_one.sh / run_httest_torchrun.sh: fail fast on empty wheel/nvcc/include resolution and verify the JIT include tree before publishing .ready; run_httest_torchrun.sh is now one-torchrun-per-node (launch with --ntasks-per-node=1) instead of spawning 8 torchruns per node. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Addresses PR review nit: the backend-status filter comprehensions (`for b, on in ((name, flag), ...) if on`) now use `is_enabled`, which reads clearer than a bare `on`. Pure rename, no behavior change. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
…evice - bench_ep_matrix.py: under EP_REUSE_PARAMS=1, refresh the cached CombineInputParams with the current dispatch output each iteration (cin is a fresh tensor per call — a new recv buffer, or a per-iter clone under EP_SEPARATE_COMBINE_BUF=1 — so reusing the first cin combined stale data). - run_httest_torchrun.sh / run_ep_matrix_one.sh: split the nvcc lookup from the export and fail fast if nvcc is not on PATH (empty NCCL_EP_JIT_NVCC otherwise fails later in the HT JIT with a harder-to-diagnose error). - moe_ep/nccl_ep/handle.py: allocate _recv_count_t on topk_idx.device instead of the default "cuda" device, avoiding a device mismatch when the active device context differs from the routing tensor's device. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Anerudhan
force-pushed
the
feat/moe_ep/enable_compute
branch
from
June 30, 2026 18:39
1de0f19 to
8cde1f6
Compare
Anerudhan
disabled auto-merge
June 30, 2026 18:39
Anerudhan
enabled auto-merge (squash)
June 30, 2026 18:40
Collaborator
|
/bot run |
bkryu
approved these changes
Jun 30, 2026
Collaborator
waynehacking8
added a commit
to waynehacking8/flashinfer
that referenced
this pull request
Jul 1, 2026
…lashinfer-ai#3547) The EP-offset double-subtraction itself was fixed on main by flashinfer-ai#3686: pack_inputs now packs GLOBAL expert ids and the kernel owns the global->local mapping (mLocalExpertsStartIdx in RoutingKernel.h). This aligns the rest of the tree with that contract: - TestTrtllmEPOffset still asserted the old pre-subtracted packing (selected_experts - local_expert_offset), which contradicts the fixed pack_inputs; assert GLOBAL ids are preserved instead - add an EP-shard forward regression test: offset>0 output must match the bit-identical offset-0 run and must not be all-zero (the gh flashinfer-ai#3547 symptom) - drop the gh flashinfer-ai#3547 entry from the fuzzer _KNOWN_FAILURES ledger and align its CI-gate comments - fix the TrtllmFp4RoutedRunner docstrings that still described the pre-subtracted packing Fixes flashinfer-ai#3547 AI-assisted (Claude Code). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_012EK1bNDvPJbBvUMk7pumrw
This was referenced Jul 1, 2026
This was referenced Jul 2, 2026
Merged
4 tasks
Merged
3 tasks
aleozlx
pushed a commit
that referenced
this pull request
Jul 10, 2026
…3547) (#3591) ## Description **Rebased 2026-07-02 — scope reduced to tests + docs.** The EP-offset double-subtraction this PR originally fixed was fixed on main by #3686 (`pack_inputs` now packs GLOBAL expert ids; the kernel owns the global→local mapping via `mLocalExpertsStartIdx` in `include/flashinfer/trtllm/fused_moe/RoutingKernel.h`). What #3686 did **not** update: - `tests/moe/test_unified_moe.py::TestTrtllmEPOffset::test_pack_inputs_applies_local_expert_offset` still asserts the **old** pre-subtracted packing (`decoded == selected_experts - local_expert_offset`), which now contradicts `pack_inputs` — for `local_expert_offset > 0` it fails wherever it runs (it is `sm100_required`). - `tests/moe/test_unified_moe_fuzz.py` still carries the gh #3547 `_KNOWN_FAILURES` xfail entry for a now-fixed bug (per the ledger's own design, a fixed bug's entry must be removed) plus stale "pending gh #3547" CI-gate comments. - The `TrtllmFp4RoutedRunner` class/`pack_inputs` docstrings still describe the pre-subtracted packing. This PR aligns all of the above with the merged contract and adds an EP-shard forward regression test (`test_ep_shard_forward_matches_offset_zero`): `offset > 0` output must match the bit-identical `offset = 0` run and must not be all-zero (the gh #3547 symptom), so the double-subtraction can't silently come back. Fixes #3547 (the code fix landed in #3686; this closes the loop with regression coverage). ## Related Issues gh #3547 · supersedes-note: code fix landed via #3686 ## Verification - `ruff check` + `ruff format --check` clean on the three touched files - `pytest tests/moe/test_unified_moe.py` on RTX PRO 6000 (SM120, CUDA 13): **99 passed, 11 skipped** (the `sm100_required` tests, including the ones this PR edits, skip on SM120 — they need SM100 CI to execute) - `pytest --collect-only` on both test files: 196 collected, no import errors AI-assisted (Claude Code). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **Bug Fixes** * Fixed TRTLLM FP4 routed MoE all-zero outputs for nonzero expert offsets by keeping global expert IDs in packed top-k IDs and applying expert offset remapping during kernel execution. * **Tests** * Updated input-packing assertions to verify global expert ID preservation and correct shard-range mapping. * Added a GPU correctness regression test comparing nonzero-offset results against the zero-offset baseline. * Removed the prior expected-failure entry for this configuration from the fuzz/ledger logic. * **Documentation** * Clarified routed MoE packed top-k ID semantics and kernel-side offset behavior. <!-- end of auto-generated comment: release notes by coderabbit.ai --> --------- Signed-off-by: Wayne Chiu <waynehacking8@gmail.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com>
Anerudhan
pushed a commit
that referenced
this pull request
Jul 14, 2026
<!-- .github/pull_request_template.md --> ## [MoE-EP] Unify split + mega under one MoEEpLayer; add DeepGEMM / CuTeDSL mega backends ### Summary Extends #3686 (moe_ep with flat split only path). This PR merges that split path and a new mega path into a single orchestrated `MoEEpLayer`, and adds three mega backends. Major changes in these two folders: `flashinfer/moe_ep` and `/tests/moe_ep` only; rest are just minor wiring. `flashinfer/moe_ep/backends/mega/kernel` additionally includes `cutedsl_backend_kernels` which are cutedsl mega_kernel implementations from Jack Yang's team. Also see [[Design doc] ](https://github.com/mhoqueanik/flashinfer-moe_ep/blob/mega_moe_integration/docs/design_docs/moe_ep_architecture.md). Please note that current split path only supports bf16 compute on trtllm fused_moe kernel. ### What's included - **Unified entry point**: `MoEEpLayer(bootstrap, fleet_params, backend=...)` dispatches by config — `SplitConfig` → `MoEEpSplitLayer`, `MegaConfig` → `MoEEpMegaLayer`. - **Mega path**: a single symmetric-memory kernel fusing expert-parallel comm with the local MoE (no separate Fleet/Handle), with weight preprocessing + workspace lifecycle and bf16-input quantization. - **Three mega backends**: - `deep_gemm_mega` (`DeepGemmMegaMoeConfig`) — DeepGEMM mega_moe, FP8/FP4 - `nvfp4_cutedsl` (`Nvfp4CutedslMegaMoeConfig`) — CuTeDSL NVFP4 - `mxfp8_cutedsl` (`Mxfp8CutedslMegaMoeConfig`) — CuTeDSL MXFP8 (e4m3/e5m2) - all Blackwell (sm_100+); CuTeDSL kernels need NVSHMEM (skip with `MEGA_NO_DIST=1`). - **Package restructure**: shared `core/` abstractions, `backends/split/**` (comm + inner kernels) and `backends/mega/kernel/**`, `modes/` (split/mega layers + config). The #3686 split backends are migrated into this layout. - **Tests + benchmarks** updated to the new API layout; existing split correctness tests merged in; new mega multirank parity + single-rank MXFP8 preprocess-vs-reference tests. - **Docs**: package architecture doc. ### Directories affected - `flashinfer/moe_ep/` — unified `layer.py` factory, `modes/` (split_layer, mega_layer, config), `core/`, restructured `backends/split/**`, new `backends/mega/kernel/**` (deep_gemm_mega, nvfp4_cutedsl, mxfp8_cutedsl, cutedsl_backend_kernels), config/tensors/weights - `tests/moe_ep/` — mega multirank tests, MXFP8 preprocess-vs-reference, migrated split tests, updated `run_tests.sh` targets - `benchmarks/` — EP bench scripts updated for the new API layout - `docs/design_docs/` — new `moe_ep_architecture.md` - `docker/install/` — EP build script now also installs mega deps (DeepGEMM, NVSHMEM, CUTLASS DSL) ### Test plan See full run-book [[here]](https://github.com/mhoqueanik/flashinfer-moe_ep/blob/mega_moe_integration/docs/design_docs/moe_ep_runbook.md) #### Create docker - `salloc` a node (Tested on 1x4GB200 GPUs) ```shell export RW=/path/to/flashinfer/repo srun --jobid="$SLURM_JOB_ID" -N1 \ --container-image=nvcr.io/nvidia/pytorch:26.05-py3 \ --container-save=$RW/flashinfer-ep-pt2605-mega_moe_ep.sqsh \ --container-mounts=$RW:/host \ bash -lc 'bash /host/flashinfer/docker/install/build_flashinfer_ep_pytorch.sh' export IMG=$RW/flashinfer-ep-pt2605-mega_moe_ep.sqsh export ROOT=/workspace srun --jobid="$SLURM_JOB_ID" \ --overlap \ --container-image="$IMG" \ --container-mounts="$ROOT:$ROOT" \ --container-workdir="$ROOT/flashinfer" \ --pty bash -l BUILD_NVEP=0 BUILD_NCCL_EP=1 BUILD_NIXL_EP=0 \ pip install --no-cache-dir --no-build-isolation -e ".[nvep]" ``` #### Run tests - `bash tests/moe_ep/run_tests.sh unit` - `bash tests/moe_ep/run_tests.sh multirank` (4 GPU, NCCL-EP; NIXL-EP if built) - `bash tests/moe_ep/run_tests.sh split_path_correctness_bf16` (4 GPU, Blackwell) - `bash tests/moe_ep/run_tests.sh mega` (4 GPU, Blackwell sm_100+; DeepGEMM + NVFP4 + MXFP8) --- ## 🔍 Related Issues #3692 #3781 ## 🧪 Tests - [✔️] Tests have been added or updated as needed. - [✔️] All tests are passing (`unittest`, etc.). --------- Co-authored-by: root <root@ptyche0307.ptyche.clusters.nvidia.com> Co-authored-by: Md Saidul Hoque Anik <mhoqueanik@login-ptyche02.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0348.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0297.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0341.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0355.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0351.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0323.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0289.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0339.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0324.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0360.ptyche.clusters.nvidia.com> Co-authored-by: root <root@ptyche0343.ptyche.clusters.nvidia.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com> Co-authored-by: root <root@ptyche0053.ptyche.clusters.nvidia.com> Co-authored-by: Claude Fable 5 <noreply@anthropic.com> Co-authored-by: Md Saidul Hoque Anik <mhoqueanik@login-preos01.a51.clusters.nvidia.com>
This was referenced Jul 21, 2026
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
4 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.
Summary
Wires FlashInfer's unified MoE compute API into expert parallelism: a new
flashinfer.moe_ep.MoEEpLayerruns one MoE layer split across ranks asdispatch → per-expert grouped GEMM → combine, over a pluggable transport
(NCCL-EP via
nccl.ep/nccl4py, and NIXL-EP). The expert GEMM reuses theunified
flashinfer.fused_moe.MoELayeras a pure per-expert grouped GEMM (routinglives in dispatch/combine).
What's included
MoEEpLayer(flashinfer/moe_ep/) with both algorithms and all three receivelayouts:
EXPERT_MAJOR(combine reweights on receive) andRANK_MAJOR(received-routing, combine sums across ranks).
FLAT(token-major received routing, unweighted combine).nccl_ep(nccl4py ≥ 0.3.1,nccl.epv0.1.0) andnixl_ep; built viathe
[nvep]extra (BUILD_NCCL_EP=1).docker/Dockerfile.flashinfer-ep-pytorch(PyTorch base, CUDA 13.2) —required for cross-node HT (GIN/GDAKI); the CUDA-13.0 image aborts multi-node HT at
nccl_ep.cc:2884(shown to be the base image's IB/GPUDirect runtime, not FlashInfer).bench_moe_ep.py(full MoE) andbench_ep_matrix.py(comm-only,ep_bench-comparable), plus a SLURM matrix harness. An opt-in host-call fast path
(
NV_FI_EP_FAST_PATH=1) and a per-step host-wall burn-down (EP_PROFILE_HOST=1).tests/moe_ep/— multi-GPU numerical correctness(
test_moe_ep_compute_correctness.py,test_moe_ep_ht_correctness.py), layout-bridgeunit tests, config/constraints, and transport smokes.
benchmarks/MoE_benchmarks.md(container/run/results vs ep_bench) anddocs/design_docs/MoE_EP_impl.md(API + call stack + correctness-test walkthrough).Correctness
8× B200, bf16: EP
dispatch→compute→combinematches the sameMoELayerkernel runnon-EP to rel-err ≈ 0.0045 for both LL layouts and ≈ 0.007 for HT FLAT
(4096/8192 tok/rank). Two multi-rank bugs found & fixed (global-vs-local expert ids;
local→global
topk_idxfor RANK_MAJOR/HT).Benchmarks (Pre-Nyx B200) vs ep_bench
FlashInfer and the upstream
contrib/nccl_ep/ep_benchreference launch byte-identicaldispatch/combine kernels with the same GPU time, single-node and to 64 GPU (verified by
Nsight Systems). The Python host-call overhead was profiled and reduced behind the opt-in
NV_FI_EP_FAST_PATH(LL dispatch host-wall 101→17 µs); validated multi-node(
--validatedispatch+combine OK at 8/16/32/64 GPU). HT is library-bound on the blockingnccl.epdispatch. Seebenchmarks/MoE_benchmarks.md.Validated bf16 end-to-end on GB200 (SM100), GB300 (SM103), GB200-NVL36, and Pre-Nyx B200,
on both
nccl_epandnixl_ep.Known limitations / follow-ups
nccl_epfails atnccl_ep.cc:1491(LL-GIN dev-comm setup) —library-side; HT 64-GPU passes on the same nodes.
NV_FI_EP_FAST_PATHstays opt-in: it currently covers LL EXPERT_MAJOR + HT andchanges the
num_tokens/buffer-aliasing contracts (see doc); to be extended + auditedbefore defaulting.
Notes
nccl4py >= 0.3.1(thenccl.epAPI); no in-tree NCCL build. NCCL-EP's GINtransport needs the GDAKI/GPUDirect stack at runtime, and NCCL ≥ 2.30.7 on B200
(with
NCCL_MNNVL_ENABLE=1for multi-node).🤖 Generated with Claude Code
Summary by CodeRabbit
Release Notes
New Features
MoEEpLayerwith an optional layout-aware dispatch → compute → combine path.EpLayout) and improved backend/quant variant validation.Documentation
Build & Infrastructure
Tests