DRAFT proof: Kimi K2.5 + Eagle3 shadow-failover stack (do not merge) - #13394
DRAFT proof: Kimi K2.5 + Eagle3 shadow-failover stack (do not merge)#13394galletas1712 wants to merge 17 commits into
Conversation
Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
Introduce the GPU Memory Service (GMS) weight-loading prototype for Dynamo's cross-engine zero-copy weight sharing: - `tensorrt_llm/_torch/memory/gpu_memory_backend.py` — `GMSBackend` wrapper around `gpu_memory_service.client.torch.allocator` providing publish/materialize primitives (`materialize_module`, `defer_finalize_write`, `move_untracked_params`) and `mem_pool_scope` for directing `torch.empty`/`torch.zeros` allocations into GMS-backed virtual memory. - `tensorrt_llm/_torch/pyexecutor/model_loader.py` — extend `LoadFormat` with `GMS` and branch on `gms_backend.is_rw` to publish weights (RW) or materialize an existing layout zero-copy (RO). - `tensorrt_llm/llmapi/llm_args.py` — new `gms_mode`, `gms_tag`, `gms_socket_path` fields; API stability YAML + test coverage. - `tensorrt_llm/_torch/modules/linear.py` — mark `Linear` modules that materialized zero-copy from a committed GMS layout with `_weights_presharded = True`; gate re-sharding in `load_weight_shard` itself via an optional `module=` kwarg so every call site (not just the three unquantized helpers) is protected. The `_weights_presharded` gate lives in `load_weight_shard` to avoid scattering the same ternary across ~90 callers (quantized scales, MoE expert loaders, fused/triton linear). Callers opt in by passing `module=module`; a presharded module forces `tensor_parallel_size=1`, `tensor_parallel_rank=0` and returns the tensor unchanged. Existing callers without `module=` retain legacy behavior. Tests: `tests/unittest/_torch/memory/test_gms_backend.py`, `tests/unittest/_torch/modules/test_load_weight_shard.py`, `tests/unittest/llmapi/test_gms_args.py`. Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
One-model speculative configs keep draft weights in a separate checkpoint directory pointed at by `SpeculativeConfig.speculative_model`. The AUTO load path calls `model.load_draft_weights` after the target load (model_loader.py:476-489), but the GMS RW branch skipped that step, publishing uninitialized draft weights into GMS. For Kimi K2.5 + Eagle3, DeepSeekV3's weight loader explicitly drops any tensor prefixed `draft_model.*` (modeling_deepseekv3.py) and `_call_load_weights` reports those as 'skipped'. The allocate-weights pass (model_loader.py:359-389) still creates CUDA tensors for the draft midlayer (Eagle3DecoderLayer / fc / norms), so the RW snapshot gets committed to GMS with garbage bytes for those ranges. RO clients then zero-copy the garbage. Correctness still holds for the target decode path, but draft acceptance collapses, so this reads as a silent perf regression rather than a crash. Mirror the AUTO path's draft-weight load inside the GMS RW branch. The draft tensors allocated earlier under `gms_backend.mem_pool_scope` already live in the GMS pool; calling `model.load_draft_weights` runs `copy_()` into them so they hold real weights at publish time. Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
18d83f4 to
8974e44
Compare
…pile Opened NVIDIA/TensorRT-LLM#13394 to produce an official CI tarball for the galletas1712/TensorRT-LLM schwinns/gms-refactor-20260423 branch. That pipeline publishes wheels to urm.nvidia.com/artifactory/sw-tensorrt-generic/llm-artifacts/LLM/main/L0_MergeRequest_PR/<PR#>/, which is orders of magnitude faster than having the Dynamo BuildKit pool compile the fork wheel on every run (the from-source path hits scheduling or OOM limits on the dynamo-builder-fallback node pool). Add a new, preferred path that downloads and extracts that tarball instead of compiling: - .github/workflows/build-on-demand.yml: expose 'trtllm_artifact_url' (amd64) and 'trtllm_artifact_gh200_url' (arm64/Grace) as workflow_dispatch inputs and forward them to trtllm-pipeline. - .github/workflows/build-flavor.yml: accept and forward the same two inputs to the build-flavor composite action. - .github/actions/build-flavor/action.yml: extend the 'Prebuild TensorRT-LLM wheel' step to prefer the per-arch artifact URL when set. It downloads the tarball, extracts tensorrt_llm-*.whl, and drops it into the same buildx build-context directory the Dockerfile expects (both at OUTPUT_DIR/ and OUTPUT_DIR/build/ for compatibility). If no URL is provided, the step falls back to the existing from-source compile via container/build_trtllm_wheel.sh. - container/build_trtllm_wheel.sh: default TRTLLM_CUDA_ARCHS to '100-real' (SM100 / B200) instead of '90-real'. nscale single-node failover targets B200; SM90 was the earlier Kimi-on-H200 assumption. Callers can still override via the env var. - container/context.yaml: bump github_trtllm_commit to 8974e4470aafe7d73f7eee07d55769c81f3c396f, the re-signed-off tip of galletas1712/TensorRT-LLM schwinns/gms-refactor-20260423 required by NVIDIA/TensorRT-LLM#13394's DCO gate. Same tree as the previous 18d83f4f commit, just with Signed-off-by trailers added. Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
|
/bot run |
📝 WalkthroughWalkthroughIntroduces GPU Memory Service (GMS) integration with virtual memory management, including a new weight loading format, backend abstraction for GPU memory operations, OOM retry mechanisms for virtual memory, KV cache sleep/wakeup control, pre-sharded weight detection in linear modules, and lifecycle management for GMS-backed weight loading. Changes
Sequence Diagram(s)sequenceDiagram
participant ML as ModelLoader
participant GMS as GMSBackend
participant GMSClient as GMS Client
participant CUDA as CUDA/Torch
ML->>GMS: connect()
GMS->>GMSClient: Initialize socket connection
GMSClient-->>GMS: Connected (RW mode)
GMS-->>ML: Connection success
rect rgba(0, 128, 255, 0.5)
Note over ML,CUDA: GMS RW Mode - Load Weights
ML->>GMS: mem_pool_scope()
GMS->>GMSClient: Activate memory pool
ML->>CUDA: Allocate & load weights
ML->>GMS: move_untracked_params()
GMS->>GMSClient: Map parameter tensors to GMS
ML->>GMS: defer_finalize_write()
GMS->>GMS: Store pending model
end
rect rgba(0, 200, 0, 0.5)
Note over ML,CUDA: Executor Pre-Start
ML->>GMS: finalize_pending_write()
GMS->>GMSClient: Register tensors & finalize write
GMS->>CUDA: Synchronize & transition to RO
GMS-->>ML: Committed GiB
end
sequenceDiagram
participant Executor as Executor
participant Worker as Worker
participant VirtualMem as Virtual Memory
participant CUDA as CUDA
rect rgba(100, 150, 255, 0.5)
Note over Executor,CUDA: Sleep Operation
Executor->>Worker: sleep(sleep_tags)
Worker->>VirtualMem: release_with_tag(*tags)
VirtualMem->>CUDA: Release tagged resources
end
rect rgba(150, 200, 100, 0.5)
Note over Executor,CUDA: Wakeup Operation
Executor->>Worker: wakeup(wakeup_tags)
Worker->>VirtualMem: materialize_with_tag(*tags)
loop OOM Retry
VirtualMem->>CUDA: Materialize resources
alt OOM Error
CUDA-->>VirtualMem: RuntimeError (OOM)
VirtualMem->>CUDA: Synchronize & cleanup
VirtualMem->>VirtualMem: sleep(1.0)
VirtualMem->>CUDA: Retry materialization
else Success
CUDA-->>VirtualMem: Success
end
end
alt KV_CACHE in tags
Worker->>CUDA: reset_prefix_cache()
end
Worker->>CUDA: Synchronize
end
sequenceDiagram
participant GMS as GMSBackend
participant GMSClient as GMS Client
participant Engine as Engine
rect rgba(200, 150, 100, 0.5)
Note over GMS,Engine: GMS RO Mode - Materialize Weights
GMS->>Engine: post_load_weights()
GMS->>Engine: prebind_shared_draft_submodules()
GMS->>GMS: materialize_module()
GMS->>GMSClient: Materialize via GMS
GMSClient-->>GMS: Module weights materialized
GMS->>Engine: post_load_weights()
GMS->>GMS: Mark _is_rw = False
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 7
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (2)
tensorrt_llm/_torch/model_config.py (1)
54-75:⚠️ Potential issue | 🟠 MajorDon’t catch caller
OSErrors as lock failures.This
try/exceptstill wraps theyield, so anyOSErrororPermissionErrorraised insidewith config_file_lock():is treated as a lock-acquisition problem. In the newOSErrorpath that can fall into the tempdir/unlocked retry logic and even attempt a secondyield, which obscures the real config-load failure and can tripcontextmanagerruntime errors. Limit the retry handling to lock acquire/release only.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/model_config.py` around lines 54 - 75, The context manager config_file_lock currently wraps the inner yield in a broad try/except, so exceptions raised by the caller inside "with config_file_lock()" (e.g., OSError or PermissionError) are being mistaken for lock acquisition failures and retried; change the logic so only the lock acquisition/release phases are guarded for filelock.Timeout/PermissionError/OSError (checking _is_shared_fs_lock_error for OSError), while the actual yielded block is executed outside that except handling (i.e., perform lock.acquire() in a try/except, on failure fall back to tmp_lock acquisition or unlocked path, but once a lock is successfully acquired use a plain try/finally around lock.release() and perform the yield without catching exceptions from the caller), referencing config_file_lock, lock, tmp_lock and _is_shared_fs_lock_error to locate and refactor the guarded sections.tensorrt_llm/_torch/pyexecutor/model_loader.py (1)
441-456:⚠️ Potential issue | 🟠 MajorAvoid calling
post_load_weights()before GMS RO materialization.When meta init succeeds, the
LoadFormat.GMSpath skips the generic meta→CUDA allocation branch, so this call can run while parameters are still meta tensors. The RO branch then runspost_load_weights()a second time aftermaterialize_module(). Several implementations intensorrt_llm/_torch/modules/linear.pyrewrite parameter storage inpost_load_weights(), so this can either fail before import or repack the imported tensors twice. This needs a separate pre-materialization hook, not a doublepost_load_weights()pass.Also applies to: 551-560
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/pyexecutor/model_loader.py` around lines 441 - 456, The current code calls post_load_weights() before GMS RO materialization which allows post_load_weights() to run on meta tensors (or causes a second run after materialize_module()), so change the logic to skip calling post_load_weights() for the LoadFormat.GMS path until after materialize_module() completes: detect the GMS branch (LoadFormat.GMS) and avoid invoking post_load_weights() in the pre-materialization phase (the branch that uses init_meta_tensor and model._apply), then invoke a dedicated post-materialize hook (or call post_load_weights() once) immediately after materialize_module() finishes so modules like those in tensorrt_llm/_torch/modules/linear.py only run their storage-rewriting code once on concrete CUDA storage.
🧹 Nitpick comments (3)
tensorrt_llm/_torch/virtual_memory.py (1)
17-19: Minor:__all__is not sorted.Static analysis indicates
__all__should be sorted alphabetically.♻️ Proposed fix
__all__ = [ - "RestoreMode", "maybe_scope", "scope", "release_with_tag", - "materialize_with_tag", "run_with_oom_retry" + "RestoreMode", "materialize_with_tag", "maybe_scope", "release_with_tag", + "run_with_oom_retry", "scope" ]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/virtual_memory.py` around lines 17 - 19, The __all__ list in this module is unsorted; update the __all__ assignment so the exported names are in alphabetical order (e.g., sort the strings "RestoreMode", "maybe_scope", "materialize_with_tag", "release_with_tag", "run_with_oom_retry", "scope") to satisfy static analysis; modify the __all__ declaration (the variable named __all__) to contain the same items sorted lexicographically while preserving the existing string names and commas.tensorrt_llm/_torch/memory/__init__.py (1)
4-6: Minor:__all__is not sorted alphabetically.♻️ Proposed fix
-from .gpu_memory_backend import GMSBackend, GPUMemoryBackend +from .gpu_memory_backend import GMSBackend, GPUMemoryBackend -__all__ = ["GPUMemoryBackend", "GMSBackend"] +__all__ = ["GMSBackend", "GPUMemoryBackend"]🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/_torch/memory/__init__.py` around lines 4 - 6, The exported names in __all__ are not alphabetized; update the __all__ list to be sorted alphabetically so "GMSBackend" comes before "GPUMemoryBackend" (i.e., change __all__ = ["GPUMemoryBackend", "GMSBackend"] to __all__ = ["GMSBackend", "GPUMemoryBackend"]) to keep exports consistent and easier to maintain.tensorrt_llm/llmapi/llm_args.py (1)
3549-3554: Constraingms_modeat the field type instead of validating a free-form string later.
gms_modeonly accepts three values, so making it aLiteral["auto", "rw", "ro"]gives you tighter schema generation and earlier Pydantic validation. The model validator can then stay focused on theload_formatcross-field logic.♻️ Suggested change
- gms_mode: Optional[str] = Field( + gms_mode: Literal["auto", "rw", "ro"] = Field( default="auto", description="GMS operating mode: 'auto', 'rw', or 'ro'. " "Only used when load_format='GMS'.", status="prototype", ) @@ `@model_validator`(mode="after") def validate_gms_config(self) -> 'TorchLlmArgs': - if self.load_format == LoadFormat.GMS and self.gms_mode not in ( - "auto", "rw", "ro"): - raise ValueError( - f"gms_mode must be 'auto', 'rw', or 'ro', got '{self.gms_mode}'." - ) - if self.gms_socket_path is not None and self.load_format != LoadFormat.GMS: logger.warning( "gms_socket_path is set but load_format is '%s', not 'GMS'. " "The gms_socket_path will be ignored. Set load_format='GMS' " "to enable GPU Memory Service.", self.load_format.name)As per coding guidelines, "Use
Literal[\"value1\", \"value2\"]instead ofstrin Python Pydantic fields when only certain values are allowed".Also applies to: 3775-3781
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@tensorrt_llm/llmapi/llm_args.py` around lines 3549 - 3554, Change the gms_mode Field from an unconstrained Optional[str] to a Literal["auto","rw","ro"] (keep the default "auto" and existing description/status) and import Literal (from typing or typing_extensions as used elsewhere) so Pydantic enforces allowed values at schema/validation time; update or remove any redundant runtime checks in validators that only validated gms_mode, and apply the same Literal refactor to the other similar field noted in the file (the second occurrence with the same three allowed values) so both use Literal for stricter typing and earlier validation.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@DYNAMO_FORK.md`:
- Around line 41-46: Update the comment for the trtllm YAML block to remove the
incorrect claim that github_trtllm_commit is "Used by install_tensorrt.sh";
instead either state where that variable is actually consumed (e.g., the
pipeline/job or script that checks out the trtllm repo or the installer that
builds the wheel) or delete the misleading usage note entirely. Locate the
trtllm block and the identifier github_trtllm_commit in DYNAMO_FORK.md and edit
the descriptive line so it accurately reflects the real consumer (or omits the
installer reference to install_tensorrt.sh).
In `@tensorrt_llm/_torch/memory/gpu_memory_backend.py`:
- Around line 134-142: Change has_committed_weights to rely on the internal
_is_rw flag instead of reading self._client.granted_lock_type: return False if
self._client is None else return not self._is_rw (or equivalent logic that
treats _is_rw==False as "committed"). Ensure finalize_write already flips
self._is_rw (no change) and update cleanup() to reset self._is_rw = False so
connection state is consistent after teardown; also replace other helpers that
currently inspect self._client.granted_lock_type with the same _is_rw-based
check to keep behavior consistent (referenced symbols: has_committed_weights,
finalize_write, cleanup, _is_rw, self._client.granted_lock_type).
In `@tensorrt_llm/_torch/virtual_memory.py`:
- Around line 150-168: run_with_oom_retry currently loops forever on OOMs; make
it accept a configurable max retry count (e.g., max_retries: int =
DEFAULT_MAX_OOM_RETRIES) or read from a module-level constant, increment retries
as now but break when retries > max_retries and raise a clear exception; keep
existing behavior for non-OOM errors via _is_oom_error, still call
_cleanup_after_oom() and sleep using _OOM_RETRY_INTERVAL_SECONDS between
attempts, and log an error when the retry limit is reached before raising so
callers know it exhausted retries (update the function signature and logs
referencing run_with_oom_retry, _is_oom_error, _cleanup_after_oom, and
_OOM_RETRY_INTERVAL_SECONDS).
In `@tensorrt_llm/executor/ray_gpu_worker.py`:
- Around line 298-299: The call to self.engine.reset_prefix_cache when
ExecutorMemoryType.KV_CACHE is in tags can raise AttributeError because some
engine implementations lack that method; mirror the sibling check in worker.py
by guarding the call with hasattr(self.engine, "reset_prefix_cache") so only
call self.engine.reset_prefix_cache() if that attribute exists, keeping the
conditional block using ExecutorMemoryType.KV_CACHE and the same behavior
otherwise.
In `@tests/unittest/_torch/misc/test_py_executor_creator_retry.py`:
- Around line 1-5: Add the standard NVIDIA copyright header to the top of this
new Python source file (above the existing imports), including the correct year
of latest meaningful modification; ensure the header matches the project's other
headers exactly and remains as the first lines of the file so symbols like the
imported virtual_memory and py_executor_creator stay unchanged.
In `@tests/unittest/_torch/misc/test_virtual_memory_retry.py`:
- Around line 1-57: Add the required NVIDIA copyright header to the top of this
test file (the module containing test_materialize_with_tag_retries_oom and
test_materialize_with_tag_does_not_retry_non_oom) so the file begins with the
standard NVIDIA license/comment block; ensure the header precedes any imports
and retains the existing imports and test code unchanged.
In `@tests/unittest/executor/test_ray_gpu_worker_sleep.py`:
- Around line 1-5: This new test module test_ray_gpu_worker_sleep.py is missing
the required NVIDIA SPDX/copyright header; add the same SPDX-License-Identifier
line and full NVIDIA copyright/license header block used in other new Python
files in the repo at the top of this file (before the imports), updating the
year to the latest meaningful modification year; ensure the header is a Python
comment block and appears above the existing imports (the file imports
tensorrt_llm.executor.ray_gpu_worker and ExecutorMemoryType from
tensorrt_llm.llmapi.llm_args so place the header before those lines).
---
Outside diff comments:
In `@tensorrt_llm/_torch/model_config.py`:
- Around line 54-75: The context manager config_file_lock currently wraps the
inner yield in a broad try/except, so exceptions raised by the caller inside
"with config_file_lock()" (e.g., OSError or PermissionError) are being mistaken
for lock acquisition failures and retried; change the logic so only the lock
acquisition/release phases are guarded for
filelock.Timeout/PermissionError/OSError (checking _is_shared_fs_lock_error for
OSError), while the actual yielded block is executed outside that except
handling (i.e., perform lock.acquire() in a try/except, on failure fall back to
tmp_lock acquisition or unlocked path, but once a lock is successfully acquired
use a plain try/finally around lock.release() and perform the yield without
catching exceptions from the caller), referencing config_file_lock, lock,
tmp_lock and _is_shared_fs_lock_error to locate and refactor the guarded
sections.
In `@tensorrt_llm/_torch/pyexecutor/model_loader.py`:
- Around line 441-456: The current code calls post_load_weights() before GMS RO
materialization which allows post_load_weights() to run on meta tensors (or
causes a second run after materialize_module()), so change the logic to skip
calling post_load_weights() for the LoadFormat.GMS path until after
materialize_module() completes: detect the GMS branch (LoadFormat.GMS) and avoid
invoking post_load_weights() in the pre-materialization phase (the branch that
uses init_meta_tensor and model._apply), then invoke a dedicated
post-materialize hook (or call post_load_weights() once) immediately after
materialize_module() finishes so modules like those in
tensorrt_llm/_torch/modules/linear.py only run their storage-rewriting code once
on concrete CUDA storage.
---
Nitpick comments:
In `@tensorrt_llm/_torch/memory/__init__.py`:
- Around line 4-6: The exported names in __all__ are not alphabetized; update
the __all__ list to be sorted alphabetically so "GMSBackend" comes before
"GPUMemoryBackend" (i.e., change __all__ = ["GPUMemoryBackend", "GMSBackend"] to
__all__ = ["GMSBackend", "GPUMemoryBackend"]) to keep exports consistent and
easier to maintain.
In `@tensorrt_llm/_torch/virtual_memory.py`:
- Around line 17-19: The __all__ list in this module is unsorted; update the
__all__ assignment so the exported names are in alphabetical order (e.g., sort
the strings "RestoreMode", "maybe_scope", "materialize_with_tag",
"release_with_tag", "run_with_oom_retry", "scope") to satisfy static analysis;
modify the __all__ declaration (the variable named __all__) to contain the same
items sorted lexicographically while preserving the existing string names and
commas.
In `@tensorrt_llm/llmapi/llm_args.py`:
- Around line 3549-3554: Change the gms_mode Field from an unconstrained
Optional[str] to a Literal["auto","rw","ro"] (keep the default "auto" and
existing description/status) and import Literal (from typing or
typing_extensions as used elsewhere) so Pydantic enforces allowed values at
schema/validation time; update or remove any redundant runtime checks in
validators that only validated gms_mode, and apply the same Literal refactor to
the other similar field noted in the file (the second occurrence with the same
three allowed values) so both use Literal for stricter typing and earlier
validation.
🪄 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: 79e51e54-db55-4139-9971-cba9cff4971f
📒 Files selected for processing (23)
DYNAMO_FORK.mdcpp/tensorrt_llm/runtime/virtualMemory.cppcpp/tests/unit_tests/runtime/virtualMemoryTest.cpptensorrt_llm/_torch/memory/__init__.pytensorrt_llm/_torch/memory/gpu_memory_backend.pytensorrt_llm/_torch/model_config.pytensorrt_llm/_torch/modules/linear.pytensorrt_llm/_torch/pyexecutor/model_loader.pytensorrt_llm/_torch/pyexecutor/py_executor_creator.pytensorrt_llm/_torch/virtual_memory.pytensorrt_llm/executor/ray_gpu_worker.pytensorrt_llm/executor/worker.pytensorrt_llm/llmapi/llm.pytensorrt_llm/llmapi/llm_args.pytests/unittest/_torch/memory/test_gms_backend.pytests/unittest/_torch/misc/test_py_executor_creator_retry.pytests/unittest/_torch/misc/test_virtual_memory_retry.pytests/unittest/_torch/modules/test_load_weight_shard.pytests/unittest/_torch/ray_orchestrator/single_gpu/test_llm_sleep.pytests/unittest/_torch/test_model_config.pytests/unittest/api_stability/references/llm.yamltests/unittest/executor/test_ray_gpu_worker_sleep.pytests/unittest/llmapi/test_gms_args.py
| ```yaml | ||
| trtllm: | ||
| pip_wheel: tensorrt-llm==<version> # PyPI pin (upstream) | ||
| trtllm_wheel_image: <wheel image> # Release image (upstream) | ||
| github_trtllm_commit: <sha-or-tag> # Used by install_tensorrt.sh | ||
| ``` |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Description: Check if install_tensorrt.sh references github_trtllm_commit
# Search for github_trtllm_commit in the install_tensorrt.sh script
rg -n 'github_trtllm_commit' docker/common/install_tensorrt.sh
# Also check if there are any other scripts that might use it
rg -n 'github_trtllm_commit' --type=shRepository: NVIDIA/TensorRT-LLM
Length of output: 45
Line 45 incorrectly claims github_trtllm_commit is used by install_tensorrt.sh.
The comment states github_trtllm_commit is "Used by install_tensorrt.sh", but this variable is not referenced in that script or any other shell script in the repository. Update the comment to clarify where and how github_trtllm_commit is actually consumed, or remove the misleading reference.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@DYNAMO_FORK.md` around lines 41 - 46, Update the comment for the trtllm YAML
block to remove the incorrect claim that github_trtllm_commit is "Used by
install_tensorrt.sh"; instead either state where that variable is actually
consumed (e.g., the pipeline/job or script that checks out the trtllm repo or
the installer that builds the wheel) or delete the misleading usage note
entirely. Locate the trtllm block and the identifier github_trtllm_commit in
DYNAMO_FORK.md and edit the descriptive line so it accurately reflects the real
consumer (or omits the installer reference to install_tensorrt.sh).
| def has_committed_weights(self) -> bool: | ||
| if self._client is None: | ||
| return False | ||
| try: | ||
| from gpu_memory_service.common.locks import GrantedLockType | ||
|
|
||
| return self._client.granted_lock_type == GrantedLockType.RO | ||
| except Exception: | ||
| return False |
There was a problem hiding this comment.
has_committed_weights() can stay false after a successful publish.
finalize_write() only flips self._is_rw, but has_committed_weights() still reads self._client.granted_lock_type. If the client object does not update that field in place, the writer will report “no committed weights” immediately after commit. Base this helper on _is_rw instead, and reset _is_rw in cleanup() so the connection state stays internally consistent.
Also applies to: 215-215, 283-285
🧰 Tools
🪛 Ruff (0.15.11)
[warning] 141-141: Do not catch blind exception: Exception
(BLE001)
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/_torch/memory/gpu_memory_backend.py` around lines 134 - 142,
Change has_committed_weights to rely on the internal _is_rw flag instead of
reading self._client.granted_lock_type: return False if self._client is None
else return not self._is_rw (or equivalent logic that treats _is_rw==False as
"committed"). Ensure finalize_write already flips self._is_rw (no change) and
update cleanup() to reset self._is_rw = False so connection state is consistent
after teardown; also replace other helpers that currently inspect
self._client.granted_lock_type with the same _is_rw-based check to keep behavior
consistent (referenced symbols: has_committed_weights, finalize_write, cleanup,
_is_rw, self._client.granted_lock_type).
| def run_with_oom_retry(action: Callable[[], T], *, description: str) -> T: | ||
| retries = 0 | ||
| while True: | ||
| try: | ||
| result = action() | ||
| if retries: | ||
| logger.info("%s succeeded after %d retries", description, | ||
| retries) | ||
| return result | ||
| except Exception as error: | ||
| if not _is_oom_error(error): | ||
| raise | ||
|
|
||
| retries += 1 | ||
| logger.warning( | ||
| "%s hit OOM, waiting for capacity before retry %d: %s", | ||
| description, retries, error) | ||
| _cleanup_after_oom() | ||
| time.sleep(_OOM_RETRY_INTERVAL_SECONDS) |
There was a problem hiding this comment.
Unbounded retry loop may cause infinite hangs under persistent OOM.
run_with_oom_retry retries indefinitely on OOM errors. If GPU memory cannot be freed (e.g., due to a memory leak or external process), this will loop forever with 1-second sleeps. Consider adding a configurable maximum retry count.
🛡️ Proposed fix to add retry limit
+_OOM_MAX_RETRIES = 10 # or expose as parameter
+
def run_with_oom_retry(action: Callable[[], T], *, description: str) -> T:
retries = 0
while True:
try:
result = action()
if retries:
logger.info("%s succeeded after %d retries", description,
retries)
return result
except Exception as error:
if not _is_oom_error(error):
raise
retries += 1
+ if retries > _OOM_MAX_RETRIES:
+ logger.error("%s failed after %d retries", description, retries)
+ raise
logger.warning(
"%s hit OOM, waiting for capacity before retry %d: %s",
description, retries, error)
_cleanup_after_oom()
time.sleep(_OOM_RETRY_INTERVAL_SECONDS)🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/_torch/virtual_memory.py` around lines 150 - 168,
run_with_oom_retry currently loops forever on OOMs; make it accept a
configurable max retry count (e.g., max_retries: int = DEFAULT_MAX_OOM_RETRIES)
or read from a module-level constant, increment retries as now but break when
retries > max_retries and raise a clear exception; keep existing behavior for
non-OOM errors via _is_oom_error, still call _cleanup_after_oom() and sleep
using _OOM_RETRY_INTERVAL_SECONDS between attempts, and log an error when the
retry limit is reached before raising so callers know it exhausted retries
(update the function signature and logs referencing run_with_oom_retry,
_is_oom_error, _cleanup_after_oom, and _OOM_RETRY_INTERVAL_SECONDS).
| if ExecutorMemoryType.KV_CACHE in tags: | ||
| self.engine.reset_prefix_cache() |
There was a problem hiding this comment.
Missing hasattr check for reset_prefix_cache.
The sibling implementation in tensorrt_llm/executor/worker.py:97-98 guards the call with hasattr(self.engine, 'reset_prefix_cache'). This check is missing here, which could raise AttributeError if the engine doesn't implement reset_prefix_cache.
🛡️ Proposed fix for consistency with worker.py
if ExecutorMemoryType.KV_CACHE in tags:
- self.engine.reset_prefix_cache()
+ if hasattr(self.engine, 'reset_prefix_cache'):
+ self.engine.reset_prefix_cache()🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tensorrt_llm/executor/ray_gpu_worker.py` around lines 298 - 299, The call to
self.engine.reset_prefix_cache when ExecutorMemoryType.KV_CACHE is in tags can
raise AttributeError because some engine implementations lack that method;
mirror the sibling check in worker.py by guarding the call with
hasattr(self.engine, "reset_prefix_cache") so only call
self.engine.reset_prefix_cache() if that attribute exists, keeping the
conditional block using ExecutorMemoryType.KV_CACHE and the same behavior
otherwise.
| import pytest | ||
|
|
||
| from tensorrt_llm._torch import virtual_memory | ||
| from tensorrt_llm._torch.pyexecutor import py_executor_creator | ||
|
|
There was a problem hiding this comment.
Add the required NVIDIA file header.
This is a new Python source file, so it should include the standard NVIDIA copyright header used elsewhere in the repo.
📝 Suggested change
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
import pytest
from tensorrt_llm._torch import virtual_memoryAs per coding guidelines, "All TensorRT-LLM source files must contain an NVIDIA copyright header with the year of latest meaningful modification" and "Include NVIDIA copyright header on all new files; update year on modified files".
📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| import pytest | |
| from tensorrt_llm._torch import virtual_memory | |
| from tensorrt_llm._torch.pyexecutor import py_executor_creator | |
| # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| import pytest | |
| from tensorrt_llm._torch import virtual_memory | |
| from tensorrt_llm._torch.pyexecutor import py_executor_creator |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unittest/_torch/misc/test_py_executor_creator_retry.py` around lines 1
- 5, Add the standard NVIDIA copyright header to the top of this new Python
source file (above the existing imports), including the correct year of latest
meaningful modification; ensure the header matches the project's other headers
exactly and remains as the first lines of the file so symbols like the imported
virtual_memory and py_executor_creator stay unchanged.
| import pytest | ||
|
|
||
| from tensorrt_llm._torch import virtual_memory | ||
|
|
||
|
|
||
| class _RetryingManager: | ||
|
|
||
| def __init__(self): | ||
| self.calls = 0 | ||
|
|
||
| def materialize_with_tag(self, tag: str) -> int: | ||
| self.calls += 1 | ||
| if self.calls == 1: | ||
| raise RuntimeError(f"{tag} out of memory") | ||
| return 1 | ||
|
|
||
|
|
||
| def test_materialize_with_tag_retries_oom(monkeypatch): | ||
| manager = _RetryingManager() | ||
| sleeps: list[float] = [] | ||
| sync_calls: list[None] = [] | ||
| empty_cache_calls: list[None] = [] | ||
| gc_calls: list[int] = [] | ||
|
|
||
| monkeypatch.setattr(virtual_memory, "get_virtual_memory_manager", | ||
| lambda: manager) | ||
| monkeypatch.setattr(virtual_memory.torch.cuda, "synchronize", | ||
| lambda: sync_calls.append(None)) | ||
| monkeypatch.setattr(virtual_memory.torch.cuda, "empty_cache", | ||
| lambda: empty_cache_calls.append(None)) | ||
| monkeypatch.setattr(virtual_memory.gc, "collect", | ||
| lambda: gc_calls.append(1) or 0) | ||
| monkeypatch.setattr(virtual_memory.time, "sleep", sleeps.append) | ||
|
|
||
| assert virtual_memory.materialize_with_tag("kv_cache") == 1 | ||
| assert manager.calls == 2 | ||
| assert sleeps == [1.0] | ||
| assert len(sync_calls) == 1 | ||
| assert len(empty_cache_calls) == 1 | ||
| assert gc_calls == [1] | ||
|
|
||
|
|
||
| def test_materialize_with_tag_does_not_retry_non_oom(monkeypatch): | ||
| class _FailingManager: | ||
|
|
||
| def materialize_with_tag(self, tag: str) -> int: | ||
| raise RuntimeError(f"{tag} permission denied") | ||
|
|
||
| sleeps: list[float] = [] | ||
| monkeypatch.setattr(virtual_memory, "get_virtual_memory_manager", | ||
| lambda: _FailingManager()) | ||
| monkeypatch.setattr(virtual_memory.time, "sleep", sleeps.append) | ||
|
|
||
| with pytest.raises(RuntimeError, match="permission denied"): | ||
| virtual_memory.materialize_with_tag("kv_cache") | ||
|
|
||
| assert not sleeps |
There was a problem hiding this comment.
Missing NVIDIA copyright header.
Per coding guidelines, all source files should include the NVIDIA copyright header.
📄 Proposed fix
+# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
import pytest
from tensorrt_llm._torch import virtual_memory📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| import pytest | |
| from tensorrt_llm._torch import virtual_memory | |
| class _RetryingManager: | |
| def __init__(self): | |
| self.calls = 0 | |
| def materialize_with_tag(self, tag: str) -> int: | |
| self.calls += 1 | |
| if self.calls == 1: | |
| raise RuntimeError(f"{tag} out of memory") | |
| return 1 | |
| def test_materialize_with_tag_retries_oom(monkeypatch): | |
| manager = _RetryingManager() | |
| sleeps: list[float] = [] | |
| sync_calls: list[None] = [] | |
| empty_cache_calls: list[None] = [] | |
| gc_calls: list[int] = [] | |
| monkeypatch.setattr(virtual_memory, "get_virtual_memory_manager", | |
| lambda: manager) | |
| monkeypatch.setattr(virtual_memory.torch.cuda, "synchronize", | |
| lambda: sync_calls.append(None)) | |
| monkeypatch.setattr(virtual_memory.torch.cuda, "empty_cache", | |
| lambda: empty_cache_calls.append(None)) | |
| monkeypatch.setattr(virtual_memory.gc, "collect", | |
| lambda: gc_calls.append(1) or 0) | |
| monkeypatch.setattr(virtual_memory.time, "sleep", sleeps.append) | |
| assert virtual_memory.materialize_with_tag("kv_cache") == 1 | |
| assert manager.calls == 2 | |
| assert sleeps == [1.0] | |
| assert len(sync_calls) == 1 | |
| assert len(empty_cache_calls) == 1 | |
| assert gc_calls == [1] | |
| def test_materialize_with_tag_does_not_retry_non_oom(monkeypatch): | |
| class _FailingManager: | |
| def materialize_with_tag(self, tag: str) -> int: | |
| raise RuntimeError(f"{tag} permission denied") | |
| sleeps: list[float] = [] | |
| monkeypatch.setattr(virtual_memory, "get_virtual_memory_manager", | |
| lambda: _FailingManager()) | |
| monkeypatch.setattr(virtual_memory.time, "sleep", sleeps.append) | |
| with pytest.raises(RuntimeError, match="permission denied"): | |
| virtual_memory.materialize_with_tag("kv_cache") | |
| assert not sleeps | |
| # SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |
| # SPDX-License-Identifier: Apache-2.0 | |
| import pytest | |
| from tensorrt_llm._torch import virtual_memory | |
| class _RetryingManager: | |
| def __init__(self): | |
| self.calls = 0 | |
| def materialize_with_tag(self, tag: str) -> int: | |
| self.calls += 1 | |
| if self.calls == 1: | |
| raise RuntimeError(f"{tag} out of memory") | |
| return 1 | |
| def test_materialize_with_tag_retries_oom(monkeypatch): | |
| manager = _RetryingManager() | |
| sleeps: list[float] = [] | |
| sync_calls: list[None] = [] | |
| empty_cache_calls: list[None] = [] | |
| gc_calls: list[int] = [] | |
| monkeypatch.setattr(virtual_memory, "get_virtual_memory_manager", | |
| lambda: manager) | |
| monkeypatch.setattr(virtual_memory.torch.cuda, "synchronize", | |
| lambda: sync_calls.append(None)) | |
| monkeypatch.setattr(virtual_memory.torch.cuda, "empty_cache", | |
| lambda: empty_cache_calls.append(None)) | |
| monkeypatch.setattr(virtual_memory.gc, "collect", | |
| lambda: gc_calls.append(1) or 0) | |
| monkeypatch.setattr(virtual_memory.time, "sleep", sleeps.append) | |
| assert virtual_memory.materialize_with_tag("kv_cache") == 1 | |
| assert manager.calls == 2 | |
| assert sleeps == [1.0] | |
| assert len(sync_calls) == 1 | |
| assert len(empty_cache_calls) == 1 | |
| assert gc_calls == [1] | |
| def test_materialize_with_tag_does_not_retry_non_oom(monkeypatch): | |
| class _FailingManager: | |
| def materialize_with_tag(self, tag: str) -> int: | |
| raise RuntimeError(f"{tag} permission denied") | |
| sleeps: list[float] = [] | |
| monkeypatch.setattr(virtual_memory, "get_virtual_memory_manager", | |
| lambda: _FailingManager()) | |
| monkeypatch.setattr(virtual_memory.time, "sleep", sleeps.append) | |
| with pytest.raises(RuntimeError, match="permission denied"): | |
| virtual_memory.materialize_with_tag("kv_cache") | |
| assert not sleeps |
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unittest/_torch/misc/test_virtual_memory_retry.py` around lines 1 - 57,
Add the required NVIDIA copyright header to the top of this test file (the
module containing test_materialize_with_tag_retries_oom and
test_materialize_with_tag_does_not_retry_non_oom) so the file begins with the
standard NVIDIA license/comment block; ensure the header precedes any imports
and retains the existing imports and test code unchanged.
| from contextlib import contextmanager | ||
|
|
||
| import tensorrt_llm.executor.ray_gpu_worker as ray_gpu_worker | ||
| from tensorrt_llm.llmapi.llm_args import ExecutorMemoryType | ||
|
|
There was a problem hiding this comment.
Add the required SPDX header to this new test file.
This file is new, but it does not include the NVIDIA copyright/license header used by other new Python files in the repo and in this PR.
As per coding guidelines, "All TensorRT-LLM source files must contain an NVIDIA copyright header with the year of latest meaningful modification".
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@tests/unittest/executor/test_ray_gpu_worker_sleep.py` around lines 1 - 5,
This new test module test_ray_gpu_worker_sleep.py is missing the required NVIDIA
SPDX/copyright header; add the same SPDX-License-Identifier line and full NVIDIA
copyright/license header block used in other new Python files in the repo at the
top of this file (before the imports), updating the year to the latest
meaningful modification year; ensure the header is a Python comment block and
appears above the existing imports (the file imports
tensorrt_llm.executor.ray_gpu_worker and ExecutorMemoryType from
tensorrt_llm.llmapi.llm_args so place the header before those lines).
|
/bot help |
GitHub Bot Help
Provide a user friendly way for developers to interact with a Jenkins server. Run See details below for each supported subcommand. Details
Launch build/test pipelines. All previously running jobs will be killed.
kill
Kill all running builds associated with pull request. skip
Skip testing for latest commit on pull request. reuse-pipeline
Reuse a previous pipeline to validate current commit. This action will also kill all currently running builds associated with the pull request. IMPORTANT NOTE: This is dangerous since lack of user care and validation can cause top of tree to break. |
|
/bot run --skip-test |
svc-trtllm-gh-bot
added
the
Community want to contribute
|
PR_Github #45252 [ run ] triggered by Bot. Commit: |
|
PR_Github #45252 [ run ] completed with state
|
Signed-off-by: Schwinn Saereesitthipitak <schwinns@nvidia.com>
Porting plan for this proof branchThis PR was useful as the working shadow-failover proof, but it should not be rebased wholesale onto current The intended upstream path is to keep this PR as the proof/artifact reference, then port the pieces into a smaller stack on top of:
Proposed stack
Validation target from this proof branchThe reference successful run was:
Observed acceptance numbers from the artifacts:
Explicit non-goals for the port
The goal is to turn this proof branch into a reviewable sequence: first establish clean GMS loading, then add reversible GMS sleep/wake, then add MPI propagation, memory calibration, VMM retry, MoE footprint reduction, and CUDA graph/autotuner guardrails. |
4 participants
DO NOT MERGE - proof build only
This draft PR exists to trigger NVIDIA/TensorRT-LLM external-contributor CI and to host the TRT-LLM-side patches needed by the Dynamo TRT-LLM shadow-engine failover proof.
Head:
galletas1712/TensorRT-LLM:schwinns/gms-refactor-20260423@9ba5ee40Base:
NVIDIA/TensorRT-LLM:mainDynamo-side orchestration and runtime pytest live in ai-dynamo/dynamo#8621.
Patch Rationale
TRTLLM_GMS_RELEASE_WEIGHTS_ON_SLEEP=all. The backing store remains resident once in GMS, but the parked process releases its local mappings so duplicate shadow footprint is lower.LoadFormat.GMS, string forms likegms, and serialized enum forms must round-trip throughTorchLlmArgsand MPI startup. Without this, TP workers can silently miss the GMS path.GMS_*,ENGINE_ID, andFAILOVER_LOCK_PATHmust reach MPI workers, not only the parent process. Otherwise GMS mode and failover-lock behavior diverge across TP ranks.TRTLLM_GMS_SHADOW_PEER_MEMORY_RESERVE_GIB/_BYTESsubtracts explicit headroom before applyingfree_gpu_memory_fraction; this keeps pre-capture from OOMing without lowering the serving memory fraction below0.85.0.2sretry interval to avoid tight polling across TP ranks while preserving the staggered block-on-OOM behavior.max_gpu_total_bytes. Once calibration produces a byte budget, the resource manager must honor it directly; otherwise the later KV manager can remeasure a different free-memory state and resize inconsistently.gms_weightsshadow-failover sleep preserves captured inference graphs. Releasing them made Qwen's first promoted-shadow response recapture graphs and jump from about1sto about13safterSIGKILL.TRTLLM_CUDA_GRAPH_MEMORY_SNAPSHOTS=1logs per-key graph-capture memory/timing snapshots around eager warmup, graph capture, postprocess, and graph storage. This was needed to prove bs128's incremental graph reserve and distinguish graph memory from GMS weights/KV VMM.memory_buffer_utils.Buffers.bytes_by_name()andclear()let sleep account for and release reusable buffers that are safe to rebuild after promotion.TRTLLM_GMS_LOG_PER_TAG_SLEEP_RELEASE=1breaks sleep cleanup down by VMM tag,TRTLLM_GMS_LOG_SLEEP_PHASES=1can snapshot memory immediately after GMS weight parking, and wakeup reports per-rank phase timings.shadow_failoverSleepConfig preset. Dynamo should not hardcode TRT-LLM's internal sleep tags. TRT-LLM owns the preset expansion for KV, model, draft model, executor extras, MoE comms, speculative resources, drafter, sampler, guided decoder state, and GMS weights.SleepConfigpickleable. Dynamo passesLLMArgsthrough MPI startup at TP=8. The previousdefaultdict(lambda: ...)restore-mode default was not pickleable, so worker startup failed before GMS restore. Default restore modes are now a plaindict.TLLM_AUTOTUNER_USE_CUDA_GRAPHfor autotuner profiling. Kimi K2.5 + EAGLE3 showed that autotuner profiling can allocate about37-38 GiB/rankin CUDA graph private pools even when inference CUDA graphs are disabled. This env gate disables only the autotuner's internaltorch.cuda.CUDAGraph()profiling wrapper; inference CUDA graphs remain controlled bycuda_graph_config.Current Behavior
The accepted path now preserves the failover invariant: warmup and CUDA graph capture happen before the standby parks, and promotion/wakeup does not run deferred warmup or graph capture.
The latest Kimi K2.5 + EAGLE3 TP=8 bs128 run used one standby,
free_gpu_memory_fraction=0.85, a pre-populated autotuner cache,TLLM_AUTOTUNER_USE_CUDA_GRAPH=0, explicit CUDA graph batch sizes[1,2,4,8,16,32,64,128],enable_padding=true, andTRTLLM_GMS_SHADOW_PEER_MEMORY_RESERVE_GIB=10. It passed with the standby already pre-captured and parked before the primary served traffic.The promoted-standby wakeup hot path was dominated by GMS weight restore/import/access setup:
restore_gms_weights=4.91-5.03s, whilematerialize_vmm=0.01-0.04s. No deferred warmup/CUDA graph bucket was present.The no-reserve control run failed before failover: the standby sized to roughly
75 GiB/rankKV, captured bs128, then OOMed starting bs64 with rank-0 driver free down to0.01 GiB. The 10 GiB peer reserve is therefore startup/capture headroom, not a serving-memory-fraction cheat.sequenceDiagram participant P as Primary TRT-LLM participant G as GMS participant S as Shadow TRT-LLM participant V as TRT-LLM VMM/KV/Graphs participant D as Dynamo Locks/Router P->>G: publish weights RW P->>V: autotune/cache-hit warmup and CUDA graph capture P->>V: sleep(shadow_failover), preserve graphs, release KV/VMM/caches S->>G: map weights RO from shared GMS store S->>V: calibrate KV with GMS weights + peer reserve S->>V: allocate KV, run warmup, capture CUDA graphs before parking S->>V: sleep(shadow_failover), preserve graphs, release KV/VMM/caches S->>D: park and wait on failover lock P->>V: wake for serving path and allocate full KV P->>D: serve traffic, then SIGKILL releases lock S->>D: acquire failover lock S->>G: reconnect/remap GMS weights RO S->>V: wake and materialize KV VMM; no warmup or graph capture S->>D: publish model metadata D->>S: route post-failover inference for manual inspectionValidation
TRT-LLM build/setup notes:
/workspace-dev/TensorRT-LLMand/usr/local/lib/python3.12/dist-packages/tensorrt_llmbefore the final run.Checks from this commit:
git diff --check- clean.python3 -m py_compilefor edited TRT-LLM runtime files - passed.TRTLLM_GMS_DEFER,run_deferred,capture_deferred,defer_warmup_once,defer_cuda_graph_warmup_once, etc.) - no matches in the edited pyexecutor/test paths.python3 -m py_compilefor edited TRT-LLM runtime files - passed.pytest; the nscale pod ran the runtime failover pytest against the synced installed code.Dynamo runtime validation with this TRT-LLM code synced into the pod:
Qwen/Qwen3-0.6B, TP=1, two shadows,free_gpu_memory_fraction=0.85-1 passed in 158.00s; failover-ready0.71s, post-restore response0.33s, kill-to-response1.04s; primary/promoted KV blocks[43485, 0]/[43589, 0].Qwen/Qwen3-0.6B, TP=8, two shadows,free_gpu_memory_fraction=0.85-1 passed in 198.33s; failover-ready3.34s, post-restore response0.50s, kill-to-response3.84s; primary/promoted KV blocks[344905, 0]/[345103, 0].1 passed in 360.40s; failover-ready5.21s, post-restore response1.95s, kill-to-response7.15s.TLLM_AUTOTUNER_USE_CUDA_GRAPH=0, peer reserve10 GiB-1 passed in 320.54s; failover-ready5.86s, post-restore response0.57s, kill-to-response6.43s; primary/promoted KV blocks58510/60232; manual response ended infailover-ok.free_gpu_memory_fraction=0.85.0Profiled runner=lines; test passed in299.07s; kill-to-response9.53s; primary/promoted KV blocks70660/72380.Artifacts/worklog are archived locally under
/Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427, especiallynew-runs-20260428-precapture-guard/INDEX.mdandtrtllm_failover_artifacts_with_precapture_guard_20260428.tar.gz.Known Follow-Ups
restore_gms_weightshot-path bucket. Current evidence points at CUDA handle import/access setup, not weight data copy or KV VMM materialization.This PR is not intended for merge as-is. It is a proof-build branch for the Dynamo shadow-failover stack and should be closed once the relevant pieces land through their intended upstream path.
Signed-off-by: Schwinn Saereesitthipitak schwinns@nvidia.com
Reproducibility: nscale Devpod Setup
Pod / container
nv-prd-dgxc.teleport.sh-dynamo-nscale-dev-clusterschwinns/trtllm-gms-failover-devpodcluster-0967a26d-pool-14bee067-prctr-s2877nvcr.io/nvidia/tensorrt-llm/devel:1.3.0rc12cpu=48,memory=384Gi; limitscpu=96,memory=768Gihf-cache, PVCshared-model-cache, mounted at/home/dynamo/.cache/huggingfaceworkspace-dev, mounted at/workspace-devemptyDir.medium=Memory, mounted at/dev/shm182633 MiBvisible per GPU, CUDA13.1.1.006, driver590.48.01, Python3.12.3, PyTorch2.11.0a0+eb65b36914.nv26.02, NGC PyTorch26.02, TensorRT-LLM package1.3.0rc11,gpu-memory-service==0.9.0.The source directories in the pod were synced into
/workspace-dev/dynamoand/workspace-dev/TensorRT-LLMwithout.gitmetadata. The PR SHAs used for the synced source/build were:ai-dynamo/dynamo:schwinns/gms-refactor-20260423atb1fb89cf2cc37be4a49c5cdad7d89a3b0058a6d6galletas1712/TensorRT-LLM:schwinns/gms-refactor-20260423at9ba5ee407b3e1fe4a48b427d6e6bcc10824e629bInstalled packages in the accepted pod:
Build / install steps
TRTLLM had C++ changes, so the wheel was rebuilt inside the pod with the native B200 arch and ccache:
Dynamo and GMS were installed editable from the synced Dynamo branch:
Detached
kubectl execruns needed the Python wheel.libsdirectories onLD_LIBRARY_PATH; without this, the run could fail to loadlibnixl-c7102487.so/libetcd-cpp-api-core-cac80504.so:Accepted TRTLLM shadow-failover pytest command
The main accepted Kimi K2.5 + EAGLE3 run used the shared HF cache offline and pre-captured CUDA graphs before parking the standby. Promotion does not defer warmup or CUDA graph capture into the failover hot path.
Notes on that configuration:
TRTLLM_GMS_SHADOW_PEER_MEMORY_RESERVE_GIB=10reserves headroom for the already parked peer while the promoted engine sizes KV and captures CUDA graphs. This is currently a conservative explicit reserve; a follow-up is to replace it with measured peer footprint where possible.TLLM_AUTOTUNER_USE_CUDA_GRAPH=0avoids the autotuner's CUDA-graph profiling pool, which was observed to create a private per-rank memory spike around 37-38 GiB in failed exploratory runs.bs=128with padding enabled. The default/expanded graph batch set put too much pressure on the parked-shadow setup.pytest.logfor thefailover-okresponse, rather than asserting semantic coherence in pytest.Timing / validation results
/Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427/qwen3_06b_tp1_final_1777080565/pytest.log1 passed in 158.00s43485/43589/Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427/qwen3_06b_tp8_accept_1777078131/pytest.log1 passed in 198.33s344905/345103/Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427/new-runs-20260428-precapture-guard/logs/kimi_precapture_cgbs128_reserve10_guard_snap_1777360695/pytest.log1 passed in 320.54s5.86s; post-restore response0.57s; SIGKILL-to-response6.43s; response manually inspected asfailover-ok58510/60232; primary approx61.27 GiB main + 3.01 GiB draft; promoted standby approx63.07 GiB main + 3.10 GiB draft/Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427/new-runs-20260428-tensor-audit/tensor_audit_1777364433/pytest.log1 passed in 355.83s5.91s; post-restore response0.55s; SIGKILL-to-response6.46s58510/60232; parked footprint approx9.5-9.7 GiB/rank; local tensor storage approx4.282 GiB/rank(3.437 GiBCUDA graph runner +0.845 GiBshape-classified local tensors)The hot wake path in the accepted Kimi run was dominated by remapping/restoring GMS weights, not KV VMM materialization:
Artifact index
Local artifact/worklog roots:
/Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427/WORKLOG.md/Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427/new-runs-20260427-cuda-graphs/CUDA_GRAPH_MEMORY_INVESTIGATION.md/Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427/new-runs-20260428-tensor-audit/TENSOR_AUDIT_SUMMARY.md/Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427/new-runs-20260428-precapture-guard/logs/kimi_precapture_cgbs128_reserve10_guard_snap_1777360695//Users/schwinns/dynamo-worktrees/trtllm-failover-artifacts-20260427/new-runs-20260428-tensor-audit/tensor_audit_1777364433/Pod artifact root for the latest tensor-audit run:
/workspace-dev/trtllmfailover-artifacts/tensor_audit_1777364433/pytest.log/workspace-dev/trtllmfailover-artifacts/tensor_audit_1777364433/run.env/workspace-dev/trtllmfailover-artifacts/tensor_audit_1777364433/tensor_audit/*.json