[None][fix] Align the budget split, spec layers and pool_ratio with derived per-layer KV windows

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).
 

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,24 +2585,42 @@ 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
     # free-memory budget; (2) draft layers live at global indices past the
     # 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

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