feat(trtllm): MXCheckpointLoader deployment support with auto-detect and RDMA validation by KavinKrishnan · Pull Request #218 · ai-dynamo/modelexpress

KavinKrishnan · 2026-04-09T19:58:26Z

Summary

Follow-up to #202. Replaces the LoadFormat.PRESHARDED monkey-patch approach with TRT-LLM's @register_checkpoint_loader("MX") architecture, updates deployment manifests for auto-detect source/target, and adds the patch_mx_loader.py script to bridge the base image's TRT-LLM with the new loader.

Validated E2E: 16 ranks × 90.75 GB RDMA at 363–506 Gbps on GCP GB200 (Kimi K2.5, TP=8).

Changes

Dockerfile (Dockerfile.ph3-gcp-gb200)

Replace old apply_patches.py / patch_model_loader.py (PRESHARDED approach) with:
- COPY checkpoints/mx/ from TRT-LLM fork — installs MXCheckpointLoader via @register_checkpoint_loader("MX")
- patch_mx_loader.py — patches base image's model_loader.py with model=model kwarg, p2p_succeeded check, _weights_presharded marking, publish_as_source() hook, and MX registry entries
Remove worker.py publish hook (publish now handled by MXCheckpointLoader.publish_as_source() in model_loader.py)
Add TODO comments referencing upstream PRs for eventual patch removal

K8s manifests

Remove --model-express-role source/target args (auto-detected by probing MX server)
Remove MODEL_EXPRESS_SOURCE env var from source DGD
Update image tags to v2.9.0
Add MX_HEARTBEAT_TIMEOUT_SECS=999999 and MX_GC_TIMEOUT_SECS=999999 to MX server to prevent metadata GC

Patch script (trtllm_patches/v1.3.0rc5/patch_mx_loader.py)

New script that applies minimal patches to the base image's TRT-LLM for MX loader compatibility
All patches are idempotent (check before patching)

How it works

The MXCheckpointLoader (from TRT-LLM PR #13045) replaces the old LoadFormat.PRESHARDED approach:

Dynamo engine.py sets checkpoint_format="MX" when --model-express-url is provided
TRT-LLM constructs MXCheckpointLoader (registered via @register_checkpoint_loader("MX"))
Source: _has_existing_sources() probe finds no sources → falls back to HF disk loading → publish_as_source() registers NIXL metadata with MX server
Target: probe finds sources (or MODEL_EXPRESS_TARGET=1) → MxLiveWeightLoader does NIXL RDMA transfer → p2p_succeeded=True → _weights_presharded set on Linear modules → skip weight mapper

Companion PRs

PR	Repo	Status	Description
#202	ModelExpress	Merged	TRT-LLM P2P client (`trtllm_live_transfer.py`)
#8037	Dynamo	In review	Engine integration with `--model-express-url`, auto-detect
#13045	TRT-LLM (NVIDIA)	In review	Full MX + GMS two-axis checkpoint loader (Chien-Chun)

Signed-off-by: Kavin Krishnan kavink@nvidia.com

- Remove --model-express-role args from K8s manifests (source/target auto-detected by probing MX server for existing sources) - Remove MODEL_EXPRESS_SOURCE env var from source DGD manifest - Update Dockerfile verification to check model_express_url - Add TODO comments for removing monkey patches once upstream PRs merge: - ai-dynamo/dynamo#8037 (--model-express-url native) - NVIDIA/TensorRT-LLM#12898 (LoadFormat.PRESHARDED native) Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

coderabbitai · 2026-04-09T20:01:14Z

Walkthrough

The changes remove --model-express-role arguments from Kubernetes deployment manifests and update a Dockerfile verification check to validate model_express_url instead of MODEL_EXPRESS_SOURCE, while adding TODO comments about pending patch removals.

Changes

Cohort / File(s)	Summary
Dockerfile Build Verification `examples/p2p_transfer_k8s/client/trtllm/Dockerfile.ph3-gcp-gb200`	Updated build-time verification to check for `model_express_url` instead of `MODEL_EXPRESS_SOURCE`; added TODO comments flagging patches for removal.
Kubernetes Manifests `examples/p2p_transfer_k8s/client/trtllm/kimi-disagg-mx-tp8-dgd.yaml`, `examples/p2p_transfer_k8s/client/trtllm/kimi-source-decode-dgd.yaml`	Removed `--model-express-role` arguments from container command specs; also removed `MODEL_EXPRESS_SOURCE` env var from source manifest.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~12 minutes

Poem

🐰 Hops with glee!
We've trimmed the roles, so sleek and lean,
Model Express now validates anew,
Arguments fade like morning dew,
Cleaner configs hop on through! ✨

🚥 Pre-merge checks | ✅ 2 | ❌ 1

❌ Failed checks (1 warning)

Check name	Status	Explanation	Resolution
Title check	⚠️ Warning	The PR title mentions 'MXCheckpointLoader deployment support' but the actual changes focus on auto-detecting source/target roles and removing explicit role flags from manifests—unrelated to checkpoint loading.	Update the title to reflect the actual changes, e.g. 'refactor: auto-detect TRT-LLM source/target role and remove explicit role flags' or similar.

✅ Passed checks (2 passed)

Check name	Status	Explanation
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Description Check	✅ Passed	Check skipped - CodeRabbit’s high-level summary is enabled.

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

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

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

coderabbitai

🧹 Nitpick comments (1)

examples/p2p_transfer_k8s/client/trtllm/Dockerfile.ph3-gcp-gb200 (1)
69-72: Update deployment docs for role auto-detection changes.

Since the rollout no longer relies on explicit source/target role flags and MODEL_EXPRESS_SOURCE, please update docs/DEPLOYMENT.md to reflect the new required arguments/env behavior and patch expectations.

Based on learnings: Update docs/DEPLOYMENT.md when making changes to configuration, environment variables, deployment (Docker, K8s, Helm, P2P).
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.

In `@examples/p2p_transfer_k8s/client/trtllm/Dockerfile.ph3-gcp-gb200` around
lines 69 - 72, Update docs/DEPLOYMENT.md to document that explicit source/target
role flags and the MODEL_EXPRESS_SOURCE env var are no longer required and that
the system now uses automatic role detection; describe the new required
command-line arguments and environment behavior for deployments (Docker, K8s,
Helm, P2P), and clarify that the patch step (copying and running
trtllm_patches/v1.3.0rc5/patch_model_loader.py as done in
Dockerfile.ph3-gcp-gb200) is still expected in current builds until upstream PR
merges, including when/why to include the patch and how it affects
model_loader/post_load_weights ordering.

🤖 Prompt for all review comments with AI agents

Verify each finding against the current code and only fix it if needed.

Nitpick comments:
In `@examples/p2p_transfer_k8s/client/trtllm/Dockerfile.ph3-gcp-gb200`:
- Around line 69-72: Update docs/DEPLOYMENT.md to document that explicit
source/target role flags and the MODEL_EXPRESS_SOURCE env var are no longer
required and that the system now uses automatic role detection; describe the new
required command-line arguments and environment behavior for deployments
(Docker, K8s, Helm, P2P), and clarify that the patch step (copying and running
trtllm_patches/v1.3.0rc5/patch_model_loader.py as done in
Dockerfile.ph3-gcp-gb200) is still expected in current builds until upstream PR
merges, including when/why to include the patch and how it affects
model_loader/post_load_weights ordering.

ℹ️ Review info

⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Pro

Run ID: 6a67e355-6bc7-4789-8741-cfa076524c61

📥 Commits

Reviewing files that changed from the base of the PR and between c272550 and d441d87.

📒 Files selected for processing (3)

examples/p2p_transfer_k8s/client/trtllm/Dockerfile.ph3-gcp-gb200
examples/p2p_transfer_k8s/client/trtllm/kimi-disagg-mx-tp8-dgd.yaml
examples/p2p_transfer_k8s/client/trtllm/kimi-source-decode-dgd.yaml

💤 Files with no reviewable changes (2)

examples/p2p_transfer_k8s/client/trtllm/kimi-disagg-mx-tp8-dgd.yaml
examples/p2p_transfer_k8s/client/trtllm/kimi-source-decode-dgd.yaml

…roach Replace PRESHARDED monkey patches with: - COPY checkpoints/mx/ from TRT-LLM fork (MxCheckpointLoader) - COPY worker.py from TRT-LLM fork (publish hook) - Remove apply_patches.py and patch_model_loader.py Now requires --build-context trtllm=../TensorRT-LLM in addition to dynamo. Signed-off-by: Kavin Krishnan <kavink@nvidia.com> Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

…e verify TRT-LLM imports require libcuda.so.1 which isn't available at build time. Signed-off-by: Kavin Krishnan <kavink@nvidia.com> Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

- Disable worker.py double-publish (publish now handled by MXCheckpointLoader.publish_as_source in model_loader.py) - Update DGD yamls to v2.9.0 image - Add MX_HEARTBEAT_TIMEOUT_SECS and MX_GC_TIMEOUT_SECS to prevent metadata garbage collection for long-lived sources Validated E2E: 16 ranks x 90.75 GB RDMA at 363-506 Gbps on GCP GB200 Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

devin-ai-integration

Devin Review found 4 potential issues.

View 4 additional findings in Devin Review.

…kerfile - Add Dockerfile.dynamo-runtime for tensorrtllm-runtime base images (replaces full-file copy with in-place patching) - Add trtllm_patches/dynamo/patch_dynamo_mx.py for Dynamo MX hooks (model_express_url CLI, engine checkpoint_format, compat guards) - Update DGD yamls to v3.2.0 image Validated E2E on gcp-dev-02: source publishes 8 workers, targets receive via RDMA, all pods 1/1 Running zero restarts. Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

Replace outdated PRESHARDED references with checkpoint_format="MX" architecture. Add pre-built image tags, companion PR links, and patch-based Dockerfile instructions. Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

…persist MPI worker processes get killed via SIGKILL when the engine init completes, which skips Python's normal logging shutdown. The buffered FileHandler loses all logged Gbps/timing metrics. Switch to line-buffered stream and explicit flush+fsync at end of load_weights(). Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

New examples/p2p_transfer_k8s/client/trtllm/hpa/ folder demonstrates HPA-driven horizontal scaling where new replicas receive weights via NIXL RDMA from the first (disk-loaded) replica. - kimi-agg-autoscale-dgd.yaml: aggregated TP=8 worker DGD that auto-detects source vs target via _has_existing_sources() - kimi-agg-autoscale-hpa.yaml: DGDSA + HPA wrapper exposing the scale subresource for HPA-driven replica count - hpa/README.md: full demo guide with prerequisites, validated results, and pre-built image references Validated on Kimi K2.5 (TP=8, GCP GB200): first replica ~22 min disk load, scale-up replicas ~5 min via RDMA at 361-583 Gbps/rank. No code changes needed — the demo works entirely with existing auto-detect in MXCheckpointLoader and the upstream DGDSA controller. Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

- patch_mx_loader.py: validate every str.replace() match, exit non-zero on unmatched patterns (catches upstream TRT-LLM API drift instead of silently writing partially-patched files) - Dockerfiles: extend verify step to grep for model=model, publish_as_source, and **kwargs additions in addition to p2p_succeeded - DGD yamls: replace hardcoded cluster-specific values (kavin namespace, customer-gpu-w0e/o7v node pools, kavink personal registry) with explicit <NAMESPACE>, <GPU_NODE_POOL>, <CPU_NODE_POOL>, <REGISTRY>/<NAME>:<TAG> placeholders so users can't accidentally deploy with the wrong values - mx-infra-decode.yaml: replace MX_HEARTBEAT_TIMEOUT_SECS=999999 with 4500s (75min) and document trade-off — covers 685B-class model load times without permanently disabling reaper-based cleanup - README.md: add placeholder reference table and pre-built image link Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

…uctions - Drop pre-built nvcr.io/nvidian/dynamo-dev/kavink:* references that pointed at a personal NVCR namespace - Replace with step-by-step "check out companion repos + buildx build + push" instructions using <YOUR_REGISTRY>/<YOUR_NAME>/<YOUR_TAG> placeholders consistent with the YAML placeholders - Document the directory layout the build context expects - Document what each step of the Dockerfile does - Note that future-merged PRs only need main; current PR branches documented for pre-merge users - Mention the tensorrtllm-runtime base image is published for both arm64 and amd64 (the build itself only documents arm64 for now) Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

zhengluo-nv · 2026-04-29T21:54:36Z

+)
+
+
+def patch_backend_args():


Why do we still add more patches? Shall we wait for upstream changes?

These are temporary deployment patches, not a long-term plan. The actual code changes live in dynamo PR #8037 and TRT-LLM PR #13531 — once those merge into a published tensorrtllm-runtime image, this whole trtllm_patches/ folder gets deleted.

For now they let users layer MX onto the released image without rebuilding Dynamo from source — three small in-place edits:

backend_args.py: add --model-express-url CLI

engine.py: set checkpoint_format="MX" when the URL is set

llm_worker.py: plumb the URL into get_llm_engine

Address Zheng's PR #218 review question ("why do we still add more patches? shall we wait for upstream changes?"): - Expand patch_dynamo_mx.py docstring to explain the script is a *temporary* shim that lets users layer MX onto the released tensorrtllm-runtime image without rebuilding Dynamo from source, and that the actual code changes live in dynamo PR #8037 (delete this script once that merges and a new image is published). - Document why we patch instead of full-file copy (rebased Dynamo files transitively import symbols not in the base image's older Dynamo, e.g. register_embedding_cache_metrics). - Add _replace_or_die() helper mirroring patch_mx_loader.py so any upstream Dynamo source drift fails the build fast instead of silently producing a partially-patched file. (Devin raised the same concern on the TRT-LLM patch script.) Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

Adds runnable Kubernetes manifests and a step-by-step run book for the ModelExpress + TRT-LLM P2P weight-loading integration on GCP GB200. This is the deployment-side companion to TRT-LLM PR #13531 (MXCheckpointLoader, merged 2026-05-06) and Dynamo PR #8037 (--model-express-url CLI, in review). What's added under examples/p2p_transfer_k8s/client/trtllm/: * mx-infra-decode.yaml — ModelExpress server + Redis deployment with placeholders for namespace, registry, CPU node pool. Tunable reaper / GC timeouts (MX_HEARTBEAT_TIMEOUT_SECS=4500s, MX_GC_TIMEOUT_SECS=4500s) sized to cover 685B-class disk loads without permanently disabling cleanup. * kimi-source-decode-dgd.yaml — Source DGD (Kimi K2.5 TP=8 across 2 nodes). Loads from disk, then publishes via publish_model_params. Auto-detected as source by probing MX server for existing entries. * kimi-disagg-mx-tp8-dgd.yaml — Target DGD (Frontend + Prefill TP=8 + Decode TP=8). Auto-detected as target. Receives weights via MxLiveWeightLoader / NIXL RDMA in ~2 seconds per rank instead of 15-20 minutes from disk. * hpa/kimi-agg-autoscale-dgd.yaml — Aggregated worker DGD where every replica uses the same spec; auto-detect handles source vs target. Designed for HPA-driven horizontal scaling. * hpa/kimi-agg-autoscale-hpa.yaml — DGDSA + HorizontalPodAutoscaler wrapping the aggregated DGD. Exposes the scale subresource so HPA can drive replica count. * hpa/README.md — End-to-end demo: first replica ~22 minutes from disk, subsequent HPA-driven replicas ~5 minutes via RDMA at 361-583 Gbps/rank (validated on Kimi K2.5, GCP GB200, 4x 400G RoCE). * README.md (overall guide) — Quick-start (5 steps), build instructions, image-stack guidance for both the post-PR-13531 recommended path (just pip install modelexpress on top of tensorrtllm-runtime once it bumps to TRT-LLM 1.3.0rc15+) and a bridge path that points at the patch-based Dockerfiles still hosted on the kavink/trtllm_clean branch (PR #218) for users on older runtime images. All yamls use placeholders (<NAMESPACE>, <REGISTRY>/<NAME>:<TAG>, <GPU_NODE_POOL>, <CPU_NODE_POOL>, <MX_INFRA_*>) so users replace cluster-specific values before applying. This PR carries the deployment-example slice of PR #218 (kavink/trtllm_clean). The patch-shim slice retires once the next tensorrtllm-runtime image cuts post-PR #13531; PR #218 will be closed without merging at that point. Validated end-to-end on GCP GB200 (dynamo-gcp-dev-02, kavin namespace): 16 target ranks x 90.75 GB transferred at 363-506 Gbps, end-to-end disaggregated serving (prefill + decode + frontend) verified. Companion PRs: * TRT-LLM #13531: MXCheckpointLoader (merged 2026-05-06) * ModelExpress #202: MxLiveWeightLoader, publish_model_params (merged) * ModelExpress #267: MX_POOL_REG allocation-based registration (merged) * Dynamo #8037: --model-express-url CLI integration (open) Signed-off-by: Kavin Krishnan <kavink@nvidia.com>

pull-request-size Bot added the size/S label Apr 9, 2026

KavinKrishnan temporarily deployed to GITLAB April 9, 2026 19:58 — with GitHub Actions Inactive

github-actions Bot added the refactor label Apr 9, 2026

coderabbitai Bot reviewed Apr 9, 2026

View reviewed changes

pull-request-size Bot added size/M and removed size/S labels Apr 14, 2026

KavinKrishnan temporarily deployed to GITLAB April 14, 2026 17:27 — with GitHub Actions Inactive

fix: use file existence checks instead of Python imports in Dockerfil…

296510e

…e verify TRT-LLM imports require libcuda.so.1 which isn't available at build time. Signed-off-by: Kavin Krishnan <kavink@nvidia.com> Signed-off-by: Kavin Krishnan <kavinkrishnan@gmail.com> Made-with: Cursor

KavinKrishnan temporarily deployed to GITLAB April 14, 2026 17:59 — with GitHub Actions Inactive

pull-request-size Bot added size/L and removed size/M labels Apr 17, 2026

KavinKrishnan temporarily deployed to GITLAB April 17, 2026 17:06 — with GitHub Actions Inactive

devin-ai-integration Bot reviewed Apr 17, 2026

View reviewed changes

KavinKrishnan changed the title ~~refactor: auto-detect TRT-LLM source/target role, add upstream PR TODOs~~ feat(trtllm): MXCheckpointLoader deployment support with auto-detect and RDMA validation Apr 23, 2026

github-actions Bot added feat and removed refactor labels Apr 23, 2026

zhengluo-nv temporarily deployed to GITLAB April 23, 2026 21:50 — with GitHub Actions Inactive

KavinKrishnan force-pushed the kavink/trtllm_clean branch from a5c5138 to af2ffb2 Compare April 28, 2026 20:01

pull-request-size Bot added size/XL and removed size/L labels Apr 28, 2026

KavinKrishnan temporarily deployed to GITLAB April 28, 2026 20:01 — with GitHub Actions Inactive

KavinKrishnan temporarily deployed to GITLAB April 28, 2026 20:45 — with GitHub Actions Inactive

KavinKrishnan temporarily deployed to GITLAB April 29, 2026 03:28 — with GitHub Actions Inactive

pull-request-size Bot added size/XXL and removed size/XL labels Apr 29, 2026

KavinKrishnan temporarily deployed to GITLAB April 29, 2026 17:44 — with GitHub Actions Inactive

KavinKrishnan temporarily deployed to GITLAB April 29, 2026 19:44 — with GitHub Actions Inactive

KavinKrishnan temporarily deployed to GITLAB April 29, 2026 21:41 — with GitHub Actions Inactive

zhengluo-nv reviewed Apr 29, 2026

View reviewed changes

KavinKrishnan temporarily deployed to GITLAB April 30, 2026 01:26 — with GitHub Actions Inactive

KavinKrishnan closed this May 7, 2026

This was referenced May 8, 2026

docs(trtllm): deployment example yamls + HPA autoscale demo for MX P2P #271

Closed

fix(trtllm): flush rank log handler so MPI workers' transfer metrics … #272

Closed

Navigation Menu

Uh oh!

feat(trtllm): MXCheckpointLoader deployment support with auto-detect and RDMA validation - #218

feat(trtllm): MXCheckpointLoader deployment support with auto-detect and RDMA validation#218
KavinKrishnan wants to merge 11 commits into
mainfrom
kavink/trtllm_clean

KavinKrishnan commented Apr 9, 2026 •

edited

Loading

Uh oh!

coderabbitai Bot commented Apr 9, 2026 •

edited

Loading

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Uh oh!

devin-ai-integration Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhengluo-nv Apr 29, 2026

Uh oh!

KavinKrishnan Apr 30, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		)


		def patch_backend_args():

Uh oh!

Conversation

KavinKrishnan commented Apr 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Changes

How it works

Companion PRs

Uh oh!

coderabbitai Bot commented Apr 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Walkthrough

Changes

Estimated code review effort

Poem

❌ Failed checks (1 warning)

Uh oh!

coderabbitai Bot left a comment

Choose a reason for hiding this comment

Uh oh!

devin-ai-integration Bot left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

zhengluo-nv Apr 29, 2026

Choose a reason for hiding this comment

Uh oh!

KavinKrishnan Apr 30, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

KavinKrishnan commented Apr 9, 2026 •

edited

Loading

coderabbitai Bot commented Apr 9, 2026 •

edited

Loading