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

[chienchunhung]

Summary

Integrate optional ModelExpress (MX) checkpoint loading for the TensorRT-LLM PyTorch backend.

MX is treated as a separately deployed shared service. Operators start and manage the MX server and its Redis metadata backend, then configure each TensorRT-LLM launch with the MX server URL. TensorRT-LLM does not start or own MX/Redis processes or containers.

Deployment Model

Configuration	Behavior
mx_config.server_url is set	Use that MX server for source discovery, P2P receive, and publish.
MODEL_EXPRESS_URL is set and no explicit URL is configured	Resolve it into mx_config.server_url and use that server.
No MX URL is configured	Skip MX transfer and use Hugging Face checkpoint loading.
MX discovery or transfer fails	Fall back to Hugging Face checkpoint loading.
checkpoint_format is not MX	Existing checkpoint behavior is unchanged.

One MX server can be shared by multiple TensorRT-LLM launches that can reach the endpoint. A worker that does not find a compatible source loads from storage and publishes its weights; later compatible workers can receive them through MX.

Changes

• Add the optional mx package extra with modelexpress==0.4.1:

  pip install "tensorrt_llm[mx]"

• Pass mx_config.server_url, MODEL_EXPRESS_URL, query timeout, and model identity through the MX checkpoint loader.
• Support the released ModelExpress 0.4.1 client source-query shape and keep polling for a compatible donor when none is registered yet.
• Attach TRT-LLM SourceIdentity, weight-layout, and transform-protocol metadata so incompatible publishers are rejected before transfer.
• Serialize temporary process-wide MX environment and identity-builder state across load and publish calls.
• Synchronize pending CUDA writes before exposing post-transform source buffers through MX.
• Preserve Hugging Face fallback when source discovery, transfer, or optional client cleanup fails.
• Preserve mixed-transfer fallback behavior for pre-transform sources and reject unsafe partial post-transform transfers.
• Document independent MX/Redis deployment, shared-server configuration, and source-process lifetime requirements.
• Add focused unit coverage for URL propagation, source discovery, identity compatibility, timeout behavior, transfer, publish ordering, and fallback.

Current Scope

• Post-transform MX reception is qualified for LlamaForCausalLM with transform protocol version 1.
• Other model families fall back to Hugging Face loading until their staged loading hooks and runtime layouts are qualified.
• Loads requiring a separately loaded draft model also fall back.
• Artifact identity is not yet bound to immutable checkpoint contents or a resolved revision; that remains follow-up work before general content-safe deployment.

Validation

Completed on reviewed head 3388328514:

• All scoped pre-commit hooks passed, including formatting, lint, codespell, test-list validation, private-key checks, and DCO.
• Python syntax compilation passed for the modified MX implementation and tests.
• git diff --check passed.
• Focused pytest could not run in the local macOS environment because the TensorRT-LLM runtime dependencies (torch and transformers) are not installed.
• No /bot run CI rerun was requested for the reviewed head.

Related Work

• Staged model hook migration landed in #15432.

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #55912 [ run ] triggered by Bot. Commit: 0bbe292 Link to invocation

[tensorrt-cicd]

PR_Github #55912 [ run ] completed with state FAILURE. Commit: 0bbe292
/LLM/main/L0_MergeRequest_PR pipeline #44791 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 #56487 [ run ] triggered by Bot. Commit: f52b100 Link to invocation

[tensorrt-cicd]

PR_Github #56487 [ run ] completed with state FAILURE. Commit: f52b100
/LLM/main/L0_MergeRequest_PR pipeline #45328 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 #57252 [ run ] triggered by Bot. Commit: 22c40e0 Link to invocation

[tensorrt-cicd]

PR_Github #57252 [ run ] completed with state FAILURE. Commit: 22c40e0
/LLM/main/L0_MergeRequest_PR pipeline #46019 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]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #57855 [ run ] triggered by Bot. Commit: 5f2516e Link to invocation

[tensorrt-cicd]

PR_Github #57855 [ run ] completed with state FAILURE. Commit: 5f2516e
/LLM/main/L0_MergeRequest_PR pipeline #46551 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 #57950 [ run ] triggered by Bot. Commit: 6803176 Link to invocation

[tensorrt-cicd]

PR_Github #57950 [ run ] completed with state FAILURE. Commit: 6803176
/LLM/main/L0_MergeRequest_PR pipeline #46631 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 #58039 [ run ] triggered by Bot. Commit: 552ff8a Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #58073 [ run ] triggered by Bot. Commit: 552ff8a Link to invocation

[tensorrt-cicd]

PR_Github #58039 [ run ] completed with state ABORTED. Commit: 552ff8a

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59598 [ run ] triggered by Bot. Commit: 7abab2b Link to invocation

[tensorrt-cicd]

PR_Github #59598 [ run ] completed with state SUCCESS. Commit: 7abab2b
/LLM/main/L0_MergeRequest_PR pipeline #48041 completed with status: 'SUCCESS'

CI Report

Link to invocation

[chienchunhung]

/bot run --disable-fail-fast

[chienchunhung]

/bot run --disable-fail-fast

[tensorrt-cicd]

PR_Github #59831 [ run ] triggered by Bot. Commit: 705436e Link to invocation

[tensorrt-cicd]

PR_Github #59831 [ run ] completed with state SUCCESS. Commit: 705436e
/LLM/main/L0_MergeRequest_PR pipeline #48238 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 "Build-Docker-Images"

[tensorrt-cicd]

PR_Github #60014 [ run ] triggered by Bot. Commit: 705436e Link to invocation

[tensorrt-cicd]

PR_Github #60014 [ run ] completed with state FAILURE. Commit: 705436e
/LLM/main/L0_MergeRequest_PR pipeline #48405 (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]

/bot skip --comment "Current-head CI failures were unrelated test/infra flakes. The targeted Docker run verified all PR-relevant x86_64 and SBSA release images, including successful modelexpress==0.4.1 installation; its only failure was an SBSA tritondevel Jenkins-agent DNS/pod deletion."

[tensorrt-cicd]

PR_Github #60069 [ skip ] triggered by Bot. Commit: 705436e Link to invocation

[tensorrt-cicd]

PR_Github #60069 [ skip ] completed with state SUCCESS. Commit: 705436e
Skipping testing for commit 705436e

Link to invocation