diff --git a/docs/source/features/kvcache.md b/docs/source/features/kvcache.md index d77b9fdebaac..c494cd335bde 100644 --- a/docs/source/features/kvcache.md +++ b/docs/source/features/kvcache.md @@ -94,6 +94,19 @@ Separately, Gemma4 hybrid attention and sparse-attention models are routed to V2 unconditionally: their per-layer buffer layouts cannot be represented by V1's unified pool, so `use_kv_cache_manager_v2` does not apply to them. +For a model whose `layer_types` mixes sliding-window and full-attention layers +and that publishes a single `sliding_window` (GPT-OSS, Gemma3), the V2 manager +derives one attention window per layer from `layer_types` when +`max_attention_window` is not set: sliding layers get `sliding_window`, full +layers get `max_seq_len`, and the two window sizes form two layer groups whose +pools are sized independently. The derived list is logged at startup. Set +`max_attention_window` explicitly to override the derivation; a single entry +restores one full-context pool for every layer. With derived windows, +`pool_ratio` must carry one entry per layer group (two for such a model). If a +configured `pool_ratio` does not match the derived group count, the manager +logs a warning and keeps the single-window default, so existing configurations +continue to run. + For the native V2 cold-storage representation and codec extension contract, see [KVCacheManagerV2 Cold-Page Codec Design](../developer-guide/kv-cache-cold-page-codec.md). diff --git a/tensorrt_llm/_torch/pyexecutor/_util.py b/tensorrt_llm/_torch/pyexecutor/_util.py index cfcb6e849212..bf79860387ab 100644 --- a/tensorrt_llm/_torch/pyexecutor/_util.py +++ b/tensorrt_llm/_torch/pyexecutor/_util.py @@ -506,6 +506,56 @@ def _derive_layer_type_attention_windows( return resolved +def _derive_v2_layer_type_attention_windows( + kv_cache_config: KvCacheConfig, + kv_cache_manager_cls: type, + model_config: ModelConfig, + max_seq_len: int, +) -> Optional[List[int]]: + """Per-layer windows a KVCacheManagerV2 for `model_config` is built with. + + Shared by the static per-token cost model + (`KvCacheCreator._per_manager_cache_cost`) and `_create_kv_cache_manager`, + so the budget split sizes a manager from the same windows the manager + receives. Returns `None`, leaving `kv_cache_config` as is, when the user + supplied `max_attention_window`, for any class other than KVCacheManagerV2 + (KVCacheManager keeps the single-window default), when + `_derive_layer_type_attention_windows` has nothing to derive, for hybrid + linear-attention configs (their attention layers are interleaved with + recurrent layers, and the static cost model indexes a window list by + attention-layer position while the manager indexes it by global layer id, + so the two would read a per-layer list differently), and when a + user-supplied `pool_ratio` does not carry one entry per derived layer + group (the manager would reject that arity at construction; a warning + names the fix and the configuration keeps the single pool it was written + for). + """ + if kv_cache_config.max_attention_window is not None: + return None + if not (isinstance(kv_cache_manager_cls, type) + and issubclass(kv_cache_manager_cls, KVCacheManagerV2)): + return None + config = model_config.pretrained_config + derived_windows = _derive_layer_type_attention_windows(config, max_seq_len) + if derived_windows is None: + return None + if is_hybrid_linear(config): + return None + pool_ratio = kv_cache_config.pool_ratio + num_layer_groups = len(set(derived_windows)) + if pool_ratio is not None and len(pool_ratio) != num_layer_groups: + logger.warning_once( + f"kv_cache_config.pool_ratio has {len(pool_ratio)} entries, but the " + "per-layer attention windows derived from layer_types form " + f"{num_layer_groups} layer groups; keeping the single-window " + "default. Provide one pool_ratio entry per layer group to size the " + "sliding-window and full-attention pools separately, or set " + "max_attention_window explicitly.", + key="derived_attention_windows_pool_ratio_arity") + return None + return derived_windows + + def _get_num_pool_groups_for_estimation( model_config: object, max_seq_len: int, @@ -794,6 +844,16 @@ def _per_manager_cache_cost(self, **extra_kwargs) -> CacheCost: kv_cache_config = (kv_cache_config if kv_cache_config is not None else self._kv_cache_config) + if not is_draft: + derived_windows = _derive_v2_layer_type_attention_windows( + kv_cache_config, manager_cls, model_config, self._max_seq_len) + if derived_windows is not None: + # Cost the manager with the per-layer windows + # `_create_kv_cache_manager` builds it with, so the budget + # split between the target and draft managers matches their + # pools. The draft manager never derives (see there). + kv_cache_config = kv_cache_config.model_copy( + update={"max_attention_window": derived_windows}) return CacheCost.from_raw( manager_cls.get_cache_size_per_token( model_config, @@ -2525,7 +2585,10 @@ def _create_kv_cache_manager( # picks up the smaller per-pool block count automatically. Gemma4 hybrid # always resolves to KVCacheManagerV2 (see _non_hybrid_kv_cache_manager_cls), # so the V2 guard never excludes it. A user-supplied max_attention_window - # always wins. + # always wins. `_derive_v2_layer_type_attention_windows` applies the same + # rules in the creator's static cost model, so the budget split sizes the + # manager from the windows it is built with; it also keeps the default + # when a configured `pool_ratio` does not match the derived layer groups. # Skip derivation for the one-model draft manager: (1) in one-model # spec-decode the KV memory budget is already split between target and # draft, so VSWA sizing here would size each window pool from the unsplit @@ -2533,16 +2596,31 @@ def _create_kv_cache_manager( # target's num_hidden_layers, so _project_max_attention_window_vec would # wrap them back onto pattern[0]. The draft config sets # max_attention_window=None to opt out, not to request derivation. - if (kv_cache_config.max_attention_window is None and not is_draft - and issubclass(kv_cache_manager_cls, KVCacheManagerV2)): - derived_windows = _derive_layer_type_attention_windows( - config, max_seq_len) - if derived_windows is not None: - assert uses_vswa_kv_cache_layout(derived_windows), ( - "derived per-layer windows must select a VSWA layout; a non-VSWA " - "vector would reshape the single pool instead of splitting it") - kv_cache_config = copy.copy(kv_cache_config) - kv_cache_config.max_attention_window = derived_windows + # The cross-attention pool holds encoder-side KV that the decoder's + # `layer_types` do not describe, so it keeps the default too. + derived_windows = None + if (not is_draft and kv_cache_type + == tensorrt_llm.bindings.internal.batch_manager.CacheType.SELF): + derived_windows = _derive_v2_layer_type_attention_windows( + kv_cache_config, kv_cache_manager_cls, _model_config, max_seq_len) + if derived_windows is not None: + if layer_mask is None and spec_config is not None: + # `get_pp_layers` appends the one-model speculative layers after + # the decoder layers when they share this manager, and V2 resolves + # each layer's window by global layer id: give them the full + # context the single-window default gave them instead of wrapping + # them onto the decoder pattern. + derived_windows = derived_windows + [int( + max_seq_len)] * get_num_spec_layers(spec_config) + assert uses_vswa_kv_cache_layout(derived_windows), ( + "derived per-layer windows must select a VSWA layout; a non-VSWA " + "vector would reshape the single pool instead of splitting it") + logger.info( + "Derived per-layer max_attention_window from layer_types for " + f"{kv_cache_manager_cls.__name__}: {derived_windows} " + f"({len(set(derived_windows))} distinct windows)") + kv_cache_config = copy.copy(kv_cache_config) + kv_cache_config.max_attention_window = derived_windows # Note: Gemma4 KV sharing is handled at the model level — shared layers # use cache_layer_idx to read from the target layer's cache slot via diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py index 5a6220330125..4f9554364d1c 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_budget_split.py @@ -260,6 +260,204 @@ def get_cache_size_per_token(model_config, *args, **kwargs): else: get_manager_cls.assert_not_called() + def test_target_cost_uses_derived_layer_type_windows(self) -> None: + """A target with a mixed sliding/full `layer_types` schedule on + KVCacheManagerV2 is costed from the same derived per-layer windows + `_create_kv_cache_manager` builds it with: its three sliding layers + become a fixed per-request cost and only the full layer is charged per + token, so the split matches the manager's pools. Without the derivation + the target counted four full layers per token and no fixed cost.""" + + class TargetModelConfig: + quant_config = None + is_encoder_decoder = False + pretrained_config = SimpleNamespace( + num_hidden_layers=4, + hidden_size=1024, + num_attention_heads=8, + num_key_value_heads=8, + sliding_window=512, + layer_types=[ + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + ], + ) + + def get_num_attention_layers(self) -> int: + return 4 + + class DraftModelConfig: + """A one-layer full-attention draft head without window metadata.""" + + quant_config = None + pretrained_config = SimpleNamespace( + num_hidden_layers=1, + hidden_size=1024, + num_attention_heads=8, + num_key_value_heads=8, + ) + + def get_num_attention_layers(self) -> int: + return 1 + + target_model_config = TargetModelConfig() + draft_model_config = DraftModelConfig() + # Block reuse off: with it on, the estimator also extends every sliding + # window by the one-model draft's prompt lookahead, a separate charge + # covered in test_kv_cache_estimation.py; this test is about the windows. + target_kv_config = KvCacheConfig(enable_block_reuse=False) + mode = Mock() + mode.is_external_drafter.return_value = False + # A Mock's truthy return would send get_num_extra_kv_tokens down the + # one-engine arm; this test's config is a two-model drafter. + mode.use_one_engine.return_value = False + costed_windows: list[tuple[object, list[int] | None]] = [] + + class RecordingKVCacheManager(KVCacheManagerV2): + @staticmethod + def get_cache_size_per_token( + model_config: object, *args: object, **kwargs: object + ) -> tuple[int, int]: + costed_windows.append( + (model_config, kwargs["kv_cache_config"].max_attention_window) + ) + return KVCacheManagerV2.get_cache_size_per_token(model_config, *args, **kwargs) + + max_batch_size = 2 + creator = object.__new__(KvCacheCreator) + creator._kv_cache_config = target_kv_config + creator._tokens_per_block = 64 + creator._max_seq_len = 16384 + creator._max_batch_size = max_batch_size + creator._max_num_tokens = 128 + creator._mapping = Mock(enable_attention_dp=False, tp_size=1) + creator._mapping.pp_layers.return_value = [0, 1, 2, 3] + creator._mapping.is_last_pp_rank.return_value = True + # Neutral speculative fields: _get_generation_kv_capacity reads them + # to size the generation headroom; these values keep it at the + # non-speculative baseline of one token so the window math below stays + # the point of the test. + creator._speculative_config = SimpleNamespace( + spec_dec_mode=mode, + max_total_draft_tokens=0, + max_draft_len=0, + tokens_per_gen_step=1, + ) + creator._model_engine = SimpleNamespace( + model=SimpleNamespace(model_config=target_model_config) + ) + creator._draft_model_engine = None + creator._draft_config = draft_model_config + creator._kv_cache_manager_cls = RecordingKVCacheManager + creator._is_disagg = False + creator._should_create_separate_draft_kv_cache = Mock(return_value=True) + creator._get_effective_draft_config = Mock(return_value=draft_model_config) + creator._get_num_draft_layers = Mock(return_value=1) + + target_kv, draft_kv = creator._get_target_and_draft_cache_costs() + + # K and V, 8 heads x 128 dims, bf16. + layer_bytes_per_token = 2 * 8 * 128 * 2 + # Each sliding layer retains page-granular window blocks per request: + # a 512-token window with one headroom token spans + # ceil((512 + 1 - 2) / 64) + 1 = 9 blocks = 576 tokens. Context + # additionally retains each sliding layer for the in-flight token + # batch (max_num_tokens). + window_tokens = (math.ceil((512 + 1 - 2) / 64) + 1) * 64 + sliding_bytes_per_request = 3 * window_tokens * layer_bytes_per_token + context_batch_bytes = 3 * layer_bytes_per_token * creator._max_num_tokens + assert target_kv == CacheCost( + slope=layer_bytes_per_token, + intercept=sliding_bytes_per_request * max_batch_size + context_batch_bytes, + ) + assert draft_kv == CacheCost(slope=layer_bytes_per_token, intercept=0) + target_windows = [ + windows + for model_config, windows in costed_windows + if model_config is target_model_config + ] + assert target_windows and all( + windows == [512, 512, 16384, 512] for windows in target_windows + ) + draft_windows = [ + windows + for model_config, windows in costed_windows + if model_config is draft_model_config + ] + assert draft_windows == [None] + # The creator's own config is left untouched. + assert target_kv_config.max_attention_window is None + + def test_target_cost_projects_derived_windows_onto_the_rank_layers(self) -> None: + """On a pipeline rank other than the first, the static cost model reads + the derived per-layer windows through the rank's global layer ids, as + the runtime manager does. Global windows [S, S, F, S] on the second of + two PP ranks (layers 2 and 3) cost one full layer per token and one + sliding window per request; a local 0..N-1 read would have costed the + rank as two sliding layers and no per-token bytes.""" + + class TargetModelConfig: + quant_config = None + is_encoder_decoder = False + pretrained_config = SimpleNamespace( + num_hidden_layers=4, + hidden_size=1024, + num_attention_heads=8, + num_key_value_heads=8, + sliding_window=512, + layer_types=[ + "sliding_attention", + "sliding_attention", + "full_attention", + "sliding_attention", + ], + ) + + def get_num_attention_layers(self) -> int: + return 4 + + target_model_config = TargetModelConfig() + costed_windows: list[list[int] | None] = [] + + class RecordingKVCacheManager(KVCacheManagerV2): + @staticmethod + def get_cache_size_per_token( + model_config: object, *args: object, **kwargs: object + ) -> tuple[int, int]: + costed_windows.append(kwargs["kv_cache_config"].max_attention_window) + return KVCacheManagerV2.get_cache_size_per_token(model_config, *args, **kwargs) + + max_batch_size = 2 + creator = object.__new__(KvCacheCreator) + creator._kv_cache_config = KvCacheConfig() + creator._tokens_per_block = 64 + creator._max_seq_len = 16384 + creator._max_batch_size = max_batch_size + creator._max_num_tokens = 128 + creator._speculative_config = None + # The second of two pipeline ranks: it holds global layers 2 and 3. + creator._mapping = Mock(enable_attention_dp=False, tp_size=1) + creator._mapping.pp_layers.return_value = [2, 3] + + target_kv = creator._per_manager_cache_cost(RecordingKVCacheManager, target_model_config) + + # K and V, 8 heads x 128 dims, bf16. + layer_bytes_per_token = 2 * 8 * 128 * 2 + # The rank's one sliding layer retains ceil((512 + 1 - 2) / 64) + 1 = 9 + # blocks = 576 tokens per request, plus the context charge for the + # in-flight token batch (max_num_tokens). + window_tokens = (math.ceil((512 + 1 - 2) / 64) + 1) * 64 + assert target_kv == CacheCost( + slope=layer_bytes_per_token, + intercept=window_tokens * layer_bytes_per_token * max_batch_size + + layer_bytes_per_token * creator._max_num_tokens, + ) + # The cost model receives the global list; the projection onto the + # rank's layers happens inside the manager's static estimator. + assert costed_windows == [[512, 512, 16384, 512]] + def test_v1_mixed_draft_build_uses_original_max_seq_len(self, mocker): c = _make_creator(max_gpu_total_bytes=10 * GB) original_max_seq_len = 16384 diff --git a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py index 6c819f101fb0..b7604a635275 100644 --- a/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py +++ b/tests/unittest/_torch/executor/kv_cache/test_kv_cache_estimation.py @@ -17,11 +17,19 @@ import pytest import torch +import tensorrt_llm.bindings.internal.batch_manager as batch_manager from tensorrt_llm._torch.model_config import ModelConfig from tensorrt_llm._torch.models.modeling_multimodal_mixin import MultimodalModelMixin -from tensorrt_llm._torch.pyexecutor._util import CacheCost, KvCacheCreator +from tensorrt_llm._torch.pyexecutor._util import ( + CacheCost, + KvCacheCreator, + _create_kv_cache_manager, + _derive_v2_layer_type_attention_windows, +) from tensorrt_llm._torch.pyexecutor.config_utils import get_layer_attention_window from tensorrt_llm._torch.pyexecutor.kv_cache.kv_cache_manager_v2 import KVCacheManagerV2 +from tensorrt_llm._torch.pyexecutor.resource_manager import KVCacheManager +from tensorrt_llm._torch.speculative.interface import SpeculativeDecodingMode from tensorrt_llm.inputs.multimodal import MultimodalParams from tensorrt_llm.llmapi.llm_args import ( KvCacheConfig, @@ -917,8 +925,6 @@ def test_mla_branch_forwards_max_num_tokens_to_manager() -> None: and OOM risk during KV cache estimation). """ - from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager - captured_kwargs = {} class _RecordingManager: @@ -1051,8 +1057,6 @@ def test_manager_estimation_clamps_only_temporary_avg_seq_len( ) -> None: import torch - from tensorrt_llm._torch.pyexecutor._util import _create_kv_cache_manager - captured_configs = [] class _RecordingKVCacheManagerV2(KVCacheManagerV2): @@ -1158,3 +1162,200 @@ def test_separate_one_model_draft_normalizes_target_pool_ratio() -> None: assert draft_config.pool_ratio == [1.0] assert create_manager.call_args.kwargs["cold_page_codec_provider"] is codec_provider assert creator._kv_cache_config.pool_ratio == target_pool_ratio + + +# --------------------------------------------------------------------------- +# Per-layer attention windows derived from a mixed layer_types schedule +# --------------------------------------------------------------------------- +# +# `_derive_layer_type_attention_windows` (covered in +# test_layer_type_attention_windows.py) turns a mixed sliding/full +# `layer_types` schedule into one window per layer for KVCacheManagerV2. +# `_derive_v2_layer_type_attention_windows` applies it in +# `_create_kv_cache_manager` (tested here) and in the static per-token cost +# model the creator splits the GPU budget with (tested in +# test_kv_cache_budget_split.py), and keeps the single-window default for +# hybrid linear-attention configs and for a `pool_ratio` that does not match +# the derived layer groups. The creator also gives one-model speculative +# layers appended after the decoder stack the full context, and leaves the +# cross-attention pool alone. + +_SLIDING = "sliding_attention" +_FULL = "full_attention" + + +def _create_manager_and_capture_config( + pretrained: SimpleNamespace, + kv_cache_config: KvCacheConfig, + max_seq_len: int, + manager_cls: type[KVCacheManager] | type[KVCacheManagerV2] = KVCacheManagerV2, + **factory_overrides: object, +) -> KvCacheConfig: + """Run `_create_kv_cache_manager` with a recording subclass of `manager_cls` + and return the `KvCacheConfig` the manager was constructed with. + `factory_overrides` replace the default keyword arguments of the factory + call (for example `spec_config`, `layer_mask` or `kv_cache_type`).""" + captured = [] + + class _RecordingManager(manager_cls): + def __init__(self, kv_cache_config: KvCacheConfig, *args: object, **kwargs: object) -> None: + captured.append(kv_cache_config) + + model_config = Mock() + model_config.pretrained_config = pretrained + model_config.quant_config = None + + factory_kwargs: dict[str, object] = dict( + model_engine=None, + kv_cache_manager_cls=_RecordingManager, + mapping=Mock(), + kv_cache_config=kv_cache_config, + tokens_per_block=32, + max_seq_len=max_seq_len, + max_batch_size=4, + spec_config=None, + sparse_attention_config=None, + max_num_tokens=1024, + max_beam_width=1, + kv_connector_manager=None, + model_config=model_config, + dtype=torch.bfloat16, + is_draft=False, + ) + factory_kwargs.update(factory_overrides) + _create_kv_cache_manager(**factory_kwargs) + + assert len(captured) == 1 + return captured[0] + + +def _mixed_schedule_pretrained(layer_types: list[str], **overrides: object) -> SimpleNamespace: + fields = dict( + hidden_size=1024, + num_attention_heads=8, + num_key_value_heads=8, + num_hidden_layers=len(layer_types), + vocab_size=32000, + layer_types=layer_types, + sliding_window=512, + ) + fields.update(overrides) + return SimpleNamespace(**fields) + + +def _eagle3_one_model_spec_config(num_draft_hidden_layers: int | None = None) -> SimpleNamespace: + """The fields `get_num_spec_layers` and `should_use_separate_draft_kv_cache` + read from an Eagle3 one-model speculative config.""" + return SimpleNamespace( + spec_dec_mode=SpeculativeDecodingMode.EAGLE3_ONE_MODEL, + _use_shared_kv_cache=False, + _allow_separate_draft_kv_cache=True, + _num_draft_hidden_layers=num_draft_hidden_layers, + ) + + +def test_create_kv_cache_manager_pool_ratio_arity_mismatch_keeps_single_window() -> None: + """A `pool_ratio` written for the single pool (one entry) does not match the + two layer groups the derived windows would create. Rather than failing the + manager's arity check at startup, the derivation is skipped with a warning + and the configuration keeps the single-window default it was written for.""" + kv_cache_config = KvCacheConfig(pool_ratio=[1.0]) + + with patch("tensorrt_llm._torch.pyexecutor._util.logger") as mock_logger: + manager_config = _create_manager_and_capture_config( + _mixed_schedule_pretrained([_SLIDING, _SLIDING, _FULL, _SLIDING]), + kv_cache_config, + max_seq_len=2048, + ) + + assert manager_config is kv_cache_config + assert manager_config.max_attention_window is None + mock_logger.warning_once.assert_called_once() + message = mock_logger.warning_once.call_args.args[0] + assert "pool_ratio has 1 entries" in message + assert "2 layer groups" in message + + +def test_create_kv_cache_manager_pool_ratio_per_layer_group_keeps_derivation() -> None: + """One `pool_ratio` entry per derived layer group is the intended pairing.""" + kv_cache_config = KvCacheConfig(pool_ratio=[0.5, 0.5]) + + manager_config = _create_manager_and_capture_config( + _mixed_schedule_pretrained([_SLIDING, _SLIDING, _FULL, _SLIDING]), + kv_cache_config, + max_seq_len=2048, + ) + + assert manager_config.max_attention_window == [512, 512, 2048, 512] + assert manager_config.pool_ratio == [0.5, 0.5] + + +def test_derive_v2_windows_skip_hybrid_linear_configs() -> None: + """A hybrid linear-attention model (recurrent layers interleaved with + attention layers) keeps the single-window default even when its + `layer_types` and `sliding_window` would derive a per-layer list: the + static cost model indexes windows by attention-layer position while the + manager indexes them by global layer id, so the two would disagree.""" + model_config = Mock() + model_config.pretrained_config = _mixed_schedule_pretrained( + [_SLIDING, _FULL, _SLIDING, _FULL], hybrid_override_pattern="M*M*" + ) + + assert ( + _derive_v2_layer_type_attention_windows( + KvCacheConfig(), KVCacheManagerV2, model_config, max_seq_len=2048 + ) + is None + ) + + # Without the hybrid marker the same schedule derives a per-layer list. + model_config.pretrained_config = _mixed_schedule_pretrained([_SLIDING, _FULL, _SLIDING, _FULL]) + assert _derive_v2_layer_type_attention_windows( + KvCacheConfig(), KVCacheManagerV2, model_config, max_seq_len=2048 + ) == [512, 2048, 512, 2048] + + +@pytest.mark.parametrize( + ("layer_mask", "expected_windows"), + [ + (None, [512, 512, 2048, 512, 2048, 2048]), + ([True] * 4, [512, 512, 2048, 512]), + ], + ids=["spec_layers_appended", "target_only_mask"], +) +def test_create_kv_cache_manager_keeps_appended_spec_layers_full_context( + layer_mask: list[bool] | None, + expected_windows: list[int], +) -> None: + """Without a `layer_mask`, `get_pp_layers` appends the one-model speculative + layers after the decoder layers and V2 reads the window list modulo its + length, so the derived list gets one full-context entry per appended layer. + A target-only mask holds the decoder layers alone and needs no extra entry.""" + manager_config = _create_manager_and_capture_config( + _mixed_schedule_pretrained([_SLIDING, _SLIDING, _FULL, _SLIDING]), + KvCacheConfig(), + max_seq_len=2048, + spec_config=_eagle3_one_model_spec_config(num_draft_hidden_layers=2), + layer_mask=layer_mask, + ) + + assert manager_config.max_attention_window == expected_windows + + +def test_create_kv_cache_manager_cross_pool_keeps_single_window_default() -> None: + """The cross-attention pool stores encoder-side KV, which the decoder's + `layer_types` do not describe.""" + kv_cache_config = KvCacheConfig() + + manager_config = _create_manager_and_capture_config( + _mixed_schedule_pretrained([_SLIDING, _SLIDING, _FULL, _SLIDING]), + kv_cache_config, + max_seq_len=2048, + kv_cache_type=batch_manager.CacheType.CROSS, + num_layers=4, + num_kv_heads=8, + head_dim=128, + ) + + assert manager_config is kv_cache_config + assert manager_config.max_attention_window is None