[None][feat] DSv4 prep: runtime cache foundations#15378
Conversation
📝 WalkthroughWalkthroughThis PR adds V2 KV cache statistics tracking, event management, and configurable block-hash algorithms to TensorRT-LLM. It introduces new stats delta data structures, a pending-stats buffering system inside ChangesKV Cache V2 Statistics, Events, and Hash Support
Sequence Diagram(s)sequenceDiagram
participant PyExecutor
participant KVCacheManagerV2
participant KVCacheManager as KVCacheManager (runtime)
participant _KVCache
participant KVCacheEventManager
PyExecutor->>KVCacheManagerV2: prepare_resources(scheduled_batch)
KVCacheManagerV2->>KVCacheManager: create_kv_cache(expected_prompt_length)
KVCacheManager-->>_KVCache: init with _PendingStats
KVCacheManagerV2->>KVCacheManagerV2: update_context_resources(scheduled_batch)
_KVCache->>_KVCache: resize / record_allocation_range / record_reuse
_KVCache->>KVCacheEventManager: add_stored_block_event_from_block()
PyExecutor->>KVCacheManagerV2: _commit_kv_cache_stats(scheduled_batch)
KVCacheManagerV2->>_KVCache: commit_pending_stats()
_KVCache-->>KVCacheManager: KVCacheStatsDelta
KVCacheManager->>KVCacheManager: commit_stats(delta, iteration_stats)
KVCacheManagerV2->>KVCacheEventManager: flush_iteration_events()
KVCacheManagerV2->>KVCacheManagerV2: build KVCacheV2IterationStatsReport
PyExecutor->>KVCacheManagerV2: get_iteration_stats()
KVCacheManagerV2-->>PyExecutor: KVCacheV2IterationStatsReport
PyExecutor->>PyExecutor: append_kv_cache_iteration_stats(stats_dict, report)
Estimated code review effort🎯 5 (Critical) | ⏱️ ~120 minutes Suggested labels
Suggested reviewers
✨ Finishing Touches🧪 Generate unit tests (beta)
|
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 (1)
tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py (1)
79-89:⚠️ Potential issue | 🟡 MinorAlign mypyc package metadata with the repository's Python 3.10+ target.
The
setup_mypyc.pycurrently declarespython_requires=">=3.8"(line 139) while the mainsetup.pydeclares>=3.10, <4. Since newly added modules to mypyc use Python 3.10+ syntax, this version constraint should match.- python_requires=">=3.8", + python_requires=">=3.10,<4",🤖 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 `@tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py` around lines 79 - 89, Update the python_requires parameter in setup_mypyc.py to match the main setup.py requirement. Change the python_requires value from ">=3.8" to ">=3.10, <4" to align with the repository's Python 3.10+ target, since the newly added modules in the mypyc configuration use Python 3.10+ syntax features.Sources: Coding guidelines, Learnings
🧹 Nitpick comments (5)
tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py (1)
22-37: ⚡ Quick winType-annotate
_StatsDeltaMixinoperation signatures.
add,subtract, andcopyare missing full type annotations for operands/return type in this new API surface, which weakens static checks in downstream stats plumbing.♻️ Proposed fix
from dataclasses import dataclass, fields +from typing import TypeVar +_StatsDeltaT = TypeVar("_StatsDeltaT", bound="_StatsDeltaMixin") + class _StatsDeltaMixin: __slots__ = () - def add(self, other) -> None: + def add(self: _StatsDeltaT, other: _StatsDeltaT) -> None: for field in fields(self): name = field.name setattr(self, name, getattr(self, name) + getattr(other, name)) - def subtract(self, other) -> None: + def subtract(self: _StatsDeltaT, other: _StatsDeltaT) -> None: for field in fields(self): name = field.name setattr(self, name, getattr(self, name) - getattr(other, name)) - def copy(self): + def copy(self: _StatsDeltaT) -> _StatsDeltaT: return type(self)(**{field.name: getattr(self, field.name) for field in fields(self)})As per coding guidelines, Python functions should always be annotated, including explicit return types and typed parameters.
🤖 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 `@tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py` around lines 22 - 37, The methods `add`, `subtract`, and `copy` in the `_StatsDeltaMixin` class are missing type annotations. Add type annotation for the `other` parameter in both the `add` and `subtract` methods (should be typed as the same class type since they operate on instances of the same type), and add a return type annotation to the `copy` method (should return the same class type). This will ensure static type checkers can properly validate usage of these methods throughout the codebase.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py (1)
1995-2005: 💤 Low valueUse explicit
| Nonetype hint instead of implicitOptional.Per PEP 484 and the static analysis hint,
kv_cache_dtype_byte_size: float = Noneis an implicit Optional. Use the explicit union syntax for clarity.Suggested fix
def update_resources( self, scheduled_batch: ScheduledRequests, attn_metadata: "AttentionMetadata" = None, - kv_cache_dtype_byte_size: float = None, + kv_cache_dtype_byte_size: float | None = None, ):🤖 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 `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py` around lines 1995 - 2005, In the update_resources method, change the type hints for the parameters attn_metadata and kv_cache_dtype_byte_size to use explicit union syntax with None. Modify attn_metadata from "AttentionMetadata" to "AttentionMetadata" | None, and modify kv_cache_dtype_byte_size from float to float | None, to explicitly declare that these parameters accept None values instead of relying on implicit Optional inference.Source: Linters/SAST tools
tensorrt_llm/_torch/pyexecutor/_util.py (1)
270-273: ⚡ Quick winAnnotate the new fallback helper signature.
kv_cache_manager_clsand the method return are currently untyped, so this new helper weakens static checking around manager-class routing.Proposed fix
def _fallback_if_unsupported_kv_cache_manager_v2( self, - kv_cache_manager_cls, - model_config: Optional[ModelConfig] = None): + kv_cache_manager_cls: type, + model_config: Optional[ModelConfig] = None) -> type:As per coding guidelines, “Always annotate functions. Make the return type
Noneif the function does not return anything.”🤖 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 `@tensorrt_llm/_torch/pyexecutor/_util.py` around lines 270 - 273, Add missing type annotations to the `_fallback_if_unsupported_kv_cache_manager_v2` method signature. The parameter `kv_cache_manager_cls` is currently untyped and needs a type annotation (likely a class type or type object representing a KV cache manager class). Additionally, add an explicit return type annotation to the method; based on the method name and purpose, this should be `None` if the method does not return a value. This will strengthen static type checking around the manager-class routing logic.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/resource_manager.py (1)
2353-2353: ⚡ Quick winAdd a return annotation to the module
__getattr__.This new function is part of the module import surface but omits its return type.
Proposed fix
-def __getattr__(name: str): +def __getattr__(name: str) -> object:As per coding guidelines, “Always annotate functions. Make the return type
Noneif the function does not return anything.”🤖 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 `@tensorrt_llm/_torch/pyexecutor/resource_manager.py` at line 2353, Add a return type annotation to the module-level __getattr__ function definition at line 2353. The function currently has a parameter annotation for the name parameter but is missing the return type annotation after the closing parenthesis. Add the appropriate return type annotation using the arrow syntax (e.g., -> ReturnType) to indicate what the __getattr__ function returns, following the project's coding guidelines that require all functions to have explicit return type annotations.Source: Coding guidelines
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py (1)
387-488: Coverage is strong for this module; add one multi-rank follow-up test outside this PR.This file already covers serialization/hash/coalescing/lifecycle and DP-gather semantics well. Follow-up suggestion: add a real multi-rank integration case in
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.pyto validate attention-DP gathering across processes/devices (not only mocked gather behavior).As per coding guidelines, reviews under
tests/**should state coverage sufficiency and provide concrete file-level follow-up actions.🤖 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 `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py` around lines 387 - 488, This is a follow-up action for future work, not a fix to current code. Add a new integration test in tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py that validates real multi-rank attention-DP gathering behavior across multiple processes or devices (not using mocked gather functions like the existing tests do). The test should verify that KVCacheEventManager correctly handles actual distributed gathering when attention_dp_rank varies across multiple ranks and ensure event coalescence and capacity management work correctly in a true multi-rank scenario.Source: Coding guidelines
🤖 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 `@tensorrt_llm/_torch/pyexecutor/py_executor.py`:
- Line 2105: The _commit_kv_cache_stats() method calls are timing stats commits
incorrectly: they commit the newly scheduled batch before the older executed
batch is processed, causing newer batch deltas to attach to older IterationStats
records through the draining behavior of _update_iter_stats(). Move the
_commit_kv_cache_stats() call at the anchor location (line 2105) to execute
after the older batch has been processed rather than before the new batch is
scheduled. Apply the same fix at the sibling location (line 3336). For the
overlap batch handling at lines 3404-3407, ensure that only the previous batch
about to be consumed is committed for stats in the PP/overlap scenario, not
newer scheduled batches, while keeping the non-overlap commit logic intact.
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Around line 92-100: The helper function
_warn_if_unsupported_v1_kv_cache_event_hash_algo is defined but is never
invoked, causing V1 manager runs with unsupported kv_cache_event_hash_algo
values to proceed silently without warning. Add a call to
_warn_if_unsupported_v1_kv_cache_event_hash_algo in the KVCacheManager.__init__
method, passing the kv_cache_event_hash_algo parameter, before the V1 event
manager is constructed to ensure users are warned when using an unsupported hash
algorithm configuration.
In `@tensorrt_llm/runtime/kv_cache_hash.py`:
- Around line 31-32: The truncate_sha256_hash_to_int64 function accepts any byte
length without validation, even though SHA-256 digests are always exactly 32
bytes. Add a validation check at the start of the function that verifies the
block_hash parameter has a length of exactly 32 bytes, and raise a ValueError
with a descriptive message if the length is incorrect. This ensures the
function's contract matches its name and prevents silent truncation of incorrect
inputs that could mask upstream errors.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py`:
- Around line 419-425: The __slots__ tuple in the BlockRadixTree class needs to
be sorted alphabetically to comply with Ruff's RUF023 rule. Reorder the slot
names ("_life_cycles", "_tokens_per_block", "_event_manager", "next",
"__rawref__") in alphabetical order within the tuple definition to satisfy the
linter and avoid CI-only failures.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py`:
- Around line 802-806: In the is_enough function, modify the zip() call that
pairs get_num_slots(num_blocks * tokens_per_block) with remaining_slots by
adding the strict=True parameter. This ensures both iterables have the same
length and makes the invariant explicit to satisfy the Ruff B905 linting rule
for paired comparisons.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py`:
- Line 63: The __init__ method in the test class has a default argument that
calls CacheLevel(0) at function definition time, violating Ruff B008. Remove the
function call from the default argument by changing cache_level=CacheLevel(0) to
cache_level=None, then add logic inside the __init__ method body to initialize
cache_level to CacheLevel(0) when it is None.
In `@tests/unittest/llmapi/test_llm_kv_cache_events.py`:
- Around line 920-925: The list comprehensions that filter created_events and
stored_events unconditionally index into event["data"]["type"] without checking
if event is None or if the event has the required data structure. Add guards to
filter out None events and verify the event contains a "data" key before
accessing event["data"]["type"]. This will prevent TypeError exceptions and make
the test more robust when llm.get_kv_cache_events() returns optional or
incomplete event entries.
---
Outside diff comments:
In `@tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py`:
- Around line 79-89: Update the python_requires parameter in setup_mypyc.py to
match the main setup.py requirement. Change the python_requires value from
">=3.8" to ">=3.10, <4" to align with the repository's Python 3.10+ target,
since the newly added modules in the mypyc configuration use Python 3.10+ syntax
features.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/_util.py`:
- Around line 270-273: Add missing type annotations to the
`_fallback_if_unsupported_kv_cache_manager_v2` method signature. The parameter
`kv_cache_manager_cls` is currently untyped and needs a type annotation (likely
a class type or type object representing a KV cache manager class).
Additionally, add an explicit return type annotation to the method; based on the
method name and purpose, this should be `None` if the method does not return a
value. This will strengthen static type checking around the manager-class
routing logic.
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 1995-2005: In the update_resources method, change the type hints
for the parameters attn_metadata and kv_cache_dtype_byte_size to use explicit
union syntax with None. Modify attn_metadata from "AttentionMetadata" to
"AttentionMetadata" | None, and modify kv_cache_dtype_byte_size from float to
float | None, to explicitly declare that these parameters accept None values
instead of relying on implicit Optional inference.
In `@tensorrt_llm/_torch/pyexecutor/resource_manager.py`:
- Line 2353: Add a return type annotation to the module-level __getattr__
function definition at line 2353. The function currently has a parameter
annotation for the name parameter but is missing the return type annotation
after the closing parenthesis. Add the appropriate return type annotation using
the arrow syntax (e.g., -> ReturnType) to indicate what the __getattr__ function
returns, following the project's coding guidelines that require all functions to
have explicit return type annotations.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py`:
- Around line 22-37: The methods `add`, `subtract`, and `copy` in the
`_StatsDeltaMixin` class are missing type annotations. Add type annotation for
the `other` parameter in both the `add` and `subtract` methods (should be typed
as the same class type since they operate on instances of the same type), and
add a return type annotation to the `copy` method (should return the same class
type). This will ensure static type checkers can properly validate usage of
these methods throughout the codebase.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py`:
- Around line 387-488: This is a follow-up action for future work, not a fix to
current code. Add a new integration test in
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py that
validates real multi-rank attention-DP gathering behavior across multiple
processes or devices (not using mocked gather functions like the existing tests
do). The test should verify that KVCacheEventManager correctly handles actual
distributed gathering when attention_dp_rank varies across multiple ranks and
ensure event coalescence and capacity management work correctly in a true
multi-rank scenario.
🪄 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: 679682d4-6f2a-4627-b048-f73f56101672
📒 Files selected for processing (33)
cpp/include/tensorrt_llm/batch_manager/llmRequest.hcpp/tensorrt_llm/nanobind/batch_manager/bindings.cpptensorrt_llm/_torch/pyexecutor/_util.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/_torch/pyexecutor/kv_cache_stats.pytensorrt_llm/_torch/pyexecutor/model_engine.pytensorrt_llm/_torch/pyexecutor/py_executor.pytensorrt_llm/_torch/pyexecutor/resource_manager.pytensorrt_llm/_utils.pytensorrt_llm/executor/base_worker.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/metrics/collector.pytensorrt_llm/runtime/kv_cache_hash.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pytensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyitensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.pytensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.pytensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytensorrt_llm/runtime/kv_cache_manager_v2/_stats.pytensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.pytests/unittest/_torch/executor/test_resource_manager.pytests/unittest/bindings/test_bindings_ut.pytests/unittest/executor/test_stats_serializer.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.pytests/unittest/llmapi/test_llm_args.pytests/unittest/llmapi/test_llm_kv_cache_events.pytests/unittest/metrics/test_collector.py
|
/bot run --disable-fail-fast |
|
PR_Github #54311 [ run ] triggered by Bot. Commit: |
|
PR_Github #54311 [ run ] completed with state
|
49c5348 to
da71066
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #54410 [ run ] triggered by Bot. Commit: |
|
PR_Github #54410 [ run ] completed with state
|
|
/bot run --disable-fail-fast |
|
PR_Github #54600 [ run ] triggered by Bot. Commit: |
|
PR_Github #54600 [ run ] completed with state
|
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com>
Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com>
81517ec to
70f0917
Compare
|
/bot run --disable-fail-fast |
|
PR_Github #55381 [ run ] triggered by Bot. Commit: |
|
PR_Github #55381 [ run ] completed with state
|
|
/bot run |
1 similar comment
|
/bot run |
|
PR_Github #55467 [ run ] triggered by Bot. Commit: |
|
PR_Github #55467 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55477 [ run ] triggered by Bot. Commit: |
|
PR_Github #55477 [ run ] completed with state
|
|
/bot run |
|
PR_Github #55485 [ run ] triggered by Bot. Commit: |
galagam
left a comment
There was a problem hiding this comment.
AutoDeploy change looks OK. Did not review the entire PR.
|
PR_Github #55485 [ run ] completed with state
|
|
/bot run --stage-list "GB200-4_GPUs-PyTorch-PerfSanity-1, GB200-4_GPUs-PyTorch-PerfSanity-2" |
|
PR_Github #55486 [ run ] triggered by Bot. Commit: |
|
PR_Github #55486 [ run ] completed with state |
|
/bot reuse-pipeline |
|
PR_Github #55511 [ reuse-pipeline ] triggered by Bot. Commit: |
|
PR_Github #55511 [ reuse-pipeline ] completed with state |
Signed-off-by: Fanrong Li <lfr-0531@users.noreply.github.com> Signed-off-by: Jiagan Cheng <jiaganc@nvidia.com> Signed-off-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com> Co-authored-by: Fanrong Li <lfr-0531@users.noreply.github.com> Co-authored-by: Jiagan Cheng <jiaganc@nvidia.com> Co-authored-by: Yi Zhang <187001205+yizhang-nv@users.noreply.github.com>
This is PR-1 from the DSv4 umbrella split of #14751. It lands generic runtime/KV foundations only, without compressor, sparse attention, MoE, model, or tokenizer files.
Scope:
Verification on B300 GPU 0:
python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --cuda_architectures "90-real;100-real" --configure_cmake..venv-3.12withpip install --force-reinstall --no-deps build/tensorrt_llm-1.3.0rc18-cp312-cp312-linux_x86_64.whl./home/scratch.fanrongl_coreai/TRT_LLM_dsv4_split/tensorrt_llm/...and wheel imports from.venv-3.12/lib/python3.12/site-packages/tensorrt_llm/....CUDA_VISIBLE_DEVICES=0 timeout 30m .venv-3.12/bin/python -m pytest -q --tb=short -ra tests/unittest/kv_cache_manager_v2_tests: 103 passed, 12 skipped, 5 warnings.CUDA_VISIBLE_DEVICES=0 timeout 35m .venv-3.12/bin/python -m pytest -q --tb=short -ra tests/unittest/llmapi/test_llm_kv_cache_events.py: 34 passed, 1 skipped, 3 warnings.CUDA_VISIBLE_DEVICES=0 timeout 30m .venv-3.12/bin/python -m pytest -q --tb=short -ra tests/unittest/metrics/test_collector.py tests/unittest/executor/test_stats_serializer.py tests/unittest/_torch/executor/test_resource_manager.py tests/unittest/bindings/test_bindings_ut.py: 123 passed, 2 warnings, 7 subtests passed.python3 -m pre_commit run --files $(git diff --cached --name-only): passed before commit; commit hooks also passed.git diff --name-only github/main...HEAD | rg 'compressor|mhc|IndexerTopK|indexerTopK|attention_backend/sparse/deepseek_v4|fused_moe|modeling_deepseekv4|tokenizer/deepseek_v4'returned no output.Notes:
use_kv_cache_manager_v2=Trueandevent_buffer_max_size > 0; V1 event tests continue to pass.Summary by CodeRabbit
New Features
Improvements