Repository navigation
[TRTLLM-13617][feat] Inclusive host KV-cache tier (shadow reuse) for KVCacheManager V2 - #15828
reasonsolo wants to merge 1 commit into
Conversation
de63646 to
b75c19e
Compare
…anager V2 Add an optional inclusive host cache tier to KVCacheManager V2. When a full, immutable committed page is recalled host->GPU, its clean host copy is retained as a "shadow" instead of being freed; a later eviction of that page back to host reuses the shadow and skips the redundant device-to-host (D2H) copy. This eliminates redundant D2H rewrites of hot prefixes that repeatedly bounce GPU<->host, which dominate host-tier write traffic on agentic multi-turn workloads. Gated by a new KvCacheConfig flag enable_inclusive_host_cache (default False; exclusive paging unchanged). - _config.py / llm_args.py / kv_cache_manager_v2.py / _kv_cache_manager.py: plumb enable_inclusive_host_cache into StorageManager. - _page.py: add shadow_slot/shadow_level/shadow_node to CommittedPage; release the shadow on drop/rebase. - _storage_manager.py: retain a shadow on inclusive recall, reuse it copy-free on eviction, and reclaim shadows oldest-first under host memory pressure (per-(level,pool-group) FIFO), on shrink, and on destroy. Also fix take_finish_event() to run after the TemporaryCudaStream `with` block so the offload path does not raise on an empty copy batch. - test: behavioral test asserting a shadow is attached on recall and reused on re-eviction (host free-slot count and shadow FIFO tracked via _num_shadow_slots), on top of the existing StorageManager suspend/resume harness. Signed-off-by: Lizhi Zhou <1432185+reasonsolo@users.noreply.github.com>
b75c19e to
873e7e0
Compare
📝 WalkthroughWalkthroughThis PR adds an ChangesInclusive host cache feature
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant Test as Caller
participant SM as StorageManager
participant Page as CommittedPage
Test->>SM: _prepare_free_slots(level, pool group)
SM->>SM: _reclaim_shadows(deficit)
Test->>SM: _batched_migrate(pages)
SM->>SM: _is_inclusive_recall(page)
SM->>SM: _shadow_reuse_slot(page)
alt shadow reused
SM->>Page: rebind to retained shadow slot
else copy required
SM->>SM: batched_copy(copy_pages)
SM->>Page: _attach_shadow(page, slot)
end
Page->>SM: drop_shadow(page) on __del__
Suggested reviewers: 🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✨ Finishing Touches🧪 Generate unit tests (beta)
Comment |
There was a problem hiding this comment.
Actionable comments posted: 4
🤖 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/runtime/kv_cache_manager_v2/_storage_manager.py`:
- Around line 697-723: The exception cleanup in the shadow-slot rebinding path
is releasing slots that have already been rebound to src in the update_src and
not defrag branch of the storage manager logic. Update the try/except handling
in _storage_manager.py so the cleanup only releases shadow slots that are still
detached and not yet reassigned, and skips any reused_shadow entry that has
already been attached back to a src page. Use the existing src.set_slot,
reused_shadow, and _levels[dst_level].storage.release flow to distinguish live
rebound slots from still-detached ones before releasing.
- Around line 560-573: The _is_inclusive_recall predicate in _storage_manager.py
is too broad because it matches any migration to a warmer tier, which lets
disk→host prefetch attach a shadow and blocks the intended host→GPU reuse path.
Tighten the condition so it only returns true for host→GPU recalls (not
disk→host), while keeping the existing checks for enablement, CommittedPage, and
missing shadow_slot; update the logic around _is_inclusive_recall and any
callers in the migration flow accordingly.
- Around line 465-472: The reclaim logic in _storage_manager.py currently only
calls _reclaim_shadows inside the deficit path, which can cause real evictable
pages to be chosen before cheaper shadow slots are freed. Update the allocation
flow in the storage manager method that computes
shadow_count/old_free_cnt/evictable_cnt so shadow slots are reclaimed first to
cover the free-slot shortfall, then recompute the available counts and only
after that calculate num_to_evict for real pages.
In `@tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py`:
- Around line 550-563: The shadow reuse test is asserting the wrong host
free-space behavior after recall: since the suspended request already holds the
host slots and they are retained as shadows, `host_free()` should remain
unchanged rather than drop. Update the assertion in
`test_kv_cache_manager_v2.py` around `host_free_before`, `total_shadows()`, and
`hot.kv_cache.resume(stream)` so it checks for no change in free host slots
after recall while keeping the existing shadow count verification.
🪄 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: a523a39e-6552-40d4-9162-ce0b2cdcb388
📒 Files selected for processing (7)
tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.pytensorrt_llm/llmapi/llm_args.pytensorrt_llm/runtime/kv_cache_manager_v2/_config.pytensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache_manager.pytensorrt_llm/runtime/kv_cache_manager_v2/_page.pytensorrt_llm/runtime/kv_cache_manager_v2/_storage_manager.pytests/unittest/kv_cache_manager_v2_tests/test_kv_cache_manager_v2.py
| # Shadows are reclaimable for free (GPU holds the real copy), so reclaim them oldest | ||
| # -first to cover any shortfall before evicting real pages or raising. | ||
| shadow_cnt = self._num_shadow_slots(lvl_id, pg_idx) | ||
| if shadow_cnt > 0: | ||
| deficit = (goal + fallen) - (old_free_cnt + evictable_cnt) | ||
| if deficit > 0: | ||
| self._reclaim_shadows(lvl_id, pg_idx, min(deficit, shadow_cnt)) | ||
| old_free_cnt = storage.get_num_free_slots(pg_idx) |
There was a problem hiding this comment.
🚀 Performance & Scalability | 🟠 Major | ⚡ Quick win
Reclaim shadow slots before evicting real pages.
Line 469 only reclaims shadows when free + evictable cannot satisfy the request, so real evictable pages can be evicted while cheaper shadow slots remain occupied. Reclaim against the free-slot shortfall first, then compute num_to_evict.
Proposed fix
- shadow_cnt = self._num_shadow_slots(lvl_id, pg_idx)
- if shadow_cnt > 0:
- deficit = (goal + fallen) - (old_free_cnt + evictable_cnt)
- if deficit > 0:
- self._reclaim_shadows(lvl_id, pg_idx, min(deficit, shadow_cnt))
- old_free_cnt = storage.get_num_free_slots(pg_idx)
+ shadow_cnt = self._num_shadow_slots(lvl_id, pg_idx)
+ if shadow_cnt > 0:
+ deficit = (goal + fallen) - old_free_cnt
+ if deficit > 0:
+ self._reclaim_shadows(lvl_id, pg_idx, min(deficit, shadow_cnt))
+ old_free_cnt = storage.get_num_free_slots(pg_idx)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| # Shadows are reclaimable for free (GPU holds the real copy), so reclaim them oldest | |
| # -first to cover any shortfall before evicting real pages or raising. | |
| shadow_cnt = self._num_shadow_slots(lvl_id, pg_idx) | |
| if shadow_cnt > 0: | |
| deficit = (goal + fallen) - (old_free_cnt + evictable_cnt) | |
| if deficit > 0: | |
| self._reclaim_shadows(lvl_id, pg_idx, min(deficit, shadow_cnt)) | |
| old_free_cnt = storage.get_num_free_slots(pg_idx) | |
| # Shadows are reclaimable for free (GPU holds the real copy), so reclaim them oldest | |
| # -first to cover any shortfall before evicting real pages or raising. | |
| shadow_cnt = self._num_shadow_slots(lvl_id, pg_idx) | |
| if shadow_cnt > 0: | |
| deficit = (goal + fallen) - old_free_cnt | |
| if deficit > 0: | |
| self._reclaim_shadows(lvl_id, pg_idx, min(deficit, shadow_cnt)) | |
| old_free_cnt = storage.get_num_free_slots(pg_idx) |
🤖 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/_storage_manager.py` around lines
465 - 472, The reclaim logic in _storage_manager.py currently only calls
_reclaim_shadows inside the deficit path, which can cause real evictable pages
to be chosen before cheaper shadow slots are freed. Update the allocation flow
in the storage manager method that computes
shadow_count/old_free_cnt/evictable_cnt so shadow slots are reclaimed first to
cover the free-slot shortfall, then recompute the available counts and only
after that calculate num_to_evict for real pages.
| def _is_inclusive_recall( | ||
| self, page: Page, src_level: CacheLevel, dst_level: CacheLevel | ||
| ) -> bool: | ||
| """A migration that should retain the source slot as a clean host shadow. | ||
|
|
||
| True only for the inclusive-host-cache fast path: a committed (immutable) page being | ||
| recalled from a lower tier toward GPU. The retained source slot becomes the page's shadow. | ||
| """ | ||
| return ( | ||
| self._enable_inclusive_host_cache | ||
| and dst_level < src_level | ||
| and isinstance(page, CommittedPage) | ||
| and page.shadow_slot is None | ||
| ) |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Restrict shadow attachment to host→GPU recalls.
The predicate currently matches any warmer-tier migration. With a disk tier, disk→host prefetch can attach a disk shadow, then page.shadow_slot is not None prevents the later host→GPU recall from retaining the host copy this feature is meant to reuse.
Proposed fix
return (
self._enable_inclusive_host_cache
- and dst_level < src_level
+ and dst_level == GPU_LEVEL
+ and self._levels[src_level].cache_tier == CacheTier.HOST_MEM
and isinstance(page, CommittedPage)
and page.shadow_slot is None
)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| def _is_inclusive_recall( | |
| self, page: Page, src_level: CacheLevel, dst_level: CacheLevel | |
| ) -> bool: | |
| """A migration that should retain the source slot as a clean host shadow. | |
| True only for the inclusive-host-cache fast path: a committed (immutable) page being | |
| recalled from a lower tier toward GPU. The retained source slot becomes the page's shadow. | |
| """ | |
| return ( | |
| self._enable_inclusive_host_cache | |
| and dst_level < src_level | |
| and isinstance(page, CommittedPage) | |
| and page.shadow_slot is None | |
| ) | |
| def _is_inclusive_recall( | |
| self, page: Page, src_level: CacheLevel, dst_level: CacheLevel | |
| ) -> bool: | |
| """A migration that should retain the source slot as a clean host shadow. | |
| True only for the inclusive-host-cache fast path: a committed (immutable) page being | |
| recalled from a lower tier toward GPU. The retained source slot becomes the page's shadow. | |
| """ | |
| return ( | |
| self._enable_inclusive_host_cache | |
| and dst_level == GPU_LEVEL | |
| and self._levels[src_level].cache_tier == CacheTier.HOST_MEM | |
| and isinstance(page, CommittedPage) | |
| and page.shadow_slot is 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/runtime/kv_cache_manager_v2/_storage_manager.py` around lines
560 - 573, The _is_inclusive_recall predicate in _storage_manager.py is too
broad because it matches any migration to a warmer tier, which lets disk→host
prefetch attach a shadow and blocks the intended host→GPU reuse path. Tighten
the condition so it only returns true for host→GPU recalls (not disk→host),
while keeping the existing checks for enablement, CommittedPage, and missing
shadow_slot; update the logic around _is_inclusive_recall and any callers in the
migration flow accordingly.
| if update_src and not defrag: | ||
| # Rebind shadow-reuse pages to their retained shadow slots (no copy happened). | ||
| for i, shadow in enumerate(reused_shadow): | ||
| if shadow is None: | ||
| continue | ||
| src = src_pages[i] | ||
| scheduled_for_eviction = src.scheduled_for_eviction | ||
| if scheduled_for_eviction: | ||
| self.exclude_from_eviction(src) | ||
| src_pool_group.release(src) | ||
| src.set_slot(shadow) | ||
| src.cache_level = dst_level | ||
| if emit_cache_level_updates: | ||
| self._emit_cache_level_updated_event( | ||
| src, src_level, dst_level, emitted_update_keys | ||
| ) | ||
| if scheduled_for_eviction: | ||
| self.schedule_for_eviction(src) | ||
| return None if update_src else dst_slots | ||
| except Exception: | ||
| for s in dst_slots: | ||
| dst_pool_group.release(s) | ||
| # Release any detached-but-not-yet-rebound shadow slots so they are not leaked. | ||
| for shadow in reused_shadow: | ||
| if shadow is not None and shadow.has_valid_slot: | ||
| self._levels[dst_level].storage.release(pool_group_index, shadow) | ||
| raise |
There was a problem hiding this comment.
🗄️ Data Integrity & Integration | 🟠 Major | ⚡ Quick win
Do not release rebound shadow slots during exception cleanup.
After Line 707, the shadow slot is owned by src again, but reused_shadow still points at it. If a later operation in the try-block raises, Lines 720-722 can release that live slot back to storage.
Proposed fix
src_pool_group.release(src)
src.set_slot(shadow)
+ reused_shadow[i] = None
src.cache_level = dst_level📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| if update_src and not defrag: | |
| # Rebind shadow-reuse pages to their retained shadow slots (no copy happened). | |
| for i, shadow in enumerate(reused_shadow): | |
| if shadow is None: | |
| continue | |
| src = src_pages[i] | |
| scheduled_for_eviction = src.scheduled_for_eviction | |
| if scheduled_for_eviction: | |
| self.exclude_from_eviction(src) | |
| src_pool_group.release(src) | |
| src.set_slot(shadow) | |
| src.cache_level = dst_level | |
| if emit_cache_level_updates: | |
| self._emit_cache_level_updated_event( | |
| src, src_level, dst_level, emitted_update_keys | |
| ) | |
| if scheduled_for_eviction: | |
| self.schedule_for_eviction(src) | |
| return None if update_src else dst_slots | |
| except Exception: | |
| for s in dst_slots: | |
| dst_pool_group.release(s) | |
| # Release any detached-but-not-yet-rebound shadow slots so they are not leaked. | |
| for shadow in reused_shadow: | |
| if shadow is not None and shadow.has_valid_slot: | |
| self._levels[dst_level].storage.release(pool_group_index, shadow) | |
| raise | |
| if update_src and not defrag: | |
| # Rebind shadow-reuse pages to their retained shadow slots (no copy happened). | |
| for i, shadow in enumerate(reused_shadow): | |
| if shadow is None: | |
| continue | |
| src = src_pages[i] | |
| scheduled_for_eviction = src.scheduled_for_eviction | |
| if scheduled_for_eviction: | |
| self.exclude_from_eviction(src) | |
| src_pool_group.release(src) | |
| src.set_slot(shadow) | |
| reused_shadow[i] = None | |
| src.cache_level = dst_level | |
| if emit_cache_level_updates: | |
| self._emit_cache_level_updated_event( | |
| src, src_level, dst_level, emitted_update_keys | |
| ) | |
| if scheduled_for_eviction: | |
| self.schedule_for_eviction(src) | |
| return None if update_src else dst_slots | |
| except Exception: | |
| for s in dst_slots: | |
| dst_pool_group.release(s) | |
| # Release any detached-but-not-yet-rebound shadow slots so they are not leaked. | |
| for shadow in reused_shadow: | |
| if shadow is not None and shadow.has_valid_slot: | |
| self._levels[dst_level].storage.release(pool_group_index, shadow) | |
| raise |
🤖 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/_storage_manager.py` around lines
697 - 723, The exception cleanup in the shadow-slot rebinding path is releasing
slots that have already been rebound to src in the update_src and not defrag
branch of the storage manager logic. Update the try/except handling in
_storage_manager.py so the cleanup only releases shadow slots that are still
detached and not yet reassigned, and skips any reused_shadow entry that has
already been attached back to a src page. Use the existing src.set_slot,
reused_shadow, and _levels[dst_level].storage.release flow to distinguish live
rebound slots from still-detached ones before releasing.
| host_free_before = host_free() | ||
| with TemporaryCudaStream([]) as s: | ||
| stream = cast(CudaStream, s.handle) | ||
| assert hot.kv_cache.resume(stream) | ||
| s.take_finish_event() | ||
| # After recall, immutable committed blocks retained their host copy as shadows. | ||
| shadows_after_recall = total_shadows() | ||
| self.assertGreater(shadows_after_recall, 0) | ||
| # Shadows occupy host slots, so free host slots dropped versus before the recall. | ||
| self.assertLess(host_free(), host_free_before) | ||
|
|
||
| # Suspend again: the hot blocks evict GPU->host and reuse their shadows (no fresh host slot, | ||
| # no D2H copy). The shadow FIFO drains as shadows are rebound to the evicted pages. | ||
| hot.kv_cache.suspend() |
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Fix the host-free assertion in the shadow reuse test.
host_free_before is captured while the suspended request already owns the host slots. Retaining those slots as shadows should keep the free count unchanged after recall, not lower it. Coverage is sufficient with this assertion corrected; no follow-up outside this PR needed. As per path instructions, tests/** feedback should assess coverage and stay actionable.
Proposed fix
- # Shadows occupy host slots, so free host slots dropped versus before the recall.
- self.assertLess(host_free(), host_free_before)
+ # The recalled pages retain the same host slots as shadows, so free host slots stay stable.
+ self.assertEqual(host_free(), host_free_before)
# Suspend again: the hot blocks evict GPU->host and reuse their shadows (no fresh host slot,
# no D2H copy). The shadow FIFO drains as shadows are rebound to the evicted pages.
hot.kv_cache.suspend()
self.assertLess(total_shadows(), shadows_after_recall)
+ self.assertEqual(host_free(), host_free_before)📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| host_free_before = host_free() | |
| with TemporaryCudaStream([]) as s: | |
| stream = cast(CudaStream, s.handle) | |
| assert hot.kv_cache.resume(stream) | |
| s.take_finish_event() | |
| # After recall, immutable committed blocks retained their host copy as shadows. | |
| shadows_after_recall = total_shadows() | |
| self.assertGreater(shadows_after_recall, 0) | |
| # Shadows occupy host slots, so free host slots dropped versus before the recall. | |
| self.assertLess(host_free(), host_free_before) | |
| # Suspend again: the hot blocks evict GPU->host and reuse their shadows (no fresh host slot, | |
| # no D2H copy). The shadow FIFO drains as shadows are rebound to the evicted pages. | |
| hot.kv_cache.suspend() | |
| host_free_before = host_free() | |
| with TemporaryCudaStream([]) as s: | |
| stream = cast(CudaStream, s.handle) | |
| assert hot.kv_cache.resume(stream) | |
| s.take_finish_event() | |
| # After recall, immutable committed blocks retained their host copy as shadows. | |
| shadows_after_recall = total_shadows() | |
| self.assertGreater(shadows_after_recall, 0) | |
| # The recalled pages retain the same host slots as shadows, so free host slots stay stable. | |
| self.assertEqual(host_free(), host_free_before) | |
| # Suspend again: the hot blocks evict GPU->host and reuse their shadows (no fresh host slot, | |
| # no D2H copy). The shadow FIFO drains as shadows are rebound to the evicted pages. | |
| hot.kv_cache.suspend() | |
| self.assertLess(total_shadows(), shadows_after_recall) | |
| self.assertEqual(host_free(), host_free_before) |
🤖 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 550 - 563, The shadow reuse test is asserting the wrong host free-space
behavior after recall: since the suspended request already holds the host slots
and they are retained as shadows, `host_free()` should remain unchanged rather
than drop. Update the assertion in `test_kv_cache_manager_v2.py` around
`host_free_before`, `total_shadows()`, and `hot.kv_cache.resume(stream)` so it
checks for no change in free host slots after recall while keeping the existing
shadow count verification.
Source: Path instructions
|
Please hold on until bug reporter's response. It's a design choice. It uses more D2H bandwidth but D2H bandwidth is under-utilized anyway, and it improves cache hit rate because it avoids duplication and can hold more data. |
|
[by Codex] @lowsfer Friendly reminder to review this PR as the primary KV-cache-manager reviewer when you have a chance. Thanks! |
|
[by Codex] @lowsfer Could you please review this PR? Thank you! |
|
[by Codex] @lowsfer Friendly review reminder: this PR is awaiting your review. Thanks! |
Add an optional inclusive host cache tier to KVCacheManager V2. When a full, immutable committed page is recalled host->GPU, its clean host copy is retained as a "shadow" instead of being freed; a later eviction of that page back to host reuses the shadow and skips the redundant device-to-host (D2H) copy. This eliminates redundant D2H rewrites of hot prefixes that repeatedly bounce GPU<->host, which dominate host-tier write traffic on agentic multi-turn workloads. Gated by a new KvCacheConfig flag enable_inclusive_host_cache (default False; exclusive paging unchanged).
withblock so the offload path does not raise on an empty copy batch.This feature is able to reduce offloaded KvCache, but no E2E perf gain was observed:
Summary by CodeRabbit
Description
Test Coverage
PR Checklist
Please review the following before submitting your PR:
PR description clearly explains what and why. If using CodeRabbit's summary, please make sure it makes sense.
PR Follows TRT-LLM CODING GUIDELINES to the best of your knowledge.
Test cases are provided for new code paths (see test instructions)
If PR introduces API changes, an appropriate PR label is added - either
api-compatibleorapi-breaking. Forapi-breaking, includeBREAKINGin the PR title.Any new dependencies have been scanned for license and vulnerabilities
CODEOWNERS updated if ownership changes
Documentation updated as needed
Update tava architecture diagram if there is a significant design change in PR.
The reviewers assigned automatically/manually are appropriate for the PR.
Please check this after reviewing the above items as appropriate for this PR.
GitHub Bot Help
To see a list of available CI bot commands, please comment
/bot help.