diff --git a/docs/source/developer-guide/telemetry.md b/docs/source/developer-guide/telemetry.md index 90d9f3b9c7d8..b6b26dcd4ec1 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_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` | `` | `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/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py b/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py index e8577cb6ceaa..c4ebd8c6460c 100644 --- a/examples/disaggregated/slurm/cache_transceiver_test/run_cache_transceiver_test.py +++ b/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 diff --git a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py b/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py index 518c7b711162..006f4056be24 100644 --- a/tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py +++ b/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. diff --git a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py b/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py index e7cab986c2ef..dfafee2f8662 100644 --- a/tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py +++ b/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__( diff --git a/tensorrt_llm/llmapi/__init__.py b/tensorrt_llm/llmapi/__init__.py index 1bd895dbd59b..0d5b0667f8b9 100644 --- a/tensorrt_llm/llmapi/__init__.py +++ b/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', diff --git a/tensorrt_llm/llmapi/llm_args.py b/tensorrt_llm/llmapi/llm_args.py index 68d9fdbeb207..5404fed86b19 100644 --- a/tensorrt_llm/llmapi/llm_args.py +++ b/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}) diff --git a/tensorrt_llm/usage/llm_args_golden_manifest.json b/tensorrt_llm/usage/llm_args_golden_manifest.json index f20e02169d62..e530f74aeb6d 100644 --- a/tensorrt_llm/usage/llm_args_golden_manifest.json +++ b/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": "", + "converter": "", + "kind": "value", + "path": "kv_cache_config.block_reuse_config.max_num_turns" }, { "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 b625cc974dee..4a103840160c 100644 --- a/tests/unittest/_torch/executor/test_kv_cache_manager_v2.py +++ b/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: diff --git a/tests/unittest/_torch/executor/test_mamba_cache_manager.py b/tests/unittest/_torch/executor/test_mamba_cache_manager.py index e7d61248e7b4..09d4a9a152ee 100644 --- a/tests/unittest/_torch/executor/test_mamba_cache_manager.py +++ b/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, diff --git a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py b/tests/unittest/disaggregated/test_cache_transceiver_single_process.py index 26b9e2e9199d..00eaf14a10da 100644 --- a/tests/unittest/disaggregated/test_cache_transceiver_single_process.py +++ b/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 diff --git a/tests/unittest/disaggregated/test_kv_transfer.py b/tests/unittest/disaggregated/test_kv_transfer.py index c2cfc1382e37..3001edddbb74 100644 --- a/tests/unittest/disaggregated/test_kv_transfer.py +++ b/tests/unittest/disaggregated/test_kv_transfer.py @@ -18,7 +18,7 @@ # progress thread is enough here: these tests verify transfer logic, not # transfer-engine threading. os.environ["TRTLLM_NIXL_NUM_THREADS"] = "1" -from dataclasses import dataclass +from dataclasses import dataclass, field from typing import List, Optional import numpy as np @@ -47,7 +47,7 @@ from tensorrt_llm.bindings import DataType from tensorrt_llm.bindings import LayerType as LayerTypeCpp from tensorrt_llm.bindings import ModelConfig as ModelConfigCpp -from tensorrt_llm.llmapi.llm_args import KvCacheConfig +from tensorrt_llm.llmapi.llm_args import BlockReuseConfig, KvCacheConfig from tensorrt_llm.logger import logger # Default to 4 worker threads for all KV transfer tests in this module. @@ -83,7 +83,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 # V2 specific field max_util_for_resume: float = 0.95 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 9c130f12977e..9df2f238d8ba 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 @@ -25,7 +25,7 @@ from tensorrt_llm.bindings import DataType, SamplingConfig from tensorrt_llm.bindings.internal.batch_manager import CacheType from tensorrt_llm.bindings.internal.testing import simulate_prefill_completion_only_use_for_testing -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 from tensorrt_llm.sampling_params import SamplingParams @@ -126,7 +126,7 @@ def _create_manager( max_gpu_total_bytes=gpu_bytes, max_util_for_resume=1.0, max_attention_window=max_attention_window, - block_reuse_policy=block_reuse_policy, + block_reuse_config=BlockReuseConfig(block_reuse_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 183481805aed..49cc8393a313 100644 --- a/tests/unittest/llmapi/test_llm_args.py +++ b/tests/unittest/llmapi/test_llm_args.py @@ -29,8 +29,9 @@ from tensorrt_llm.commands.serve import get_llm_args, is_non_default_or_required from tensorrt_llm.llmapi import CapacitySchedulerPolicy, SchedulerConfig # fmt: off -from tensorrt_llm.llmapi.llm_args import (BaseLlmArgs, CacheTransceiverConfig, - CalibConfig, ContextChunkingPolicy, +from tensorrt_llm.llmapi.llm_args import (BaseLlmArgs, BlockReuseConfig, + CacheTransceiverConfig, CalibConfig, + ContextChunkingPolicy, CudaGraphConfig, DecodeCudaGraphConfig, DecodingBaseConfig, @@ -664,7 +665,7 @@ def test_KvCacheConfig_declaration(): assert KvCacheConfig().mamba_state_cache_interval is None assert KvCacheConfig().mamba_state_config.periodic_snapshot_interval == 0 assert KvCacheConfig().kv_cache_event_hash_algo == "auto" - assert KvCacheConfig().block_reuse_policy == "all_reusable" + assert KvCacheConfig().block_reuse_config == BlockReuseConfig() assert KvCacheConfig().enable_swa_scratch_reuse is False assert KvCacheConfig().use_kv_cache_manager_v2 == "auto" assert KvCacheConfig( @@ -697,7 +698,9 @@ def test_KvCacheConfig_declaration(): ), pool_ratio=[0.25, 0.75], avg_seq_len=2048, - block_reuse_policy="per_request", + block_reuse_config=BlockReuseConfig( + block_reuse_policy="per_request", + max_num_turns=2), attention_dp_events_gather_period_ms=10) pybind_config = config._to_pybind() @@ -716,7 +719,8 @@ 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_policy == "per_request" + assert config.block_reuse_config.block_reuse_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 == [ 128 @@ -726,7 +730,7 @@ def test_KvCacheConfig_declaration(): ] assert not hasattr(pybind_config, "pool_ratio") assert not hasattr(pybind_config, "avg_seq_len") - assert not hasattr(pybind_config, "block_reuse_policy") + assert not hasattr(pybind_config, "block_reuse_config") assert not hasattr(pybind_config, "enable_swa_scratch_reuse") assert KvCacheConfig( kv_cache_event_hash_algo="auto").kv_cache_event_hash_algo == "auto" @@ -739,10 +743,12 @@ 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 (KvCacheConfig(block_reuse_policy="per_conversation"). + assert (BlockReuseConfig(block_reuse_policy="per_conversation"). block_reuse_policy == "per_conversation") with pytest.raises(ValidationError): - KvCacheConfig(block_reuse_policy="invalid") + BlockReuseConfig(block_reuse_policy="invalid") + with pytest.raises(ValidationError): + BlockReuseConfig(max_num_turns=0) def test_MambaStateConfig_defaults_use_independent_lists(): @@ -823,7 +829,8 @@ def test_KvCacheConfig_warns_when_disabling_periodic_conversation_snapshots( lambda message: warnings_seen.append(message)) config = KvCacheConfig( - block_reuse_policy="per_conversation", + block_reuse_config=BlockReuseConfig( + block_reuse_policy="per_conversation"), mamba_state_config=MambaStateConfig( periodic_snapshot_interval=64, additional_snapshot_offsets_from_end=[0], @@ -834,11 +841,13 @@ 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_policy=per_conversation" in warnings_seen[0] + assert ("block_reuse_config.block_reuse_policy=per_conversation" + in warnings_seen[0]) assert "setting it to 0" in warnings_seen[0] warnings_seen.clear() - KvCacheConfig(block_reuse_policy="per_conversation") + KvCacheConfig(block_reuse_config=BlockReuseConfig( + block_reuse_policy="per_conversation")) assert warnings_seen == [] @@ -3420,7 +3429,8 @@ def _capture_warnings(monkeypatch): ), ( KvCacheConfig( - block_reuse_policy="per_conversation", + block_reuse_config=BlockReuseConfig( + block_reuse_policy="per_conversation"), mamba_state_config=MambaStateConfig( periodic_snapshot_interval=64), use_kv_cache_manager_v2=True,