[None][chore] KVCacheManagerV2: Python and test preparation for a C++ backend#16218
Open
lowsfer wants to merge 2 commits into
Open
[None][chore] KVCacheManagerV2: Python and test preparation for a C++ backend#16218lowsfer wants to merge 2 commits into
lowsfer wants to merge 2 commits into
Conversation
Contributor
📝 WalkthroughWalkthroughThis change adds Python/C++ KV-cache V2 backend selection and introspection, exposes pool layout descriptors, updates lifecycle, statistics, storage, radix-tree, and event handling, and rewrites disaggregation page-table construction. Tests and configuration manifests are updated for the new APIs and Blake3 hash algorithms. ChangesKV cache V2 platform and storage
Estimated code review effort: 5 (Critical) | ~120 minutes Sequence Diagram(s)sequenceDiagram
participant KVCacheManager
participant Introspection
participant CppBackend
participant Disaggregation
participant EventManager
KVCacheManager->>Introspection: request lifecycle and storage metadata
Introspection->>CppBackend: delegate when native backend is active
CppBackend-->>Introspection: return statistics and descriptors
Introspection-->>KVCacheManager: provide backend-independent metadata
Disaggregation->>KVCacheManager: read pool_group_descs and init_config
Disaggregation-->>Disaggregation: build page tables from buffer offsets
EventManager->>KVCacheManager: resolve page references and event keys
Possibly related PRs
Suggested reviewers: 🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 5
🧹 Nitpick comments (2)
tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py (1)
100-128: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winAdd type annotations to the new fallback functions.
_load_cpp_module,_RawRef.__init__,_RawRef.__call__, and_RawRef.__class_getitem__omit parameter and/or return annotations. As per coding guidelines, Python functions must always be annotated.Also applies to: 205-217
🤖 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/__init__.py` around lines 100 - 128, Add complete parameter and return type annotations to _load_cpp_module, _RawRef.__init__, _RawRef.__call__, and _RawRef.__class_getitem__, including appropriate self, argument, and return types consistent with their existing behavior and surrounding typing conventions.Source: Coding guidelines
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py (1)
1245-1251: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueUse an f-string conversion flag.
Ruff reports RUF010 for
str(role)inside the f-string.Proposed fix
- f"role={str(role)}, lifecycle_id={int(lifecycle_id)}, " + f"role={role!s}, lifecycle_id={int(lifecycle_id)}, "🤖 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 1245 - 1251, Replace str(role) in the return expression of the relevant KV cache manager method with the f-string conversion flag !s, preserving the existing output while resolving Ruff RUF010.Source: Linters/SAST tools
🤖 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/kv_cache_manager_v2.py`:
- Around line 860-870: Preserve the configured disk cache tier when host
registration falls back to GPU-only tiers. Update the retry/filter logic in the
KV cache manager initialization, near the host registration handling, to retain
both GPU and disk tiers while excluding only the failed host tier; ensure the
existing DiskCacheTierConfig created from disk_cache_size and disk_cache_path
remains available for suspended request resumption.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py`:
- Around line 119-128: In the backend-loading function containing the
kv_cache_manager_v2 import, replace the assert validating
find_spec("kv_cache_manager_v2") with an explicit check for a missing spec or
origin, and raise a descriptive ImportError before constructing
Path(spec.origin). Preserve the existing path setup and cleanup for valid specs.
- Line 166: Guard the stats helpers that use `_cpp_introspection` so native
`KVCacheManager` does not fall back to the private `self.impl._storage`
contract. Update the relevant methods in `KVCacheManager` to gate storage-based
paths by backend or fail fast when native introspection is unavailable, while
preserving the Python manager fallback.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 2648-2686: Make
test_excess_scratch_slot_waits_for_ready_event_on_new_stream deterministic by
gating producer completion after its work is enqueued, then assert
consumer_marker does not complete while that gate is held and does complete
after releasing it. Replace the current producer_marker.query_complete assertion
with explicit producer-side synchronization/control that proves the consumer
waits for the producer’s ready event across streams.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.py`:
- Around line 90-117: Extend test_manager_stats_config_and_api to exercise both
collection modes: after creating the cache, resume it, resize it, and commit it,
then assert enable_stats=True produces non-empty allocation/reuse counters in
the committed or iteration stats, while enable_stats=False keeps those stats
empty. Retain the existing configuration and dirty/excluded-state assertions,
and use the cache lifecycle methods exposed by the test fixture.
---
Nitpick comments:
In `@tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py`:
- Around line 1245-1251: Replace str(role) in the return expression of the
relevant KV cache manager method with the f-string conversion flag !s,
preserving the existing output while resolving Ruff RUF010.
In `@tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py`:
- Around line 100-128: Add complete parameter and return type annotations to
_load_cpp_module, _RawRef.__init__, _RawRef.__call__, and
_RawRef.__class_getitem__, including appropriate self, argument, and return
types consistent with their existing behavior and surrounding typing
conventions.
🪄 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: 97de4c46-d2af-46d1-b2e9-4872b9a0986f
📒 Files selected for processing (36)
tensorrt_llm/_torch/attention_backend/sparse/deepseek_v4/cache_manager.pytensorrt_llm/_torch/disaggregation/resource/kv_extractor.pytensorrt_llm/_torch/disaggregation/resource/utils.pytensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/runtime/kv_cache_hash.pytensorrt_llm/runtime/kv_cache_manager_v2/AGENTS.mdtensorrt_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/_common.pytensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_copy_engine.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.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/_event_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_eviction_controller/_eviction_controller.pytensorrt_llm/runtime/kv_cache_manager_v2/_introspection.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytensorrt_llm/runtime/kv_cache_manager_v2/_stats.pytensorrt_llm/runtime/kv_cache_manager_v2/_storage/_core.pytensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_utils.pytensorrt_llm/runtime/kv_cache_manager_v2/setup_mypyc.pytensorrt_llm/usage/llm_args_golden_manifest.jsontests/integration/defs/accuracy/test_kv_pool_rebalance_accuracy.pytests/integration/defs/perf/create_perf_comparison_report.pytests/unittest/_torch/attention/sparse/deepseek_v4/test_deepseek_v4_cache_manager.pytests/unittest/disaggregated/test_kv_transfer.pytests/unittest/kv_cache_manager_v2_tests/fake_engine.pytests/unittest/kv_cache_manager_v2_tests/test_branch_reuse.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_event_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_api.pytests/unittest/llmapi/test_llm_args.py
💤 Files with no reviewable changes (1)
- tensorrt_llm/runtime/kv_cache_manager_v2/_config.py
Comment on lines
+860
to
+870
| cache_tiers.append(HostCacheTierConfig(quota=int(host_quota))) | ||
| logger.info( | ||
| f"KV cache manager v2 host cache quota set to {host_quota / (1 << 30):.2f}GiB" | ||
| ) | ||
| disk_cache_size = kv_cache_config.disk_cache_size | ||
| if disk_cache_size is not None and disk_cache_size > 0: | ||
| disk_cache_path = kv_cache_config.disk_cache_path | ||
| assert disk_cache_path is not None | ||
| cache_tiers.append(DiskCacheTierConfig(quota=disk_cache_size, path=disk_cache_path)) | ||
| cache_tiers.append( | ||
| DiskCacheTierConfig(quota=int(disk_cache_size), path=disk_cache_path) | ||
| ) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Preserve the configured disk tier during host-tier fallback.
If host registration fails, the existing retry at Line 895 filters to GPU-only tiers. A configured disk tier is therefore silently removed too, leaving no secondary tier and potentially preventing suspended requests from resuming.
Proposed fix
- cache_tiers_gpu_only = [t for t in cache_tiers if isinstance(t, GpuCacheTierConfig)]
+ cache_tiers_without_host = [
+ tier
+ for tier in cache_tiers
+ if not isinstance(tier, HostCacheTierConfig)
+ ]
config = self._build_cache_config(
kv_cache_config,
tokens_per_block=tokens_per_block,
vocab_size=vocab_size,
- cache_tiers=cache_tiers_gpu_only,
+ cache_tiers=cache_tiers_without_host,
)
- cache_tiers = cache_tiers_gpu_only
+ cache_tiers = cache_tiers_without_host🤖 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 860 -
870, Preserve the configured disk cache tier when host registration falls back
to GPU-only tiers. Update the retry/filter logic in the KV cache manager
initialization, near the host registration handling, to retain both GPU and disk
tiers while excluding only the failed host tier; ensure the existing
DiskCacheTierConfig created from disk_cache_size and disk_cache_path remains
available for suspended request resumption.
Comment on lines
+119
to
+128
| spec = find_spec("kv_cache_manager_v2") | ||
| assert spec is not None and spec.origin is not None | ||
| trtllm_root = str(Path(spec.origin).parent.parent.parent) | ||
| sys.path.insert(0, trtllm_root) | ||
| try: | ||
| from bindings.internal.batch_manager import kv_cache_manager_v2 | ||
|
|
||
| return kv_cache_manager_v2 | ||
| finally: | ||
| sys.path.remove(trtllm_root) |
Contributor
There was a problem hiding this comment.
🩺 Stability & Availability | 🟡 Minor | ⚡ Quick win
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
set -euo pipefail
file="tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py"
echo "== file outline =="
ast-grep outline "$file" --view expanded || true
echo
echo "== relevant lines =="
sed -n '1,220p' "$file" | cat -n
echo
echo "== search for related import/error handling =="
rg -n "find_spec\\(|ImportError|assert spec is not None|kv_cache_manager_v2|bindings.internal.batch_manager" "$file" . || trueRepository: NVIDIA/TensorRT-LLM
Length of output: 50376
Raise ImportError when the native backend can’t be resolved. If find_spec() returns None or has no origin, raise ImportError before Path(spec.origin); assert disappears under -O and the import fails later with a less useful AttributeError.
🤖 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/__init__.py` around lines 119 - 128,
In the backend-loading function containing the kv_cache_manager_v2 import,
replace the assert validating find_spec("kv_cache_manager_v2") with an explicit
check for a missing spec or origin, and raise a descriptive ImportError before
constructing Path(spec.origin). Preserve the existing path setup and cleanup for
valid specs.
Comment on lines
+2648
to
+2686
| def test_excess_scratch_slot_waits_for_ready_event_on_new_stream(self): | ||
| num_layers = 512 | ||
| self._prepare_scratch( | ||
| num_layers=num_layers, | ||
| window_size=32, | ||
| tokens_per_block=32, | ||
| gpu_quota=16 << 20, | ||
| ) | ||
| producer_prompt = [self.next_token() for _ in range(64)] | ||
| consumer_prompt = [self.next_token() for _ in range(256)] | ||
| producer = self.manager.create_kv_cache(None, producer_prompt) | ||
| consumer = self.manager.create_kv_cache(None, consumer_prompt) | ||
| producer_stream_holder = CachedCudaStream() | ||
| consumer_stream_holder = CachedCudaStream() | ||
| producer_stream = cast(CudaStream, producer_stream_holder.handle) | ||
| consumer_stream = cast(CudaStream, consumer_stream_holder.handle) | ||
| cached_cuda_event = get_cached_cuda_event_type() | ||
| producer_marker = None | ||
|
|
||
| try: | ||
| self.assertTrue(producer.resume(producer_stream)) | ||
| self.assertTrue(producer.resize(64)) | ||
| with enable_kernel_delay(): | ||
| for _ in range(8): | ||
| self.engine.execute([Step(producer, producer_prompt, [])], producer_stream) | ||
| producer_marker = cached_cuda_event(producer_stream) | ||
| producer.close() | ||
|
|
||
| self.assertTrue(consumer.resume(producer_stream)) | ||
| self.assertTrue(consumer.resize(256)) | ||
| self.assertTrue(consumer.has_scratch_slots) | ||
|
|
||
| consumer.cuda_stream = consumer_stream | ||
| self.assertTrue(consumer.resize(288, 256)) | ||
| self.assertFalse(consumer.has_scratch_slots) | ||
|
|
||
| consumer_marker = cached_cuda_event(consumer_stream) | ||
| consumer_marker.synchronize() | ||
| self.assertTrue(producer_marker.query_complete()) |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Make the cross-stream ordering assertion deterministic.
Synchronizing consumer_marker before querying producer_marker does not prove the consumer waited: the producer may complete independently. Add a controlled producer-side gate and assert the consumer cannot complete before it is released. Coverage is currently insufficient for the regression named by this test.
🤖 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_manager_v2.py` around
lines 2648 - 2686, Make
test_excess_scratch_slot_waits_for_ready_event_on_new_stream deterministic by
gating producer completion after its work is enqueued, then assert
consumer_marker does not complete while that gate is held and does complete
after releasing it. Replace the current producer_marker.query_complete assertion
with explicit producer-side synchronization/control that proves the consumer
waits for the producer’s ready event across streams.
Source: Path instructions
Comment on lines
+90
to
+117
| @pytest.mark.parametrize("enable_stats", [False, True]) | ||
| def test_manager_stats_config_and_api(enable_stats: bool) -> None: | ||
| manager = KVCacheManager(_make_config(enable_stats=enable_stats)) | ||
| cache = None | ||
| try: | ||
| assert manager.init_config.enable_stats is enable_stats | ||
| assert manager.get_committed_stats() == KVCacheStatsDelta() | ||
| assert manager.get_and_reset_iteration_stats() == {} | ||
| peak_stats = manager.get_and_reset_iteration_peak_block_stats(GPU_LEVEL) | ||
| assert len(peak_stats) == 1 | ||
| assert peak_stats[0].available >= 0 | ||
| assert peak_stats[0].unavailable >= 0 | ||
| assert peak_stats[0].evictable >= 0 | ||
|
|
||
| manager.mark_stats_dirty(11) | ||
| manager.mark_stats_dirty(None) | ||
| assert manager.get_dirty_stats_kv_cache_ids() == {11} | ||
| manager.mark_stats_excluded(11) | ||
| assert manager.is_stats_excluded(11) | ||
| assert manager.get_dirty_stats_kv_cache_ids() == set() | ||
| manager.clear_stats_excluded(11) | ||
| assert not manager.is_stats_excluded(11) | ||
|
|
||
| cache = manager.create_kv_cache(id=17, expected_prompt_length=8) | ||
| manager.mark_stats_dirty(17) | ||
| assert cache.commit_pending_stats() == KVCacheStatsDelta() | ||
| assert manager.get_dirty_stats_kv_cache_ids() == set() | ||
| cache.discard_pending_stats() |
Contributor
There was a problem hiding this comment.
🎯 Functional Correctness | 🟡 Minor | ⚡ Quick win
Exercise the enabled-versus-disabled collection path.
Both variants only assert empty state, so this passes even if enable_stats is ignored. Resume, resize, and commit a cache, then assert enabled stats contain allocation/reuse counters while disabled stats remain empty. Coverage is insufficient for the config contract.
🤖 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_stats_api.py` around
lines 90 - 117, Extend test_manager_stats_config_and_api to exercise both
collection modes: after creating the cache, resume it, resize it, and commit it,
then assert enable_stats=True produces non-empty allocation/reuse counters in
the committed or iteration stats, while enable_stats=False keeps those stats
empty. Retain the existing configuration and dirty/excluded-state assertions,
and use the cache lifecycle methods exposed by the test fixture.
Source: Path instructions
lowsfer
force-pushed
the
kvCacheManagerV2-py-prep
branch
from
acb5620 to
fa5ef6d
Compare
lowsfer
changed the title
lowsfer
added
the
api-compatible
… backend Backend-neutral Python groundwork for the upcoming C++ implementation of KVCacheManagerV2 (introduced in a follow-up PR): - Expand the public API of tensorrt_llm.runtime.kv_cache_manager_v2: export layout descriptors (pool_group_descs: PoolDesc, PoolGroupDesc, SlotDesc, CoalescedBuffer, ExpandedBuffer), stats/eventing types, and a backend-neutral _introspection helper module. - Use the public layout/introspection APIs in the disaggregation kv_extractor and the DSv4/MiniMax sparse cache managers instead of reaching into implementation internals. - Keep vocab_size in the _build_cache_config() virtual-method contract but drop it from the generic KVCacheManagerConfig. - Add v2_blake3 / v2_blake3_64 KV cache event hash options to llm_args and regenerate the LLM args golden manifest. - Fix StagingBuffer wrap-around when the ring tail cannot satisfy min_size; misc small fixes surfaced while translating the code. - Extend unit test coverage (stats API, hash options); route tests through the public API where possible. - Fix create_perf_comparison_report.py crash when the only perf test is waived and no CSV is produced (pre-existing main issue). - Rename debug assertion env var to TLLM_DEBUG_MODE. Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
Signed-off-by: Yao Yao <lowsfer@users.noreply.github.com>
lowsfer
force-pushed
the
kvCacheManagerV2-py-prep
branch
from
8645d9f to
d9b0aa9
Compare
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58657 [ run ] triggered by Bot. Commit: |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Backend-neutral Python groundwork for the upcoming C++ implementation of KVCacheManagerV2 (the C++ backend and its dispatcher come in a follow-up PR stacked on top). This PR contains no backend-switching code — the pure-Python implementation remains the only backend.
tensorrt_llm.runtime.kv_cache_manager_v2: export layout descriptors (pool_group_descs:PoolDesc,PoolGroupDesc,SlotDesc,CoalescedBuffer,ExpandedBuffer), stats/eventing types, and a backend-neutral_introspectionhelper module.kv_extractorand the DSv4/MiniMax sparse cache managers instead of reaching into implementation internals.vocab_sizein the_build_cache_config()virtual-method contract (subclasses override it) but drop it from the genericKVCacheManagerConfig.v2_blake3/v2_blake3_64KV cache event hash options toKvCacheConfigand regenerate the LLM args golden manifest.StagingBufferwrap-around when the ring tail cannot satisfymin_size; misc small fixes surfaced while translating the code.create_perf_comparison_report.pycrash when the only perf test is waived and no CSV is produced (pre-existing main issue).TLLM_DEBUG_MODE.Test Coverage
tests/unittest/kv_cache_manager_v2_tests/(full directory): 138 passed, 12 skipped.tests/unittest/llmapi/test_llm_args.pyhash-algo tests pass.PR Checklist
🤖 Generated with Claude Code