[None][feat] Add per-conversation KV cache block reuse

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -17,6 +17,7 @@
 import os
 import sys
 from collections import OrderedDict, defaultdict
+from dataclasses import dataclass
 from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union
 
 import numpy as np
@@ -58,6 +59,7 @@
     LayerId,
     LifeCycleId,
     PageIndexMode,
+    PlannedDropHandle,
     PoolGroupPeakBlockStats,
     ReuseScope,
     SwaScratchReuseConfig,
@@ -130,6 +132,90 @@ class Role:
 class BlockReusePolicy(StrEnum):
     ALL_REUSABLE = "all_reusable"
     PER_REQUEST = "per_request"
+    PER_CONVERSATION = "per_conversation"
+
+
+def _request_conversation_id(request: LlmRequest) -> Optional[str]:
+    if request.is_dummy_request:
+        return None
+    conversation_params = request.py_conversation_params
+    if conversation_params is None:
+        return None
+    conversation_id = conversation_params.conversation_id.strip()
+    return conversation_id or None
+
+
+@dataclass(slots=True)
+class _ConversationState:
+    current_request_id: Optional[int] = None
+    planned_drop_handle: Optional[PlannedDropHandle] = None
+
+
+class ConversationManager:
+    """Track the current request and drop plan for each conversation."""
+
+    def __init__(self) -> None:
+        self._conversation_states: Dict[str, _ConversationState] = {}
+
+    def save_drop_plan(self, request: LlmRequest, kv_cache: _KVCache) -> None:
+        """Save a completed context's drop plan and apply the preceding plan on success."""
+        request_id = request.py_request_id
+        conversation_id = _request_conversation_id(request)
+        if conversation_id is None:
+            return
+
+        state = self._conversation_states[conversation_id]
+        if state.current_request_id != request_id:
+            return
+
+        drop_handle = kv_cache.plan_committed_block_drop()
+        if drop_handle is None:
+            logger.warning(
+                f"Committed blocks for request {request_id} in conversation "
+                f"{conversation_id} have been dropped."
+            )
+        else:
+            previous_handle = state.planned_drop_handle
+            state.planned_drop_handle = drop_handle
+            if previous_handle is not None:
+                previous_handle.drop()
+
+        self.finish_request(request)
+
+    def prepare_request(self, request: LlmRequest) -> None:
+        """Register a context request unless its conversation has another active one."""
+        conversation_id = _request_conversation_id(request)
+        if conversation_id is None:
+            return
+        request_id = request.py_request_id
+        state = self._conversation_states.setdefault(conversation_id, _ConversationState())
+        current_request_id = state.current_request_id
+        if current_request_id is not None and current_request_id != request_id:
+            logger.warning(
+                f"Conversation {conversation_id} already has current request "
+                f"{current_request_id}. Request {request_id} will ignore "
+                "conversation params."
+            )
+            return
+
+        state.current_request_id = request_id
+
+    def finish_request(self, request: LlmRequest) -> None:
+        """Clear a request as active while preserving any saved drop plan."""
+        conversation_id = _request_conversation_id(request)
+        if conversation_id is None:
+            return
+        state = self._conversation_states.get(conversation_id)
+        if state is None or state.current_request_id != request.py_request_id:
+            return
+
+        state.current_request_id = None
+        if state.planned_drop_handle is None:
+            self._conversation_states.pop(conversation_id)
+
+    def clear(self) -> None:
+        """Clear state after reusable KV-cache blocks have been cleared."""
+        self._conversation_states.clear()
 
 
 def _estimate_full_attn_size_per_token(
@@ -990,6 +1076,12 @@ def append_to_kv_heads_per_layer(
         self.enable_block_reuse = kv_cache_config.enable_block_reuse
         self.enable_partial_reuse = kv_cache_config.enable_partial_reuse
         self.disk_prefetch_num_reqs = kv_cache_config.disk_prefetch_num_reqs
+        enable_conversation_manager = (
+            self.enable_block_reuse
+            and self.block_reuse_policy == BlockReusePolicy.PER_CONVERSATION
+            and not self.is_draft
+        )
+        self.conversation_manager = ConversationManager() if enable_conversation_manager else None
 
         # With pipeline parallelism, multiple microbatches can be in-flight
         # simultaneously, so we need slots for all concurrent sequences.
@@ -1954,6 +2046,8 @@ def prepare_context(self, req: LlmRequest) -> bool:
 
     def _prepare_context_impl(self, req: LlmRequest) -> bool:
         if req.is_first_context_chunk:
+            if self.conversation_manager is not None:
+                self.conversation_manager.prepare_request(req)
             kv_cache = self.kv_cache_map.get(req.py_request_id)
             if kv_cache is None:
                 all_tokens = req.get_tokens(DEFAULT_BEAM_INDEX)
@@ -2825,6 +2919,8 @@ def release_index_slot(self, request_id: int) -> None:
         self._early_freed_index_requests.add(request_id)
 
     def free_resources(self, request: LlmRequest, pin_on_release: bool = False):
+        if self.conversation_manager is not None:
+            self.conversation_manager.finish_request(request)
         self._allocated_draft_lens.pop(request.py_request_id, None)
         kv_cache = self.kv_cache_map.pop(request.py_request_id, None)
         if kv_cache is None:
@@ -3065,6 +3161,8 @@ def shutdown(self):
             kv_cache.close()
         self.kv_cache_map.clear()
         self.impl.shutdown()
+        if self.conversation_manager is not None:
+            self.conversation_manager.clear()
 
     def get_max_resource_count(self) -> int:
         # TODO: implement this
@@ -3146,6 +3244,8 @@ def update_context_resources(self, scheduled_batch: ScheduledRequests):
             if should_commit:
                 self.try_commit_blocks(req)
             if req.context_remaining_length == 0:
+                if self.conversation_manager is not None:
+                    self.conversation_manager.save_drop_plan(req, kv_cache)
                 # Scratch blocks are only for prefill chunks. Disable them at
                 # the context/generation boundary so generation uses normal KV
                 # pages before the first generation allocation.
@@ -3328,3 +3428,5 @@ def prefetch_for_context_tokens(self, requests: list) -> bool:
 
     def reset_reuse_state(self):
         self.impl.clear_reusable_blocks()
+        if self.conversation_manager is not None:
+            self.conversation_manager.clear()

tensorrt_llm/llmapi/llm_args.py

@@ -3666,12 +3666,19 @@ class KvCacheConfig(StrictBaseModel, PybindMirror):
         "pool_ratio is set.")
 
     # This is a pure python field, not a pybind field. It is only for the Pytorch backend.
-    block_reuse_policy: Literal["all_reusable", "per_request"] = Field(
-        default="all_reusable",
-        status="prototype",
-        description="KV cache manager v2 block reuse policy. "
-        "With SWA scratch reuse and 'all_reusable', only non-scratch "
-        "blocks are saved for reuse.")
+    block_reuse_policy: Literal[
+        "all_reusable", "per_request", "per_conversation"] = Field(
+            default="all_reusable",
+            status="prototype",
+            description="KV cache manager v2 block reuse policy. "
+            "'all_reusable' commits reusable blocks after every context chunk; "
+            "'per_request' commits them only after the final context chunk; "
+            "'per_conversation' uses 'per_request' commits and drops the previous "
+            "turn's committed SWA-window blocks after the current turn's final context "
+            "chunk. All reusable blocks remain subject to normal cache eviction. "
+            "Requests without conversation params use 'per_request' behavior. When "
+            "'all_reusable' and SWA scratch reuse are both enabled, only non-scratch "
+            "blocks are committed for reuse.")
 
     def _to_pybind(self):
         config = _KvCacheConfig(

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.py

@@ -53,6 +53,7 @@
     ExpandedBuffer,
     KVCacheManager,
     PageIndexConverter,
+    PlannedDropHandle,
     PoolDesc,
     PoolGroupDesc,
     PoolGroupPeakBlockStats,
@@ -118,6 +119,7 @@
     "KVCacheUpdatedData",
     "KvCacheStatus",
     "LayerGroupId",
+    "PlannedDropHandle",
     "LayerId",
     "LifeCycleId",
     "MemAddress",

tensorrt_llm/runtime/kv_cache_manager_v2/__init__.pyi

@@ -51,6 +51,9 @@ CacheLevel = NewType("CacheLevel", int)
 TokenId = NewType("TokenId", int)
 TokenIdExt = Union[TokenId, bytes]
 
+class PlannedDropHandle:
+    def drop(self) -> None: ...
+
 class ReuseScope(NamedTuple):
     lora_id: int | None = None
     salt: int | None = None
@@ -349,6 +352,7 @@ class _KVCache:
     def committed_tokens(self) -> list[TokenIdExt]: ...
     @property
     def reuse_scope(self) -> ReuseScope: ...
+    def plan_committed_block_drop(self) -> PlannedDropHandle | None: ...
     def stop_committing(self) -> None: ...
     def suspend(self) -> None: ...
     def resume(self, cuda_stream: CudaStream | None = None) -> bool: ...

tensorrt_llm/runtime/kv_cache_manager_v2/_core/__init__.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 
 from .._common import DEFAULT_BEAM_INDEX, BeamIndex
-from ._kv_cache import _KVCache
+from ._kv_cache import PlannedDropHandle, _KVCache
 from ._kv_cache_manager import (
     AggregatedPageDesc,
     ExpandedBuffer,
@@ -29,6 +29,7 @@
 __all__ = [
     "KVCacheManager",
     "_KVCache",
+    "PlannedDropHandle",
     "BeamIndex",
     "DEFAULT_BEAM_INDEX",
     "AggregatedPageDesc",

tensorrt_llm/runtime/kv_cache_manager_v2/_core/_kv_cache.py

@@ -16,7 +16,7 @@
 import array
 import enum
 import math
-from collections.abc import Sequence
+from collections.abc import Iterable, Sequence
 from contextlib import contextmanager
 from dataclasses import dataclass
 from itertools import chain
@@ -37,6 +37,7 @@
     CudaStream,
     PageIndex,
     PageIndexMode,
+    PageStatus,
     Priority,
     TokenIdExt,
 )
@@ -129,6 +130,58 @@ def __del__(self) -> None:
         self.pages.clear()
 
 
+class PlannedDropHandle:
+    """Track committed pages planned for dropping without owning them.
+
+    The handle stores weak references and does not keep pages alive. Dropping it
+    decrements each live page's planned-drop count and removes an already-droppable
+    page from eviction tracking when no plans remain.
+    """
+
+    __slots__ = ("_page_refs",)
+
+    _page_refs: tuple[rawref.ref[CommittedPage], ...] | None
+
+    def __init__(self, pages: Iterable[CommittedPage]) -> None:
+        planned_pages = tuple({id(page): page for page in pages}.values())
+        self._page_refs = tuple(rawref.ref(page) for page in planned_pages)
+        for page in planned_pages:
+            page.planned_drop_count += 1
+
+    def drop(self) -> None:
+        """Apply this drop plan and invalidate the handle.
+
+        A live page is removed from eviction tracking only when this is its final
+        plan and it is already droppable and queued for eviction. Calling this
+        method twice is invalid.
+        """
+        page_refs = self._page_refs
+        if page_refs is None:
+            raise ValueError("Planned drop handle has already been dropped")
+
+        pages = list[CommittedPage]()
+        for page_ref in page_refs:
+            page = page_ref()
+            if page is not None:
+                if page.planned_drop_count <= 0:
+                    raise ValueError("Committed page has no planned drop")
+                pages.append(page)
+
+        self._page_refs = None
+        for page in pages:
+            page.planned_drop_count -= 1
+            if (
+                page.planned_drop_count == 0
+                and page.status == PageStatus.DROPPABLE
+                and page.scheduled_for_eviction
+            ):
+                page.manager.exclude_from_eviction(page)
+
+    def __del__(self) -> None:
+        if self._page_refs is not None:
+            self.drop()
+
+
 class _Status(enum.Enum):
     ACTIVE = enum.auto()
     SUSPENDED = enum.auto()
@@ -993,6 +1046,43 @@ def committed_tokens(self) -> list[TokenIdExt]:
     def reuse_scope(self) -> ReuseScope:
         return self._reuse_scope
 
+    def plan_committed_block_drop(self) -> PlannedDropHandle | None:
+        """Plan dropping SWA blocks needed only by the next conversation turn.
+
+        The plan covers committed pages in each SWA life cycle's current
+        attention window. Full-attention and attention-sink blocks are excluded
+        because later turns may still need them. SSM state is not yet supported.
+        This must be called after stop_committing(). Returns None without
+        creating a plan if any required SWA page is unavailable.
+        """
+        if self._commit_state != self.CommitState.USER_STOP:
+            raise LogicError("plan_committed_block_drop() requires stop_committing()")
+
+        end = self._num_committed_blocks
+        pages_to_drop: list[CommittedPage] = []
+        for lc_idx, lc in self.manager._life_cycles.items():
+            if isinstance(lc, SsmLifeCycle):
+                # TODO: Support recording reusable SSM state pages.
+                continue
+            if lc.window_size is None:
+                continue
+            stale_range = _KVCache._get_stale_range(
+                self.tokens_per_block, self.num_committed_tokens, lc
+            )
+            window_start = min(stale_range.end, end)
+            for ordinal in typed_range(window_start, end):
+                tree_block = self._blocks[ordinal].tree_block
+                if tree_block is None:
+                    return None
+                page_ref = tree_block.storage[lc_idx]
+                if page_ref is None:
+                    return None
+                page = page_ref()
+                if page is None:
+                    return None
+                pages_to_drop.append(page)
+        return PlannedDropHandle(pages_to_drop)
+
     # Users promise to not commit any more tokens. For cases where we shouldn't reuse generated tokens
     # (eg. CoT), this helps us drop (instead of evict) out-of-window blocks for SWA layers.
     # If there is a uncommitted block containing committed tokens, we will commit the block immediately.

tensorrt_llm/runtime/kv_cache_manager_v2/_page.py

@@ -241,6 +241,7 @@ class CommittedPage(Page):
     """
 
     block: rawref.ref["Block"]
+    planned_drop_count: int
     __rawref__: rawref.ref["CommittedPage"]
 
     def is_committed(self) -> bool:
@@ -256,6 +257,7 @@ def __init__(
         priority: Priority,
     ):
         self.block = rawref.ref(block)
+        self.planned_drop_count = 0
         self.__rawref__ = rawref.NULL
         Page.__init__(
             self,

tensorrt_llm/usage/llm_args_golden_manifest.json

@@ -561,9 +561,10 @@
     {
       "allowed_values": [
         "all_reusable",
-        "per_request"
+        "per_request",
+        "per_conversation"
       ],
-      "annotation": "Literal['all_reusable', 'per_request']",
+      "annotation": "Literal['all_reusable', 'per_request', 'per_conversation']",
       "converter": "",
       "kind": "categorical",
       "path": "kv_cache_config.block_reuse_policy"