Skip to content

[MX-299][feat] Delegate MX loading to ModelExpress strategies #17029

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
zhengluo-nv wants to merge 2 commits into
base: main
Choose a base branch
from
Open

[MX-299][feat] Delegate MX loading to ModelExpress strategies #17029

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
9d49310
fix: rebase ModelExpress integration on main
zhengluo-nv Aug 27, 2026
08eb2dd
fix: harden ModelExpress fallback lifecycle
zhengluo-nv Aug 27, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
34 changes: 18 additions & 16 deletions docs/source/features/model-express.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -120,7 +120,7 @@ uses a metadata-only view of the donor's canonical snapshot and contains no
weight shards. A positive result therefore requires direct transfer; disk
fallback cannot accidentally satisfy the test.

Run the TP=1 smoke test against an isolated ModelExpress 0.4.1 service with
Run the TP=1 smoke test against an isolated ModelExpress 0.5.1 service with
NIXL enabled:

```bash
Expand All @@ -144,7 +144,7 @@ row. `TRTLLM_MX_E2E_TIMEOUT_S` controls the 1200-second timeout used for the
baseline worker, receiver worker, and donor-readiness wait; increase it for
slow model storage or startup.

The dedicated H100 CI stages own isolated Redis and ModelExpress 0.4.1
The dedicated H100 CI stages own isolated Redis and ModelExpress 0.5.1
sidecars. The two-GPU TP=1 stage is classified as multi-GPU: it runs
automatically in post-merge pipelines or when a multi-GPU file changes, while
direct pre-merge dispatch requires the `ci: full pre-merge approved` label.
Expand Down Expand Up @@ -190,11 +190,14 @@ When adding an ABI ID:

## Installation

The official TensorRT LLM release container includes the MX Python client. No
additional Python package installation is required in that container. MX
remains opt-in at runtime: TensorRT LLM uses the client only when the MX
checkpoint-loading path and a server URL are configured. Installing the client
does not expand the model support scope described above.
TensorRT LLM release containers that include this feature already install a
compatible MX Python client; no additional client installation is needed for
P2P transfer. For earlier TensorRT LLM releases, use a release or container
built with this feature; upgrading the MX package alone does not add the
missing TensorRT LLM integration. MX remains opt-in at runtime: TensorRT LLM
uses the client only when the MX checkpoint-loading path and a server URL are
configured. Installing the client does not expand the model support scope
described above.

For pip installations outside the official release container, install the MX
Python client through the optional `mx` extra:
Expand All @@ -203,15 +206,14 @@ Python client through the optional `mx` extra:
pip install "tensorrt-llm[mx]"
```

The extra accepts ModelExpress client versions `>=0.4.1,<0.6.0`. Version
`0.4.1` is the minimum client API qualified by this integration, while the
upper bound prevents resolving unqualified `0.6.0` or newer client APIs.
Deploy a compatible MX server version.
The extra accepts ModelExpress client versions `>=0.5.1,<0.6.0`. Version
`0.5.1` is the minimum client release that provides the TensorRT LLM adapter,
while the upper bound prevents resolving unqualified `0.6.0` or newer client
APIs. Deploy a compatible MX server version.
Comment thread
moraxu marked this conversation as resolved.
The extra can be added to an existing TensorRT LLM installation. If the MX
loading path is configured but the client cannot be imported, TensorRT LLM
fails with an actionable installation message instead of silently loading from
the Hugging Face checkpoint. Source discovery and transfer failures continue to
use the Hugging Face fallback described above.
logs a warning and uses the Hugging Face fallback described above. Source
discovery and transfer failures use the same fallback.

## Deploy the MX Service

Expand All @@ -236,7 +238,7 @@ docker run -d --name modelexpress-server \
-e MODEL_EXPRESS_LOG_LEVEL=info \
-e MX_METADATA_BACKEND=redis \
-e REDIS_URL=redis://modelexpress-redis:6379 \
nvcr.io/nvidia/ai-dynamo/modelexpress-server:0.4.1
nvcr.io/nvidia/ai-dynamo/modelexpress-server:0.5.1
```

## Configure TensorRT LLM
Expand Down Expand Up @@ -272,7 +274,7 @@ path.
| Field | Default | Description |
|-------|---------|-------------|
| `mx_config.server_url` | `null` | URL of the separately managed MX server. |
| `mx_config.server_query_timeout_s` | `null` | Timeout for MX source discovery. When unset, TensorRT LLM uses a short fallback cap when no source exists and otherwise lets MX wait for long donor loads. |
| `mx_config.server_query_timeout_s` | `null` | Deprecated and ignored. MX checks once for a compatible source, then falls back to native checkpoint loading. |

## Notes and Limitations

Expand Down
6 changes: 3 additions & 3 deletions jenkins/L0_Test.groovy
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ ARTIFACTORY_CREDENTIALS_ID = "trtllm-artifactory-credentials"
// DLFW torch image
DLFW_IMAGE = "urm.nvidia.com/docker/nvidia/pytorch:26.05-py3"

MODEL_EXPRESS_VERSION = "0.4.1"
MODEL_EXPRESS_VERSION = "0.5.1"
MODEL_EXPRESS_NIXL_VERSION = "1.4.0"
MODEL_EXPRESS_SERVER_IMAGE = "urm.nvidia.com/docker/nvidia/ai-dynamo/modelexpress-server:${MODEL_EXPRESS_VERSION}"
MODEL_EXPRESS_REDIS_IMAGE = "urm.nvidia.com/docker/redis:7-alpine"
Expand Down Expand Up @@ -3704,7 +3704,7 @@ def createKubernetesPodConfig(image, type, arch = "amd64", gpuCount = 1, perfMod
- name: TRTLLM_MX_E2E_REQUIRED
value: "1"
"""
// Mirrors the ModelExpress v0.4.1 Redis deployment and image contract.
// Mirrors the ModelExpress Redis deployment and image contract.
// The image exposes /app/modelexpress-server and accepts the port/backend settings below.
// Use regular containers because the Jenkins Kubernetes launcher does not
// reliably attach to pods containing restartable init-container sidecars.
Expand Down Expand Up @@ -4993,7 +4993,7 @@ def runLLMTestlistOnPlatformImpl(pipeline, platform, testList, config=VANILLA_CO
}
if (stageName.contains("-ModelExpress-")) {
trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install modelexpress==${MODEL_EXPRESS_VERSION}")
// ModelExpress 0.4.1 imports nixl._api, while requirements-dev.txt
// ModelExpress imports nixl._api, while requirements-dev.txt
// installs only the nixl-cu13 backend. Install the matching
// namespace shim without pulling the unused CUDA 12 backend.
trtllm_utils.llmExecStepWithRetry(pipeline, script: "pip3 install --no-deps nixl==${MODEL_EXPRESS_NIXL_VERSION}")
Expand Down
2 changes: 1 addition & 1 deletion setup.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -148,7 +148,7 @@ def has_ext_modules(self):
Path("requirements-dev-windows.txt"
if on_windows else "requirements-dev.txt"))
openengine_deps, _ = parse_requirements(Path("requirements-openengine.txt"))
mx_deps = ["modelexpress>=0.4.1,<0.6.0"]
mx_deps = ["modelexpress>=0.5.1,<0.6.0"]
Comment thread
moraxu marked this conversation as resolved.
# Gateway protocol adapters are opt-in extras: the default installation must
# not carry any gateway protobuf package. Each gateway owns a dedicated
# requirements-<gateway>.txt as the single source of truth for its pins; CI
Expand Down
Loading
Loading