[Bugfix] Take the sliding window from the layer, not the KV cache group (#51756)
Signed-off-by: Nick Hill <nickhill123@gmail.com>
Nick Hill committed
Aug 11, 2026 at 09:10 UTC
513f83e7ecee5e39be727fcb7b7de660bd851cf2
4 files changed
+116
-26
tests/v1/attention/test_group_sliding_window.py
new
+48
@@ -0,0 +1,48 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""A KV cache group may hold both windowed and global layers — e.g. Gemma-3 with
4
+the hybrid KV cache manager disabled, where the sliding-window layers are
5
+promoted to full-attention *storage* and merged into one group whose spec still
6
+records the window. Backends must take the window from the layers, never from
7
+that group spec, or the global layers get windowed too and silently lose access
8
+to everything older than the window.
9
+"""
10
+
11
+from types import SimpleNamespace
12
+from unittest.mock import MagicMock, patch
13
+
14
+import pytest
15
+
16
+from vllm.v1.attention.backends.cpu_attn import (
17
+ CPUAttentionBackendImpl,
18
+ CPUAttentionMetadataBuilder,
19
+)
20
+
21
+
22
+def _layers(layer_windows: list[int]):
23
+ """Stand-in attention layers, one per window, as one KV cache group."""
24
+ return {
25
+ f"layer_{i}": SimpleNamespace(
26
+ impl=MagicMock(spec=CPUAttentionBackendImpl, sliding_window=window)
27
+ )
28
+ for i, window in enumerate(layer_windows)
29
+ }
30
+
31
+
32
+@pytest.mark.parametrize(
33
+ "layer_windows,expected",
34
+ [
35
+ ([512, 512], 512), # uniform sliding window -> shared by the group
36
+ ([-1, -1], -1), # all global -> no window
37
+ ([512, -1], -1), # mixed -> the group cannot assume either window
38
+ ],
39
+)
40
+def test_cpu_group_sliding_window(layer_windows, expected):
41
+ layers = _layers(layer_windows)
42
+ builder = SimpleNamespace(vllm_config=None, layer_names=list(layers))
43
+ with patch(
44
+ "vllm.v1.attention.backends.cpu_attn.get_layers_from_vllm_config",
45
+ return_value=layers,
46
+ ):
47
+ window = CPUAttentionMetadataBuilder._group_sliding_window(builder)
48
+ assert window == expected
tests/v1/e2e/general/test_correctness_sliding_window.py
+28
@@ -83,6 +83,34 @@ def test_sliding_window_retrieval(
83
)
84
85
86
+@pytest.mark.parametrize("model", ["google/gemma-3-1b-it"])
87
+def test_hybrid_kv_cache_manager_output_equivalence(model, vllm_runner):
88
+ """Disabling the hybrid KV cache manager must not change what the model
89
+ computes: it promotes the sliding-window layers to full-attention *storage*
90
+ only, so the global layers must keep attending globally.
91
+
92
+ Guards against the group's KV cache spec, which then covers windowed and
93
+ global layers alike, imposing its window on the global layers too.
94
+ """
95
+ prompts, _, _ = prep_prompts(2, ln_range=model_config[model].ln_range)
96
+ sampling_params = SamplingParams(temperature=0.0, max_tokens=32)
97
+ enforce_eager = current_platform.is_rocm()
98
+
99
+ outputs = []
100
+ for disable_hybrid in (False, True):
101
+ with vllm_runner(
102
+ model,
103
+ max_model_len=None,
104
+ enable_chunked_prefill=None,
105
+ disable_hybrid_kv_cache_manager=disable_hybrid,
106
+ enforce_eager=enforce_eager,
107
+ ) as runner:
108
+ responses = runner.get_llm().generate(prompts, sampling_params)
109
+ outputs.append([response.outputs[0].text for response in responses])
110
+
111
+ assert outputs[0] == outputs[1]
112
+
113
+
114
def check_length(prompts: list[str], llm: LLM, sliding_window: int):
115
"""
116
Check if the prompt length is valid, i.e., longer than the sliding window
vllm/v1/attention/backends/cpu_attn.py
+32
-4
@@ -11,8 +11,13 @@ import torch
11
12
from vllm import _custom_ops as ops
13
from vllm import envs
14
-from vllm.config import VllmConfig, get_current_vllm_config
14
+from vllm.config import (
15
+ VllmConfig,
16
+ get_current_vllm_config,
17
+ get_layers_from_vllm_config,
18
+)
19
from vllm.logger import init_logger
20
+from vllm.model_executor.layers.attention import Attention
21
from vllm.platforms import CpuArchEnum, current_platform
22
from vllm.utils.torch_utils import is_quantized_kv_cache
23
from vllm.v1.attention.backend import (
@@ -151,9 +156,8 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
156
)
157
self.head_dim = kv_cache_spec.head_size
158
self.dtype = vllm_config.model_config.dtype
154
- self.window_size = getattr(kv_cache_spec, "sliding_window", -1)
155
- if self.window_size is None:
156
- self.window_size = -1
159
+ # Resolved from the layers on the first build(), once they exist.
160
+ self.window_size: int | None = None
161
self.block_size = vllm_config.cache_config.block_size
162
self.kv_cache_dtype = vllm_config.cache_config.cache_dtype
163
self.isa = _get_attn_isa(
@@ -167,12 +171,36 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata]
171
kv_cache_spec, EncoderOnlyAttentionSpec
172
)
173
174
+ def _group_sliding_window(self) -> int:
175
+ """The window shared by every layer in this group, else -1 (no window).
176
+
177
+ Taken from the layers rather than the group spec: one KV cache group can
178
+ hold both windowed and global layers (e.g. Gemma-3 with the hybrid KV
179
+ cache manager disabled), and the scheduler metadata built here is shared
180
+ by the whole group, so it may only assume a window all of them agree on.
181
+ """
182
+ layers = get_layers_from_vllm_config(
183
+ self.vllm_config, Attention, self.layer_names
184
+ )
185
+ windows = {
186
+ layer.impl.sliding_window
187
+ for layer in layers.values()
188
+ if isinstance(layer.impl, CPUAttentionBackendImpl)
189
+ }
190
+ if len(windows) != 1:
191
+ return -1
192
+ window = windows.pop()
193
+ return -1 if window is None else window
194
+
195
def build(
196
self,
197
common_prefix_len: int,
198
common_attn_metadata: CommonAttentionMetadata,
199
fast_build: bool = False,
200
) -> CPUAttentionMetadata:
201
+ if self.window_size is None:
202
+ self.window_size = self._group_sliding_window()
203
+
204
num_reqs = common_attn_metadata.num_reqs
205
num_actual_tokens = common_attn_metadata.num_actual_tokens
206
max_query_len = common_attn_metadata.max_query_len
vllm/v1/attention/backends/flash_attn.py
+8
-22
@@ -287,8 +287,6 @@ class FlashAttentionMetadata:
287
288
causal: bool | torch.Tensor = True
289
290
- sliding_window: tuple[int, int] | None = None
291
-
290
# PrefixLM bidirectional range containing each scheduled query token.
291
# Shape: (num_actual_tokens, 2) int32, absolute [start, end] bounds;
292
# (-1, -1) for query tokens outside every multimodal range.
@@ -713,13 +711,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
711
if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32:
712
causal = causal.to(torch.int32)
713
716
- # Symmetrize the spec's sliding_window for non-causal attention
717
- group_sliding_window = getattr(self.kv_cache_spec, "sliding_window", None)
718
- base_window = (
719
- (-1, -1) if group_sliding_window is None else (group_sliding_window - 1, 0)
720
- )
721
- effective_sliding_window = _maybe_symmetrize_window(base_window, causal)
722
-
714
attn_metadata = FlashAttentionMetadata(
715
num_actual_tokens=num_actual_tokens,
716
max_query_len=max_query_len,
@@ -743,7 +734,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
734
prefix_scheduler_metadata=prefix_scheduler_metadata,
735
max_num_splits=max_num_splits,
736
causal=causal,
746
- sliding_window=effective_sliding_window,
737
)
738
739
# Compute mm_prefix range tensor if the batch contains
@@ -1047,18 +1037,18 @@ class FlashAttentionImpl(AttentionImpl):
1037
)
1038
return output
1039
else:
1050
- window = (
1051
- attn_metadata.sliding_window
1052
- if attn_metadata.sliding_window is not None
1053
- else self.sliding_window
1054
- )
1040
+ causal = attn_metadata.causal
1041
+ is_dynamic_causal = isinstance(causal, torch.Tensor)
1042
+
1043
+ # The layer's own window wins over the group's: one KV cache
1044
+ # group can hold both windowed and global layers (e.g. Gemma-3
1045
+ # with the hybrid KV cache manager disabled), and the group spec
1046
+ # cannot describe both.
1047
+ window = _maybe_symmetrize_window(self.sliding_window, causal)
1048
sliding_window_size: list[int] | None = (
1049
list(window) if window is not None else None
1050
)
1051
1059
- causal = attn_metadata.causal
1060
- is_dynamic_causal = isinstance(causal, torch.Tensor)
1061
-
1052
mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor
1053
mm_mask_mod = None
1054
mm_aux = None
@@ -1068,10 +1058,6 @@ class FlashAttentionImpl(AttentionImpl):
1058
and causal is True
1059
and self.vllm_flash_attn_version == 4
1060
):
1071
- # Use the layer impl's window, not attn_metadata's. The
1072
- # metadata field comes from kv_cache_spec (model-wide, e.g.
1073
- # Gemma4's 512), while the impl holds the per-layer window
1074
- # the mask_mod must encode (e.g. a test override).
1061
# Triton convention: 1 + window_size[0]. Global layers store
1062
# (-1, -1) → sw stays None.
1063
layer_window = self.sliding_window