diff --git a/tests/v1/attention/test_group_sliding_window.py b/tests/v1/attention/test_group_sliding_window.py new file mode 100644 index 0000000000..436e54a9b9 --- /dev/null +++ b/tests/v1/attention/test_group_sliding_window.py @@ -0,0 +1,48 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""A KV cache group may hold both windowed and global layers — e.g. Gemma-3 with +the hybrid KV cache manager disabled, where the sliding-window layers are +promoted to full-attention *storage* and merged into one group whose spec still +records the window. Backends must take the window from the layers, never from +that group spec, or the global layers get windowed too and silently lose access +to everything older than the window. +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.v1.attention.backends.cpu_attn import ( + CPUAttentionBackendImpl, + CPUAttentionMetadataBuilder, +) + + +def _layers(layer_windows: list[int]): + """Stand-in attention layers, one per window, as one KV cache group.""" + return { + f"layer_{i}": SimpleNamespace( + impl=MagicMock(spec=CPUAttentionBackendImpl, sliding_window=window) + ) + for i, window in enumerate(layer_windows) + } + + +@pytest.mark.parametrize( + "layer_windows,expected", + [ + ([512, 512], 512), # uniform sliding window -> shared by the group + ([-1, -1], -1), # all global -> no window + ([512, -1], -1), # mixed -> the group cannot assume either window + ], +) +def test_cpu_group_sliding_window(layer_windows, expected): + layers = _layers(layer_windows) + builder = SimpleNamespace(vllm_config=None, layer_names=list(layers)) + with patch( + "vllm.v1.attention.backends.cpu_attn.get_layers_from_vllm_config", + return_value=layers, + ): + window = CPUAttentionMetadataBuilder._group_sliding_window(builder) + assert window == expected diff --git a/tests/v1/e2e/general/test_correctness_sliding_window.py b/tests/v1/e2e/general/test_correctness_sliding_window.py index a8a29203d9..5209d1894d 100644 --- a/tests/v1/e2e/general/test_correctness_sliding_window.py +++ b/tests/v1/e2e/general/test_correctness_sliding_window.py @@ -83,6 +83,34 @@ def test_sliding_window_retrieval( ) +@pytest.mark.parametrize("model", ["google/gemma-3-1b-it"]) +def test_hybrid_kv_cache_manager_output_equivalence(model, vllm_runner): + """Disabling the hybrid KV cache manager must not change what the model + computes: it promotes the sliding-window layers to full-attention *storage* + only, so the global layers must keep attending globally. + + Guards against the group's KV cache spec, which then covers windowed and + global layers alike, imposing its window on the global layers too. + """ + prompts, _, _ = prep_prompts(2, ln_range=model_config[model].ln_range) + sampling_params = SamplingParams(temperature=0.0, max_tokens=32) + enforce_eager = current_platform.is_rocm() + + outputs = [] + for disable_hybrid in (False, True): + with vllm_runner( + model, + max_model_len=None, + enable_chunked_prefill=None, + disable_hybrid_kv_cache_manager=disable_hybrid, + enforce_eager=enforce_eager, + ) as runner: + responses = runner.get_llm().generate(prompts, sampling_params) + outputs.append([response.outputs[0].text for response in responses]) + + assert outputs[0] == outputs[1] + + def check_length(prompts: list[str], llm: LLM, sliding_window: int): """ Check if the prompt length is valid, i.e., longer than the sliding window diff --git a/vllm/v1/attention/backends/cpu_attn.py b/vllm/v1/attention/backends/cpu_attn.py index a1b71f036c..3d09d789f3 100644 --- a/vllm/v1/attention/backends/cpu_attn.py +++ b/vllm/v1/attention/backends/cpu_attn.py @@ -11,8 +11,13 @@ import torch from vllm import _custom_ops as ops from vllm import envs -from vllm.config import VllmConfig, get_current_vllm_config +from vllm.config import ( + VllmConfig, + get_current_vllm_config, + get_layers_from_vllm_config, +) from vllm.logger import init_logger +from vllm.model_executor.layers.attention import Attention from vllm.platforms import CpuArchEnum, current_platform from vllm.utils.torch_utils import is_quantized_kv_cache from vllm.v1.attention.backend import ( @@ -151,9 +156,8 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata] ) self.head_dim = kv_cache_spec.head_size self.dtype = vllm_config.model_config.dtype - self.window_size = getattr(kv_cache_spec, "sliding_window", -1) - if self.window_size is None: - self.window_size = -1 + # Resolved from the layers on the first build(), once they exist. + self.window_size: int | None = None self.block_size = vllm_config.cache_config.block_size self.kv_cache_dtype = vllm_config.cache_config.cache_dtype self.isa = _get_attn_isa( @@ -167,12 +171,36 @@ class CPUAttentionMetadataBuilder(AttentionMetadataBuilder[CPUAttentionMetadata] kv_cache_spec, EncoderOnlyAttentionSpec ) + def _group_sliding_window(self) -> int: + """The window shared by every layer in this group, else -1 (no window). + + Taken from the layers rather than the group spec: one KV cache group can + hold both windowed and global layers (e.g. Gemma-3 with the hybrid KV + cache manager disabled), and the scheduler metadata built here is shared + by the whole group, so it may only assume a window all of them agree on. + """ + layers = get_layers_from_vllm_config( + self.vllm_config, Attention, self.layer_names + ) + windows = { + layer.impl.sliding_window + for layer in layers.values() + if isinstance(layer.impl, CPUAttentionBackendImpl) + } + if len(windows) != 1: + return -1 + window = windows.pop() + return -1 if window is None else window + def build( self, common_prefix_len: int, common_attn_metadata: CommonAttentionMetadata, fast_build: bool = False, ) -> CPUAttentionMetadata: + if self.window_size is None: + self.window_size = self._group_sliding_window() + num_reqs = common_attn_metadata.num_reqs num_actual_tokens = common_attn_metadata.num_actual_tokens max_query_len = common_attn_metadata.max_query_len diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index acdb1661dc..e7ef03d99c 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -287,8 +287,6 @@ class FlashAttentionMetadata: causal: bool | torch.Tensor = True - sliding_window: tuple[int, int] | None = None - # PrefixLM bidirectional range containing each scheduled query token. # Shape: (num_actual_tokens, 2) int32, absolute [start, end] bounds; # (-1, -1) for query tokens outside every multimodal range. @@ -713,13 +711,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad if isinstance(causal, torch.Tensor) and causal.dtype != torch.int32: causal = causal.to(torch.int32) - # Symmetrize the spec's sliding_window for non-causal attention - group_sliding_window = getattr(self.kv_cache_spec, "sliding_window", None) - base_window = ( - (-1, -1) if group_sliding_window is None else (group_sliding_window - 1, 0) - ) - effective_sliding_window = _maybe_symmetrize_window(base_window, causal) - attn_metadata = FlashAttentionMetadata( num_actual_tokens=num_actual_tokens, max_query_len=max_query_len, @@ -743,7 +734,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad prefix_scheduler_metadata=prefix_scheduler_metadata, max_num_splits=max_num_splits, causal=causal, - sliding_window=effective_sliding_window, ) # Compute mm_prefix range tensor if the batch contains @@ -1047,18 +1037,18 @@ class FlashAttentionImpl(AttentionImpl): ) return output else: - window = ( - attn_metadata.sliding_window - if attn_metadata.sliding_window is not None - else self.sliding_window - ) + causal = attn_metadata.causal + is_dynamic_causal = isinstance(causal, torch.Tensor) + + # The layer's own window wins over the group's: one KV cache + # group can hold both windowed and global layers (e.g. Gemma-3 + # with the hybrid KV cache manager disabled), and the group spec + # cannot describe both. + window = _maybe_symmetrize_window(self.sliding_window, causal) sliding_window_size: list[int] | None = ( list(window) if window is not None else None ) - causal = attn_metadata.causal - is_dynamic_causal = isinstance(causal, torch.Tensor) - mm_prefix_query_ranges = attn_metadata.mm_prefix_query_range_tensor mm_mask_mod = None mm_aux = None @@ -1068,10 +1058,6 @@ class FlashAttentionImpl(AttentionImpl): and causal is True and self.vllm_flash_attn_version == 4 ): - # Use the layer impl's window, not attn_metadata's. The - # metadata field comes from kv_cache_spec (model-wide, e.g. - # Gemma4's 512), while the impl holds the per-layer window - # the mask_mod must encode (e.g. a test override). # Triton convention: 1 + window_size[0]. Global layers store # (-1, -1) → sw stays None. layer_window = self.sliding_window