[TRTLLM-9389][chore] Refactor AlltoallMethodType.

[bobboli]

Summary by CodeRabbit

• Refactor

  • Refactored Mixture of Experts all-to-all communication method selection from string-based backend flags to dedicated method type enumerations, improving code clarity and consistency.
  • Removed public moe_alltoall_backend property; method selection now uses dedicated enum types.
  • Updated default behavior for all-to-all method selection.

• Tests

  • Updated tests to cover new all-to-all method type variants.

_{✏️ Tip: You can customize this high-level summary in your review settings.}

Description

Test Coverage

PR Checklist

Please review the following before submitting your PR:

• PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.

• PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.

• Test cases are provided for new code paths (see test instructions)

• Any new dependencies have been scanned for license and vulnerabilities

• CODEOWNERS updated if ownership changes

• Documentation updated as needed

• Update tava architecture diagram if there is a significant design change in PR.

• The reviewers assigned automatically/manually are appropriate for the PR.

• Please check this after reviewing the above items as appropriate for this PR.

GitHub Bot Help

/bot [-h] ['run', 'kill', 'skip', 'reuse-pipeline'] ...

Provide a user friendly way for developers to interact with a Jenkins server.

Run /bot [-h|--help] to print this help message.

See details below for each supported subcommand.

Details

run [--reuse-test (optional)pipeline-id --disable-fail-fast --skip-test --stage-list "A10-PyTorch-1, xxx" --gpu-type "A30, H100_PCIe" --test-backend "pytorch, cpp" --add-multi-gpu-test --only-multi-gpu-test --disable-multi-gpu-test --post-merge --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" --detailed-log --debug(experimental)]

Launch build/test pipelines. All previously running jobs will be killed.

--reuse-test (optional)pipeline-id (OPTIONAL) : Allow the new pipeline to reuse build artifacts and skip successful test stages from a specified pipeline or the last pipeline if no pipeline-id is indicated. If the Git commit ID has changed, this option will be always ignored. The DEFAULT behavior of the bot is to reuse build artifacts and successful test results from the last pipeline.

--disable-reuse-test (OPTIONAL) : Explicitly prevent the pipeline from reusing build artifacts and skipping successful test stages from a previous pipeline. Ensure that all builds and tests are run regardless of previous successes.

--disable-fail-fast (OPTIONAL) : Disable fail fast on build/tests/infra failures.

--skip-test (OPTIONAL) : Skip all test stages, but still run build stages, package stages and sanity check stages. Note: Does NOT update GitHub check status.

--stage-list "A10-PyTorch-1, xxx" (OPTIONAL) : Only run the specified test stages. Examples: "A10-PyTorch-1, xxx". Note: Does NOT update GitHub check status.

--gpu-type "A30, H100_PCIe" (OPTIONAL) : Only run the test stages on the specified GPU types. Examples: "A30, H100_PCIe". Note: Does NOT update GitHub check status.

--test-backend "pytorch, cpp" (OPTIONAL) : Skip test stages which don't match the specified backends. Only support [pytorch, cpp, tensorrt, triton]. Examples: "pytorch, cpp" (does not run test stages with tensorrt or triton backend). Note: Does NOT update GitHub pipeline status.

--only-multi-gpu-test (OPTIONAL) : Only run the multi-GPU tests. Note: Does NOT update GitHub check status.

--disable-multi-gpu-test (OPTIONAL) : Disable the multi-GPU tests. Note: Does NOT update GitHub check status.

--add-multi-gpu-test (OPTIONAL) : Force run the multi-GPU tests in addition to running L0 pre-merge pipeline.

--post-merge (OPTIONAL) : Run the L0 post-merge pipeline instead of the ordinary L0 pre-merge pipeline.

--extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx" (OPTIONAL) : Run the ordinary L0 pre-merge pipeline and specified test stages. Examples: --extra-stage "H100_PCIe-TensorRT-Post-Merge-1, xxx".

--detailed-log (OPTIONAL) : Enable flushing out all logs to the Jenkins console. This will significantly increase the log volume and may slow down the job.

--debug (OPTIONAL) : Experimental feature. Enable access to the CI container for debugging purpose. Note: Specify exactly one stage in the stage-list parameter to access the appropriate container environment. Note: Does NOT update GitHub check status.

For guidance on mapping tests to stage names, see docs/source/reference/ci-overview.md
and the scripts/test_to_stage_mapping.py helper.

kill

kill

Kill all running builds associated with pull request.

skip

skip --comment COMMENT

Skip testing for latest commit on pull request. --comment "Reason for skipping build/test" is required. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

reuse-pipeline

reuse-pipeline

Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break.

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #25488 [ run ] triggered by Bot. Commit: eee5304

[coderabbitai]

📝 Walkthrough

Walkthrough

This change refactors the MOE all-to-all backend selection from string-based checks to an enumerated type system. The MNNVL enum value is replaced with NVLinkOneSided and NVLinkTwoSided, the moe_alltoall_backend cached property is removed, and all conditional logic across multiple implementations is updated to use AlltoallMethodType enum comparisons instead of backend string checks.

Changes

Cohort / File(s)	Change Summary
Enum definition tensorrt_llm/_torch/modules/fused_moe/interface.py	Refactored AlltoallMethodType enum: removed MNNVL = 1, added NVLinkOneSided = 1 and NVLinkTwoSided = 2, renumbered DeepEP to 3 and DeepEPLowLatency to 4
Cutlass implementation tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py	Replaced moe_alltoall_backend string checks with AlltoallMethodType enum comparisons; removed cached_property; updated select_alltoall_method_type to default to NVLinkOneSided; refactored control flow for alltoall enablement, workspace selection, and payload handling to use enum-based routing
TRTLLMGen implementation tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py	Replaced moe_alltoall_backend string checks with AlltoallMethodType enum comparisons; removed cached_property; updated default selection logic from MNNVL to NVLinkOneSided; refactored initialization, quantization, and forward logic to consistently route on alltoall_method_type
WideEP implementation tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py	Replaced MNNVL references with NVLinkTwoSided across feasibility checks, dispatch paths, and post-quantization handling; removed moe_alltoall_backend cached_property; removed TRTLLM_MOE_DISABLE_ALLTOALLV environment variable gate
Test updates tests/unittest/_torch/modules/test_fused_moe.py	Updated parametrization to use NVLinkOneSided and NVLinkTwoSided instead of MNNVL; relaxed output-shape conditions to apply sum(dim=1) for all method types when output.ndim == 3

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

• Enum value reordering: Verify enum value changes (MNNVL = 1 → NVLinkOneSided = 1; DeepEP = 2 → 3; DeepEPLowLatency = 3 → 4) do not break serialization, deserialization, or stored configurations
• Default behavior change: Confirm NVLinkOneSided default (replacing MNNVL) maintains expected behavior and does not introduce performance regressions
• Removed moe_alltoall_backend property: Ensure no external code or tests depend on this public cached_property
• TRTLLM_MOE_DISABLE_ALLTOALLV removal: Verify removal of this environment variable gate is intentional and does not break existing deployments that may rely on it
• Enum conditional logic: Check all AlltoallMethodType comparisons in each implementation (Cutlass, TRTLLMGen, WideEP) correctly handle the new enum variants and default cases

Pre-merge checks and finishing touches

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The PR description is essentially empty—it only contains the repository template and CodeRabbit prompts with no actual implementation details, rationale, or test coverage information.	Add a concrete description explaining the refactoring objectives (replacing MNNVL with NVLinkOneSided/NVLinkTwoSided), affected modules, and relevant test coverage.
Docstring Coverage	⚠️ Warning	Docstring coverage is 30.00% which is insufficient. The required threshold is 80.00%.	You can run @coderabbitai generate docstrings to improve docstring coverage.

✅ Passed checks (1 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly identifies the main change: refactoring the AlltoallMethodType enum, which aligns with the primary modifications across multiple files.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment

Tip

📝 Customizable high-level summaries are now available in beta!

You can now customize how CodeRabbit generates the high-level summary in your pull requests — including its content, structure, tone, and formatting.

• Provide your own instructions using the high_level_summary_instructions setting.
• Format the summary however you like (bullet lists, tables, multi-section layouts, contributor stats, etc.).
• Use high_level_summary_in_walkthrough to move the summary from the description to the walkthrough section.

Example instruction:

"Divide the high-level summary into five sections:

1. 📝 Description — Summarize the main change in 50–60 words, explaining what was done.
2. 📓 References — List relevant issues, discussions, documentation, or related PRs.
3. 📦 Dependencies & Requirements — Mention any new/updated dependencies, environment variable changes, or configuration updates.
4. 📊 Contributor Summary — Include a Markdown table showing contributions:
   | Contributor | Lines Added | Lines Removed | Files Changed |
5. ✔️ Additional Notes — Add any extra reviewer context.
   Keep each section concise (under 200 words) and use bullet or numbered lists for clarity."

Note: This feature is currently in beta for Pro-tier users, and pricing will be announced later.

---

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

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

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

[coderabbitai]

Actionable comments posted: 0

Caution

Some comments are outside the diff and can’t be posted inline due to platform limitations.

⚠️ Outside diff range comments (1)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (1)

163-172: Minor: error message references the wrong class name

The NotImplementedError for DeepEP/DeepEPLowLatency currently says "not supported for CutlassFusedMoE yet" even though this is inside TRTLLMGenFusedMoE.select_alltoall_method_type.

A tiny wording fix will avoid confusion:

-                raise NotImplementedError(
-                    "DeepEP and DeepEPLowLatency are not supported for CutlassFusedMoE yet"
-                )
+                raise NotImplementedError(
+                    "DeepEP and DeepEPLowLatency are not supported for TRTLLMGenFusedMoE yet"
+                )

🧹 Nitpick comments (1)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (1)

696-700: Chunked DP overlap is disabled for NVLink alltoall (by design?)

For the multi‑chunk path, the overlap between compute and reducescatter_or_allreduce is now skipped when alltoall_method_type is NVLinkOneSided or NVLinkTwoSided, while it remains enabled for the non‑alltoall case.

That’s reasonable given that alltoall already dominates communication cost, but if you later add more alltoall types it may be cleaner to express this as a dedicated flag (e.g., supports_chunk_overlap) rather than hard‑coding specific enum values.

Also applies to: 717-721

📜 Review details

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between c045e35 and eee5304.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (12 hunks)
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (10 hunks)
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (9 hunks)
• tensorrt_llm/_torch/modules/fused_moe/interface.py (1 hunks)
• tests/unittest/_torch/modules/test_fused_moe.py (4 hunks)

🧰 Additional context used

🧠 Learnings (13)

📓 Common learnings

Learnt from: nzmora-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 9163
File: tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py:107-113
Timestamp: 2025-11-14T11:22:03.729Z
Learning: In TensorRT-LLM AutoDeploy custom ops, when adding hardware capability checks to select between kernel implementations (e.g., cuBLAS vs. CUDA kernel), use descriptive variable names that identify the specific GPU architectures or families being targeted (e.g., `is_blackwell_geforce_or_ada`) rather than generic names like `enable_cuda_core`. This makes it clear that the code is selecting an implementation path based on hardware capabilities, not enabling/disabling hardware features.

📚 Learning: 2025-08-14T06:36:40.701Z

Learnt from: timlee0212
Repo: NVIDIA/TensorRT-LLM PR: 6886
File: tensorrt_llm/_torch/models/modeling_deepseekv3.py:0-0
Timestamp: 2025-08-14T06:36:40.701Z
Learning: In DeepSeek V3 model (tensorrt_llm/_torch/models/modeling_deepseekv3.py), the disagreement between AllReduce.__init__ guard and _compute_mlp_tp_size logic for MNNVL usage is expected by design. The AllReduce component and MLP TP-size computation intentionally use different criteria for MNNVL availability decisions.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/interface.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

📚 Learning: 2025-08-21T02:39:12.009Z

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 7104
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1475-1480
Timestamp: 2025-08-21T02:39:12.009Z
Learning: The min latency mode functionality in TensorRT-LLM MOE kernels (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu) is deprecated and no longer being maintained/updated, as confirmed by djns99. Bug reports and optimization suggestions for the computeStridesTmaWarpSpecializedLowLatencyKernel and related min latency code paths should be deprioritized.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/interface.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

📚 Learning: 2025-08-14T23:23:27.449Z

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4010-4012
Timestamp: 2025-08-14T23:23:27.449Z
Learning: For MOE (Mixture of Experts) code reviews in TensorRT-LLM, avoid repeatedly suggesting finalize fusion validation checks and safety assertions. The user djns99 has indicated these suggestions are repetitive and unwanted across multiple MOE-related changes.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
• tests/unittest/_torch/modules/test_fused_moe.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

📚 Learning: 2025-08-19T03:35:20.866Z

Learnt from: djns99
Repo: NVIDIA/TensorRT-LLM PR: 6915
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:4616-4626
Timestamp: 2025-08-19T03:35:20.866Z
Learning: In the MOE profiler TMA workspace preparation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu), the overlapping of TMA WS regions for NONE and FINALIZE variants is deliberate design to save memory space, as confirmed by djns99. The comment "reuse the same pointers to save space" reflects this intentional behavior.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

📚 Learning: 2025-09-23T15:12:38.312Z

Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: cpp/tensorrt_llm/thop/allreduceOp.cpp:352-446
Timestamp: 2025-09-23T15:12:38.312Z
Learning: In TensorRT-LLM NCCL device allreduce implementation (cpp/tensorrt_llm/thop/allreduceOp.cpp), the goto pattern in runNCCLAllReduceDeviceFusion is intentionally used for future extensibility, allowing multiple switch cases to fallback to the default handler. While not aesthetically ideal, this pattern supports adding more fusion cases later that can reuse the same fallback logic.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

📚 Learning: 2025-08-09T20:57:04.084Z

Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu:118-127
Timestamp: 2025-08-09T20:57:04.084Z
Learning: In the CUTLASS MoE finalize fusion implementation (cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_gemm_tma_warp_specialized_input.cu), when setting `fused_finalize_epilogue.stride_final_output` with shape `(hidden_size, num_output_tokens, 1)`, the `num_rows_in_final_output` should be set to `num_output_tokens` (not `hidden_size`) because of a swap+transpose operation that maps rows of the output tensor to `hidden_size` and columns to `num_output_tokens`.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
• tests/unittest/_torch/modules/test_fused_moe.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

📚 Learning: 2025-08-19T12:45:11.997Z

Learnt from: amitz-nv
Repo: NVIDIA/TensorRT-LLM PR: 7033
File: tensorrt_llm/_torch/pyexecutor/model_engine.py:0-0
Timestamp: 2025-08-19T12:45:11.997Z
Learning: In tensorrt_llm/_torch/pyexecutor/model_engine.py, DoRA (Delta Orthogonal Rank Adaptation) functionality was removed from the PyTorch flow to eliminate issues with inverted DoRA detection logic. The original is_dora condition was checking if scaling_vec_pointer == 0, which was potentially incorrect.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

📚 Learning: 2025-08-08T22:03:40.707Z

Learnt from: sklevtsov-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 3294
File: cpp/tensorrt_llm/kernels/cutlass_kernels/moe_gemm/moe_kernels.cu:1198-1209
Timestamp: 2025-08-08T22:03:40.707Z
Learning: In the CUTLASS MoE kernels (cpp/tensorrt_llm/cutlass_extensions), when `layout_info.fusion` is set to `TmaWarpSpecializedGroupedGemmInput::EpilogueFusion::FINALIZE`, the `router_scales` parameter must be non-null by design. The fused finalize kernel epilogue does not perform nullptr checks and requires valid router scales to function correctly. This is an implicit contract that callers must satisfy when enabling the FINALIZE fusion mode.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

📚 Learning: 2025-09-19T21:28:13.751Z

Learnt from: jhaotingc
Repo: NVIDIA/TensorRT-LLM PR: 7856
File: cpp/tensorrt_llm/thop/fp8BlockScaleMoe.cpp:159-166
Timestamp: 2025-09-19T21:28:13.751Z
Learning: In TensorRT-LLM blockScaleMoe routing (cpp/tensorrt_llm/kernels/trtllmGenKernels/blockScaleMoe/runner.cu), the DeepSeek routing method performs reinterpret_cast<float*>(routingLogits) at line 89, which could cause issues if routing_logits are BF16. However, Qwen3-FP8 models use RenormalizeNaive routing method and are not affected by this dtype casting issue.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

📚 Learning: 2025-11-14T11:22:03.729Z

Learnt from: nzmora-nvidia
Repo: NVIDIA/TensorRT-LLM PR: 9163
File: tensorrt_llm/_torch/auto_deploy/custom_ops/quant.py:107-113
Timestamp: 2025-11-14T11:22:03.729Z
Learning: In TensorRT-LLM AutoDeploy custom ops, when adding hardware capability checks to select between kernel implementations (e.g., cuBLAS vs. CUDA kernel), use descriptive variable names that identify the specific GPU architectures or families being targeted (e.g., `is_blackwell_geforce_or_ada`) rather than generic names like `enable_cuda_core`. This makes it clear that the code is selecting an implementation path based on hardware capabilities, not enabling/disabling hardware features.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

📚 Learning: 2025-08-20T07:43:36.447Z

Learnt from: ChristinaZ
Repo: NVIDIA/TensorRT-LLM PR: 7068
File: cpp/tensorrt_llm/kernels/moeTopKFuncs.cuh:169-172
Timestamp: 2025-08-20T07:43:36.447Z
Learning: In TensorRT-LLM MOE kernels, when processing up to 128 experts across 32 threads, each thread handles at most 4 experts (N < 5 constraint), where N represents candidates per thread rather than total system capacity.

Applied to files:

• tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py

📚 Learning: 2025-10-13T19:45:03.518Z

Learnt from: nv-lschneider
Repo: NVIDIA/TensorRT-LLM PR: 7910
File: tests/unittest/_torch/multi_gpu/test_nccl_device.py:138-149
Timestamp: 2025-10-13T19:45:03.518Z
Learning: In test_nccl_device.py, the NCCL device AllReduce implementation compares the entire residual tensor on each rank, unlike the UB implementation which compares per-rank chunks. The residual chunking calculations in the test are intentionally overridden to reflect this design difference.

Applied to files:

• tests/unittest/_torch/modules/test_fused_moe.py
• tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py

🧬 Code graph analysis (4)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (1)

tensorrt_llm/_torch/modules/fused_moe/interface.py (1)

• AlltoallMethodType (26-36)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (1)

tensorrt_llm/_torch/modules/fused_moe/interface.py (1)

• AlltoallMethodType (26-36)

tests/unittest/_torch/modules/test_fused_moe.py (1)

tensorrt_llm/_torch/modules/fused_moe/interface.py (1)

• AlltoallMethodType (26-36)

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (2)

tensorrt_llm/_torch/modules/fused_moe/interface.py (1)

• AlltoallMethodType (26-36)

tensorrt_llm/_torch/distributed/moe_alltoall.py (1)

• MoeAlltoAll (26-235)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (1)

• GitHub Check: Pre-commit Check

🔇 Additional comments (11)

tensorrt_llm/_torch/modules/fused_moe/interface.py (1)

25-36: Enum value changes – confirm there are no external numeric dependencies

AlltoallMethodType now renumbers members (NVLinkOneSided=1, NVLinkTwoSided=2, DeepEP=3, DeepEPLowLatency=4). This is fine as long as all call sites use the symbolic names or env strings, but it would break any persisted/int-based contracts (e.g. configs, logs parsed as raw ints, cross-language bindings). Please double‑check that nothing outside this repo relies on the previous numeric values.

tests/unittest/_torch/modules/test_fused_moe.py (2)

210-215: Alltoall test parametrization correctly tracks new enum values

The parametrized alltoall_method_type sets now cover NVLinkOneSided, NVLinkTwoSided, DeepEP, DeepEPLowLatency, and NotEnabled where appropriate, matching the new AlltoallMethodType contract used in the implementations. This keeps coverage in sync with the refactor.

Also applies to: 323-329, 685-689

---

305-307: Shape‑based reduction for 3D outputs looks appropriate

Switching to if output.ndim == 3: output = output.sum(dim=1) decouples the test from specific alltoall methods and relies only on the presence of an extra top_k dimension. That matches how the WideEP/Cutlass paths structure their outputs today and keeps the test robust as new alltoall methods are added.

If any backend ever returns 3D output for reasons other than a top_k dimension, you may want a more explicit condition here.

tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py (3)

145-168: NVLinkTwoSided vs NVLinkOneSided initialization looks consistent

The split between NVLinkTwoSided (Mnnvl workspaces) and NVLinkOneSided (Python MoeAlltoAll with num_experts=self.num_slots and invalid‑expert id self.num_slots) matches the WideEPMoE usage pattern and the MoeAlltoAll contract. Workspace sizing via TRTLLM_MOE_A2A_WORKSPACE_MB is clear and bounded per rank.

---

233-237: Stronger default to NVLinkOneSided – consider external expectations

select_alltoall_method_type now always returns NVLinkOneSided when DP>1, moe_tp_size==1, and MnnvlMemory.supports_mnnvl(), instead of honoring previous heuristics like ep_size > top_k or an explicit disable flag. This is likely desired for performance, but it does change behavior for existing deployments that might have relied on the old heuristic or TRTLLM_MOE_DISABLE_ALLTOALLV.

If you have workloads that were intentionally running without alltoall despite MNNVL being available, you may want to:

• Document the new default, and/or
• Recommend TRTLLM_FORCE_ALLTOALL_METHOD=NotEnabled where disabling is still needed.

Also applies to: 239-243

---

314-317: Load‑balancer stats and NVLink alltoall flows are wired correctly

• Using ignore_allreduce = (alltoall_method_type == NVLinkTwoSided) delegates global statistics aggregation to the NVLink prepare/combine path, avoiding double allreduce.
• The NVLinkTwoSided path prepares with mnnvl_moe_alltoallv_prepare_without_allgather, updates global stats via update_statistic_with_gathered_statistic, and uses memset_expert_ids(..., invalid_id=self.num_slots) for padding, which matches the existing MNNVL semantics.
• The NVLinkOneSided path uses MoeAlltoAll.dispatch/combine with consistent payload ordering and view/reshape logic, and correctly reuses a workspace‑backed output when available.

Overall, the control flow and tensor reshaping look sound.

Also applies to: 427-460, 571-597

tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py (3)

113-153: WideEPMoE alltoall initialization for NVLink/DeepEP variants looks correct

The constructor cleanly maps:

• NVLinkTwoSided → MnnvlMemory.initialize() plus MnnvlMoe.get_moe_workspaces/get_moe_prepare_workspace.
• DeepEP → shared buffer_pool with reserve(hidden_size, dtype).
• DeepEPLowLatency → low‑latency buffer with reserve(self.deep_ep_max_num_tokens, hidden_size, self.num_slots) and NVSHMEM_QP_DEPTH tuning.

Unknown/NotEnabled method types are rejected early. This is consistent with the intended separation between NVLink and DeepEP backends.

---

185-223: Alltoall selection and gating – behavior changes look intentional, please confirm

• select_alltoall_method_type now:

  • Requires enable_attention_dp and mapping.tp_size > 1.
  • Disables alltoall when moe_ep_size <= top_k.
  • Prefers NVLinkTwoSided if MnnvlMemory.supports_mnnvl().
  • Otherwise, optionally selects DeepEP/DeepEPLowLatency based on TRTLLM_CAN_USE_DEEP_EP and BF16.

• can_use_alltoall always returns True for NVLinkTwoSided, even when multiple MoE chunks are required, while still disabling alltoall under chunking for DeepEP/LowLatency.

• is_post_quant_all2all_supported now treats NVLinkTwoSided as always supporting post‑quant alltoall, and narrows DeepEP/LowLatency support to specific quant algorithms.

This aligns with the new NVLink‑centric design, but it does slightly broaden the scenarios where NVLinkTwoSided alltoall is used (e.g., with chunking). Please verify this matches the expected capability/perf envelope of the NVLink path.

Also applies to: 245-263, 367-377

---

400-401: NVLinkTwoSided load‑balancer integration and combine flow look consistent

• Using ignore_allreduce = (alltoall_method_type == NVLinkTwoSided) delegates global load‑balancer stats to the NVLink prepare path via alltoall_prepare and _load_balancer_update_statistic_with_gathered_statistic, avoiding redundant allreduces.
• alltoall_prepare / alltoall_dispatch / alltoall_combine wrap the Mnnvl NVLink calls cleanly and centrally, with expert‑id padding handled via memset_expert_ids.
• The combine branch for NVLinkTwoSided reuses these helpers and respects alltoall_result_do_sum, while DeepEP/LowLatency paths continue to use their respective buffer‑based combine routines.

Overall, the NVLinkTwoSided flow through prepare → dispatch → MoEOp → combine is structured and consistent with the other backends.

Also applies to: 427-433, 463-476, 554-559, 651-687

tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py (2)

111-119: TRTLLMGenFusedMoE alltoall selection and init mirror Cutlass logic

• The class now chooses an alltoall_method_type and logs it once, then:

  • For NVLinkTwoSided, initializes Mnnvl workspaces via MnnvlMemory/MnnvlMoe.
  • For NVLinkOneSided, sets up a Python MoeAlltoAll instance with num_experts=self.num_slots and user‑configurable workspace size.
  • Rejects DeepEP/LowLatency as not yet supported.

• select_alltoall_method_type follows the same pattern as Cutlass: require DP>1, moe_tp_size == 1, MNNVL support, optional TRTLLM_FORCE_ALLTOALL_METHOD override, otherwise default to NVLinkOneSided.

• enable_alltoall is derived purely from the enum (non‑NotEnabled), which keeps downstream logic simple.

This brings TRTLLMGenFusedMoE in line with the other MoE backends.

Also applies to: 122-138, 151-179, 184-189

---

339-369: NVLink alltoall + MoeAlltoAll integration in forward path looks sound

• For post‑quant communication, load‑balancer stats are updated with ignore_allreduce=True when using NVLinkTwoSided, relying on the subsequent Mnnvl gather to provide global statistics.

• The NVLinkTwoSided branch:

  • Prepares alltoall via mnnvl_moe_alltoallv_prepare_without_allgather, optionally passing LB local stats and feeding back gathered stats.
  • Does a multi‑payload alltoall of [x, x_sf, token_selected_experts, token_final_scales].
  • Uses memset_expert_ids(..., invalid_id=-1) consistent with TRTLLM‑Gen semantics.

• The NVLinkOneSided branch:

  • Uses MoeAlltoAll.dispatch/combine with payloads [x, (optional x_sf), token_selected_experts, token_final_scales] and invalid‑id -1, and correctly reshapes the 3D recv tensors back to 2D.
  • Optionally obtains a workspace‑backed output tensor via get_combine_payload_tensor_in_workspace when using W4A8_MXFP4_MXFP8, avoiding an extra copy on combine.

• The final combine block dispatches to either Mnnvl combine or MoeAlltoAll.combine based on the enum, with clear error handling for unsupported types.

Overall, the control flow and tensor reshaping/dtypes for both NVLink variants match the intended semantics.

Also applies to: 350-355, 370-475, 495-502, 761-799

[tensorrt-cicd]

PR_Github #25488 [ run ] completed with state SUCCESS. Commit: eee5304
/LLM/main/L0_MergeRequest_PR pipeline #19299 completed with status: 'FAILURE'

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #25495 [ run ] triggered by Bot. Commit: eee5304

[bobboli]

/bot kill

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #25507 [ run ] triggered by Bot. Commit: beea101

[tensorrt-cicd]

PR_Github #25495 [ run ] completed with state ABORTED. Commit: eee5304
LLM/main/L0_MergeRequest_PR #19306 (Blue Ocean) completed with status: ABORTED

[tensorrt-cicd]

PR_Github #25508 [ kill ] triggered by Bot. Commit: beea101

[tensorrt-cicd]

PR_Github #25507 [ run ] completed with state ABORTED. Commit: beea101

[tensorrt-cicd]

PR_Github #25508 [ kill ] completed with state SUCCESS. Commit: beea101
Successfully killed previous jobs for commit beea101

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #25524 [ run ] triggered by Bot. Commit: beea101

[tensorrt-cicd]

PR_Github #25524 [ run ] completed with state SUCCESS. Commit: beea101
/LLM/main/L0_MergeRequest_PR pipeline #19330 completed with status: 'FAILURE'

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #25549 [ run ] triggered by Bot. Commit: beea101

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #25665 [ run ] triggered by Bot. Commit: 018a80a

[tensorrt-cicd]

PR_Github #25665 [ run ] completed with state SUCCESS. Commit: 018a80a
/LLM/main/L0_MergeRequest_PR pipeline #19451 completed with status: 'FAILURE'

[bobboli]

/bot run --reuse-test

[tensorrt-cicd]

PR_Github #25725 [ run ] triggered by Bot. Commit: 018a80a

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #25728 [ run ] triggered by Bot. Commit: f85b11f

[tensorrt-cicd]

PR_Github #25725 [ run ] completed with state ABORTED. Commit: 018a80a
LLM/main/L0_MergeRequest_PR #19507 (Blue Ocean) completed with status: ABORTED

[bobboli]

/bot run

[tensorrt-cicd]

PR_Github #25739 [ run ] triggered by Bot. Commit: f85b11f

[tensorrt-cicd]

PR_Github #25728 [ run ] completed with state ABORTED. Commit: f85b11f
LLM/main/L0_MergeRequest_PR #19510 (Blue Ocean) completed with status: ABORTED

[tensorrt-cicd]

PR_Github #25739 [ run ] completed with state SUCCESS. Commit: f85b11f
/LLM/main/L0_MergeRequest_PR pipeline #19518 completed with status: 'FAILURE'

[bobboli]

/bot run --reuse-test

[tensorrt-cicd]

PR_Github #25769 [ run ] triggered by Bot. Commit: f85b11f

[tensorrt-cicd]

PR_Github #25769 [ run ] completed with state FAILURE. Commit: f85b11f
/LLM/main/L0_MergeRequest_PR pipeline #19543 completed with status: 'FAILURE'

[bobboli]

/bot run --reuse-test

[bobboli]

/bot run --reuse-test

[bobboli]

/bot run --reuse-test

[tensorrt-cicd]

PR_Github #25946 [ run ] triggered by Bot. Commit: f85b11f

[tensorrt-cicd]

PR_Github #25946 [ run ] completed with state SUCCESS. Commit: f85b11f
/LLM/main/L0_MergeRequest_PR pipeline #19677 completed with status: 'SUCCESS'
Pipeline passed with automatic retried tests. Check the rerun report for details.