| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from vllm.config.vllm import VllmConfig |
| 5 | from vllm.model_executor.layers.attention import Attention |
| 6 | from vllm.v1.kv_cache_interface import KVCacheSpec, RSWASpec, get_kv_quant_mode |
| 7 | |
| 8 | |
| 9 | class RSWAAttention(Attention): |
| 10 | """Attention layer that reports ``RSWASpec`` as its KV cache spec. |
| 11 | |
| 12 | Drop-in replacement for the standard ``Attention`` layer when the model is |
| 13 | configured with Reference Sliding Window Attention (R-SWA, |
| 14 | ``rswa_window > 0``). The actual masking logic lives in the attention |
| 15 | backend (FlexAttention or FA4 mask_mod); this layer only overrides |
| 16 | ``get_kv_cache_spec`` so the KV cache manager instantiates ``RSWAManager`` |
| 17 | (instead of ``FullAttentionManager``) and can therefore evict "gap" blocks |
| 18 | to keep per-request KV memory bounded at O(prefix + window). |
| 19 | """ |
| 20 | |
| 21 | def __init__(self, *args, rswa_window: int, **kwargs) -> None: |
| 22 | super().__init__(*args, **kwargs) |
| 23 | self._rswa_window = rswa_window |
| 24 | |
| 25 | def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None: |
| 26 | spec = super().get_kv_cache_spec(vllm_config) |
| 27 | if spec is None: |
| 28 | return None |
| 29 | return RSWASpec( |
| 30 | block_size=vllm_config.cache_config.block_size, |
| 31 | num_kv_heads=self.num_kv_heads, |
| 32 | head_size=self.head_size, |
| 33 | head_size_v=self.head_size_v, |
| 34 | dtype=self.kv_cache_torch_dtype, |
| 35 | kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype), |
| 36 | rswa_window=self._rswa_window, |
| 37 | ) |