[TRTLLM-12352][feat] complete MX post-transform qualification foundation

[chienchunhung]

Summary

Complete the post-transform qualification foundation needed before adding MX support for another model family.

This PR keeps LlamaForCausalLM as the only enabled profile. It replaces broad capability assumptions with an exact, test-backed contract for deciding whether a producer may publish and a receiver may consume post-transform weights.

Dependencies

• Builds on merged [TRTLLM-12352][feat] integrate ModelExpress checkpoint loading #15641 for the ModelExpress checkpoint-loading integration and packaging.
• Remains stacked on open [None][feat] Bind SourceIdentity to checkpoint artifacts #16159 for ArtifactIdentity and content-bound SourceIdentity; this branch was rebased onto [None][feat] Bind SourceIdentity to checkpoint artifacts #16159 head 3df2ef3777.
• Includes the post-transform profile review follow-up from commit 3f2eee5133, replayed on the current dependency stack.

The PR targets main, so its visible diff still includes #16159 while that dependency remains open. It should remain draft and be rebased after #16159 merges.

Changes

Foundation A: exact capability profiles

• Match the exact root model class, architecture, model type, transfer scope, protocol, speculative mode, and feature constraints.
• Capture architecture and model type from the resolved config before model construction can normalize them.
• Resolve qualification once per load and use the same decision for SourceIdentity, publication, reception, and staged finalization.
• Validate the public support table against the runtime registry so stale or missing profile rows fail tests.

Foundation B: reusable qualification harness

• Parameterize full-load versus staged-load lifecycle equivalence.
• Compare parameter and buffer names, shapes, dtypes, layouts, strides, storage offsets, and exact values.
• Verify aliases, transform guards, derived state, and deterministic output probes.
• Require an unregistered exact-root negative control.
• Preserve explicit no-disk staged-receiver and unqualified disk-fallback coverage.

Foundation C: transform-layout ABI

• Define the initial Llama ABI as trtllm-llama-target-layout-v1.
• Bind the ABI into SourceIdentity format version 3 and MX source metadata.
• Reject old identities and missing or mismatched ABI metadata before P2P transfer.
• Keep format and ABI compatibility mandatory even under the ENFORCE identity policy.
• Document ABI bump rules and the exact currently supported profile.

Scope

• No new model family is enabled.
• Separately loaded draft models remain unsupported for post-transform transfer.
• Real donor/receiver GPU qualification remains required before making a readiness claim.

Validation

• Rebase range-diff: the two Foundation commits were preserved; the only conflict combined [None][feat] Bind SourceIdentity to checkpoint artifacts #16159's newer fallback fixture with the ABI assertion.
• Iterative self-review: no actionable findings after the rebase.
• Changed-file pre-commit hooks: passed.
• Python syntax compilation for all changed Python files: passed.
• git diff --check: passed.
• DCO signatures: preserved and passing on GitHub.
• Focused pytest collection is unavailable locally because this checkout has no TRT-LLM runtime environment (torch and transformers are not installed); full CI is required.

Dev Engineer Review

• Added exact post-transform qualification across model identity, architecture, model type, transfer scope, protocol, speculative mode, feature constraints, and transform-layout ABI.
• Qualification is resolved once per load and reused across identity creation, publication, reception, and finalization.
• Added Llama ABI trtllm-llama-target-layout-v1, bound to SourceIdentity format version 3. ABI mismatches or missing metadata fail closed before P2P transfer.
• Existing behavior remains limited to qualified LlamaForCausalLM; unsupported models and speculative/draft-model cases fall back safely.
• Source identity validation, serialization, fingerprinting, and enforced matching consistently include the ABI.
• Documentation defines the qualification workflow and immutable ABI rules.
• No config or test-list changes were identified. git diff --check, syntax compilation, hooks, and DCO validation passed. Focused pytest execution was unavailable because torch and transformers are absent; full CI remains required.

QA Engineer Review

Test code changes include:

• Added and updated MX/GMS ModelLoader qualification, fallback, ABI propagation, normalization, speculative-mode, and output-behavior tests.
• Added MX checkpoint metadata publication and ABI compatibility tests.
• Added post-transform profile validation and configuration-identity tests.
• Added SourceIdentity ABI binding, serialization, fingerprint, mismatch, and format-version tests.
• Expanded the lifecycle harness with unqualified-model negative controls and output probes.

No corresponding coverage entries were found in test-db/, qa/, or tests/integration/test_lists/ for these focused tests. Verdict: needs follow-up.

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #60831 [ run ] triggered by Bot. Commit: 63503c8 Link to invocation

[tensorrt-cicd]

PR_Github #60831 [ run ] completed with state SUCCESS. Commit: 63503c8
/LLM/main/L0_MergeRequest_PR pipeline #49107 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 --stage-list "A100X-PyTorch-1,B300-PyTorch-1,DGX_B200-PyTorch-5,DGX_H100-PyTorch-5"

[tensorrt-cicd]

PR_Github #61117 [ run ] triggered by Bot. Commit: 63503c8 Link to invocation

[tensorrt-cicd]

PR_Github #61117 [ run ] completed with state SUCCESS. Commit: 63503c8
/LLM/main/L0_MergeRequest_PR pipeline #49371 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61352 [ run ] triggered by Bot. Commit: 63503c8 Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #61352 [ run ] completed with state SUCCESS. Commit: 63503c8
/LLM/main/L0_MergeRequest_PR pipeline #49580 completed with status: 'SUCCESS'

CI Report

Link to invocation

[tensorrt-cicd]

PR_Github #61354 [ run ] triggered by Bot. Commit: c0c0de5 Link to invocation

[tensorrt-cicd]

PR_Github #61354 [ run ] completed with state FAILURE. Commit: c0c0de5
/LLM/main/L0_MergeRequest_PR pipeline #49582 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 #61986 [ run ] triggered by Bot. Commit: c0c0de5 Link to invocation

[tensorrt-cicd]

PR_Github #61986 [ run ] completed with state SUCCESS. Commit: c0c0de5
/LLM/main/L0_MergeRequest_PR pipeline #50175 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 --stage-list "DGX_H100-PyTorch-1"

[tensorrt-cicd]

PR_Github #62044 [ run ] triggered by Bot. Commit: c0c0de5 Link to invocation

[tensorrt-cicd]

PR_Github #62044 [ run ] completed with state SUCCESS. Commit: c0c0de5
/LLM/main/L0_MergeRequest_PR pipeline #50228 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #62268 [ run ] triggered by Bot. Commit: c0c0de5 Link to invocation

[tensorrt-cicd]

PR_Github #62268 [ run ] completed with state SUCCESS. Commit: c0c0de5
/LLM/main/L0_MergeRequest_PR pipeline #50436 completed with status: 'SUCCESS'

CI Report

Link to invocation

[coderabbitai]

Walkthrough

Post-transform ModelExpress support now uses explicit capability profiles and immutable transform-layout ABI identifiers. Source identities, MX metadata, loader qualification, fallback behavior, tests, and documentation were updated to enforce exact ABI compatibility before P2P transfers.

Changes

Post-transform ABI qualification

Layer / File(s)	Summary
Identity and qualification contracts tensorrt_llm/_torch/weight_sharing/..., tests/unittest/_torch/weight_sharing/...	Post-transform profiles and SourceIdentity now carry validated ABI identifiers through construction, serialization, fingerprints, matching, and compatibility checks.
Loader qualification and propagation tensorrt_llm/_torch/pyexecutor/model_loader.py, tests/unittest/_torch/executor/...	ModelLoader captures configuration identity before normalization, computes one qualification decision, propagates its ABI into SourceIdentity, and applies it across MX/GMS staging and publishing.
MX ABI metadata and transfer gates tensorrt_llm/_torch/models/checkpoints/mx/..., tests/unittest/_torch/models/checkpoints/mx/...	MX publishing records the ABI, while receiving requires a non-empty exact ABI match and otherwise falls back to disk loading.
Qualification harness and documentation tests/unittest/utils/post_transform_qualification.py, docs/source/features/model-express.md	Qualification tests add unregistered-root and output-equivalence checks; documentation defines exact profiles, compatibility gates, and immutable ABI rules.

Estimated code review effort: 4 (Complex) | ~60 minutes

Sequence Diagram(s)

sequenceDiagram
  participant ModelLoader
  participant MXCheckpointLoader
  participant SourceIdentity
  ModelLoader->>SourceIdentity: bind qualified transform ABI
  ModelLoader->>MXCheckpointLoader: publish or receive post-transform weights
  MXCheckpointLoader->>SourceIdentity: compare source and receiver ABI
  SourceIdentity-->>MXCheckpointLoader: compatibility result
  MXCheckpointLoader-->>ModelLoader: use MX transfer or fall back to disk

Loading

Suggested labels: api-breaking

Suggested reviewers: bowenfu

🚥 Pre-merge checks | ✅ 3 | ❌ 2

❌ Failed checks (2 warnings)

Check name	Status	Explanation	Resolution
Description check	⚠️ Warning	The description is informative but doesn't follow the required template and omits explicit Test Coverage and PR Checklist sections.	Rewrite it using the repository template, adding Description, Test Coverage, and PR Checklist sections with concrete details.
Docstring Coverage	⚠️ Warning	Docstring coverage is 12.50% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (3 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title matches the PR's main change and follows the required [ticket][type] format.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py (1)

994-1012: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low value

Remove the unused _metadata_is_post_transform helper. No call sites remain in the repo, so this wrapper can be dropped to avoid dead code.

🤖 Prompt for AI Agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py` around lines
994 - 1012, Remove the unused `_metadata_is_post_transform` helper entirely;
retain `_metadata_weight_layout_status` and its callers unchanged.

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Nitpick comments:
In `@tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py`:
- Around line 994-1012: Remove the unused `_metadata_is_post_transform` helper
entirely; retain `_metadata_weight_layout_status` and its callers unchanged.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 09ee05a3-92e7-479d-b15a-7fab33f49529

📥 Commits

Reviewing files that changed from the base of the PR and between 5a69240 and c0c0de5.

📒 Files selected for processing (13)

• docs/source/features/model-express.md
• tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py
• tensorrt_llm/_torch/pyexecutor/model_loader.py
• tensorrt_llm/_torch/weight_sharing/__init__.py
• tensorrt_llm/_torch/weight_sharing/post_transform_profiles.py
• tensorrt_llm/_torch/weight_sharing/source_identity.py
• tests/unittest/_torch/executor/test_model_loader_gms.py
• tests/unittest/_torch/executor/test_model_loader_mx.py
• tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py
• tests/unittest/_torch/weight_sharing/_source_identity_fakes.py
• tests/unittest/_torch/weight_sharing/test_post_transform_profiles.py
• tests/unittest/_torch/weight_sharing/test_source_identity.py
• tests/unittest/utils/post_transform_qualification.py

[2ez4bz]

Approving changes to on the modeling side.

[schetlur-nv]

rubber stamp for trt-llm-devs