diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index b6b26dcd4ec1..c95b0674be6a 100644 --- a/docs/source/developer-guide/telemetry.md +++ b/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` | `` | `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` | `` | `value` | | | | `kv_cache_config.cross_kv_cache_fraction` | `Optional[float]` | `value` | | | | `kv_cache_config.disk_cache_size` | `Optional[Annotated[int, Ge(ge=0)]]` | `value` | | | diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index b8b395abdc4d..3f3da5058d0e 100644 --- a/docs/source/features/kvcache.md +++ b/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] diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 967c7d413f68..8d31e943788e 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index 75eec8a8b5aa..7000b2bdaf0f 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/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__( diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 00f2d2c4f678..1f7e855280c3 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/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}) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index c317ee871293..f7b26ee0a1ff 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/tensorrt_llm/usage/llm_args_golden_manifest.json @@ -609,6 +609,13 @@ "kind": "value", "path": "kv_cache_config.avg_seq_len" }, + { + "allowed_values": [], + "annotation": "", + "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": "", - "converter": "", - "kind": "value", - "path": "kv_cache_config.block_reuse_config.max_num_turns" + "path": "kv_cache_config.block_reuse_config.policy" }, { "allowed_values": [], diff --git a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py b/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py index 4a103840160c..89a8b81fdc99 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/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, ), ), diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index aab465e0a475..8ceddf7e03a3 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/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 diff --git a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py b/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py index c24be87516ab..d848244d16f5 100644 --- a/tests/unittest/kv_cache_manager_v2_tests/test_kv_cache_stats_behavior.py +++ b/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, diff --git a/tests/unittest/llmapi/test_llm_args.py b/tests/unittest/llmapi/test_llm_args.py index 1fb188c84563..09a7ffb167c6 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -719,8 +719,7 @@ def test_KvCacheConfig_declaration(): pool_ratio=[0.25, 0.75], avg_seq_len=2048, block_reuse_config=BlockReuseConfig( - block_reuse_policy="per_request", - max_num_turns=2), + policy="per_request", max_num_turns=2), attention_dp_events_gather_period_ms=10) pybind_config = config._to_pybind() @@ -739,7 +738,7 @@ def test_KvCacheConfig_declaration(): assert config.kv_cache_event_hash_algo == "v2_sha256_64" assert config.pool_ratio == [0.25, 0.75] assert config.avg_seq_len == 2048 - assert config.block_reuse_config.block_reuse_policy == "per_request" + assert config.block_reuse_config.policy == "per_request" assert config.block_reuse_config.max_num_turns == 2 assert config.mamba_state_config.periodic_snapshot_interval == 0 assert config.mamba_state_config.additional_snapshot_offsets_from_start == [ @@ -763,14 +762,23 @@ def test_KvCacheConfig_declaration(): assert pybind_config.enable_partial_reuse == True assert pybind_config.copy_on_partial_reuse == True assert pybind_config.attention_dp_events_gather_period_ms == 10 - assert (BlockReuseConfig(block_reuse_policy="per_conversation"). - block_reuse_policy == "per_conversation") + assert BlockReuseConfig( + policy="per_conversation").policy == "per_conversation" with pytest.raises(ValidationError): - BlockReuseConfig(block_reuse_policy="invalid") + BlockReuseConfig(policy="invalid") with pytest.raises(ValidationError): BlockReuseConfig(max_num_turns=0) +@pytest.mark.cpu_only +def test_BlockReuseConfig_reports_renamed_policy_field(): + with pytest.raises(ValidationError, match="block_reuse_config\\.policy"): + KvCacheConfig.model_validate( + {"block_reuse_config": { + "block_reuse_policy": "per_request" + }}) + + @pytest.mark.cpu_only def test_MambaStateConfig_defaults_use_independent_lists(): first = MambaStateConfig() @@ -855,8 +863,7 @@ def test_KvCacheConfig_warns_when_disabling_periodic_conversation_snapshots( lambda message: warnings_seen.append(message)) config = KvCacheConfig( - block_reuse_config=BlockReuseConfig( - block_reuse_policy="per_conversation"), + block_reuse_config=BlockReuseConfig(policy="per_conversation"), mamba_state_config=MambaStateConfig( periodic_snapshot_interval=64, additional_snapshot_offsets_from_end=[0], @@ -867,13 +874,12 @@ def test_KvCacheConfig_warns_when_disabling_periodic_conversation_snapshots( assert config.mamba_state_config.additional_snapshot_offsets_from_end == [0] assert len(warnings_seen) == 1 assert "periodic_snapshot_interval=64" in warnings_seen[0] - assert ("block_reuse_config.block_reuse_policy=per_conversation" - in warnings_seen[0]) + assert ("block_reuse_config.policy=per_conversation" in warnings_seen[0]) assert "setting it to 0" in warnings_seen[0] warnings_seen.clear() KvCacheConfig(block_reuse_config=BlockReuseConfig( - block_reuse_policy="per_conversation")) + policy="per_conversation")) assert warnings_seen == [] @@ -3530,7 +3536,7 @@ def _capture_warnings(monkeypatch): ( KvCacheConfig( block_reuse_config=BlockReuseConfig( - block_reuse_policy="per_conversation"), + policy="per_conversation"), mamba_state_config=MambaStateConfig( periodic_snapshot_interval=64), use_kv_cache_manager_v2=True,