[None][chore] KVCacheManagerV2: Python and test preparation for a C++ backend - #16218
Merged
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
lowsfer
force-pushed
the
kvCacheManagerV2-py-prep
branch
from
July 10, 2026 06:52
acb5620 to
fa5ef6d
Compare
lowsfer
changed the title
lowsfer
added
the
api-compatible
lowsfer
force-pushed
the
kvCacheManagerV2-py-prep
branch
from
July 10, 2026 10:10
8645d9f to
d9b0aa9
Compare
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58657 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58657 [ run ] completed with state
|
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58713 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58713 [ run ] completed with state
|
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #58747 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #58747 [ run ] completed with state |
2 tasks
Collaborator
Merge conflict — how to resolveThis PR is currently in a File: What conflicts:
Both symbols are genuinely in use, so there is no logical conflict — just keep both sides' imports:
After rebasing / merging |
… 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>
lowsfer
force-pushed
the
kvCacheManagerV2-py-prep
branch
from
July 14, 2026 08:39
d9b0aa9 to
5588af1
Compare
lowsfer
requested review from
BowenFu,
YihuiLu512,
chienchunhung,
dpitman-nvda,
nv-xtf,
schetlur-nv and
yiqingy0
July 14, 2026 08:40
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59222 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59222 [ run ] completed with state
|
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59350 [ run ] triggered by Bot. Commit: |
yizhang-nv
approved these changes
Jul 15, 2026
Collaborator
|
PR_Github #59350 [ run ] completed with state
|
Member
Author
|
/bot run --disable-fail-fast |
Collaborator
|
PR_Github #59400 [ run ] triggered by Bot. Commit: |
Collaborator
|
PR_Github #59400 [ run ] completed with state |
3 tasks
This was referenced Jul 16, 2026
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
9 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