Skip to content

[None][feat] DSv4 prep: runtime cache foundations#15378

Merged
jiaganc merged 16 commits into
NVIDIA:mainfrom
lfr-0531:user/fanrongl/dsv4-runtime-kv
Jun 24, 2026
Merged

[None][feat] DSv4 prep: runtime cache foundations#15378
jiaganc merged 16 commits into
NVIDIA:mainfrom
lfr-0531:user/fanrongl/dsv4-runtime-kv

Conversation

@lfr-0531

@lfr-0531 lfr-0531 commented Jun 15, 2026

Copy link
Copy Markdown
Collaborator

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:

  • KV cache manager V2 runtime events/stats/hash plumbing and pyexecutor/resource-manager integration.
  • KV cache event/stat serialization and metrics collector coverage.
  • Minimal C++ request/nanobind fields needed by KV cache transfer metrics.
  • Focused unit coverage for V2 KV cache events/stats, resource-manager behavior, metrics, serializer, LLM args, and bindings.

Verification on B300 GPU 0:

  • Built C++/nanobind changes with python3 ./scripts/build_wheel.py --trt_root /usr/local/tensorrt --benchmarks --use_ccache --cuda_architectures "90-real;100-real" --configure_cmake.
  • Installed wheel into .venv-3.12 with pip install --force-reinstall --no-deps build/tensorrt_llm-1.3.0rc18-cp312-cp312-linux_x86_64.whl.
  • Import/path smoke confirmed repo source imports from /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.
  • Scope check 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:

  • V2 KV events remain opt-in via use_kv_cache_manager_v2=True and event_buffer_max_size > 0; V1 event tests continue to pass.
  • One multi-GPU event test was skipped locally because it requires at least 2 GPUs.

Summary by CodeRabbit

  • New Features

    • Added KV cache performance statistics tracking and reporting, including allocation, reuse, and miss counters.
    • Introduced KV cache event management system for detailed lifecycle tracking.
    • Added configurable hash algorithm selection for KV cache events.
    • Extended iteration statistics reporting with per-pool-group and per-lifecycle breakdowns.
  • Improvements

    • Enhanced KV cache manager with detailed performance metrics visibility.
    • Added KV cache transfer timestamp tracking.
    • Improved metrics aggregation and collection across executor workflows.

@lfr-0531
lfr-0531 requested review from a team as code owners June 15, 2026 13:27
@lfr-0531
lfr-0531 requested a review from schetlur-nv June 15, 2026 13:27
@coderabbitai

coderabbitai Bot commented Jun 15, 2026

Copy link
Copy Markdown
Contributor

Review Change Stack

📝 Walkthrough

Walkthrough

This 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 _KVCache, a KVCacheEventManager for lifecycle event emission, migration recording through StorageManager and BlockRadixTree, extended KVCacheManager stats APIs, and wires all of this through the pyexecutor loop and Prometheus metrics export.

Changes

KV Cache V2 Statistics, Events, and Hash Support

Layer / File(s) Summary
Stats delta types, hash constants, and config fields
tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py, tensorrt_llm/runtime/kv_cache_hash.py, tensorrt_llm/runtime/kv_cache_manager_v2/_config.py, tensorrt_llm/llmapi/llm_args.py, tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py, tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi
Adds KVCacheStatsDelta/KVCacheIterationStatsDelta with an arithmetic mixin; introduces hash algorithm constants and get_effective_kv_cache_event_hash_algo/truncate_sha256_hash_to_int64; adds enable_stats to KVCacheManagerConfig and kv_cache_event_hash_algo to KvCacheConfig; re-exports new types from the package root with updated stubs.
Pending stats buffering inside _KVCache
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py, tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
Introduces _PendingAllocationSegment, _PendingStatsDelta, and _PendingStats for buffered block-range and reuse accounting; wires these into _KVCache resize (allocation recording/subtraction, generation alloc readiness), resume (migration recorder, intra-device copy stats), reuse setup (full/partial reuse classification), and commit/close paths.
KVCacheEventManager and event payload types
tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py, tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py
Adds a new thread-safe event manager with event coalescing, batched removal, attention-DP gather, and get_latest_events; defines immutable event payload dataclasses (KVCacheCreatedData, KVCacheStoredData, KVCacheRemovedData, KVCacheUpdatedData, KVCacheEvent); registers new modules for mypyc compilation.
Migration recording and cache-level updated events
tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py, tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
Adds MigrationRecorder callable type; threads migration callback through all StorageManager slot allocation and eviction paths into _batched_migrate; adds _emit_cache_level_updated_event and copy-stream sync; extends batched_lock_to_gpu with the migration recorder parameter.
Event hooks in BlockRadixTree
tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
Wires KVCacheEventManager into BlockRadixTree; emits removed-block events during remove_subtree, sibling pruning, and unset_page predecessor pruning; emits stored-block and stored-lifecycle events from _commit_block; relaxes RootBlock.make_key to accept nullable reuse_scope.
KVCacheManager stats API and event manager wiring
tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
Adds optional event_manager parameter, stats fields (dirty/excluded sets, committed stats, per-lifecycle iteration stats), commit_stats, get_committed_stats, get_and_reset_iteration_stats, dirty/excluded tracking methods, event_manager property, expected_prompt_length parameter on create_kv_cache; changes clamp_max_seq_len_for_mem to return 0 gracefully instead of asserting.
KVCacheManagerV2 event/stats init and iteration reporting
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
Adds enable_stats flag, conditional _event_manager construction, threads enable_stats into config; constructs KVCacheManagerPy with event manager; configures window sizes and emits created/iteration events; reworks get_kv_cache_stats from committed storage, adds report-building helpers, updates commit_scheduled_kv_cache_stats/get_iteration_stats/event flush; adds update_context_resources.
V2 iteration stats serialization
tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py
Adds KVCacheV2PoolGroupIterationStats, KVCacheV2LifeCycleIterationStats, KVCacheV2IterationStatsReport; adds serialize_kv_cache_iteration_stats and append_kv_cache_iteration_stats for injecting stats into output dicts under up to three keys.
Executor loop, serializer, and metrics integration
tensorrt_llm/_torch/pyexecutor/py_executor.py, tensorrt_llm/executor/base_worker.py, tensorrt_llm/metrics/collector.py, tensorrt_llm/_utils.py, tensorrt_llm/_torch/pyexecutor/resource_manager.py, tensorrt_llm/_torch/pyexecutor/_util.py, tensorrt_llm/_torch/pyexecutor/model_engine.py
Adds _commit_kv_cache_stats helper and inserts it in PP/non-PP/overlap executor paths; delegates kvCacheIterationStats injection to append_kv_cache_iteration_stats in executor and worker; updates MetricsCollector to aggregate from V2 lifecycle/pool-group fields with legacy fallback; extends event JSON serialization with hash_algo/layer_group_id; adds hash-algo warning helper and lazy __getattr__; centralizes V2 fallback logic; fixes V2 scheduler capacity to use max_batch_size.
C++ helper and Python bindings
cpp/include/tensorrt_llm/batch_manager/llmRequest.h, cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
Adds updateKvCachePerfMetrics batch helper to GenericLlmRequest; exposes kv_cache_transfer_start/end properties, transfer timestamp setters, size mutators, and update_kv_cache_perf_metrics in Python bindings.
Tests
tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py, tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py, tests/unittest/executor/test_stats_serializer.py, tests/unittest/bindings/test_bindings_ut.py, tests/unittest/_torch/executor/test_resource_manager.py, tests/unittest/llmapi/test_llm_args.py, tests/unittest/llmapi/test_llm_kv_cache_events.py, tests/unittest/metrics/test_collector.py
Covers event manager serialization/coalescing/hash modes/DP gather/lifecycle events (including CUDA-gated block chain tests), stats correctness across enable/disable/reuse/partial-reuse/SWA/pool-group scenarios, V2 pool-group serializer, binding transfer metrics, hash-algo warning behavior, kv_cache_event_hash_algo field, V2 event hash validation, and V2 metrics aggregation.

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)
Loading

Estimated code review effort

🎯 5 (Critical) | ⏱️ ~120 minutes

Suggested labels

deepseek-v4

Suggested reviewers

  • niukuo
  • zeroepoch
  • yizhang-nv
  • yuxianq
  • dongxuy04
  • schetlur-nv
  • chang-l
  • hchings
  • thorjohnsen
  • Superjomn
  • tburt-nv
✨ Finishing Touches
🧪 Generate unit tests (beta)
  • Create PR with unit tests

@coderabbitai coderabbitai Bot left a comment

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Actionable comments posted: 7

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 | 🟡 Minor

Align mypyc package metadata with the repository's Python 3.10+ target.

The setup_mypyc.py currently declares python_requires=">=3.8" (line 139) while the main setup.py declares >=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 win

Type-annotate _StatsDeltaMixin operation signatures.

add, subtract, and copy are 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 value

Use explicit | None type hint instead of implicit Optional.

Per PEP 484 and the static analysis hint, kv_cache_dtype_byte_size: float = None is 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 win

Annotate the new fallback helper signature.

kv_cache_manager_cls and 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 None if 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 win

Add 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 None if 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.py to 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

📥 Commits

Reviewing files that changed from the base of the PR and between 20b6068 and da71066.

📒 Files selected for processing (33)
  • cpp/include/tensorrt_llm/batch_manager/llmRequest.h
  • cpp/tensorrt_llm/nanobind/batch_manager/bindings.cpp
  • tensorrt_llm/_torch/pyexecutor/_util.py
  • tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py
  • tensorrt_llm/_torch/pyexecutor/kv_cache_stats.py
  • tensorrt_llm/_torch/pyexecutor/model_engine.py
  • tensorrt_llm/_torch/pyexecutor/py_executor.py
  • tensorrt_llm/_torch/pyexecutor/resource_manager.py
  • tensorrt_llm/_utils.py
  • tensorrt_llm/executor/base_worker.py
  • tensorrt_llm/llmapi/llm_args.py
  • tensorrt_llm/metrics/collector.py
  • tensorrt_llm/runtime/kv_cache_hash.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi
  • tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/_config.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/_core/_pending_stats.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/_event_manager.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/_page.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/_stats.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.py
  • tensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.py
  • tests/unittest/_torch/executor/test_resource_manager.py
  • tests/unittest/bindings/test_bindings_ut.py
  • tests/unittest/executor/test_stats_serializer.py
  • tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py
  • tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py
  • tests/unittest/llmapi/test_llm_args.py
  • tests/unittest/llmapi/test_llm_kv_cache_events.py
  • tests/unittest/metrics/test_collector.py

Comment thread tensorrt_llm/_torch/pyexecutor/py_executor.py Outdated
Comment thread tensorrt_llm/_torch/pyexecutor/resource_manager.py Outdated
Comment thread tensorrt_llm/runtime/kv_cache_hash.py
Comment thread tensorrt_llm/runtime/kv_cache_manager_v2/_block_radix_tree.py
Comment thread tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.py
Comment thread tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.py Outdated
Comment thread tests/unittest/llmapi/test_llm_kv_cache_events.py Outdated
@lfr-0531

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #54311 [ run ] triggered by Bot. Commit: da71066 Link to invocation

@lfr-0531
lfr-0531 requested a review from jiaganc June 15, 2026 16:02
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #54311 [ run ] completed with state SUCCESS. Commit: da71066
/LLM/main/L0_MergeRequest_PR pipeline #43382 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@lfr-0531
lfr-0531 force-pushed the user/fanrongl/dsv4-runtime-kv branch from 49c5348 to da71066 Compare June 16, 2026 01:41
@lfr-0531

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #54410 [ run ] triggered by Bot. Commit: 3d4e8bd Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #54410 [ run ] completed with state FAILURE. Commit: 3d4e8bd
/LLM/main/L0_MergeRequest_PR pipeline #43479 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@lfr-0531

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #54600 [ run ] triggered by Bot. Commit: 81517ec Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #54600 [ run ] completed with state FAILURE. Commit: 81517ec
/LLM/main/L0_MergeRequest_PR pipeline #43642 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

Link to invocation

lfr-0531 and others added 5 commits June 17, 2026 01:05
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>
@lfr-0531
lfr-0531 force-pushed the user/fanrongl/dsv4-runtime-kv branch from 81517ec to 70f0917 Compare June 17, 2026 02:52
@lfr-0531

Copy link
Copy Markdown
Collaborator Author

/bot run --disable-fail-fast

@lfr-0531 lfr-0531 closed this Jun 17, 2026
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55381 [ run ] triggered by Bot. Commit: b5798b6 Link to invocation

@jiaganc
jiaganc requested a review from lowsfer June 24, 2026 03:16
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55381 [ run ] completed with state FAILURE. Commit: b5798b6
/LLM/main/L0_MergeRequest_PR pipeline #44331 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@jiaganc

jiaganc commented Jun 24, 2026

Copy link
Copy Markdown
Collaborator

/bot run

1 similar comment
@jiaganc

jiaganc commented Jun 24, 2026

Copy link
Copy Markdown
Collaborator

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55467 [ run ] triggered by Bot. Commit: b5798b6 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55467 [ run ] completed with state FAILURE. Commit: b5798b6
/LLM/main/L0_MergeRequest_PR pipeline #44395 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@Superjomn Superjomn left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

LGTM

@Superjomn

Copy link
Copy Markdown
Collaborator

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55477 [ run ] triggered by Bot. Commit: b5798b6 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55477 [ run ] completed with state FAILURE. Commit: b5798b6
/LLM/main/L0_MergeRequest_PR pipeline #44404 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@jiaganc

jiaganc commented Jun 24, 2026

Copy link
Copy Markdown
Collaborator

/bot run

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55485 [ run ] triggered by Bot. Commit: b5798b6 Link to invocation

@jiaganc
jiaganc enabled auto-merge (squash) June 24, 2026 10:40

@galagam galagam left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

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

AutoDeploy change looks OK. Did not review the entire PR.

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55485 [ run ] completed with state FAILURE. Commit: b5798b6
/LLM/main/L0_MergeRequest_PR pipeline #44410 completed with status: 'FAILURE'

CI Report

⚠️ Action Required:

  • Please check the failed tests and fix your PR
  • If you cannot view the failures, ask the CI triggerer to share details
  • Once fixed, request an NVIDIA team member to trigger CI again

CI Agent Failure Analysis

Link to invocation

@jiaganc

jiaganc commented Jun 24, 2026

Copy link
Copy Markdown
Collaborator

/bot run --stage-list "GB200-4_GPUs-PyTorch-PerfSanity-1, GB200-4_GPUs-PyTorch-PerfSanity-2"

@jiaganc
jiaganc disabled auto-merge June 24, 2026 11:11
@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55486 [ run ] triggered by Bot. Commit: cd8c038 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55486 [ run ] completed with state SUCCESS. Commit: cd8c038
/LLM/main/L0_MergeRequest_PR pipeline #44411 (Partly Tested) completed with status: 'SUCCESS'

CI Report

Link to invocation

jiaganc commented Jun 24, 2026

Copy link
Copy Markdown
Collaborator

/bot reuse-pipeline

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55511 [ reuse-pipeline ] triggered by Bot. Commit: cd8c038 Link to invocation

@tensorrt-cicd

Copy link
Copy Markdown
Collaborator

PR_Github #55511 [ reuse-pipeline ] completed with state SUCCESS. Commit: cd8c038
Reusing PR_Github #55486 (Partly Tested) for commit cd8c038

Link to invocation

@jiaganc
jiaganc merged commit eb0cbdb into NVIDIA:main Jun 24, 2026
7 checks passed
BrianLi23 pushed a commit to BrianLi23/TensorRT-LLM that referenced this pull request Jul 9, 2026
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>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

9 participants