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

[chienchunhung]

Summary

Wave 4 of the staged post-load hooks rollout, stacked on #15386.

This adds the dormant MX receiver-side cutover infrastructure for consuming post-transform weights once Wave 5 teaches MX publishers to advertise that layout. Production behavior remains unchanged today: the MX post-transform signal defaults false and the per-model allow-list ships empty.

What Changed

• Added a checkpoint-loader signal for direct preloads that are already post-transform, keeping the base default false.
• Added MX receiver gating for post-transform preloads: MX format, successful preload, compatible SourceIdentity, and (model_class, transform_protocol_version) allow-list entry are all required before skipping transform_weights().
• Added the staged receiver path in ModelLoader: setup_aliases() -> mark participating modules as _weights_transformed=True -> cache_derived_state(), while preserving the existing full post_load_weights() path for all non-allow-listed cases.
• Kept MX production behavior dormant with _source_metadata_is_post_transform() returning false until Wave 5 wires explicit MX metadata.
• Added a Wave 5 code reminder that non-allow-listed post-transform bytes must not fall through into the full post-load path once real post-transform publishing is enabled.
• Added focused unit coverage for the synthetic allow-list path, fallback path, transformed-flag marking, and post-transform partial-transfer fallback to a full disk load.

Dependency / prerequisite stack

This PR is Wave 4 in the staged post-load hooks rollout. The foundation PRs #14770 and #14878 and upstream wave PRs #15014 and #15288 are already 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 (open)"]
    PR15387["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15387'>#15387</a>: Wave 4 MX receiver cutover (this PR, open)"]
    PR15432["<a href='https://github.com/NVIDIA/TensorRT-LLM/pull/15432'>#15432</a>: Wave 5 MX post-transform Llama receiver (open)"]
    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 PR15386,PR15432 inflight;
    class PR15387 current;
    class VERIFY downstream;

Loading

Immediate merge dependency for this PR: #15386 must land first; after it lands, rebase this branch onto main so the PR diff collapses to the Wave 4 delta.

Test Plan

• git diff --check
• python -m py_compile tensorrt_llm/_torch/models/checkpoints/base_checkpoint_loader.py 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_mx.py
• python -m pytest tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py tests/unittest/_torch/pyexecutor/test_model_loader_mx.py -q attempted, but blocked in this local shell because transformers is not installed.
• 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.

Next Steps

• Wave 5 should wire explicit MX post-transform metadata and ensure non-allow-listed post-transform sources fail or fall back before full post-load transforms can run on transformed bytes.
• After [TRTLLM-13248][feat] Wave 3: migrate MoE staged hooks #15386 lands or rebases, rebase this branch so the PR diff collapses to the Wave 4 commit only.

Summary by CodeRabbit

Release Notes

• New Features

  • Introduced staged weight loading with support for recognizing and utilizing pre-transformed weights from external sources (e.g., peer-to-peer transfers).

• Improvements

  • Weight transformation processing is now more robust, with idempotent transformations ensuring consistent behavior across multiple invocations.
  • Enhanced efficiency when pre-transformed weights are available, reducing redundant processing.

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54387 [ run ] triggered by Bot. Commit: 6d5feca Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54483 [ run ] triggered by Bot. Commit: 0b3b58d Link to invocation

[tensorrt-cicd]

PR_Github #54387 [ run ] completed with state ABORTED. Commit: 6d5feca

Link to invocation

[tensorrt-cicd]

PR_Github #54483 [ run ] completed with state SUCCESS. Commit: 0b3b58d
/LLM/main/L0_MergeRequest_PR pipeline #43546 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

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast --stage-list "DGX_H100-PyTorch-2,H100_PCIe-PyTorch-Ray-1,DGX_B200-4_GPUs-PyTorch-Ray-1"

[tensorrt-cicd]

PR_Github #54639 [ run ] triggered by Bot. Commit: 0b3b58d Link to invocation

[tensorrt-cicd]

PR_Github #54639 [ run ] completed with state SUCCESS. Commit: 0b3b58d
/LLM/main/L0_MergeRequest_PR pipeline #43671 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #54670 [ run ] triggered by Bot. Commit: 0b3b58d Link to invocation

[tensorrt-cicd]

PR_Github #54670 [ run ] completed with state SUCCESS. Commit: 0b3b58d
/LLM/main/L0_MergeRequest_PR pipeline #43702 completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57118 [ run ] triggered by Bot. Commit: ef8cf45 Link to invocation

[tensorrt-cicd]

PR_Github #57061 [ run ] completed with state ABORTED. Commit: ef8cf45

Link to invocation

[tensorrt-cicd]

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

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57214 [ run ] triggered by Bot. Commit: 45e24ee Link to invocation

[tensorrt-cicd]

PR_Github #57214 [ run ] completed with state SUCCESS. Commit: 45e24ee
/LLM/main/L0_MergeRequest_PR pipeline #45986 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

[brb-nv]

Approving changes under tensorrt_llm/_torch/models with a comment.

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57268 [ run ] triggered by Bot. Commit: 9d4a470 Link to invocation

[tensorrt-cicd]

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

/bot run --disable-fail-fast

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57632 [ run ] triggered by Bot. Commit: 6c70685 Link to invocation

[tensorrt-cicd]

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

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57902 [ run ] triggered by Bot. Commit: 7f7d95c Link to invocation

[tensorrt-cicd]

PR_Github #57902 [ run ] completed with state SUCCESS. Commit: 7f7d95c
/LLM/main/L0_MergeRequest_PR pipeline #46596 completed with status: 'SUCCESS'

CI Report

Link to invocation