[None][feat] BREAKING Support saving last N turns in per-conversation policy

docs/source/developer-guide/telemetry.md

@@ -106,7 +106,7 @@ unset or when the safety sanitizer rejects the runtime value.
 | `iter_stats_max_iterations` | `Optional[int]` | `value` |  |  |
 | `kv_cache_config.attention_dp_events_gather_period_ms` | `<class 'int'>` | `value` |  |  |
 | `kv_cache_config.avg_seq_len` | `Optional[Annotated[int, Gt(gt=0)]]` | `value` |  |  |
-| `kv_cache_config.block_reuse_policy` | `Literal['all_reusable', 'per_request', 'per_conversation']` | `categorical` |  | `all_reusable`, `per_request`, `per_conversation` |
+| `kv_cache_config.block_reuse_config.block_reuse_policy` | `Literal['all_reusable', 'per_request', 'per_conversation']` | `categorical` |  | `all_reusable`, `per_request`, `per_conversation` |
 | `kv_cache_config.copy_on_partial_reuse` | `<class 'bool'>` | `value` |  |  |
 | `kv_cache_config.cross_kv_cache_fraction` | `Optional[float]` | `value` |  |  |
 | `kv_cache_config.disk_cache_size` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` |  |  |

examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py

@@ -40,7 +40,7 @@
 import pickle
 import signal
 import sys
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import List, Optional
 
 import torch
@@ -58,7 +58,7 @@
 from tensorrt_llm._torch.pyexecutor.kv_cache_transceiver import create_kv_cache_transceiver
 from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest, LlmRequestState, LlmRequestType
 from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager
-from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig
+from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, CacheTransceiverConfig
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.sampling_params import SamplingParams
 
@@ -105,7 +105,7 @@ class KvCacheConfigV2:
     dtype: str = "auto"
     pool_ratio: Optional[List[float]] = None
     avg_seq_len: Optional[int] = None
-    block_reuse_policy: str = "all_reusable"
+    block_reuse_config: BlockReuseConfig = field(default_factory=BlockReuseConfig)
     enable_swa_scratch_reuse: bool = False
     disk_prefetch_num_reqs: int = 4
     max_util_for_resume: float = 0.95

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -16,8 +16,8 @@
 import math
 import os
 import sys
-from collections import OrderedDict, defaultdict
-from dataclasses import dataclass, replace
+from collections import OrderedDict, defaultdict, deque
+from dataclasses import dataclass, field, replace
 from typing import TYPE_CHECKING, Dict, Iterable, List, NamedTuple, Optional, Sequence, Tuple, Union
 
 import numpy as np
@@ -153,13 +153,14 @@ def _request_conversation_id(request: LlmRequest) -> Optional[str]:
 @dataclass(slots=True)
 class _ConversationState:
     current_request_id: Optional[int] = None
-    planned_drop_handle: Optional[PlannedDropHandle] = None
+    planned_drop_handles: deque[PlannedDropHandle] = field(default_factory=deque)
 
 
 class ConversationManager:
     """Track the current request and drop plan for each conversation."""
 
-    def __init__(self) -> None:
+    def __init__(self, max_num_turns: int) -> None:
+        self._max_num_turns = max_num_turns
         self._conversation_states: Dict[str, _ConversationState] = {}
 
     def save_drop_plan(self, request: LlmRequest, kv_cache: _KVCache) -> None:
@@ -180,10 +181,9 @@ def save_drop_plan(self, request: LlmRequest, kv_cache: _KVCache) -> None:
                 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()
+            state.planned_drop_handles.append(drop_handle)
+            if len(state.planned_drop_handles) > self._max_num_turns:
+                state.planned_drop_handles.popleft().drop()
 
         self.finish_request(request)
 
@@ -215,7 +215,7 @@ def finish_request(self, request: LlmRequest) -> None:
             return
 
         state.current_request_id = None
-        if state.planned_drop_handle is None:
+        if not state.planned_drop_handles:
             self._conversation_states.pop(conversation_id)
 
     def clear(self) -> None:
@@ -794,7 +794,8 @@ def __init__(
         self.enable_swa_scratch_reuse = (
             kv_cache_config.enable_swa_scratch_reuse and not self.is_draft
         )
-        self.block_reuse_policy = BlockReusePolicy(kv_cache_config.block_reuse_policy)
+        block_reuse_config = kv_cache_config.block_reuse_config
+        self.block_reuse_policy = BlockReusePolicy(block_reuse_config.block_reuse_policy)
         self.num_local_layers = len(self.pp_layers)
         self.layer_offsets = {idx: offset for offset, idx in enumerate(self.pp_layers)}
         self.max_beam_width = max_beam_width
@@ -1136,7 +1137,11 @@ def append_to_kv_heads_per_layer(
             and self.block_reuse_policy == BlockReusePolicy.PER_CONVERSATION
             and not self.is_draft
         )
-        self.conversation_manager = ConversationManager() if enable_conversation_manager else None
+        self.conversation_manager = (
+            ConversationManager(block_reuse_config.max_num_turns)
+            if enable_conversation_manager
+            else None
+        )
 
         # With pipeline parallelism, multiple microbatches can be in-flight
         # simultaneously, so we need slots for all concurrent sequences.

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -2629,12 +2629,15 @@ def __init__(
 
         kv_cache_config = kv_cache_config.model_copy(deep=True)
         if any(mamba_layer_mask) and kv_cache_config.enable_block_reuse:
+            block_reuse_config = kv_cache_config.block_reuse_config
             block_reuse_policy = BlockReusePolicy(
-                kv_cache_config.block_reuse_policy)
+                block_reuse_config.block_reuse_policy)
             if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE:
                 # SSM reuse is valid only at explicit snapshot boundaries.
-                kv_cache_config.block_reuse_policy = (
-                    BlockReusePolicy.PER_REQUEST.value)
+                kv_cache_config.block_reuse_config = block_reuse_config.model_copy(
+                    update={
+                        "block_reuse_policy": BlockReusePolicy.PER_REQUEST.value
+                    })
         self.kv_cache_config = kv_cache_config
 
         super().__init__(

tensorrt_llm/llmapi/__init__.py

@@ -7,7 +7,7 @@
 from .llm import LLM, RequestOutput
 # yapf: disable
 from .llm_args import (AttentionDpConfig, AutoDecodingConfig, BatchingType,
-                       CacheTransceiverConfig, CalibConfig,
+                       BlockReuseConfig, CacheTransceiverConfig, CalibConfig,
                        CapacitySchedulerPolicy, ContextChunkingPolicy,
                        CudaGraphConfig, DecodeCudaGraphConfig,
                        DeepSeekSparseAttentionConfig,
@@ -42,6 +42,7 @@
     'DisaggregatedParams',
     'ConversationParams',
     'DisaggScheduleStyle',
+    'BlockReuseConfig',
     'KvCacheConfig',
     'MambaStateConfig',
     'KvCacheRetentionConfig',

tensorrt_llm/llmapi/llm_args.py

@@ -3618,6 +3618,35 @@ class MambaStateConfig(StrictBaseModel):
         "snapshots require KV cache manager V2.")
 
 
+class BlockReuseConfig(StrictBaseModel):
+    """Configuration for KV cache block reuse policies."""
+
+    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 retains committed "
+            "SWA-window blocks and Mamba stable-boundary state for up to "
+            "`max_num_turns` completed turns. Periodic Mamba state snapshots "
+            "are disabled with 'per_conversation'. 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.")
+
+    max_num_turns: PositiveInt = Field(
+        default=1,
+        status="prototype",
+        description=
+        "Maximum number of completed conversation turns whose committed SWA-window "
+        "blocks and Mamba stable-boundary state are retained by KV cache manager v2. "
+        "Only used when "
+        "`block_reuse_policy` is 'per_conversation'.")
+
+
 @PybindMirror.mirror_pybind_fields(_KvCacheConfig)
 class KvCacheConfig(StrictBaseModel, PybindMirror):
     """Configuration for the KV cache."""
@@ -3840,21 +3869,11 @@ class KvCacheConfig(StrictBaseModel, PybindMirror):
         "unset. This does not take effect when 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", "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 and Mamba stable-boundary state "
-            "after the current turn's final context chunk. Periodic Mamba state "
-            "snapshots are disabled with 'per_conversation'. 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.")
+    block_reuse_config: BlockReuseConfig = Field(
+        default_factory=BlockReuseConfig,
+        status="prototype",
+        description="KV cache manager v2 configuration for block reuse policies."
+    )
 
     def _to_pybind(self):
         config = _KvCacheConfig(
@@ -3940,13 +3959,14 @@ def migrate_legacy_mamba_interval(self) -> 'KvCacheConfig':
     def disable_periodic_mamba_snapshots_for_conversations(
             self) -> 'KvCacheConfig':
         """Use only explicit stable boundaries for conversation reuse."""
-        if (self.block_reuse_policy == "per_conversation"
+        if (self.block_reuse_config.block_reuse_policy == "per_conversation"
                 and self.mamba_state_config.periodic_snapshot_interval != 0):
             interval = self.mamba_state_config.periodic_snapshot_interval
             logger.warning(
                 f"'kv_cache_config.mamba_state_config.periodic_snapshot_interval={interval}' "
                 "is ignored because "
-                "'kv_cache_config.block_reuse_policy=per_conversation' disables "
+                "'kv_cache_config.block_reuse_config."
+                "block_reuse_policy=per_conversation' disables "
                 "periodic Mamba snapshots; setting it to 0.")
             self.mamba_state_config = self.mamba_state_config.model_copy(
                 update={"periodic_snapshot_interval": 0})

tensorrt_llm/usage/llm_args_golden_manifest.json

@@ -577,7 +577,14 @@
       "annotation": "Literal['all_reusable', 'per_request', 'per_conversation']",
       "converter": "",
       "kind": "categorical",
-      "path": "kv_cache_config.block_reuse_policy"
+      "path": "kv_cache_config.block_reuse_config.block_reuse_policy"
+    },
+    {
+      "allowed_values": [],
+      "annotation": "<class 'int'>",
+      "converter": "",
+      "kind": "value",
+      "path": "kv_cache_config.block_reuse_config.max_num_turns"
     },
     {
       "allowed_values": [],

tests/unittest/_torch/executor/test_kv_cache_manager_v2.py

@@ -25,7 +25,7 @@
 from tensorrt_llm.bindings import DataType
 from tensorrt_llm.bindings.internal.batch_manager import CacheType
 from tensorrt_llm.conversation_params import ConversationParams
-from tensorrt_llm.llmapi.llm_args import KvCacheConfig
+from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, KvCacheConfig
 from tensorrt_llm.mapping import Mapping
 from tensorrt_llm.runtime.kv_cache_manager_v2 import (
     DEFAULT_BEAM_INDEX,
@@ -71,7 +71,9 @@ def _make_cache_config_for_test(
     cache_manager.enable_swa_scratch_reuse = False
     cache_manager.num_extra_kv_tokens = num_extra_kv_tokens
     cache_manager.enable_stats = False
-    cache_manager.block_reuse_policy = BlockReusePolicy(kv_cache_config.block_reuse_policy)
+    cache_manager.block_reuse_policy = BlockReusePolicy(
+        kv_cache_config.block_reuse_config.block_reuse_policy
+    )
     cache_manager.is_draft = is_draft
     cache_manager.num_local_layers = 1
     cache_manager.pp_layers = [0]
@@ -107,7 +109,7 @@ def test_commit_min_snapshot_follows_block_reuse_policy(
     config = _make_cache_config_for_test(
         KvCacheConfig(
             enable_block_reuse=enable_block_reuse,
-            block_reuse_policy=block_reuse_policy,
+            block_reuse_config=BlockReuseConfig(block_reuse_policy=block_reuse_policy),
             enable_partial_reuse=True,
         ),
         is_draft=is_draft,
@@ -285,7 +287,12 @@ def set_prepopulated_prompt_len(self, length: int, tokens_per_block: int) -> Non
 
 
 @pytest.fixture
-def manager() -> KVCacheManagerV2:
+def max_num_turns() -> int:
+    return 1
+
+
+@pytest.fixture
+def manager(max_num_turns: int) -> KVCacheManagerV2:
     if not torch.cuda.is_available():
         pytest.skip("requires CUDA")
     init_cuda_once()
@@ -296,7 +303,10 @@ def manager() -> KVCacheManagerV2:
             max_gpu_total_bytes=16 << 20,
             max_attention_window=[MAX_SEQ_LEN, TOKENS_PER_BLOCK],
             max_util_for_resume=1.0,
-            block_reuse_policy="per_conversation",
+            block_reuse_config=BlockReuseConfig(
+                block_reuse_policy="per_conversation",
+                max_num_turns=max_num_turns,
+            ),
         ),
         CacheType.SELF,
         num_layers=2,
@@ -474,6 +484,39 @@ def test_per_conversation_policy_drops_previous_divergent_blocks(
         _free_if_active(manager, request_a)
 
 
+@pytest.mark.parametrize("max_num_turns", [2])
+def test_per_conversation_policy_retains_configured_number_of_turns(
+    manager: KVCacheManagerV2,
+) -> None:
+    request_a = _ContextRequest(1, list(range(8)), 8, "conv-1")
+    request_b = _ContextRequest(2, list(range(100, 108)), 8, "conv-1")
+    request_a_probe = _ContextRequest(3, list(range(8)), 8, "conv-2")
+    request_c = _ContextRequest(4, list(range(200, 208)), 8, "conv-1")
+    request_a_after_eviction = _ContextRequest(5, list(range(8)), 8, "conv-3")
+
+    try:
+        _run_context(manager, request_a)
+        _free_if_active(manager, request_a)
+        _run_context(manager, request_b)
+        _free_if_active(manager, request_b)
+
+        assert manager.prepare_context(request_a_probe)
+        assert request_a_probe.prepopulated_prompt_len == request_a_probe.prompt_len - 1
+        _free_if_active(manager, request_a_probe)
+
+        _run_context(manager, request_c)
+        _free_if_active(manager, request_c)
+
+        assert manager.prepare_context(request_a_after_eviction)
+        assert request_a_after_eviction.prepopulated_prompt_len == 0
+    finally:
+        _free_if_active(manager, request_a_after_eviction)
+        _free_if_active(manager, request_c)
+        _free_if_active(manager, request_a_probe)
+        _free_if_active(manager, request_b)
+        _free_if_active(manager, request_a)
+
+
 def test_per_conversation_policy_ignores_overlapping_request(
     manager: KVCacheManagerV2,
 ) -> None:

tests/unittest/_torch/executor/test_mamba_cache_manager.py

@@ -53,6 +53,7 @@
 from tensorrt_llm._utils import torch_dtype_to_binding
 from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType
 from tensorrt_llm.llmapi.llm_args import (
+    BlockReuseConfig,
     CacheTransceiverConfig,
     KvCacheConfig,
     MambaStateConfig,
@@ -1680,7 +1681,7 @@ def _build_v2_hybrid_with_mamba_layer(
         max_tokens=512,
         enable_block_reuse=enable_block_reuse,
         enable_partial_reuse=enable_partial_reuse,
-        block_reuse_policy=block_reuse_policy,
+        block_reuse_config=BlockReuseConfig(block_reuse_policy=block_reuse_policy),
         enable_swa_scratch_reuse=enable_swa_scratch_reuse,
         mamba_state_config=MambaStateConfig(
             periodic_snapshot_interval=periodic_snapshot_interval,

tests/unittest/disaggregated/test_cache_transceiver_single_process.py

@@ -32,7 +32,7 @@
 # injects; see test_kv_transfer.py for the full rationale.
 os.environ["UCX_TLS"] = "^ib,gdr_copy"
 os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1"
-from dataclasses import dataclass
+from dataclasses import dataclass, field
 from typing import Dict, List, Optional
 
 import pytest
@@ -54,7 +54,7 @@
 from tensorrt_llm.bindings import LayerType as LayerTypeCpp
 from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp
 from tensorrt_llm.bindings.internal.batch_manager import CacheType as CacheTypeCpp
-from tensorrt_llm.llmapi.llm_args import CacheTransceiverConfig, KvCacheConfig
+from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, CacheTransceiverConfig, KvCacheConfig
 
 AttentionTypeCpp = tensorrt_llm.bindings.internal.batch_manager.AttentionType
 
@@ -174,7 +174,7 @@ class KvCacheConfigV2:
     disk_prefetch_num_reqs: int = 4
     pool_ratio: Optional[List[float]] = None
     avg_seq_len: Optional[int] = None
-    block_reuse_policy: str = "all_reusable"
+    block_reuse_config: BlockReuseConfig = field(default_factory=BlockReuseConfig)
     enable_swa_scratch_reuse: bool = False
     max_util_for_resume: float = 0.95