Skip to content

docs: add docstring for llm.rs - #267

Merged
nvda-mesharma merged 3 commits into
ai-dynamo:mainfrom
jcoleharrison:main
Apr 11, 2025
Merged

docs: add docstring for llm.rs#267
nvda-mesharma merged 3 commits into
ai-dynamo:mainfrom
jcoleharrison:main

Conversation

@jcoleharrison

@jcoleharrison jcoleharrison commented Mar 18, 2025

Copy link
Copy Markdown
Contributor

Overview:

Details:

• Added a comprehensive module-level docstring explaining the purpose of each sub-module and that type annotations are applied for better user experience.

Where should the reviewer start?

dynamo/lib/bindings/python/rust/llm.rs to see the updated documentation

Related Issues: (use one of the action keywords Closes / Fixes / Resolves / Relates to)

None

@copy-pr-bot

copy-pr-bot Bot commented Mar 18, 2025

Copy link
Copy Markdown

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

@github-actions

github-actions Bot commented Mar 18, 2025

Copy link
Copy Markdown
Contributor

👋 Hi jcoleharrison! Thank you for contributing to ai-dynamo/dynamo.

Just a reminder: The NVIDIA Test Github Validation CI runs an essential subset of the testing framework to quickly catch errors.Your PR reviewers may elect to test the changes comprehensively before approving your changes.

🚀

@jcoleharrison jcoleharrison changed the title add docstring for llm.rs docs: add docstring for llm.rs Mar 18, 2025
@grahamking

grahamking commented Mar 18, 2025

Copy link
Copy Markdown
Contributor

Could you add a bit of detail to the comments?

dmitry-tokarev-nv
dmitry-tokarev-nv requested changes Mar 18, 2025
View reviewed changes

@dmitry-tokarev-nv dmitry-tokarev-nv 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.

Please address Graham's request

@jcoleharrison

jcoleharrison commented Mar 18, 2025

Copy link
Copy Markdown
Contributor Author

Please address Graham's request

I've addressed this and added more specific/detailed comments ✔️ @grahamking @dmitry-tokarev-nv

@nvda-mesharma nvda-mesharma added the external-contribution Pull request is from an external contributor label Mar 19, 2025
jcoleharrison
jcoleharrison commented Mar 24, 2025
View reviewed changes

@jcoleharrison jcoleharrison left a comment •
edited
Loading

Copy link
Copy Markdown
Contributor Author

Choose a reason for hiding this comment

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

Resolved, please re-review @dmitry-tokarev-nv

@jcoleharrison
jcoleharrison requested a review from a team April 4, 2025 17:28
@harryskim

harryskim commented Apr 10, 2025

Copy link
Copy Markdown
Contributor

@dmitry-tokarev-nv please review and target for 0.2.0 release.

@nvda-mesharma
nvda-mesharma merged this pull request into ai-dynamo:main Apr 11, 2025
saturley-hall pushed a commit that referenced this pull request Apr 11, 2025
@jcoleharrison @saturley-hall
docs: add docstring for llm.rs (#267)
531366c
kylehh pushed a commit to kylehh/dynamo that referenced this pull request Apr 11, 2025
@jcoleharrison @saturley-hall
docs: add docstring for llm.rs (ai-dynamo#267)
447840c
oandreeva-nv pushed a commit that referenced this pull request Apr 22, 2025
@jcoleharrison @oandreeva-nv
docs: add docstring for llm.rs (#267)
b6bd20a
KavinKrishnan added a commit that referenced this pull request May 8, 2026
@KavinKrishnan
feat(trtllm): ModelExpress P2P weight transfer with auto-detect sourc…
2b7b9c3 
…e/target

Add ModelExpress P2P weight transfer support to the Dynamo TRT-LLM
engine. When --model-express-url is set, the engine auto-detects whether
to act as source or target by probing the MX server for existing
sources -- no explicit role flag needed.

Auto-detect behavior:
- Sources found  -> target mode: TRT-LLM uses checkpoint_format="MX",
                    MxLiveWeightLoader does NIXL RDMA receive
- No sources     -> source mode: standard HF disk load,
                    publish_from_worker registers via MX after load

Changes:

* components/src/dynamo/trtllm/backend_args.py:
  - Add --model-express-url CLI argument (env MODEL_EXPRESS_URL).

* components/src/dynamo/trtllm/engine.py:
  - Plumb model_express_url through TensorRTLLMEngine and
    get_llm_engine(); when set, configure checkpoint_format="MX" and
    seed MODEL_EXPRESS_URL env so the upstream MXCheckpointLoader
    (TRT-LLM PR #13531) and modelexpress.publish_from_worker take over.

* components/src/dynamo/trtllm/workers/llm_worker.py:
  - Pass model_express_url from backend args to engine constructor.
  - Use config.exclude_tools_when_tool_choice_none directly (drop the
    defensive getattr/hasattr guard; the field is part of the
    DynamoRuntimeConfig contract; .ai/python-guidelines.md flags
    defensive access on known types).
  - Construct RequestHandlerConfig directly (drop the
    inspect.signature filter; the inspect import was inside the
    function body and the silent kwarg filtering hid contract
    mismatches; .ai/python-guidelines.md requires fail-fast).

* container/{context.yaml, templates/args.Dockerfile,
  templates/trtllm_runtime.Dockerfile}:
  - Add ENABLE_MODELEXPRESS_P2P + MODELEXPRESS_REF build args; install
    modelexpress from git when enabled.

* recipes/deepseek-v3.2/trtllm/disagg/dep8x2/deploy.yaml:
  - DeepSeek-V3.2 disaggregated TP=8 prefill + TP=8 decode recipe with
    --model-express-url wired into both components.

* recipes/deepseek-v3.2/README.md (new):
  - Documents the deviation from recipes/CONTRIBUTING.md standard
    structure -- this recipe intentionally has no model-cache/ since MX
    P2P bypasses local cache for fast scale-up. Points operators at
    sibling DeepSeek recipes if they need the standard model-cache flow.

Validated end-to-end on GCP GB200 (Kimi K2.5 TP=8, 2 nodes per
component): 16 target ranks x 90.75 GB transferred at 365-509 Gbps.

Companion PRs:
* TRT-LLM PR #13531: MXCheckpointLoader (merged 2026-05-06)
* ModelExpress PR #202: MxLiveWeightLoader, publish_from_worker (merged)
* ModelExpress PR #267: MX_POOL_REG allocation-based registration (merged)

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
KavinKrishnan added a commit that referenced this pull request May 8, 2026
@KavinKrishnan
feat(trtllm): ModelExpress P2P weight transfer with auto-detect sourc…
517869e 
…e/target

Add ModelExpress P2P weight transfer support to the Dynamo TRT-LLM
engine. When --model-express-url is set, the engine auto-detects whether
to act as source or target by probing the MX server for existing
sources -- no explicit role flag needed.

Auto-detect behavior:
- Sources found  -> target mode: TRT-LLM uses checkpoint_format="MX",
                    MxLiveWeightLoader does NIXL RDMA receive
- No sources     -> source mode: standard HF disk load,
                    publish_from_worker registers via MX after load

Changes:

* components/src/dynamo/trtllm/backend_args.py:
  - Add --model-express-url CLI argument (env MODEL_EXPRESS_URL).

* components/src/dynamo/trtllm/engine.py:
  - Plumb model_express_url through TensorRTLLMEngine and
    get_llm_engine(); when set, configure checkpoint_format="MX" and
    seed MODEL_EXPRESS_URL env so the upstream MXCheckpointLoader
    (TRT-LLM PR #13531) and modelexpress.publish_from_worker take over.

* components/src/dynamo/trtllm/workers/llm_worker.py:
  - Pass model_express_url from backend args to engine constructor.
  - Use config.exclude_tools_when_tool_choice_none directly (drop the
    defensive getattr/hasattr guard; the field is part of the
    DynamoRuntimeConfig contract; .ai/python-guidelines.md flags
    defensive access on known types).
  - Construct RequestHandlerConfig directly (drop the
    inspect.signature filter; the inspect import was inside the
    function body and the silent kwarg filtering hid contract
    mismatches; .ai/python-guidelines.md requires fail-fast).

* container/{context.yaml, templates/args.Dockerfile,
  templates/trtllm_runtime.Dockerfile}:
  - Add ENABLE_MODELEXPRESS_P2P + MODELEXPRESS_REF build args; install
    modelexpress from git when enabled.

* recipes/deepseek-v3.2/trtllm/disagg/dep8x2/deploy.yaml:
  - DeepSeek-V3.2 disaggregated TP=8 prefill + TP=8 decode recipe with
    --model-express-url wired into both components.

* recipes/deepseek-v3.2/README.md (new):
  - Documents the deviation from recipes/CONTRIBUTING.md standard
    structure -- this recipe intentionally has no model-cache/ since MX
    P2P bypasses local cache for fast scale-up. Points operators at
    sibling DeepSeek recipes if they need the standard model-cache flow.

Validated end-to-end on GCP GB200 (Kimi K2.5 TP=8, 2 nodes per
component): 16 target ranks x 90.75 GB transferred at 365-509 Gbps.

Companion PRs:
* TRT-LLM PR #13531: MXCheckpointLoader (merged 2026-05-06)
* ModelExpress PR #202: MxLiveWeightLoader, publish_from_worker (merged)
* ModelExpress PR #267: MX_POOL_REG allocation-based registration (merged)

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
KavinKrishnan added a commit that referenced this pull request May 8, 2026
@KavinKrishnan
feat(trtllm): ModelExpress P2P weight transfer with auto-detect sourc…
0f1041e 
…e/target

Add ModelExpress P2P weight transfer support to the Dynamo TRT-LLM
engine. When --model-express-url is set, the engine auto-detects whether
to act as source or target by probing the MX server for existing
sources -- no explicit role flag needed.

Auto-detect behavior:
- Sources found  -> target mode: TRT-LLM uses checkpoint_format="MX",
                    MxLiveWeightLoader does NIXL RDMA receive
- No sources     -> source mode: standard HF disk load,
                    publish_from_worker registers via MX after load

Changes:

* components/src/dynamo/trtllm/backend_args.py:
  - Add --model-express-url CLI argument (env MODEL_EXPRESS_URL).

* components/src/dynamo/trtllm/engine.py:
  - Plumb model_express_url through TensorRTLLMEngine and
    get_llm_engine(); when set, configure checkpoint_format="MX" and
    seed MODEL_EXPRESS_URL env so the upstream MXCheckpointLoader
    (TRT-LLM PR #13531) and modelexpress.publish_from_worker take over.

* components/src/dynamo/trtllm/workers/llm_worker.py:
  - Pass model_express_url from backend args to engine constructor.
  - Use config.exclude_tools_when_tool_choice_none directly (drop the
    defensive getattr/hasattr guard; the field is part of the
    DynamoRuntimeConfig contract; .ai/python-guidelines.md flags
    defensive access on known types).
  - Construct RequestHandlerConfig directly (drop the
    inspect.signature filter; the inspect import was inside the
    function body and the silent kwarg filtering hid contract
    mismatches; .ai/python-guidelines.md requires fail-fast).

* container/{context.yaml, templates/args.Dockerfile,
  templates/trtllm_runtime.Dockerfile}:
  - Add ENABLE_MODELEXPRESS_P2P + MODELEXPRESS_REF build args; install
    modelexpress from git when enabled.

* recipes/deepseek-v3.2/trtllm/disagg/dep8x2/deploy.yaml:
  - DeepSeek-V3.2 disaggregated TP=8 prefill + TP=8 decode recipe with
    --model-express-url wired into both components.

* recipes/deepseek-v3.2/README.md (new):
  - Documents the deviation from recipes/CONTRIBUTING.md standard
    structure -- this recipe intentionally has no model-cache/ since MX
    P2P bypasses local cache for fast scale-up. Points operators at
    sibling DeepSeek recipes if they need the standard model-cache flow.

Validated end-to-end on GCP GB200 (Kimi K2.5 TP=8, 2 nodes per
component): 16 target ranks x 90.75 GB transferred at 365-509 Gbps.

Companion PRs:
* TRT-LLM PR #13531: MXCheckpointLoader (merged 2026-05-06)
* ModelExpress PR #202: MxLiveWeightLoader, publish_from_worker (merged)
* ModelExpress PR #267: MX_POOL_REG allocation-based registration (merged)

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
KavinKrishnan added a commit that referenced this pull request May 21, 2026
@KavinKrishnan
feat(trtllm): ModelExpress P2P weight transfer with auto-detect sourc…
4d58d8c 
…e/target

## What this is and why we want it

Cold-starting a large MoE model in TRT-LLM is dominated by weight loading,
not by engine init. On Kimi K2.5 (TP=8, 685 GB) we see ~15-20 minutes per
worker to load weights from a `shared-model-cache` PVC, with all GPUs idle
during the load. For DeepSeek-V3.2, GLM-5, and other ~1 TB MoE models the
penalty is even worse, and it's paid every time a replica scales up,
restarts, or migrates. That stall is the bottleneck for HPA-driven
autoscale, multi-tenant serving, and any deployment where the model size
exceeds local-disk read throughput.

ModelExpress (MX) is the upstream P2P weight transfer system in
`ai-dynamo/modelexpress`. It runs a small gRPC server + Redis sidecar in
the cluster, advertises which workers hold which model versions, and lets
cold-starting workers pull weights via NIXL RDMA directly from a peer
that's already serving. In practice that turns a ~20 minute disk load
into ~2 seconds per rank at 360-500 Gbps (validated on Kimi K2.5 TP=8,
GCP GB200, 4x 400G RoCE).

Two upstream pieces landed in the last six weeks to make this work end-
to-end:
  - **TRT-LLM PR #13531** (merged 2026-05-06, ships in 1.3.0rc15+): adds
    native `MXCheckpointLoader` for `checkpoint_format="MX"` in TRT-LLM's
    PyTorch backend. Does the actual RDMA receive on the target side,
    publishes via `publish_model_params(model)` on the source side, and
    falls back to HF disk loading when MX is unavailable.
  - **ModelExpress PR #202 + #267** (merged): provides `MxClient`,
    `MxLiveWeightLoader`, `publish_from_worker`, and `MX_POOL_REG`
    allocation-based NIXL pool registration on the client side.

This PR is the Dynamo-side glue that lets users opt in from the command
line. It adds `--model-express-url` to the TRT-LLM backend; when set, the
engine probes the MX server at startup and auto-detects whether to act as
the source (no peers yet, load from disk and publish) or target (peers
already serving, pull via RDMA). No `--mx-role source|target` flag — same
DGD spec works for both, which is what unlocks the HPA-driven scale-up
case (every replica uses identical config; the auto-detect handles which
role each one plays).

There's no Linear/JIRA ticket; the workstream is tracked in the
companion PRs below.

## How auto-detect works

1. Engine calls `_has_existing_sources()` which probes the MX server via
   `list_sources()`
2. **Sources found** -> target mode: sets `checkpoint_format="MX"` so
   TRT-LLM's upstream `MXCheckpointLoader` does the NIXL RDMA receive
3. **No sources** -> source mode: sets `MODEL_EXPRESS_URL` env var,
   workers auto-publish via `publish_from_worker()` after disk load

## Changes

* **components/src/dynamo/trtllm/backend_args.py**: Add
  `--model-express-url` CLI argument (env `MODEL_EXPRESS_URL`).

* **components/src/dynamo/trtllm/engine.py**: Plumb `model_express_url`
  through `TensorRTLLMEngine` and `get_llm_engine()`; when set,
  configure `checkpoint_format="MX"` and seed `MODEL_EXPRESS_URL` env
  so the upstream `MXCheckpointLoader` and
  `modelexpress.publish_from_worker` take over.

* **components/src/dynamo/trtllm/workers/llm_worker.py**: Pass
  `model_express_url` from backend args to engine constructor.
  Use `config.exclude_tools_when_tool_choice_none` directly (drop the
  defensive getattr/hasattr guard; the field is part of the
  DynamoRuntimeConfig contract; .ai/python-guidelines.md flags
  defensive access on known types). Construct `RequestHandlerConfig`
  directly (drop the inspect.signature filter; the inspect import was
  inside the function body and the silent kwarg filtering hid contract
  mismatches; .ai/python-guidelines.md requires fail-fast).

* **container/{context.yaml, templates/args.Dockerfile,
  templates/trtllm_runtime.Dockerfile}**: Add `ENABLE_MODELEXPRESS_P2P`
  + `MODELEXPRESS_REF` build args; install modelexpress from git when
  enabled. The new `RUN` step uses the same `--mount=type=cache,target=
  /home/dynamo/.cache/uv,uid=1000,gid=0,mode=0775,sharing=shared` cache
  mount as the rest of the file.

* **recipes/deepseek-v3.2/trtllm/disagg/dep8x2/deploy.yaml**: DeepSeek-
  V3.2 disaggregated TP=8 prefill + TP=8 decode recipe with
  `--model-express-url` wired into both components.

* **recipes/deepseek-v3.2/README.md** (new): Documents the deviation
  from `recipes/CONTRIBUTING.md` standard structure - this recipe
  intentionally has no `model-cache/` since MX P2P bypasses local
  cache for fast scale-up. Points operators at sibling DeepSeek
  recipes if they need the standard model-cache flow.

## Validation

Validated end-to-end on GCP GB200 (Kimi K2.5 TP=8, 2 nodes per
component): 16 target ranks x 90.75 GB transferred at 365-509 Gbps,
end-to-end disaggregated inference verified. Same workload on a single-
replica HPA-aggregated deploy: first replica ~22 minutes from disk,
subsequent HPA-driven replicas ~5 minutes via RDMA at 361-583 Gbps/rank.

## Companion PRs

* TRT-LLM PR #13531: `MXCheckpointLoader` (merged 2026-05-06,
  NVIDIA/TensorRT-LLM#13531)
* ModelExpress PR #202: `MxLiveWeightLoader`, `publish_from_worker`
  (merged, ai-dynamo/modelexpress#202)
* ModelExpress PR #267: `MX_POOL_REG` allocation-based registration
  (merged, ai-dynamo/modelexpress#267)
* ModelExpress PR #271 / #272: deployment examples + MPI worker log
  flush fix (open, in review, do not block this PR)

Signed-off-by: Kavin Krishnan <kavink@nvidia.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@biswapanda biswapanda biswapanda approved these changes
@ishandhanani ishandhanani ishandhanani approved these changes
@ryanolson ryanolson Awaiting requested review from ryanolson
@grahamking grahamking Awaiting requested review from grahamking
@paulhendricks paulhendricks Awaiting requested review from paulhendricks
@tmonty12 tmonty12 Awaiting requested review from tmonty12
@GuanLuo GuanLuo Awaiting requested review from GuanLuo
@rmccorm4 rmccorm4 Awaiting requested review from rmccorm4
@alec-flowers alec-flowers Awaiting requested review from alec-flowers
@kkranen kkranen Awaiting requested review from kkranen
@dmitry-tokarev-nv dmitry-tokarev-nv Awaiting requested review from dmitry-tokarev-nv

Assignees

No one assigned

Labels

external-contribution Pull request is from an external contributor

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

7 participants

@jcoleharrison @grahamking @harryskim @biswapanda @ishandhanani @dmitry-tokarev-nv @nvda-mesharma