[TRTLLM-13248][feat] Wave 3: migrate MoE staged hooks by chienchunhung · Pull Request #15386 · NVIDIA/TensorRT-LLM

chienchunhung · 2026-06-15T19:52:25Z

Summary

Wave 3 of the staged post-load hooks rollout, built on merged Wave 2 / #15288 and rebased onto main.

This migrates the remaining MoE, Mamba, min-latency, and sparse attention post-load work into the staged hook protocol. Tensor-layout and irreversible MoE weight work now lives in transform_weights() behind _weights_transformed, while repeatable Python-side state refresh lives in cache_derived_state(). Existing full-load paths continue to use post_load_weights() as the backward-compatible shim.

What Changed

Added staged MoE hooks in MoE and ConfigurableMoE, with transform_weights() guarded by _weights_transformed and cache_derived_state() delegated to the quant method/backend.
Split fused-MoE quant-method work so shared-weight finalization and backend-specific transforms run in transform_weights(), while quant scales and load-balancer assignment refresh run in cache_derived_state().
Routed CUTLASS, Triton, TRTLLM-Gen, DenseGEMM, WideEP, and MegaMoE backend post-load shims through the staged module hooks.
Moved Mamba2, Llama min-latency, and DSA derived-state recomputation into cache_derived_state() while keeping post_load_weights() compatibility shims.
Added MoE hook idempotency coverage plus focused DSA and Mamba derived-state tests.

Dependency / prerequisite stack

This PR is Wave 3 in the staged post-load hooks rollout. The foundation PRs and Waves 1-2 (#14770, #14878, #15014, and #15288) are merged. The remaining wave PRs should merge in sequence; after each upstream wave lands, rebase the next wave onto main so review and CI focus on that wave's delta.

Arrows point from prerequisite to dependent. PR numbers in graph nodes are clickable.

graph TD
    PR14770["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/14770'>#14770</a>: staged-hook contract (merged)"]
    PR14878["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/14878'>#14878</a>: GMS SourceIdentity gate (merged)"]
    PR15014["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15014'>#15014</a>: Wave 1 aliases + GMS RO load (merged)"]
    PR15288["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15288'>#15288</a>: Wave 2 Linear/Attention transforms (merged)"]
    PR15386["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15386'>#15386</a>: Wave 3 MoE/Mamba staged hooks (this PR, ready for review)"]
    PR15387["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15387'>#15387</a>: Wave 4 MX receiver cutover (open, ready for review)"]
    PR15432["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15432'>#15432</a>: Wave 5 MX publisher + Llama receiver (open, ready for review)"]
    VERIFY["post-migration verification / demo (planned)"]

    PR14770 -->|satisfied| PR15014
    PR14878 -->|satisfied| PR15014
    PR15014 -->|satisfied| PR15288
    PR15288 -->|satisfied| PR15386
    PR15386 -->|blocking| PR15387
    PR15387 -->|blocking| PR15432
    PR15432 -.->|planned| VERIFY

    classDef merged fill:#dcfce7,stroke:#16a34a,color:#14532d;
    classDef inflight fill:#dbeafe,stroke:#2563eb,color:#1e3a8a;
    classDef current fill:#ede9fe,stroke:#7c3aed,color:#3b0764,stroke-width:3px;
    classDef downstream fill:#f3f4f6,stroke:#6b7280,color:#374151,stroke-dasharray:5 5;
    linkStyle 0,1,2,3 stroke:#16a34a,stroke-width:2px;
    linkStyle 4,5 stroke:#ea580c,stroke-width:3px;
    linkStyle 6 stroke:#6b7280,stroke-width:2px,stroke-dasharray:5 5;

    class PR14770,PR14878,PR15014,PR15288 merged;
    class PR15387,PR15432 inflight;
    class PR15386 current;
    class VERIFY downstream;

Immediate merge dependency satisfied: #15288 merged on 2026-06-26, and this branch is rebased onto main so the PR diff is the Wave 3 delta.

Test Plan

git diff --check origin/codex/staged-hooks-wave2-transform-weights..HEAD
python -m py_compile tensorrt_llm/_torch/attention_backend/sparse/dsa.py tensorrt_llm/_torch/models/modeling_llama_min_latency.py tensorrt_llm/_torch/modules/fused_moe/configurable_moe.py tensorrt_llm/_torch/modules/fused_moe/fused_moe_cute_dsl_b12x.py tensorrt_llm/_torch/modules/fused_moe/fused_moe_cutlass.py tensorrt_llm/_torch/modules/fused_moe/fused_moe_densegemm.py tensorrt_llm/_torch/modules/fused_moe/fused_moe_triton.py tensorrt_llm/_torch/modules/fused_moe/fused_moe_trtllm_gen.py tensorrt_llm/_torch/modules/fused_moe/fused_moe_wide_ep.py tensorrt_llm/_torch/modules/fused_moe/interface.py tensorrt_llm/_torch/modules/fused_moe/mega_moe/mega_moe_deepgemm.py tensorrt_llm/_torch/modules/fused_moe/quantization.py tensorrt_llm/_torch/modules/mamba/mamba2_mixer.py tests/unittest/_torch/attention/sparse/test_dsa_indexer.py tests/unittest/_torch/modules/moe/test_moe_backend.py tests/unittest/_torch/modules/mamba/test_mamba2_mixer.py
Pre-commit on commit: isort, yapf, ruff, ruff-format, codespell, DCO, and other applicable hooks passed. waive list check and validate-test-lists were skipped because the hook runtime fails in scripts/check_test_list.py on str | None before inspecting this diff.
Focused pytest attempted but blocked in this macOS shell because transformers is not installed and the local torch package lacks nn / distributed.

Next Steps

Complete Wave 3 review and CI, then merge [TRTLLM-13248][feat] Wave 3: migrate MoE staged hooks #15386.
After [TRTLLM-13248][feat] Wave 3: migrate MoE staged hooks #15386 lands, rebase [TRTLLM-13249][feat] Wave 4: add MX staged receiver cutover #15387 onto main; after [TRTLLM-13249][feat] Wave 4: add MX staged receiver cutover #15387 lands, repeat for [TRTLLM-13250][feat] Wave 5: Enable MX post-transform Llama receiver #15432.

Summary by CodeRabbit

Refactor
- Restructured internal weight loading and model initialization lifecycle management across model backends for improved code organization and maintainability.
Tests
- Added unit tests validating weight transformation and state caching behavior.

chienchunhung · 2026-06-15T19:54:21Z

/bot run

tensorrt-cicd · 2026-06-15T20:01:44Z

PR_Github #54355 [ run ] triggered by Bot. Commit: e378d6d Link to invocation

tensorrt-cicd · 2026-06-15T20:51:02Z

PR_Github #54355 [ run ] completed with state SUCCESS. Commit: e378d6d
/LLM/main/L0_MergeRequest_PR pipeline #43426 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

chienchunhung · 2026-06-15T22:13:28Z

/bot run --disable-fail-fast --stage-list "L0_Test-x86_64-Single-GPU, L0_Test-SBSA-Single-GPU"

tensorrt-cicd · 2026-06-15T22:21:45Z

PR_Github #54368 [ run ] triggered by Bot. Commit: e378d6d Link to invocation

tensorrt-cicd · 2026-06-15T22:58:44Z

PR_Github #54368 [ run ] completed with state FAILURE. Commit: e378d6d
/LLM/main/L0_MergeRequest_PR pipeline #43438 (Partly Tested) 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

chienchunhung · 2026-06-15T23:23:23Z

/bot run

tensorrt-cicd · 2026-06-15T23:29:00Z

PR_Github #54372 [ run ] triggered by Bot. Commit: e378d6d Link to invocation

tensorrt-cicd · 2026-06-16T03:07:58Z

PR_Github #54372 [ run ] completed with state FAILURE. Commit: e378d6d
/LLM/main/L0_MergeRequest_PR pipeline #43442 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

chienchunhung · 2026-06-16T05:40:35Z

/bot run --disable-fail-fast

chienchunhung · 2026-06-16T05:47:01Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-16T05:53:04Z

PR_Github #54488 [ run ] triggered by Bot. Commit: 32646d9 Link to invocation

tensorrt-cicd · 2026-06-16T13:52:38Z

PR_Github #54488 [ run ] completed with state SUCCESS. Commit: 32646d9
/LLM/main/L0_MergeRequest_PR pipeline #43551 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

chienchunhung · 2026-06-16T16:45:59Z

/bot run --disable-fail-fast --stage-list "A100X-PyTorch-1,DGX_B200-PyTorch-2,H100_PCIe-AutoDeploy-1,DGX_B200-8_GPUs-PyTorch-1,DGX_H100-2_GPUs-PyTorch-Others-2"

tensorrt-cicd · 2026-06-16T16:51:57Z

PR_Github #54633 [ run ] triggered by Bot. Commit: 32646d9 Link to invocation

tensorrt-cicd · 2026-06-16T19:27:08Z

PR_Github #54633 [ run ] completed with state SUCCESS. Commit: 32646d9
/LLM/main/L0_MergeRequest_PR pipeline #43665 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

chienchunhung · 2026-06-16T20:41:35Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-16T20:47:44Z

PR_Github #54669 [ run ] triggered by Bot. Commit: 32646d9 Link to invocation

tensorrt-cicd · 2026-06-16T21:14:27Z

PR_Github #54669 [ run ] completed with state SUCCESS. Commit: 32646d9
/LLM/main/L0_MergeRequest_PR pipeline #43701 completed with status: 'SUCCESS'

CI Report

Link to invocation

chienchunhung · 2026-07-02T16:11:08Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-02T16:17:31Z

PR_Github #57215 [ run ] triggered by Bot. Commit: 9a30e43 Link to invocation

tensorrt-cicd · 2026-07-02T16:21:14Z

PR_Github #57211 [ run ] completed with state ABORTED. Commit: 9a30e43

Link to invocation

tensorrt-cicd · 2026-07-02T22:03:44Z

PR_Github #57215 [ run ] completed with state SUCCESS. Commit: 9a30e43
/LLM/main/L0_MergeRequest_PR pipeline #45987 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

chienchunhung · 2026-07-02T22:22:42Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-02T22:29:45Z

PR_Github #57267 [ run ] triggered by Bot. Commit: 9a30e43 Link to invocation

tensorrt-cicd · 2026-07-02T22:30:50Z

PR_Github #57267 [ run ] completed with state FAILURE. Commit: 9a30e43

Link to invocation

chienchunhung · 2026-07-02T22:47:46Z

/bot run --disable-fail-fast

chienchunhung · 2026-07-02T23:10:20Z

/bot run --disable-fail-fast

chienchunhung · 2026-07-02T23:22:23Z

/bot run --disable-fail-fast

chienchunhung · 2026-07-02T23:48:55Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-02T23:54:45Z

PR_Github #57275 [ run ] triggered by Bot. Commit: 9a30e43 Link to invocation

tensorrt-cicd · 2026-07-03T00:06:18Z

PR_Github #57275 [ run ] completed with state FAILURE. Commit: 9a30e43
/LLM/main/L0_MergeRequest_PR pipeline #46037 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

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

chienchunhung · 2026-07-04T17:38:31Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-04T17:44:12Z

PR_Github #57571 [ run ] triggered by Bot. Commit: 7963f1a Link to invocation

tensorrt-cicd · 2026-07-05T00:29:44Z

PR_Github #57571 [ run ] completed with state SUCCESS. Commit: 7963f1a
/LLM/main/L0_MergeRequest_PR pipeline #46299 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

chienchunhung · 2026-07-06T03:39:24Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-06T03:45:14Z

PR_Github #57683 [ run ] triggered by Bot. Commit: 7963f1a Link to invocation

tensorrt-cicd · 2026-07-06T05:42:13Z

PR_Github #57683 [ run ] completed with state FAILURE. Commit: 7963f1a
/LLM/main/L0_MergeRequest_PR pipeline #46403 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

chienchunhung · 2026-07-06T18:17:16Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-06T18:23:02Z

PR_Github #57801 [ run ] triggered by Bot. Commit: 7963f1a Link to invocation

tensorrt-cicd · 2026-07-06T18:53:29Z

PR_Github #57801 [ run ] completed with state SUCCESS. Commit: 7963f1a
/LLM/main/L0_MergeRequest_PR pipeline #46498 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

chienchunhung · 2026-07-06T19:16:52Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-06T19:22:51Z

PR_Github #57809 [ run ] triggered by Bot. Commit: 7963f1a Link to invocation

tensorrt-cicd · 2026-07-06T20:41:19Z

PR_Github #57809 [ run ] completed with state SUCCESS. Commit: 7963f1a
/LLM/main/L0_MergeRequest_PR pipeline #46506 completed with status: 'SUCCESS'

CI Report

Link to invocation

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

github-actions Bot assigned chienchunhung Jun 15, 2026

chienchunhung mentioned this pull request Jun 16, 2026

[TRTLLM-13249][feat] Wave 4: add MX staged receiver cutover #15387

Merged

xxi-nv mentioned this pull request Jun 16, 2026

[TRTLLM-12950][feat] Add MegaMoECuteDsl NVFP4 MoE backend #14608

Merged

chienchunhung force-pushed the codex/staged-hooks-wave3-moe-mamba branch from e378d6d to 32646d9 Compare June 16, 2026 05:40

chienchunhung requested a review from Funatiq June 16, 2026 05:44

chienchunhung marked this pull request as ready for review June 22, 2026 18:03

chienchunhung requested review from a team as code owners June 22, 2026 18:03

[TRTLLM-13248][feat] Wave 3: migrate MoE staged hooks

7963f1a

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

chienchunhung force-pushed the codex/staged-hooks-wave3-moe-mamba branch from 9a30e43 to 7963f1a Compare July 4, 2026 17:38

chienchunhung merged commit 9f689ec into NVIDIA:main Jul 6, 2026
7 checks passed

chienchunhung mentioned this pull request Jul 8, 2026

[https://nvbugs/6428057][fix] Propagate host stop criteria across PP ranks #16131

Closed

BrianLi23 pushed a commit to BrianLi23/TensorRT-LLM that referenced this pull request Jul 9, 2026

[TRTLLM-13248][feat] Wave 3: migrate MoE staged hooks (NVIDIA#15386)

21690de

Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

Uh oh!

Conversation

chienchunhung commented Jun 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

What Changed

Dependency / prerequisite stack

Test Plan

Next Steps

Summary by CodeRabbit

Uh oh!

chienchunhung commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

chienchunhung commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

chienchunhung commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 15, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

chienchunhung commented Jun 16, 2026

Uh oh!

chienchunhung commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

chienchunhung commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

chienchunhung commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

chienchunhung commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

chienchunhung commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

chienchunhung commented Jul 2, 2026

Uh oh!

chienchunhung commented Jul 2, 2026

Uh oh!

chienchunhung commented Jul 2, 2026

Uh oh!

chienchunhung commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 2, 2026

Uh oh!

tensorrt-cicd commented Jul 3, 2026

Uh oh!

chienchunhung commented Jul 4, 2026

Uh oh!

tensorrt-cicd commented Jul 4, 2026

Uh oh!

tensorrt-cicd commented Jul 5, 2026

chienchunhung commented Jun 15, 2026 •

edited

Loading