[None][refactor] Delegate MX checkpoint loading to ModelExpress

[zhengluo-nv]

Description

This PR narrows the TensorRT-LLM side of the ModelExpress checkpoint-loading integration.

TensorRT-LLM now stays responsible for lifecycle orchestration only:

• register the modelexpress checkpoint loader name;
• pass the live model object into the loader so ModelExpress can transfer directly into TRT-LLM weight buffers;
• skip normal weight mapping only when the ModelExpress loader reports that weights were already preloaded;
• call post_load_publish() after TRT-LLM finishes weight loading, post-load processing, and CUDA stream synchronization.

The canonical user-facing checkpoint/load format is modelexpress, not mx, to avoid confusion with MXFP4/MXFP8 formats. The in-tree shim keeps the existing MXCheckpointLoader class name only as an internal compatibility name for the ModelExpress-owned Python loader.

The low-level ModelExpress behavior is delegated to the modelexpress Python package instead of living in TensorRT-LLM. Source discovery, identity construction, RDMA transfer, fallback behavior, source publication, and metadata lifecycle are owned by ModelExpress.

The post-load publish hook is intentionally ordered after TRT-LLM's post-load path. That prevents publishing tensor addresses/content before TRT-LLM has finished applying fallback weights or model-specific post-processing, while still allowing successful P2P receivers to become sources for later replicas.

Test Coverage

• Targeted TensorRT-LLM Python syntax validation on changed files:
  • tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py
  • tensorrt_llm/_torch/models/checkpoints/auto_mapper.py
  • tensorrt_llm/_torch/models/checkpoints/hf/config_loader.py
  • tensorrt_llm/_torch/models/checkpoints/hf/weight_loader.py
  • tensorrt_llm/_torch/models/checkpoints/hf/weight_mapper.py
  • tensorrt_llm/_torch/pyexecutor/model_loader.py
  • tensorrt_llm/llmapi/llm_args.py
  • tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py
  • tests/unittest/_torch/pyexecutor/test_model_loader_mx.py
• Unit tests updated to cover:
  • ModelExpress loader construction through the TRT-LLM checkpoint-loader registry;
  • P2P success skipping native weight mapping;
  • fallback weights still going through native TRT-LLM mapping;
  • post_load_publish() running after CUDA stream synchronization.
• QA list updates are not required for this PR. The test-list guideline applies to integration QA lists under tests/integration/test_lists/qa/*, while this PR only updates unit tests such as test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works, test_mx_partial_fallback_merges_returned_weights, and test_mx_fallback_runs_standard_weight_mapping in tests/unittest/_torch/pyexecutor/test_model_loader_mx.py.
• End-to-end ModelExpress + TRT-LLM validation:
  • direct trtllm-serve with nvidia/Kimi-K2.5-NVFP4;
  • TP=4 source loaded from disk and published 4 Ready ModelExpress metadata records;
  • TP=4 target found rank-matched source workers and used RDMA instead of disk fallback;
  • each target rank transferred 2188 tensors / 151.11 GB in about 5.55s at about 217-218 Gbps;
  • target published 4 additional Ready metadata records after post-load publish;
  • target serving checks returned HTTP 200 for /health and /v1/models.

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.

[coderabbitai]

📝 Walkthrough

Walkthrough

This PR migrates the in-tree MX checkpoint loader to an external modelexpress package dependency. The loader now delegates all functionality upstream, the post-processing pipeline is reordered to apply changes before publishing, and tests are simplified to validate the new shim-based architecture and callback ordering.

Changes

ModelExpress Integration Refactoring

Layer / File(s)	Summary
MX Checkpoint Loader Shim Implementation tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py	MXCheckpointLoader is converted from a full in-tree implementation extending HfCheckpointLoader to a thin shim class whose __new__ dynamically imports and instantiates modelexpress.engines.trtllm.loader.MXCheckpointLoader, raising ImportError when the upstream package is unavailable.
Checkpoint Post-Processing Pipeline Reordering tensorrt_llm/_torch/pyexecutor/model_loader.py	checkpoint_loader.post_load_apply(model, weights_preloaded=...) is called earlier in the load() pipeline, followed later by checkpoint_loader.post_load_publish(model, checkpoint_dir=..., weights_preloaded=...) before module finalization, splitting and reordering these two post-processing stages.
Configuration Documentation Update tensorrt_llm/llmapi/llm_args.py	ModelExpressConfig.server_query_timeout_s field description is updated to reflect the timeout's role in "TRT-LLM source discovery" and fallback behavior.
MX Loader Test Suite Simplification tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py	Test suite is reduced from comprehensive coverage of internal load/publish/query logic to a minimal shim-focused suite: FakeModelexpressMXCheckpointLoader and _install_fake_modelexpress_loader helper are added, and three tests verify shim instantiation via direct construction, registry selection, and ImportError on missing upstream package.
Post-Processing Event Ordering Validation Tests tests/unittest/_torch/pyexecutor/test_model_loader_mx.py	Mocks for torch.cuda.current_stream().synchronize() and checkpoint_loader.post_load_publish are updated to record event markers; three MX/fallback test cases now assert that the last two recorded events are ["sync", "post_load_publish"] to validate the post-processing order.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~25 minutes

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Docstring Coverage	⚠️ Warning	Docstring coverage is 9.09% which is insufficient. The required threshold is 80.00%.	Write docstrings for the functions missing them to satisfy the coverage threshold.

✅ Passed checks (4 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically describes the main refactor: delegating MX checkpoint loading to the ModelExpress package instead of in-tree implementation.
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.
Description check	✅ Passed	The PR description comprehensively explains the rationale, implementation details, test coverage, and design decisions for narrowing TRT-LLM's ModelExpress integration.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

✨ Finishing Touches

🧪 Generate unit tests (beta)

• Create PR with unit tests

---

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

[coderabbitai]

🧹 Nitpick comments (1)

tests/unittest/_torch/pyexecutor/test_model_loader_mx.py (1)

92-169: QA list updates are unnecessary for this PR scope.

These changes are unit-test-only and do not add or modify integration test definitions, so no update to tests/integration/test_lists/qa/* is needed in this PR.

As per coding guidelines "If the PR only touches unittest/ or narrow unit scope, say explicitly whether QA list updates are unnecessary or optional."

🤖 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 `@tests/unittest/_torch/pyexecutor/test_model_loader_mx.py` around lines 92 -
169, The reviewer notes QA list updates are unnecessary because these changes
are unit-test-only; update the PR description (or add a brief comment in the
PR/commit message) to explicitly state that QA list updates are unnecessary per
the guideline, referencing that this PR only touches unit tests such as
test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works,
test_mx_partial_fallback_merges_returned_weights, and
test_mx_fallback_runs_standard_weight_mapping so maintainers know no changes to
tests/integration/test_lists/qa/* are required.

🤖 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 `@tests/unittest/_torch/pyexecutor/test_model_loader_mx.py`:
- Around line 92-169: The reviewer notes QA list updates are unnecessary because
these changes are unit-test-only; update the PR description (or add a brief
comment in the PR/commit message) to explicitly state that QA list updates are
unnecessary per the guideline, referencing that this PR only touches unit tests
such as
test_mx_success_initializes_mapper_skips_weight_mapping_and_reload_works,
test_mx_partial_fallback_merges_returned_weights, and
test_mx_fallback_runs_standard_weight_mapping so maintainers know no changes to
tests/integration/test_lists/qa/* are required.

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: 6c2df948-8a9c-4bbd-aeec-274b73e739a8

📥 Commits

Reviewing files that changed from the base of the PR and between 7021547 and 976bc76.

📒 Files selected for processing (5)

• tensorrt_llm/_torch/models/checkpoints/mx/checkpoint_loader.py
• tensorrt_llm/_torch/pyexecutor/model_loader.py
• tensorrt_llm/llmapi/llm_args.py
• tests/unittest/_torch/models/checkpoints/mx/test_mx_checkpoint_loader.py
• tests/unittest/_torch/pyexecutor/test_model_loader_mx.py

[chienchunhung]

A couple of things to flag initially:

1. In TRTLLM we would like to preserve the existing fallback class MXCheckpointLoader(HfCheckpointLoader): if MX checkpoint loading fails, we can still fallback to the default Hf checkpoint loader. The better approach here is MX defines an interface that downstream (e.g., TRTLLM and other frameworks) can inherit and implement custom logic. This allows upstream to unify the core logic, while the downstream can still customize instead of running as blackbox. PS: I think that's what the offline discussion concluded.
2. In PR#13531, the writer publishes the weights first, then both reader and writer transform. While in this PR the order is reversed: the writer transforms the weights first before publishing. Wondering if there is any reason why the order is reversed?

[zhengluo-nv]

1. Yes that's exactly I want to do. We implemented an interface EngineAdapter in the MX repo and TRTLLM has implemented it, see PR. MX wants to define and control the fallback order, and our loading strategy is not only just RDMA transfer and default loader: we could add GDS loader and model streamer loader to the load strategy chain. We also want customer to arbitrarily enable and disable and re-order load strategies. So all these are moved into the MX repo.

2. Do you mean I changed the post_load_publish order in tensorrt_llm/_torch/pyexecutor/model_loader.py? I think TRT-LLM should publish after post_load_weights()/MoE finalization so that it publishes the real final weight tensors (assuming post_load_weights could introduce new tensors and mutate existing tensors). In vLLM we did this similarly that we publish tensors after _process_weights_after_loading(). The end to end testing validates the inference correctness after the change so I think it could be the right order.

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #48630 [ run ] triggered by Bot. Commit: 976bc76 Link to invocation

[tensorrt-cicd]

PR_Github #48630 [ run ] completed with state SUCCESS. Commit: 976bc76
/LLM/main/L0_MergeRequest_PR pipeline #38412 completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

Wanted to raise a few concerns around: what need to be hidden in the upstream vs. what can be customized in TRTLLM.

The "publish before or after weight loading" ordering is exactly one critical example. Imagine a production environment with:

• 4 replicas of Qwen2.5-72B-Instruct, all TP=8, all on the same model checkpoint.
• Replica A is configured with attn_backend="TRTLLM", nvfp4_allowed_backends=["cutlass"].
• Replica B comes up later with attn_backend="FlashInfer", nvfp4_allowed_backends=["cutedsl"] (operator-configured choice).

IIUC B's 8 ranks each find rank-matched A donors, pull A's post-transform weights (laid out for TRTLLM + cutlass). B skips its own post_load_weights (assuming the skip will be in place which seems not currently). B is now running with A's backend choices even though it's configured for FlashInfer + cutedsl. Either B silently does the wrong thing, or some hidden mismatch detection has to fail the boot.

Another example is EP. If post_load_weights is skipped on the receiver side, MoE finalize must still needs to run because the routing state is per-process GPU state. So the skip needs to be specifically module.post_load_weights() but not MoeLoadBalancer.finalize_model(). The non-uniformity is what makes things complicated

The receiver's pipeline becomes:

  receiver:                                                                                                                                                                                 
      receive(post-transform bytes from MX)                 
      # SKIP module.post_load_weights according to this PR   
      rebuild_aliases(model)              # but alias rebuilding is still required
      if isinstance(moe_load_balancer, MoeLoadBalancer):                                                                                                                                    
          moe_load_balancer.finalize_model()   # AND MoE finalize is also required                                                                                                               
      torch.cuda.current_stream().synchronize()    

It seems like the current mental model assumes "universally homogeneous" deployment configuration, which is generally common in practice. However TRTLLM needs to cover all possible setups and, even if not supported, fails out gracefully rather than silently failing.

Overall, I think there are multiple edge cases (attn_backend, quant backend list, FP8 scale fusion strategy, future quant schemes, etc) that make "transform-once-then-publish" hard to adapt. Either MX guarantees to incorporate all sources of customization into identity building / matching (which is unbounded), or the safer approach is to allow downstream to handle these customization.

PS: Even if we later decide to go the publish-post-transform route, it requires factoring module.post_load_weights() into transform vs setup steps so receivers can selectively skip

---

A couple of suggestions for what can be abstracted in MX:

1. Stable, public low-level client surface: keep MxClient, MxLiveWeightLoader, publish_model_params as public APIs.
2. Make LoadStrategyChain as a composable library, not a CheckpointLoader. I think MX providing a chain of fallback strategy across RDMA / GDS / ModelStreamer is a nice feature, but the chain should be a library the engine
   assembles, not a loader class that owns the engine's lifecycle. For example:

# MX provides:                                            
  chain = LoadStrategyChain([                                                                                                                                                               
      RdmaStrategy(mx_client=...),
      GdsStrategy(...),                                                                                                                                                                     
      ModelStreamerStrategy(...),                                                                                                                                                           
  ])
                                                                                                                                                                                            
  # Engine uses (this lives in tensorrt_llm or vllm):                                                                                                                                       
  class MXCheckpointLoader(HfCheckpointLoader):     # ← engine owns this class
      def load_weights(self, checkpoint_dir, ...):                                                                                                                                          
          result = chain.try_load(model, identity, ...)                                                                                                                                     
          if result.success:                                                                                                                                                                
              self._weights_preloaded = True                                                                                                                                                
              return result.fallback_weights  # if partial                                                                                                                                  
          # Engine-owned fallback decision
          return super().load_weights(checkpoint_dir, ...)  # disk fallback

3. Separate publish_pre_transform() and publish_post_transform() APIs. Don't bake the specific ordering into MX functions that downstream can't adapt.
4. Extensible identity schema so that engines can specify custom identity factors.

[zhengluo-nv]

1. If 2 replicas have different attn_backend, we have a source identity mechanism so they will never be matched for p2p transfer. MX must guarantee to incorporate all sources of customization into identity building / matching. We surely cannot cover 100% HF models but the CI should validate inference correctness on most of popular models in the future.

2. publish-post-transform is mainly due to VRAM consideration. If we try to publish before transform, it registers some pre-transform tensors, I assume they will still occupy VRAM even if they are unused after transformation, and the only reason they live is for model express p2p transfer, which is unacceptable to most of customers IMO.

3. Refactor and moving most of code into the MX repo is due to MX is in early stage and we will add many features, so keeping a minimum interface integration with inference engines is important to avoid to open 3 PRs to each upstream engine codebase just for one feature, and it also keeps MX behavior in 3 engines consistent.

[zhengluo-nv]

Change from MX -> modelexpress to keep consistent branding across all inference engines. In vLLM and SGLang we use modelexpress because MX may confuse customers with MXFP4/MXFP8 format.