[None][fix] BREAKING Block reuse policy rename and add more tests

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_config.block_reuse_policy` | `Literal['all_reusable', 'per_request', 'per_conversation']` | `categorical` |  | `all_reusable`, `per_request`, `per_conversation` |
+| `kv_cache_config.block_reuse_config.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` |  |  |

docs/source/features/kvcache.md

@@ -97,6 +97,9 @@ kv_cache_config:
   enable_block_reuse: true
   use_kv_cache_manager_v2: true
   avg_seq_len: 2048
+  block_reuse_config:
+    policy: per_conversation
+    max_num_turns: 2
   mamba_state_config:
     periodic_snapshot_interval: 0
     additional_snapshot_offsets_from_start: [128]

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -795,7 +795,7 @@ def __init__(
             kv_cache_config.enable_swa_scratch_reuse and not self.is_draft
         )
         block_reuse_config = kv_cache_config.block_reuse_config
-        self.block_reuse_policy = BlockReusePolicy(block_reuse_config.block_reuse_policy)
+        self.block_reuse_policy = BlockReusePolicy(block_reuse_config.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

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -2630,14 +2630,11 @@ 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(
-                block_reuse_config.block_reuse_policy)
+            block_reuse_policy = BlockReusePolicy(block_reuse_config.policy)
             if block_reuse_policy == BlockReusePolicy.ALL_REUSABLE:
                 # SSM reuse is valid only at explicit snapshot boundaries.
                 kv_cache_config.block_reuse_config = block_reuse_config.model_copy(
-                    update={
-                        "block_reuse_policy": BlockReusePolicy.PER_REQUEST.value
-                    })
+                    update={"policy": BlockReusePolicy.PER_REQUEST.value})
         self.kv_cache_config = kv_cache_config
 
         super().__init__(

tensorrt_llm/llmapi/llm_args.py

@@ -3720,21 +3720,29 @@ class MambaStateConfig(StrictBaseModel):
 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.")
+    @model_validator(mode="before")
+    @classmethod
+    def _reject_renamed_block_reuse_policy(cls, data: Any) -> Any:
+        if isinstance(data, dict) and "block_reuse_policy" in data:
+            raise ValueError(
+                "'kv_cache_config.block_reuse_config.block_reuse_policy' was "
+                "renamed to 'kv_cache_config.block_reuse_config.policy'.")
+        return data
+
+    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,
@@ -3743,7 +3751,7 @@ class BlockReuseConfig(StrictBaseModel):
         "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'.")
+        "`policy` is 'per_conversation'.")
 
 
 @PybindMirror.mirror_pybind_fields(_KvCacheConfig)
@@ -4072,14 +4080,13 @@ 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_config.block_reuse_policy == "per_conversation"
+        if (self.block_reuse_config.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_config."
-                "block_reuse_policy=per_conversation' disables "
+                "'kv_cache_config.block_reuse_config.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

@@ -609,6 +609,13 @@
       "kind": "value",
       "path": "kv_cache_config.avg_seq_len"
     },
+    {
+      "allowed_values": [],
+      "annotation": "<class 'int'>",
+      "converter": "",
+      "kind": "value",
+      "path": "kv_cache_config.block_reuse_config.max_num_turns"
+    },
     {
       "allowed_values": [
         "all_reusable",
@@ -618,14 +625,7 @@
       "annotation": "Literal['all_reusable', 'per_request', 'per_conversation']",
       "converter": "",
       "kind": "categorical",
-      "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"
+      "path": "kv_cache_config.block_reuse_config.policy"
     },
     {
       "allowed_values": [],

tests/unittest/_torch/executor/test_kv_cache_manager_v2.py

@@ -71,9 +71,7 @@ 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_config.block_reuse_policy
-    )
+    cache_manager.block_reuse_policy = BlockReusePolicy(kv_cache_config.block_reuse_config.policy)
     cache_manager.is_draft = is_draft
     cache_manager.num_local_layers = 1
     cache_manager.pp_layers = [0]
@@ -109,7 +107,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_config=BlockReuseConfig(block_reuse_policy=block_reuse_policy),
+            block_reuse_config=BlockReuseConfig(policy=block_reuse_policy),
             enable_partial_reuse=True,
         ),
         is_draft=is_draft,
@@ -304,7 +302,7 @@ def manager(max_num_turns: int) -> KVCacheManagerV2:
             max_attention_window=[MAX_SEQ_LEN, TOKENS_PER_BLOCK],
             max_util_for_resume=1.0,
             block_reuse_config=BlockReuseConfig(
-                block_reuse_policy="per_conversation",
+                policy="per_conversation",
                 max_num_turns=max_num_turns,
             ),
         ),

tests/unittest/_torch/executor/test_mamba_cache_manager.py

@@ -52,6 +52,7 @@
 from tensorrt_llm._torch.pyexecutor.scheduler import ScheduledRequests
 from tensorrt_llm._utils import torch_dtype_to_binding
 from tensorrt_llm.bindings.internal.batch_manager import LinearCacheType
+from tensorrt_llm.conversation_params import ConversationParams
 from tensorrt_llm.llmapi.llm_args import (
     BlockReuseConfig,
     CacheTransceiverConfig,
@@ -1678,6 +1679,7 @@ def _build_v2_hybrid_with_mamba_layer(
     enable_block_reuse=False,
     enable_partial_reuse=True,
     block_reuse_policy="all_reusable",
+    max_num_turns=1,
     periodic_snapshot_interval=0,
     additional_snapshot_offsets_from_end=None,
     enable_attention_dp=False,
@@ -1700,7 +1702,10 @@ def _build_v2_hybrid_with_mamba_layer(
         max_tokens=512,
         enable_block_reuse=enable_block_reuse,
         enable_partial_reuse=enable_partial_reuse,
-        block_reuse_config=BlockReuseConfig(block_reuse_policy=block_reuse_policy),
+        block_reuse_config=BlockReuseConfig(
+            policy=block_reuse_policy,
+            max_num_turns=max_num_turns,
+        ),
         enable_swa_scratch_reuse=enable_swa_scratch_reuse,
         mamba_state_config=MambaStateConfig(
             periodic_snapshot_interval=periodic_snapshot_interval,
@@ -1750,6 +1755,39 @@ def _make_wide_spec_config(max_draft_len=2, tokens_per_gen_step=5):
     )
 
 
+def _make_v2_conversation_request(
+    request_id: int,
+    tokens: list[int],
+    conversation_id: str,
+) -> LlmRequest:
+    request = LlmRequest(
+        request_id=request_id,
+        max_new_tokens=1,
+        input_tokens=tokens,
+        sampling_config=SamplingConfig(),
+        is_streaming=False,
+    )
+    request.py_conversation_params = ConversationParams(conversation_id=conversation_id)
+    return request
+
+
+def _run_v2_hybrid_context(
+    manager: MambaHybridCacheManagerV2,
+    request: LlmRequest,
+) -> None:
+    manager.prepare_expect_snapshot_points([request])
+    assert manager.prepare_context(request)
+    num_tokens = request.context_remaining_length
+    assert manager.resize_context(request, num_tokens=num_tokens)
+
+    batch = ScheduledRequests()
+    batch.append_context_request(request)
+    manager.prepare_resources(batch)
+    request.context_current_position = request.prompt_len
+    assert request.context_remaining_length == 0
+    manager.update_context_resources(batch)
+
+
 def _assert_replay_layer_cache_uses_history_size(layer_cache, history_size):
     assert layer_cache.old_x is not None
     assert layer_cache.old_B is not None
@@ -2018,6 +2056,90 @@ def test_v2_hybrid_preserves_per_conversation_and_disables_periodic_snapshots():
         mgr.shutdown()
 
 
+@skip_no_cuda
+def test_v2_hybrid_retains_configured_number_of_conversation_turns():
+    mgr = _build_v2_hybrid_with_mamba_layer(
+        enable_block_reuse=True,
+        block_reuse_policy="per_conversation",
+        max_num_turns=2,
+        additional_snapshot_offsets_from_end=[0],
+    )
+    request_a = _make_v2_conversation_request(1, list(range(64)), "conv-1")
+    request_b = _make_v2_conversation_request(2, list(range(100, 164)), "conv-1")
+    # Use fresh conversation IDs so probes query the shared prefix cache
+    # without altering conv-1's retained-turn accounting.
+    # Probe one token past the exact SSM snapshots committed at token 64.
+    request_a_probe = _make_v2_conversation_request(3, list(range(65)), "conv-2")
+    request_b_probe = _make_v2_conversation_request(4, list(range(100, 165)), "conv-3")
+    request_c = _make_v2_conversation_request(5, list(range(200, 264)), "conv-1")
+    request_a_after_eviction = _make_v2_conversation_request(6, list(range(65)), "conv-4")
+    request_b_after_eviction = _make_v2_conversation_request(7, list(range(100, 165)), "conv-5")
+    requests = [
+        request_a,
+        request_b,
+        request_a_probe,
+        request_b_probe,
+        request_c,
+        request_a_after_eviction,
+        request_b_after_eviction,
+    ]
+
+    try:
+        _run_v2_hybrid_context(mgr, request_a)
+        request_a_state_index = mgr.get_state_indices([request_a.py_request_id], [False])[0]
+        mgr.free_resources(request_a)
+        _run_v2_hybrid_context(mgr, request_b)
+        request_b_state_index = mgr.get_state_indices([request_b.py_request_id], [False])[0]
+        mgr.free_resources(request_b)
+
+        mgr.prepare_expect_snapshot_points([request_a_probe])
+        assert mgr.prepare_context(request_a_probe)
+        probe_batch = ScheduledRequests()
+        probe_batch.append_context_request(request_a_probe)
+        mgr.prepare_resources(probe_batch)
+        assert request_a_probe.prepopulated_prompt_len == request_a_probe.prompt_len - 1
+        assert mgr.get_state_indices([request_a_probe.py_request_id], [False]) == [
+            request_a_state_index
+        ]
+        mgr.free_resources(request_a_probe)
+
+        mgr.prepare_expect_snapshot_points([request_b_probe])
+        assert mgr.prepare_context(request_b_probe)
+        probe_batch = ScheduledRequests()
+        probe_batch.append_context_request(request_b_probe)
+        mgr.prepare_resources(probe_batch)
+        assert request_b_probe.prepopulated_prompt_len == request_b_probe.prompt_len - 1
+        assert mgr.get_state_indices([request_b_probe.py_request_id], [False]) == [
+            request_b_state_index
+        ]
+        mgr.free_resources(request_b_probe)
+
+        _run_v2_hybrid_context(mgr, request_c)
+        mgr.free_resources(request_c)
+
+        mgr.prepare_expect_snapshot_points([request_a_after_eviction])
+        assert mgr.prepare_context(request_a_after_eviction)
+        assert request_a_after_eviction.prepopulated_prompt_len == 0
+
+        mgr.prepare_expect_snapshot_points([request_b_after_eviction])
+        assert mgr.prepare_context(request_b_after_eviction)
+        probe_batch = ScheduledRequests()
+        probe_batch.append_context_request(request_b_after_eviction)
+        mgr.prepare_resources(probe_batch)
+        assert (
+            request_b_after_eviction.prepopulated_prompt_len
+            == request_b_after_eviction.prompt_len - 1
+        )
+        assert mgr.get_state_indices([request_b_after_eviction.py_request_id], [False]) == [
+            request_b_state_index
+        ]
+    finally:
+        for request in requests:
+            if request.py_request_id in mgr.kv_cache_map:
+                mgr.free_resources(request)
+        mgr.shutdown()
+
+
 def test_v2_hybrid_saves_conversation_plan_only_after_final_context_chunk():
     mgr = object.__new__(MambaHybridCacheManagerV2)
     mgr.enable_block_reuse = True

tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py

@@ -127,7 +127,7 @@ def _create_manager(
             max_gpu_total_bytes=gpu_bytes,
             max_util_for_resume=1.0,
             max_attention_window=max_attention_window,
-            block_reuse_config=BlockReuseConfig(block_reuse_policy=block_reuse_policy),
+            block_reuse_config=BlockReuseConfig(policy=block_reuse_policy),
         ),
         CacheType.SELF,
         num_layers=num_layers,