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

examples/p2p_transfer_k8s/client/trtllm/Dockerfile.ph3-gcp-gb200

@@ -1,29 +1,25 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 # SPDX-License-Identifier: Apache-2.0
 
-# Phase 3 P2P image for GCP GB200 (ARM64)
-# Layers ModelExpress client + Dynamo P2P hooks + PRESHARDED patches
-# on karenc's Dynamo v1.0.0 image
+# TRT-LLM P2P image for GCP GB200 (ARM64)
+# Layers ModelExpress client + Chien-Chun's MX checkpoint loader + Dynamo hooks
 #
 # Base: nvcr.io/nvidian/dynamo-dev/karenc:dynamo-trtllm-v1.0.0-a9b6f95
-# - TRT-LLM 1.3.0rc5
-# - NIXL pre-installed
-# - ARM64 (aarch64)
-# - Dynamo worker (python3 -m dynamo.trtllm)
-#
-# Requires dynamo fork at ../dynamo (the Dynamo branch with P2P support)
+# Requires:
+#   - dynamo fork with --model-express-url support
+#   - TRT-LLM fork with MXCheckpointLoader (chienchun/dynamo-integration-prototype)
 #
 # Build (from modelexpress repo root):
 #   docker buildx build --platform linux/arm64 --no-cache \
-#       -f examples/p2p_transfer_trtllm/Dockerfile.ph3-gcp-gb200 \
+#       -f examples/p2p_transfer_k8s/client/trtllm/Dockerfile.ph3-gcp-gb200 \
 #       --build-context dynamo=../dynamo \
+#       --build-context trtllm=../TensorRT-LLM \
 #       -t <REGISTRY>/dynamo-trtllm-mx:<TAG> \
 #       --push .
 FROM nvcr.io/nvidian/dynamo-dev/karenc:dynamo-trtllm-v1.0.0-a9b6f95
 
 USER root
 
-# Multinode MPI: sshd requires privilege separation directory owned by root, mode 0755
 RUN mkdir -p /run/sshd && chmod 0755 /run/sshd && chown root:root /run/sshd
 
 # Install ModelExpress client (gRPC + NIXL transfer)
@@ -41,8 +37,8 @@ RUN cd /tmp/modelexpress_client && \
     pip install --no-cache-dir . && \
     rm -rf /tmp/modelexpress_client /tmp/proto
 
+# TODO: Remove once ai-dynamo/dynamo PR #8037 merges
 # Apply Dynamo P2P hooks (--model-express-url support)
-# From the Dynamo branch with P2P support: backend_args.py, engine.py, llm_worker.py
 COPY --from=dynamo components/src/dynamo/trtllm/backend_args.py /tmp/dynamo_p2p/backend_args.py
 COPY --from=dynamo components/src/dynamo/trtllm/engine.py /tmp/dynamo_p2p/engine.py
 COPY --from=dynamo components/src/dynamo/trtllm/workers/llm_worker.py /tmp/dynamo_p2p/llm_worker.py
@@ -52,27 +48,36 @@ RUN DYNAMO_PKG=/opt/dynamo/venv/lib/python3.12/site-packages/dynamo/trtllm && \
     cp /tmp/dynamo_p2p/engine.py "$DYNAMO_PKG/engine.py" && \
     cp /tmp/dynamo_p2p/llm_worker.py "$DYNAMO_PKG/workers/llm_worker.py" && \
     rm -rf /tmp/dynamo_p2p && \
-    grep -q "model.express.url\|model_express_url" "$DYNAMO_PKG/backend_args.py" && \
+    grep -q "model_express_url" "$DYNAMO_PKG/backend_args.py" && \
     echo "Dynamo P2P hooks OK"
 
-# Apply PRESHARDED patches to TRT-LLM 1.3.0rc5
-COPY trtllm_patches/v1.3.0rc5/apply_patches.py /tmp/apply_patches.py
-RUN python3 /tmp/apply_patches.py && rm /tmp/apply_patches.py
+# TODO: Remove once NVIDIA/TensorRT-LLM PR #13045 merges
+# Install MX checkpoint loader from TRT-LLM fork
+COPY --from=trtllm tensorrt_llm/_torch/models/checkpoints/mx /tmp/trtllm_mx/mx
+
+RUN TRTLLM_PKG=/opt/dynamo/venv/lib/python3.12/site-packages/tensorrt_llm && \
+    cp -r /tmp/trtllm_mx/mx "$TRTLLM_PKG/_torch/models/checkpoints/mx" && \
+    rm -rf /tmp/trtllm_mx && \
+    test -f "$TRTLLM_PKG/_torch/models/checkpoints/mx/checkpoint_loader.py" && \
+    echo "MXCheckpointLoader installed OK"
 
-# Patch tp_allgather to use chunked allgather (fixes MPI_ERR_TRUNCATE with ob1 TCP BTL)
+# Apply minimal patches to base TRT-LLM for MX loader compatibility
+COPY trtllm_patches/v1.3.0rc5/patch_mx_loader.py /tmp/patch_mx_loader.py
+RUN python3 /tmp/patch_mx_loader.py && rm /tmp/patch_mx_loader.py
+
+# TODO: Remove once TRT-LLM fixes MPI allgather truncation natively
 COPY trtllm_patches/v1.3.0rc5/patch_tp_allgather.py /tmp/patch_tp_allgather.py
 RUN python3 /tmp/patch_tp_allgather.py && rm /tmp/patch_tp_allgather.py
 
-# Patch model_loader: source publishes BEFORE post_load_weights, target runs full post_load_weights
-COPY trtllm_patches/v1.3.0rc5/patch_model_loader.py /tmp/patch_model_loader.py
-RUN python3 /tmp/patch_model_loader.py && rm /tmp/patch_model_loader.py
-
-# Verify installation (can't import full TRT-LLM at build time — no GPU)
-RUN python3 -c "from modelexpress.client import MxClient; print('ModelExpress client OK')" && \
-    python3 -c "from modelexpress.trtllm_live_transfer import publish_model_params; print('publish_model_params OK')" && \
-    grep -q "PRESHARDED" /opt/dynamo/venv/lib/python3.12/site-packages/tensorrt_llm/llmapi/llm_args.py && \
-    echo "LoadFormat.PRESHARDED OK" && \
-    grep -q "publish_model_params" /opt/dynamo/venv/lib/python3.12/site-packages/tensorrt_llm/_torch/pyexecutor/model_loader.py && \
-    echo "model_loader.py source publish hook OK" && \
-    grep -q "MODEL_EXPRESS_SOURCE" /opt/dynamo/venv/lib/python3.12/site-packages/dynamo/trtllm/engine.py && \
-    echo "engine.py MODEL_EXPRESS_SOURCE OK"
+# Verify installation
+RUN TRTLLM_PKG=/opt/dynamo/venv/lib/python3.12/site-packages/tensorrt_llm && \
+    python3 -c "from modelexpress.client import MxClient; print('ModelExpress client OK')" && \
+    grep -q "model_express_url" /opt/dynamo/venv/lib/python3.12/site-packages/dynamo/trtllm/engine.py && \
+    echo "engine.py model_express_url OK" && \
+    grep -q "p2p_succeeded" "$TRTLLM_PKG/_torch/pyexecutor/model_loader.py" && \
+    echo "model_loader.py MX patch OK" && \
+    grep -q "MXCheckpointLoader" "$TRTLLM_PKG/_torch/models/checkpoints/__init__.py" && \
+    echo "checkpoints/__init__.py patched OK" && \
+    grep -q "MODEL_EXPRESS_URL" "$TRTLLM_PKG/executor/worker.py" && \
+    echo "worker.py patched OK" && \
+    echo "All patches verified OK"

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

@@ -98,12 +98,12 @@ spec:
                 - key: cloud.google.com/gke-nodepool
                   operator: In
                   values:
-                  - <GPU_NODE_POOL>
+                  - customer-gpu-w0e
         containers: null
         imagePullSecrets:
           - name: nvcr-imagepullsecret
         mainContainer:
-          image: "<REGISTRY>/dynamo-trtllm-mx:<TAG>"
+          image: "nvcr.io/nvidian/dynamo-dev/kavink:dynamo-trtllm-mx-v2.9.0"
           command:
             - python3
           args:
@@ -122,9 +122,9 @@ spec:
             - name: HF_HOME
               value: /model-cache
             - name: NATS_SERVER
-              value: "nats://dynamo-platform-nats.default.svc.cluster.local:4222"
+              value: "nats://dynamo-platform-nats.kavin.svc.cluster.local:4222"
             - name: ETCD_ENDPOINTS
-              value: "dynamo-platform-etcd.default.svc.cluster.local:2379"
+              value: "dynamo-platform-etcd.kavin.svc.cluster.local:2379"
           volumeMounts:
             - mountPath: /model-cache
               name: model-cache
@@ -196,7 +196,7 @@ spec:
                 - key: cloud.google.com/gke-nodepool
                   operator: In
                   values:
-                  - <GPU_NODE_POOL>
+                  - customer-gpu-w0e
           podAffinity:
             requiredDuringSchedulingIgnoredDuringExecution:
             - labelSelector:
@@ -210,7 +210,7 @@ spec:
         imagePullSecrets:
           - name: nvcr-imagepullsecret
         mainContainer:
-          image: "<REGISTRY>/dynamo-trtllm-mx:<TAG>"
+          image: "nvcr.io/nvidian/dynamo-dev/kavink:dynamo-trtllm-mx-v2.9.0"
           workingDir: /workspace/
           securityContext:
             privileged: true
@@ -232,9 +232,7 @@ spec:
             - --disaggregation-mode
             - prefill
             - --model-express-url
-            - modelexpress-server-decode.default.svc.cluster.local:8001
-            - --model-express-role
-            - target
+            - modelexpress-server-decode.kavin.svc.cluster.local:8001
             - --tensor-parallel-size
             - "8"
             - --publish-events-and-metrics
@@ -258,13 +256,15 @@ spec:
             - name: HF_MODULES_CACHE
               value: /tmp/hf_modules
             - name: MODEL_EXPRESS_URL
-              value: "modelexpress-server-decode.default.svc.cluster.local:8001"
+              value: "modelexpress-server-decode.kavin.svc.cluster.local:8001"
+            - name: MODEL_EXPRESS_TARGET
+              value: "1"
             - name: MODEL_NAME
               value: "baseten-admin/Kimi-2.5-text-nvfp4-v3"
             - name: NATS_SERVER
-              value: "nats://dynamo-platform-nats.default.svc.cluster.local:4222"
+              value: "nats://dynamo-platform-nats.kavin.svc.cluster.local:4222"
             - name: ETCD_ENDPOINTS
-              value: "dynamo-platform-etcd.default.svc.cluster.local:2379"
+              value: "dynamo-platform-etcd.kavin.svc.cluster.local:2379"
             - name: NCCL_CUMEM_ENABLE
               value: "1"
             - name: NCCL_NVLS_ENABLE
@@ -346,7 +346,7 @@ spec:
             operator: Exists
         resourceClaims:
           - name: compute-domain-channel
-            resourceClaimTemplateName: my-compute-domain-channel
+            resourceClaimTemplateName: kavin-compute-domain-channel
         volumes:
           - name: shared-model-cache
             persistentVolumeClaim:
@@ -400,7 +400,7 @@ spec:
                 - key: cloud.google.com/gke-nodepool
                   operator: In
                   values:
-                  - <GPU_NODE_POOL>
+                  - customer-gpu-w0e
           podAffinity:
             requiredDuringSchedulingIgnoredDuringExecution:
             - labelSelector:
@@ -414,7 +414,7 @@ spec:
         imagePullSecrets:
           - name: nvcr-imagepullsecret
         mainContainer:
-          image: "<REGISTRY>/dynamo-trtllm-mx:<TAG>"
+          image: "nvcr.io/nvidian/dynamo-dev/kavink:dynamo-trtllm-mx-v2.9.0"
           workingDir: /workspace/
           securityContext:
             privileged: true
@@ -436,9 +436,7 @@ spec:
             - --disaggregation-mode
             - decode
             - --model-express-url
-            - modelexpress-server-decode.default.svc.cluster.local:8001
-            - --model-express-role
-            - target
+            - modelexpress-server-decode.kavin.svc.cluster.local:8001
             - --tensor-parallel-size
             - "8"
             - --publish-events-and-metrics
@@ -462,13 +460,15 @@ spec:
             - name: HF_MODULES_CACHE
               value: /tmp/hf_modules
             - name: MODEL_EXPRESS_URL
-              value: "modelexpress-server-decode.default.svc.cluster.local:8001"
+              value: "modelexpress-server-decode.kavin.svc.cluster.local:8001"
+            - name: MODEL_EXPRESS_TARGET
+              value: "1"
             - name: MODEL_NAME
               value: "baseten-admin/Kimi-2.5-text-nvfp4-v3"
             - name: NATS_SERVER
-              value: "nats://dynamo-platform-nats.default.svc.cluster.local:4222"
+              value: "nats://dynamo-platform-nats.kavin.svc.cluster.local:4222"
             - name: ETCD_ENDPOINTS
-              value: "dynamo-platform-etcd.default.svc.cluster.local:2379"
+              value: "dynamo-platform-etcd.kavin.svc.cluster.local:2379"
             - name: NCCL_CUMEM_ENABLE
               value: "1"
             - name: NCCL_NVLS_ENABLE
@@ -552,7 +552,7 @@ spec:
             operator: Exists
         resourceClaims:
           - name: compute-domain-channel
-            resourceClaimTemplateName: my-compute-domain-channel
+            resourceClaimTemplateName: kavin-compute-domain-channel
         volumes:
           - name: shared-model-cache
             persistentVolumeClaim:

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

@@ -74,7 +74,7 @@ spec:
                 - key: cloud.google.com/gke-nodepool
                   operator: In
                   values:
-                  - <GPU_NODE_POOL>
+                  - customer-gpu-o7v
           podAffinity:
             requiredDuringSchedulingIgnoredDuringExecution:
             - labelSelector:
@@ -88,7 +88,7 @@ spec:
         imagePullSecrets:
           - name: nvcr-imagepullsecret
         mainContainer:
-          image: "<REGISTRY>/dynamo-trtllm-mx:<TAG>"
+          image: "nvcr.io/nvidian/dynamo-dev/kavink:dynamo-trtllm-mx-v2.9.0"
           imagePullPolicy: Always
           workingDir: /workspace/
           securityContext:
@@ -105,9 +105,7 @@ spec:
             - --extra-engine-args
             - /config/source-decode.yaml
             - --model-express-url
-            - modelexpress-server-decode.default.svc.cluster.local:8001
-            - --model-express-role
-            - source
+            - modelexpress-server-decode.kavin.svc.cluster.local:8001
             - --tensor-parallel-size
             - "8"
           env:
@@ -116,15 +114,13 @@ spec:
                 fieldRef:
                   fieldPath: metadata.uid
             - name: NATS_SERVER
-              value: "nats://dynamo-platform-nats.default.svc.cluster.local:4222"
+              value: "nats://dynamo-platform-nats.kavin.svc.cluster.local:4222"
             - name: ETCD_ENDPOINTS
-              value: "dynamo-platform-etcd.default.svc.cluster.local:2379"
+              value: "dynamo-platform-etcd.kavin.svc.cluster.local:2379"
             - name: MODEL_EXPRESS_URL
-              value: "modelexpress-server-decode.default.svc.cluster.local:8001"
+              value: "modelexpress-server-decode.kavin.svc.cluster.local:8001"
             - name: MODEL_NAME
               value: "baseten-admin/Kimi-2.5-text-nvfp4-v3"
-            - name: MODEL_EXPRESS_SOURCE
-              value: "1"
             - name: WORLD_SIZE
               value: "8"
             - name: HOME
@@ -212,7 +208,7 @@ spec:
             operator: Exists
         resourceClaims:
           - name: compute-domain-channel
-            resourceClaimTemplateName: my-compute-domain-channel
+            resourceClaimTemplateName: kavin-compute-domain-channel
         volumes:
           - name: infiniband
             hostPath:

examples/p2p_transfer_k8s/client/trtllm/mx-infra-decode.yaml

@@ -3,7 +3,7 @@
 
 # ModelExpress server + Redis for decode source (Phase 2 mixed TP)
 # Separate from prefill MX server so each source publishes to its own store
-# Deploys on <CPU_NODE_POOL> nodes (amd64)
+# Deploys on customer-cpu nodes (amd64)
 # Namespace: default
 ---
 apiVersion: apps/v1
@@ -24,7 +24,7 @@ spec:
     spec:
       # Update node pool names for your cluster
       nodeSelector:
-        cloud.google.com/gke-nodepool: <CPU_NODE_POOL>
+        cloud.google.com/gke-nodepool: customer-cpu
       tolerations:
         - key: dedicated
           operator: Equal
@@ -69,7 +69,7 @@ spec:
     spec:
       # Update node pool names for your cluster
       nodeSelector:
-        cloud.google.com/gke-nodepool: <CPU_NODE_POOL>
+        cloud.google.com/gke-nodepool: customer-cpu
       tolerations:
         - key: dedicated
           operator: Equal
@@ -85,6 +85,10 @@ spec:
               value: "redis"
             - name: REDIS_URL
               value: "redis://redis-decode:6379"
+            - name: MX_HEARTBEAT_TIMEOUT_SECS
+              value: "999999"
+            - name: MX_GC_TIMEOUT_SECS
+              value: "999999"
           resources:
             requests:
               cpu: "1"