[TRTLLM-14843][chore] Establish _torch/moe/ and relocate MoE modules, custom ops, and communication by lori-ren · Pull Request #17952 · NVIDIA/TensorRT-LLM

lori-ren · 2026-08-19T06:03:13Z

Dev Engineer Review

Relocates MoE code, custom operations, communication modules, and tests into _torch/moe/.
Updates imports, documentation, CODEOWNERS, lint baselines, pre-commit paths, legacy-file lists, and CI path detection.
Preserves MoeAlltoAll and custom-operation re-exports.
Adds consolidated package exports and MoE communication, health, implementation-contract, registry, and utility modules.
Review must verify import consistency, generated path lists, public exports, error handling, and the undefined Map annotation in deep_ep_utils.py.
Configuration and test-list paths appear consistent. The specified thop files remain in place.
CI previously failed. The reported failures require resolution before rerunning CI.

QA Engineer Review

Relocates MoE tests under tests/unittest/_torch/moe/.
Updates test imports, patch targets, pytest paths, duration data, waivers, and CI selections.
Adds TestAllgatherEScoreCorrectionBias, TestAllgatherExpertScales, TestScatterShardsToFull, and TestAllgatherExpertScalesNonUniform.
Adds Kimi K3 MXFP4 quantization utilities and Kimi K3 MLP reference modules.
Existing communication, routing, backend, fused-MoE, load-balancer, EPLB, quantization, and multi-GPU tests retain their reported behavior.
Relocated tests are referenced by tests/integration/test_lists/test-db/ entries.
Coverage for the new DWDP fixup tests and Kimi K3 helper modules is not confirmed in test-db/ or qa/ files.
Verdict: needs follow-up.

Description

Pure relocation into a new _torch/moe/ package — no production code changed, only import
paths and the rules that follow them. 84 renames across three commits: 54 source, 30 test.

from	to
`_torch/modules/fused_moe/` (whole tree)	`_torch/moe/fused_moe/`
`_torch/modules/fused_shared_expert.py`	`_torch/moe/fused_shared_expert.py`
`_torch/expert_statistic.py`	`_torch/moe/expert_statistic.py`
`_torch/distributed/moe_alltoall.py`	`_torch/moe/communication/moe_alltoall.py`
`_torch/custom_ops/cute_dsl_megamoe_custom_op.py`	`_torch/moe/custom_ops/cute_dsl_megamoe_custom_op.py`

Two re-exports are deliberately kept and only re-pointed: custom_ops/__init__.py's import
of cute_dsl_megamoe_custom_op (it is what registers
trtllm::cute_dsl_megamoe_nvfp4_blackwell) and distributed/__init__.py's MoeAlltoAll.

The third commit moves the 30 MoE unit tests to tests/unittest/_torch/moe/, mirroring the
source layout (the non-blocked part of TRTLLM-14848; its five thop/ files stay there,
their destination depends on TRTLLM-14844). Test node IDs change, so 214 entries across
test-db/*.yml, waives.txt and .test_durations move with them. Two directory-level
entries are load-bearing and are reproduced verbatim for the new location — test_moe_a2a.py
is named by no list and reaches the 2-GPU stage only through
unittest/_torch/multi_gpu, and three cpu_only files reach the CPU stage only through
unittest/_torch/modules.

CODEOWNERS: five source rules and seven test rules each collapse into one directory rule
(/tensorrt_llm/_torch/moe, /tests/unittest/_torch/moe); no file changes owner.

Test Coverage

tests/unittest/_torch/moe/ — the whole relocated tree, including fused_moe/ and
multi_gpu/test_moe_a2a.py
tests/unittest/_torch/thop/serial/test_moe_alltoall.py (not moved)
tests/unittest/_torch/test_mnnvl_utils.py (rewritten mock.patch string targets)

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)
If PR introduces API changes, an appropriate PR label is added - either api-compatible or api-breaking. For api-breaking, include BREAKING in the PR title.
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

To see a list of available CI bot commands, please comment /bot help.

…les, custom ops and communication Signed-off-by: Lori Ren <lorir@nvidia.com>

…e references at _torch/moe/ Signed-off-by: Lori Ren <lorir@nvidia.com>

…hey cover Signed-off-by: Lori Ren <lorir@nvidia.com>

tensorrt-cicd · 2026-08-29T04:33:51Z

PR_Github #70095 [ run ] completed with state FAILURE. Commit: f22e40a
/LLM/main/L0_MergeRequest_PR pipeline #57362 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lori-ren · 2026-08-29T05:05:37Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-29T05:10:49Z

PR_Github #70116 [ run ] triggered by Bot. Commit: f22e40a Link to invocation

tensorrt-cicd · 2026-08-29T10:30:56Z

PR_Github #70116 [ run ] completed with state SUCCESS. Commit: f22e40a
/LLM/main/L0_MergeRequest_PR pipeline #57381 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lori-ren · 2026-08-29T10:38:04Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-29T10:43:33Z

PR_Github #70143 [ run ] triggered by Bot. Commit: f22e40a Link to invocation

tensorrt-cicd · 2026-08-29T13:30:08Z

PR_Github #70143 [ run ] completed with state SUCCESS. Commit: f22e40a
/LLM/main/L0_MergeRequest_PR pipeline #57407 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lori-ren · 2026-08-29T13:41:21Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-29T13:46:52Z

PR_Github #70151 [ run ] triggered by Bot. Commit: f22e40a Link to invocation

tensorrt-cicd · 2026-08-29T15:38:44Z

PR_Github #70151 [ run ] completed with state SUCCESS. Commit: f22e40a
/LLM/main/L0_MergeRequest_PR pipeline #57414 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lori-ren · 2026-08-29T17:45:48Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-29T17:52:34Z

PR_Github #70160 [ run ] triggered by Bot. Commit: f22e40a Link to invocation

tensorrt-cicd · 2026-08-29T19:39:47Z

PR_Github #70160 [ run ] completed with state SUCCESS. Commit: f22e40a
/LLM/main/L0_MergeRequest_PR pipeline #57423 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

Please check the failed tests and fix your PR
If you cannot view the failures, ask the CI triggerer to share details
Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

lori-ren · 2026-08-29T22:24:28Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-29T22:29:49Z

PR_Github #70173 [ run ] triggered by Bot. Commit: f22e40a Link to invocation

lori-ren · 2026-08-29T23:24:55Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-29T23:31:12Z

PR_Github #70174 [ run ] triggered by Bot. Commit: f22e40a Link to invocation

tensorrt-cicd · 2026-08-29T23:35:03Z

PR_Github #70173 [ run ] completed with state ABORTED. Commit: f22e40a

Link to invocation

lori-ren · 2026-08-30T00:43:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-08-30T00:49:08Z

PR_Github #70175 [ run ] triggered by Bot. Commit: f22e40a Link to invocation

tensorrt-cicd · 2026-08-30T00:52:44Z

PR_Github #70174 [ run ] completed with state ABORTED. Commit: f22e40a

Link to invocation

tensorrt-cicd · 2026-08-30T02:29:06Z

PR_Github #70175 [ run ] completed with state SUCCESS. Commit: f22e40a
/LLM/main/L0_MergeRequest_PR pipeline #57436 completed with status: 'SUCCESS'

CI Report

Link to invocation

Resolves the conflict introduced by NVIDIA#17952, which relocated tensorrt_llm/_torch/modules/fused_moe/ to tensorrt_llm/_torch/moe/fused_moe/. That turns this branch's quantization.py hunks into modify-vs-delete. Git resolves it locally via directory-rename detection, but the merge-base..main diff spans 10217 files, past the default diff.renameLimit, so GitHub skips rename detection and reports the PR as unmergeable. Merged rather than rebased so the branch needs no force-push. Signed-off-by: xxi <xxi@nvidia.com>

Resolves the MoE tree relocation from NVIDIA#17952: quantization.py and fused_moe_cutlass.py now live under tensorrt_llm/_torch/moe/fused_moe/, and the MoE unit tests under tests/unittest/_torch/moe/. The three standalone W4A16 test-list entries are re-added at the new path in l0_h100.yml and l0_b200.yml. Signed-off-by: Dorian Magasic <dorian.magasic@turintech.ai>

lori-ren added 2 commits August 19, 2026 03:50

[TRTLLM-14843][chore] Establish _torch/moe/ and relocate the MoE modu…

1f41a4c

…les, custom ops and communication Signed-off-by: Lori Ren <lorir@nvidia.com>

[TRTLLM-14843][chore] Repoint the MoE documentation and cross-languag…

2226a86

…e references at _torch/moe/ Signed-off-by: Lori Ren <lorir@nvidia.com>

lori-ren added the api-compatible Accepted LLM API contract change that is backwards-compatible label Aug 19, 2026

github-actions Bot assigned lori-ren Aug 19, 2026

[TRTLLM-14843][chore] Move the MoE unit tests alongside the modules t…

8370c17

…hey cover Signed-off-by: Lori Ren <lorir@nvidia.com>

lori-ren marked this pull request as ready for review August 20, 2026 05:29

lori-ren requested review from a team as code owners August 20, 2026 05:29

lori-ren requested review from MrGeva, YihuiLu512, chienchunhung, mikeiovine, pcastonguay, rosong11, sunnyqgg, tburt-nv and yuanjingx87 August 20, 2026 05:29

QiJune merged commit c5c985c into NVIDIA:main Aug 30, 2026
7 checks passed

xxi-nv mentioned this pull request Aug 31, 2026

[None][fix] Restore DSv4 NVFP4 routed swiglu_limit on TRTLLM-Gen #18239

Closed

zhangcl mentioned this pull request Sep 1, 2026

[None][feat] Add SM107 NVFP4 CuTe DSL fused MoE kernels and integration #18498

Merged

Dorijan10 mentioned this pull request Sep 1, 2026

[None][fix] Pad INT8-woq MoE intermediate to 64 for tensor-parallel sharding #16198

Open

Dorijan10 mentioned this pull request Sep 3, 2026

[None][feat] Add W4A16 per-channel weight-only support for Cutlass fused MoE #18313

Open

YihuiLu512 mentioned this pull request Sep 7, 2026

[TRTLLM-14558][chore] Add the forwarding modules for the retired Attention paths #17969

Merged

1 task

Uh oh!

Conversation

lori-ren commented Aug 19, 2026 • edited by coderabbitai Bot Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Dev Engineer Review

QA Engineer Review

Description

Test Coverage

PR Checklist

GitHub Bot Help

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

lori-ren commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

lori-ren commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

lori-ren commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

lori-ren commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

lori-ren commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

lori-ren commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

tensorrt-cicd commented Aug 29, 2026

Uh oh!

lori-ren commented Aug 30, 2026

Uh oh!

tensorrt-cicd commented Aug 30, 2026

Uh oh!

tensorrt-cicd commented Aug 30, 2026

Uh oh!

tensorrt-cicd commented Aug 30, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

12 participants

lori-ren commented Aug 19, 2026 •

edited by coderabbitai Bot

Loading