[None][fix] Fix window vector layer indexing

tensorrt_llm/_torch/attention_backend/flashinfer.py

@@ -587,7 +587,10 @@ def _sanitize_swa_page_indices(self, page_indices: torch.Tensor,
         """Replace evicted SWA pages with a safe in-range page index."""
         window_vec = getattr(self.kv_cache_manager, 'max_attention_window_vec',
                              None)
-        if not window_vec or window_vec[layer_idx % len(window_vec)] is None:
+        if not window_vec:
+            return
+        local_layer_idx = self.kv_cache_manager.layer_offsets[layer_idx]
+        if window_vec[local_layer_idx] is None:
             return
 
         # KVCacheManagerV2 marks evicted out-of-window pages with -1.

tensorrt_llm/_torch/attention_backend/trtllm.py

@@ -2541,12 +2541,10 @@ def forward(
         # the model does not specify one, so paged-context attention does not
         # read stale page-table entries (BAD_PAGE_INDEX).
         if forward_args.attention_window_size is None and metadata.kv_cache_manager is not None:
-            window_vec = getattr(metadata.kv_cache_manager,
-                                 'max_attention_window_vec', None)
-            if window_vec:
-                window = window_vec[self.local_layer_idx % len(window_vec)]
-                if window is not None:
-                    forward_args.attention_window_size = window
+            window = metadata.kv_cache_manager.max_attention_window_vec[
+                self.local_layer_idx]
+            if window is not None:
+                forward_args.attention_window_size = window
         if forward_args.attention_window_size is None:
             forward_args.attention_window_size = metadata.max_seq_len
 

tensorrt_llm/_torch/pyexecutor/_util.py

@@ -15,7 +15,7 @@
 import copy
 import dataclasses
 import os
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Sequence, Union
 
 import torch
 
@@ -524,6 +524,25 @@ def draft_config_defines_attention_layout(
         or bool(getattr(draft_pretrained_config, "layer_types", None)))
 
 
+def _expand_attention_window_pattern_to_global_layers(
+    max_attention_window: Optional[Sequence[int]],
+    layer_mask: Sequence[bool],
+) -> Optional[List[int]]:
+    """Expand an enabled-layer pattern into physical global-layer order."""
+    if max_attention_window is None:
+        return None
+
+    pattern = list(max_attention_window)
+    global_windows = [pattern[0]] * len(layer_mask)
+    enabled_layer_offset = 0
+    for layer_idx, enabled in enumerate(layer_mask):
+        if enabled:
+            global_windows[layer_idx] = pattern[enabled_layer_offset %
+                                                len(pattern)]
+            enabled_layer_offset += 1
+    return global_windows
+
+
 def _derive_draft_max_attention_window(
     kv_cache_config: KvCacheConfig,
     draft_pretrained_config: object,
@@ -1560,6 +1579,11 @@ def _create_one_model_draft_kv_cache_manager(
             kv_cache_config,
             max_seq_len,
             estimating_kv_cache=estimating_kv_cache)
+        draft_kv_config.max_attention_window = (
+            _expand_attention_window_pattern_to_global_layers(
+                draft_kv_config.max_attention_window,
+                spec_dec_layer_mask,
+            ))
         if (not uses_vswa_kv_cache_layout(draft_kv_config.max_attention_window)
                 and draft_kv_config.pool_ratio is not None
                 and len(draft_kv_config.pool_ratio) != 1):

tensorrt_llm/_torch/pyexecutor/kv_cache_manager_v2.py

@@ -97,10 +97,11 @@
     BaseResourceManager,
     CacheTypeCpp,
     DataType,
-    KVCacheManager,
     ModelConfigCpp,
     ModelConfigPython,
+    _clamp_max_attention_window_vec,
     _populate_dummy_mrope_config,
+    _project_max_attention_window_vec,
     get_pp_layers,
     request_context,
 )
@@ -386,6 +387,17 @@ def _get_single_swa_pool_slot_bytes(
     return sum(layer_sizes) * int(tokens_per_block)
 
 
+def _resolve_v2_max_attention_window_vec(
+    max_attention_window_vec: Optional[Sequence[int]],
+    max_seq_len: int,
+    pp_layers: Sequence[int],
+) -> List[Optional[int]]:
+    """Resolve a V2 window pattern into exact cache-local layer order."""
+    clamped_window_pattern = _clamp_max_attention_window_vec(max_attention_window_vec, max_seq_len)
+    local_windows = _project_max_attention_window_vec(clamped_window_pattern, pp_layers)
+    return [None if window == max_seq_len else window for window in local_windows]
+
+
 def _get_static_cache_size_layer_components(
     model_config: ModelConfigPython,
     mapping: Mapping,
@@ -424,24 +436,38 @@ def _get_static_cache_size_layer_components(
         )
         layer_size = cache_size_per_token * 2
 
-    num_attention_layers = KVCacheManager._resolve_num_attention_layers(
-        model_config, mapping, num_layers
-    )
+    if num_layers is None:
+        total_attention_layers = model_config.get_num_attention_layers()
+        local_layer_ids, _ = get_pp_layers(total_attention_layers, mapping)
+    else:
+        local_layer_ids = list(range(max(num_layers, 1)))
+    num_attention_layers = len(local_layer_ids)
     layer_sizes = [layer_size] * num_attention_layers
     window_pattern = kv_cache_config.max_attention_window if kv_cache_config is not None else None
 
-    def get_window_size(layer_idx: int) -> Optional[int]:
-        if window_pattern is None or not isinstance(window_pattern, (list, tuple)):
-            return None
-        window_size = window_pattern[layer_idx % len(window_pattern)]
+    # Static estimation accepts an unknown max_seq_len and treats recurrent-state
+    # sentinels as full-attention cost. Runtime resolution must preserve those
+    # sentinels for layer construction, so it cannot be reused here directly.
+    def normalize_window_size(window_size: Optional[int]) -> Optional[int]:
         if window_size is None or window_size <= 0:
             return None
         window_size = int(window_size)
+        if max_seq_len is not None:
+            window_size = min(window_size, int(max_seq_len))
         if max_seq_len is not None and window_size == int(max_seq_len):
             return None
         return window_size
 
-    attention_windows = [get_window_size(layer_idx) for layer_idx in range(num_attention_layers)]
+    if window_pattern is None or not isinstance(window_pattern, (list, tuple)):
+        attention_windows = [None] * num_attention_layers
+    else:
+        local_window_pattern = _project_max_attention_window_vec(
+            window_pattern,
+            local_layer_ids,
+        )
+        attention_windows = [
+            normalize_window_size(window_size) for window_size in local_window_pattern
+        ]
     return layer_sizes, attention_windows
 
 
@@ -740,31 +766,43 @@ def _update_kv_cache_draft_token_location(
     else:
         use_paged_kv_cache = False
     assert use_paged_kv_cache, "Only paged kv cache is supported"
-    assert len(cache_manager.max_attention_window_vec) == 1, (
+    assert len(set(cache_manager.max_attention_window_vec)) == 1, (
         "Currently, only one max attention window size is supported."
     )
+    max_attention_window = cache_manager.max_attention_window_vec[0]
+    if max_attention_window is None:
+        max_attention_window = cache_manager.max_seq_len
 
     if use_paged_kv_cache:
         assert len(set(cache_manager.num_kv_heads_per_layer)) == 1, (
             "update_kv_cache_draft_token_location requires uniform num_kv_heads across all layers, "
             f"but got {cache_manager.num_kv_heads_per_layer}"
         )
+        local_pool_ids = {
+            int(cache_manager.kv_cache_pool_mapping[layer_idx][0])
+            for layer_idx in range(cache_manager.num_local_layers)
+        }
+        assert len(local_pool_ids) == 1, (
+            "update_kv_cache_draft_token_location requires all local layers "
+            f"in one KV pool, got pools {sorted(local_pool_ids)}"
+        )
+        pool_idx = local_pool_ids.pop()
         torch.ops.tensorrt_llm.update_kv_cache_draft_token_location(
             accepted_draft_token_offsets,
             packed_accepted_draft_tokens_indices,
             past_key_value_lengths,
             True,
-            cache_manager.num_layers,
+            cache_manager.num_local_layers,
             # Use TP-sharded num_kv_heads (per-rank) instead of the unsharded
             # total so the C++ kernel computes correct strides and grid dims.
             cache_manager.num_kv_heads_per_layer[0],
             int(cache_manager.head_dim * kv_cache_dtype_byte_size),
             cache_manager.max_total_draft_tokens,
-            cache_manager.max_attention_window_vec[0],
+            max_attention_window,
             rewind_draft_token_separate_adjustments,
             None,
-            cache_manager.kv_cache_pool_pointers,
-            attn_metadata.kv_cache_block_offsets,
+            cache_manager.kv_cache_pool_pointers[pool_idx],
+            attn_metadata.kv_cache_block_offsets[pool_idx],
             cache_manager.max_blocks_per_seq,
             cache_manager.tokens_per_block,
             None,
@@ -966,24 +1004,13 @@ def __init__(
         )
         logger.info(f"[KVCacheManager] execution_stream: {self._stream}")
 
-        # Determine max_attention_window_vec
-        if kv_cache_config.max_attention_window is not None:
-            self.max_attention_window_vec = (
-                kv_cache_config.max_attention_window.copy()
-            )  # Make a copy to avoid modifying original
-            # Clamp all window sizes to max_seq_len before calculating the
-            # number of KV cache blocks. This prevents the KV cache pool from
-            # being skewed by the largest window values.
-            self.max_attention_window_vec = [
-                min(max_seq_len, w) for w in self.max_attention_window_vec
-            ]
-
-            self.max_attention_window_vec = [
-                None if w == max_seq_len else w for w in self.max_attention_window_vec
-            ]
-
-        else:
-            self.max_attention_window_vec = [None]
+        # Materialize an exact per-local-layer vector for cache and attention consumers.
+        self.max_attention_window_vec = _resolve_v2_max_attention_window_vec(
+            kv_cache_config.max_attention_window,
+            max_seq_len,
+            self.pp_layers,
+        )
+        assert len(self.max_attention_window_vec) == self.num_local_layers
 
         event_window_size = max(
             self.max_seq_len if window_size is None else int(window_size)
@@ -1604,14 +1631,11 @@ def _kv_pool_mapping_offset(
     def _get_runtime_cache_size_layer_components(self) -> tuple[List[int], List[Optional[int]]]:
         layer_sizes = []
         attention_windows = []
-        pattern_len = len(self.max_attention_window_vec)
         for local_layer_idx in range(self.num_local_layers):
             layer_sizes.append(
                 self.get_layer_bytes_per_token(local_layer_idx=local_layer_idx, data_role=Role.ALL)
             )
-            attention_windows.append(
-                self.max_attention_window_vec[self.pp_layers[local_layer_idx] % pattern_len]
-            )
+            attention_windows.append(self.max_attention_window_vec[local_layer_idx])
         return layer_sizes, attention_windows
 
     def _get_max_tokens_from_quota(self, quota: int) -> float:
@@ -2083,9 +2107,7 @@ def _build_base_config(
                 AttentionLayerConfig(
                     layer_id=layer_id,
                     buffers=buffers,
-                    sliding_window_size=self.max_attention_window_vec[
-                        self.pp_layers[layer_id] % len(self.max_attention_window_vec)
-                    ],
+                    sliding_window_size=self.max_attention_window_vec[layer_id],
                     num_sink_tokens=None,
                 )
             )

tensorrt_llm/_torch/pyexecutor/mamba_cache_manager.py

@@ -2332,12 +2332,13 @@ def __init__(
             )
             kv_cache_config.enable_partial_reuse = False
 
-        kv_cache_config.max_attention_window = []
-        for i in range(len(layer_mask)):
-            if layer_mask[i]:
-                kv_cache_config.max_attention_window.append(
-                    LinearCacheType.RECURRENT_STATES.
-                    value if mamba_layer_mask[i] else max_seq_len)
+        # Keep the vector in physical global-layer order. Disabled entries are
+        # placeholders that are never projected into this manager.
+        kv_cache_config.max_attention_window = [
+            LinearCacheType.RECURRENT_STATES.value
+            if mamba_layer_mask[layer_idx] else max_seq_len
+            for layer_idx in range(len(layer_mask))
+        ]
 
         recurrent_states_window = LinearCacheType.RECURRENT_STATES.value
         local_windows = {