[TRTLLM-13250][feat] Wave 5: Enable MX post-transform Llama receiver by chienchunhung · Pull Request #15432 · NVIDIA/TensorRT-LLM

chienchunhung · 2026-06-16T23:25:17Z

Summary

Stacked on Wave 4 / #15387.

This implements Wave 5 of the staged post-load hook rollout for MX:

publish MX sources after post-load transforms with SourceIdentity and transform-layout metadata
let compatible, allow-listed Llama receivers consume post-transform MX bytes and run only setup_aliases() + cache_derived_state()
fail closed before P2P when SourceIdentity is missing/mismatched, transform protocol metadata is unsupported, or the model is not allow-listed
add unit coverage for metadata fallback cases, publish metadata, GMS/MX post-load publish ordering, and a tiny real-Llama staged receiver equivalence check

Dependency / prerequisite stack

This PR is Wave 5 in the staged post-load hooks rollout. The foundation PRs #14770 and #14878 are already merged. The 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 (open)"]
    PR15387["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15387'>#15387</a>: Wave 4 MX receiver cutover (open)"]
    PR15432["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15432'>#15432</a>: Wave 5 MX publisher + Llama receiver (this PR, open)"]
    VERIFY["post-migration verification / demo (planned)"]

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

    classDef merged fill:#dcfce7,stroke:#16a34a,color:#14532d;
    classDef open 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 stroke:#16a34a,stroke-width:2px;
    linkStyle 2,3,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 PR15386,PR15387 open;
    class PR15432 current;
    class VERIFY downstream;

Immediate merge dependency for this PR: #15387 must land first; after Wave 5 lands, run the post-migration verification/demo for the completed staged-hook rollout.

Validation

git diff --check
python -m py_compile tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py tensorrt_llm/_torch/pyexecutor/model_loader.py tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py tests/unittest/_torch/pyexecutor/test_model_loader_gms.py tests/unittest/_torch/pyexecutor/test_model_loader_mx.py tests/unittest/_torch/weight_sharing/test_mx_source_identity_gate.py
pre-commit on commit, with waive list check and validate-test-lists skipped locally because scripts/check_test_list.py fails under this hook interpreter with TypeError: unsupported operand type(s) for |: 'type' and 'NoneType'

Focused pytest collection is blocked in this local environment by missing transformers before tests are collected.

Summary by CodeRabbit

New Features
- Added support for staged post-transform weight transfers with source identity verification in ModelExpress transfers.
- Introduced source identity gating to validate weight compatibility across distributed systems.
Improvements
- Refactored weight transformation pipeline to separate weight transformation and state caching phases for improved clarity and maintainability.

chienchunhung · 2026-06-16T23:26:48Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-16T23:33:29Z

PR_Github #54683 [ run ] triggered by Bot. Commit: 756e717 Link to invocation

tensorrt-cicd · 2026-06-17T06:25:31Z

PR_Github #54683 [ run ] completed with state FAILURE. Commit: 756e717
/LLM/main/L0_MergeRequest_PR pipeline #43714 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-18T00:37:49Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-18T00:43:01Z

PR_Github #54888 [ run ] triggered by Bot. Commit: f123c77 Link to invocation

tensorrt-cicd · 2026-06-18T10:41:14Z

PR_Github #54888 [ run ] completed with state SUCCESS. Commit: f123c77
/LLM/main/L0_MergeRequest_PR pipeline #43893 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-19T01:32:40Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-06-19T01:38:48Z

PR_Github #54955 [ run ] triggered by Bot. Commit: 14a4537 Link to invocation

tensorrt-cicd · 2026-06-19T08:05:01Z

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

CI Report

Link to invocation

tensorrt-cicd · 2026-07-07T17:21:04Z

PR_Github #58041 [ run ] triggered by Bot. Commit: c412088 Link to invocation

tensorrt-cicd · 2026-07-07T17:27:20Z

PR_Github/15432-c412088 #58038 was force-killed by a newer pipeline run.
L0 job information not available (job may not have been triggered yet).

Link to superseding invocation

chienchunhung · 2026-07-07T21:06:12Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-07T21:11:57Z

PR_Github #58080 [ run ] triggered by Bot. Commit: 66209a0 Link to invocation

tensorrt-cicd · 2026-07-07T21:16:53Z

PR_Github #58041 [ run ] completed with state ABORTED. Commit: c412088

Link to invocation

chienchunhung · 2026-07-07T21:38:51Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-07T21:45:04Z

PR_Github #58086 [ run ] triggered by Bot. Commit: 247733e Link to invocation

tensorrt-cicd · 2026-07-07T21:50:56Z

PR_Github #58080 [ run ] completed with state ABORTED. Commit: 66209a0

Link to invocation

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

chienchunhung · 2026-07-08T00:45:30Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-08T00:52:37Z

PR_Github #58103 [ run ] triggered by Bot. Commit: 62978d5 Link to invocation

tensorrt-cicd · 2026-07-08T00:57:38Z

PR_Github #58086 [ run ] completed with state ABORTED. Commit: 247733e

Link to invocation

tensorrt-cicd · 2026-07-08T14:20:02Z

PR_Github #58103 [ run ] completed with state FAILURE. Commit: 62978d5
/LLM/main/L0_MergeRequest_PR pipeline #46766 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-08T15:12:34Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-08T15:18:29Z

PR_Github #58255 [ run ] triggered by Bot. Commit: 62978d5 Link to invocation

tensorrt-cicd · 2026-07-08T18:14:12Z

PR_Github #58255 [ run ] completed with state SUCCESS. Commit: 62978d5
/LLM/main/L0_MergeRequest_PR pipeline #46894 completed with status: 'ABORTED'

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

Link to invocation

chienchunhung · 2026-07-08T21:42:15Z

/bot run --disable-fail-fast

tensorrt-cicd · 2026-07-08T21:47:55Z

PR_Github #58302 [ run ] triggered by Bot. Commit: 62978d5 Link to invocation

tensorrt-cicd · 2026-07-08T22:14:22Z

PR_Github #58302 [ run ] completed with state SUCCESS. Commit: 62978d5
/LLM/main/L0_MergeRequest_PR pipeline #46937 completed with status: 'SUCCESS'

CI Report

Link to invocation

…VIDIA#15432) Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

github-actions Bot assigned chienchunhung Jun 16, 2026

chienchunhung force-pushed the codex/staged-hooks-wave5-mx-publisher branch 2 times, most recently from ae210cb to f123c77 Compare June 18, 2026 00:37

chienchunhung changed the title ~~[TRTLLM-13250][feat] Enable MX post-transform Llama receiver~~ [TRTLLM-13250][feat] Wav3 5: Enable MX post-transform Llama receiver Jun 18, 2026

chienchunhung force-pushed the codex/staged-hooks-wave5-mx-publisher branch from f123c77 to 14a4537 Compare June 19, 2026 01:32

chienchunhung changed the title ~~[TRTLLM-13250][feat] Wav3 5: Enable MX post-transform Llama receiver~~ [TRTLLM-13250][feat] Wave 5: Enable MX post-transform Llama receiver Jun 21, 2026

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

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

chienchunhung requested review from byshiue, danielafrimi, dongjiyingdjy, hlu1 and symphonylyh June 22, 2026 18:04

mikeiovine approved these changes Jul 7, 2026

View reviewed changes

Comment thread tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py Outdated

Comment thread tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py Outdated

chienchunhung force-pushed the codex/staged-hooks-wave5-mx-publisher branch from c412088 to 66209a0 Compare July 7, 2026 21:03

chienchunhung force-pushed the codex/staged-hooks-wave5-mx-publisher branch from 66209a0 to 247733e Compare July 7, 2026 21:25

[TRTLLM-13250][feat] Enable MX post-transform Llama receiver

62978d5

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

chienchunhung force-pushed the codex/staged-hooks-wave5-mx-publisher branch from 247733e to 62978d5 Compare July 8, 2026 00:45

Funatiq approved these changes Jul 8, 2026

View reviewed changes

litaotju approved these changes Jul 8, 2026

View reviewed changes

chienchunhung merged commit 1fd24e1 into NVIDIA:main Jul 8, 2026
7 checks passed

This was referenced Jul 9, 2026

[TRTLLM-12352][feat] integrate ModelExpress checkpoint loading #15641

Merged

[None][feat] Bind SourceIdentity to checkpoint artifacts #16159

Merged

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

[TRTLLM-13250][feat] Wave 5: Enable MX post-transform Llama receiver (N…

b0d0c57

…VIDIA#15432) Signed-off-by: Chien-Chun Hung <2679986+chienchunhung@users.noreply.github.com>

coderabbitai Bot mentioned this pull request Jul 10, 2026

[TRTLLM-12352][feat] add post-transform capability profiles #16238

Merged

7 tasks

Uh oh!

Conversation

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

Uh oh!

Summary

Dependency / prerequisite stack

Validation

Summary by CodeRabbit

Uh oh!

chienchunhung commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 16, 2026

Uh oh!

tensorrt-cicd commented Jun 17, 2026

Uh oh!

chienchunhung commented Jun 18, 2026

Uh oh!

tensorrt-cicd commented Jun 18, 2026

Uh oh!

tensorrt-cicd commented Jun 18, 2026

Uh oh!

chienchunhung commented Jun 19, 2026

Uh oh!

tensorrt-cicd commented Jun 19, 2026

Uh oh!

tensorrt-cicd commented Jun 19, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

Uh oh!

Uh oh!

chienchunhung commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

chienchunhung commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

tensorrt-cicd commented Jul 7, 2026

Uh oh!

chienchunhung commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

chienchunhung commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

chienchunhung commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

tensorrt-cicd commented Jul 8, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

6 participants

chienchunhung commented Jun 16, 2026 •

edited

Loading