Skip to content

feat(client): add TensorRT-LLM engine adapter - #565

Merged
zhengluo-nv merged 1 commit into
mainfrom
zheluo/trtllm-engine-adapter
Jul 30, 2026
Merged

feat(client): add TensorRT-LLM engine adapter#565
zhengluo-nv merged 1 commit into
mainfrom
zheluo/trtllm-engine-adapter

Conversation

@zhengluo-nv

@zhengluo-nv zhengluo-nv commented Jul 28, 2026

Copy link
Copy Markdown
Contributor

Summary

  • Add modelexpress.engines.trtllm.TrtllmAdapter and MxModelLoader so
    TensorRT-LLM uses the same LoadStrategyChain as vLLM and SGLang.
  • Accept TensorRT-LLM's authoritative SourceIdentity and include the complete
    TRT compatibility fingerprint and post-transform protocol in MX discovery.
  • Delegate source selection, RDMA receive, native fallback selection,
    publication, republish, and cleanup to shared ModelExpress components.
  • Require an exact tensor catalog for qualified TRT-LLM transfers and preserve
    TensorRT-LLM's native loader return contract on fallback.
  • Remove the legacy trtllm_live_transfer integration path.
  • Support the TRT release image's bundled CUDA 13 NIXL binding without
    installing a second NIXL stack.
  • Restrict automatic UCX NIC selection to verbs devices actually exposed by a
    Kubernetes RDMA device plugin.
  • Add the native TRT-LLM production-style deployment example and document the
    current Llama-only qualification scope.

Why

TensorRT-LLM already owns model construction, compatibility identity, Llama
post-transform qualification, and post-load lifecycle. ModelExpress should own
the transport and source-strategy behavior shared across engines. Keeping two
independent implementations made fallback, cleanup, observability, and
multi-engine behavior drift over time.

This PR establishes one MX-owned strategy path while retaining TensorRT-LLM as
the authority for model-specific correctness.

The RDMA device-plugin fix addresses a production container failure found
during E2E testing: /sys/class/infiniband exposed every host NIC, but the pod
was allocated only one /dev/infiniband/uverbsN device. Auto-pinning could
select an inaccessible host NIC and make NIXL fail with Destination is unreachable.

Cross-repository dependency

The TensorRT-LLM bridge is NVIDIA/TensorRT-LLM#17029. That PR imports
modelexpress.engines.trtllm.MxModelLoader and therefore depends on this PR.

The initial supported runtime scope remains:

Area Current scope
Model family Llama
Transfer Post-transform P2P weights
Validated topology Cross-node TP=1
Storage fallback TensorRT-LLM native Hugging Face loader
CI TRT-LLM qualification remains NOT STARTED

This draft does not claim PP/EP, speculative decoding, multi-node TP, or
PD-disaggregated production readiness.

Validation

  • ModelExpress focused runtime suite: 87 passed
  • Repository-level TRT example/docs suite: 4 passed
  • TensorRT-LLM focused GPU-runtime suite using the combined image: 30 passed
  • Combined TRT-LLM + ModelExpress image build and native import checks passed
  • Cross-node TP=1 Llama 3.1 8B E2E:
    • source loaded from disk and published READY metadata
    • target selected RDMA and did not execute disk fallback
    • 195 tensors / 16.06 GB transferred in 0.584 seconds at 220.1 Gbps
    • source and target inference passed
  • git diff --check
  • Signed commit with DCO sign-off

Before marking ready, rerun the cross-node E2E from the final PR commits and
complete the planned TP=4 CI qualification.

Summary by CodeRabbit

  • New Features

    • Added TensorRT-LLM support for ModelExpress P2P weight loading through shared transfer strategies.
    • Added a container image and Kubernetes deployment configuration for TensorRT-LLM examples.
    • Improved compatibility across supported NIXL package variants and containerized networking environments.
    • Added clearer TensorRT-LLM logging, model identity handling, and transfer verification.
  • Documentation

    • Updated build, deployment, supported-scope, and verification instructions for TensorRT-LLM examples.
    • Refreshed CI test-plan guidance for stale metadata behavior and current TRT-LLM gating.

@codecov

codecov Bot commented Jul 29, 2026

Copy link
Copy Markdown

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

@zhengluo-nv
zhengluo-nv marked this pull request as ready for review July 29, 2026 23:30
@coderabbitai

coderabbitai Bot commented Jul 29, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

Walkthrough

Changes

The pull request adds a TensorRT-LLM adapter and loader using the shared strategy chain, broadens NIXL and RDMA compatibility, improves UCX device filtering, updates model feature detection, and revises the TensorRT-LLM Kubernetes image, deployment, documentation, and tests.

TensorRT-LLM runtime integration

Layer / File(s) Summary
TRT-LLM adapter and loader path
modelexpress_client/python/modelexpress/engines/trtllm/*, modelexpress_client/python/modelexpress/__init__.py, modelexpress_client/python/tests/test_trtllm_*
Adds TensorRT-LLM identity handling, canonical tensor discovery, native/RDMA lifecycle hooks, shared strategy-chain loading, logging configuration, publishing, cleanup, and focused tests.
Shared transfer and device compatibility
modelexpress_client/python/modelexpress/load_strategy/*, modelexpress_client/python/modelexpress/{_nixl,nixl_transfer,gds_transfer,transfer_safety,ucx_utils}.py, modelexpress_client/python/tests/test_{nixl_import,transfer_safety,ucx_utils}.py
Adds NIXL package fallback loading, explicit metadata server URL handling, exact catalog capability checks, defensive TRT-LLM config parsing, and container-visible verbs-device filtering.
TRT-LLM image and Kubernetes example
examples/p2p_transfer_k8s/*, examples/p2p_transfer_k8s/client/trtllm/*, modelexpress_client/python/README.md, ci/TEST_PLAN.md, modelexpress_client/python/tests/test_trtllm_native_example.py
Adds a configurable TRT-LLM release-image Dockerfile, updates NIXL setup and deployment anti-affinity, removes the timeout setting, and revises related instructions, verification checks, and CI documentation.

Estimated code review effort: 4 (Complex) | ~60 minutes

Poem

I’m a rabbit with logs in a burrow so bright,
TRT-LLM hops through the loader tonight.
NIXL finds its path, RDMA carries the load,
Kubernetes bunnies spread across the road.
Tests thump their paws: the transfer is right!

🚥 Pre-merge checks | ✅ 4 | ❌ 1

❌ Failed checks (1 warning)

Check name Status Explanation Resolution
Docstring Coverage ⚠️ Warning Docstring coverage is 29.63% 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
Description Check ✅ Passed Check skipped - CodeRabbit’s high-level summary is enabled.
Title check ✅ Passed The title clearly matches the main change: adding a TensorRT-LLM engine adapter on the client side.
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.

Comment @coderabbitai help to get the list of available commands.

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Actionable comments posted: 7

🧹 Nitpick comments (1)
modelexpress_client/python/modelexpress/engines/trtllm/adapter.py (1)

27-27: 🎯 Functional Correctness | 🔵 Trivial | 💤 Low value

Align these alias markers with the supported TRT-LLM release. next_attn and next_layer_layernorm are current TRT-LLM internal wiring names used by several PyTorch backend models, so an upstream rename would make supported checkpoints hit the fail-closed RuntimeError path. Document the TRT-LLM version/commit these names come from, or centralize the marker list with the backend compatibility check.

🤖 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 `@modelexpress_client/python/modelexpress/engines/trtllm/adapter.py` at line
27, Update _RUNTIME_ALIAS_COMPONENTS to be explicitly tied to the supported
TRT-LLM release: document the TRT-LLM version or commit that defines next_attn
and next_layer_layernorm, or reuse a centralized marker list from the backend
compatibility check. Preserve these markers for supported checkpoints so they do
not incorrectly reach the fail-closed RuntimeError path.
🤖 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.

Inline comments:
In `@examples/p2p_transfer_k8s/client/trtllm/Dockerfile`:
- Around line 4-5: Resolve the incomplete release-image TODO in the Dockerfile
by selecting and validating a TRT-LLM release image that includes the complete
upstream ModelExpress integration before presenting this as a production-style
deployment; alternatively, clearly mark and gate the example as experimental in
the associated deployment configuration and documentation.
- Around line 6-7: Update the Dockerfile after all root-only build steps to
create and select an explicit non-root user and group, then ensure that user can
access /models, /dev/shm, GPU/RDMA devices, and IPC_LOCK at runtime. Preserve
required root privileges during image setup and add the USER declaration only
for the final serving container.

In `@examples/p2p_transfer_k8s/client/trtllm/trtllm-single-node-p2p.yaml`:
- Around line 58-64: Update the workload’s pod anti-affinity from
requiredDuringSchedulingIgnoredDuringExecution to
preferredDuringSchedulingIgnoredDuringExecution so single-node rollouts can
schedule replacement pods; preserve the app label selector and hostname topology
key.

In `@modelexpress_client/python/modelexpress/_nixl.py`:
- Around line 12-18: Update load_nixl_api so it continues only when the
candidate package or its _api submodule is genuinely unavailable. Distinguish
missing-module errors for the current package from ImportError exceptions raised
by dependencies or binary loading, and re-raise the latter instead of trying
subsequent candidates.

In `@modelexpress_client/python/modelexpress/engines/trtllm/loader.py`:
- Around line 115-131: Update cleanup so unpublish_metadata(self._ctx) is
wrapped in best-effort exception handling, logging any failure with
logger.warning while allowing execution to continue. Preserve the existing NIXL
manager shutdown and MX client close cleanup paths so they always run even when
metadata unpublishing fails.
- Around line 99-113: Broaden the exception handling in publish_model to catch
exceptions from synchronization, tensor registration, and metadata publication
so best-effort publishing never propagates failures to the caller. Preserve the
existing warning message, rank context, and continuation behavior while
including the caught exception details.

In `@modelexpress_client/python/modelexpress/transfer_safety.py`:
- Around line 48-53: Update the dtype extraction in the model configuration
feature-building logic to fall back to model_config.torch_dtype when
model_config.dtype is missing or None, not only when the attribute is absent.
Ensure features["dtype"] records the resolved dtype string rather than "None",
while preserving the existing torch. prefix removal.

---

Nitpick comments:
In `@modelexpress_client/python/modelexpress/engines/trtllm/adapter.py`:
- Line 27: Update _RUNTIME_ALIAS_COMPONENTS to be explicitly tied to the
supported TRT-LLM release: document the TRT-LLM version or commit that defines
next_attn and next_layer_layernorm, or reuse a centralized marker list from the
backend compatibility check. Preserve these markers for supported checkpoints so
they do not incorrectly reach the fail-closed RuntimeError path.
🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

  • Push a commit to this branch (recommended)
  • Create a new PR with the fixes

ℹ️ Review info
⚙️ Run configuration

Configuration used: Path: .coderabbit.yaml

Review profile: CHILL

Plan: Enterprise

Run ID: e26973c3-4999-45b9-809f-20441c112a13

📥 Commits

Reviewing files that changed from the base of the PR and between 31ff74f and bc27d6b.

📒 Files selected for processing (29)
  • ci/TEST_PLAN.md
  • examples/p2p_transfer_k8s/README.md
  • examples/p2p_transfer_k8s/client/README.md
  • examples/p2p_transfer_k8s/client/trtllm/Dockerfile
  • examples/p2p_transfer_k8s/client/trtllm/README.md
  • examples/p2p_transfer_k8s/client/trtllm/trtllm-single-node-p2p.yaml
  • modelexpress_client/python/README.md
  • modelexpress_client/python/modelexpress/__init__.py
  • modelexpress_client/python/modelexpress/_nixl.py
  • modelexpress_client/python/modelexpress/adapter.py
  • modelexpress_client/python/modelexpress/engines/trtllm/__init__.py
  • modelexpress_client/python/modelexpress/engines/trtllm/adapter.py
  • modelexpress_client/python/modelexpress/engines/trtllm/loader.py
  • modelexpress_client/python/modelexpress/gds_transfer.py
  • modelexpress_client/python/modelexpress/load_strategy/base.py
  • modelexpress_client/python/modelexpress/load_strategy/context.py
  • modelexpress_client/python/modelexpress/load_strategy/default_strategy.py
  • modelexpress_client/python/modelexpress/load_strategy/rdma_strategy.py
  • modelexpress_client/python/modelexpress/nixl_transfer.py
  • modelexpress_client/python/modelexpress/transfer_safety.py
  • modelexpress_client/python/modelexpress/trtllm_live_transfer.py
  • modelexpress_client/python/modelexpress/ucx_utils.py
  • modelexpress_client/python/tests/test_nixl_import.py
  • modelexpress_client/python/tests/test_transfer_safety.py
  • modelexpress_client/python/tests/test_trtllm_adapter.py
  • modelexpress_client/python/tests/test_trtllm_live_transfer.py
  • modelexpress_client/python/tests/test_trtllm_loader.py
  • modelexpress_client/python/tests/test_trtllm_native_example.py
  • modelexpress_client/python/tests/test_ucx_utils.py
💤 Files with no reviewable changes (2)
  • modelexpress_client/python/tests/test_trtllm_live_transfer.py
  • modelexpress_client/python/modelexpress/trtllm_live_transfer.py

Comment thread examples/p2p_transfer_k8s/client/trtllm/Dockerfile
Comment thread examples/p2p_transfer_k8s/client/trtllm/Dockerfile
Comment thread examples/p2p_transfer_k8s/client/trtllm/trtllm-single-node-p2p.yaml
Comment thread modelexpress_client/python/modelexpress/_nixl.py
Comment thread modelexpress_client/python/modelexpress/engines/trtllm/loader.py
Comment thread modelexpress_client/python/modelexpress/engines/trtllm/loader.py
Comment thread modelexpress_client/python/modelexpress/transfer_safety.py Outdated
Signed-off-by: Zheng Luo <zheluo@nvidia.com>
@zhengluo-nv
zhengluo-nv force-pushed the zheluo/trtllm-engine-adapter branch from bc27d6b to b7b42d5 Compare July 29, 2026 23:54
@copy-pr-bot
copy-pr-bot Bot temporarily deployed to automated-release July 29, 2026 23:54 Inactive
@copy-pr-bot
copy-pr-bot Bot temporarily deployed to automated-release July 29, 2026 23:54 Inactive
@zhengluo-nv zhengluo-nv self-assigned this Jul 30, 2026
@zhengluo-nv
zhengluo-nv merged commit 3542dda into main Jul 30, 2026
47 checks passed
@zhengluo-nv
zhengluo-nv deleted the zheluo/trtllm-engine-adapter branch July 30, 2026 19:53
zhengluo-nv added a commit that referenced this pull request Aug 13, 2026
Signed-off-by: Zheng Luo <zheluo@nvidia.com>
(cherry picked from commit 3542dda)
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants