A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 12:15 UTC
[BugFix][Attention] Fix sliding window attention in V1 giving incorrect results (#17574)
Signed-off-by: Lucas Wilkinson <lwilkinson@neuralmagic.com>
Lucas Wilkinson committed
May 2, 2025 at 14:01 UTC
0f87d8f7b26d2f71117211d337952396b75dac50
1 file changed
+35
-1
vllm/v1/attention/backends/flash_attn.py
+35
-1
@@ -10,9 +10,11 @@ from vllm import _custom_ops as ops
10
from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
11
AttentionMetadata, AttentionType,
12
is_quantized_kv_cache)
13
+from vllm.attention.layer import Attention
14
from vllm.attention.ops.merge_attn_states import merge_attn_states
15
from vllm.attention.utils.fa_utils import (flash_attn_supports_fp8,
16
get_flash_attn_version)
17
+from vllm.config import VllmConfig, get_layers_from_vllm_config
18
from vllm.logger import init_logger
19
from vllm.platforms import current_platform
20
from vllm.utils import cdiv
@@ -273,13 +275,23 @@ def make_local_attention_virtual_batches(
275
block_table_local
276
277
278
+def _get_sliding_window_configs(
279
+ vllm_config: VllmConfig) -> set[Optional[tuple[int, int]]]:
280
+ """Get the set of all sliding window configs used in the model."""
281
+ sliding_window_configs: set[Optional[tuple[int, int]]] = set()
282
+ layers = get_layers_from_vllm_config(vllm_config, Attention)
283
+ for layer in layers.values():
284
+ assert isinstance(layer.impl, FlashAttentionImpl)
285
+ sliding_window_configs.add(layer.impl.sliding_window)
286
+ return sliding_window_configs
287
+
288
+
289
class FlashAttentionMetadataBuilder:
290
291
def __init__(self, runner: "GPUModelRunner"):
292
model_config = runner.model_config
293
294
self.runner = runner
282
- self.aot_schedule = (get_flash_attn_version() == 3)
295
self.num_heads_q = model_config.get_num_attention_heads(
296
runner.parallel_config)
297
self.num_heads_kv = model_config.get_num_kv_heads(
@@ -287,6 +299,11 @@ class FlashAttentionMetadataBuilder:
299
self.headdim = model_config.get_head_size()
300
self.page_size = self.runner.block_size
301
302
+ self.aot_schedule = (get_flash_attn_version() == 3)
303
+ # Sliding window size to be used with the AOT scheduler will be
304
+ # populated on first build() call.
305
+ self.aot_sliding_window: Optional[tuple[int, int]] = None
306
+
307
def reorder_batch(self, input_batch: "InputBatch",
308
scheduler_output: "SchedulerOutput") -> bool:
309
return False
@@ -304,6 +321,22 @@ class FlashAttentionMetadataBuilder:
321
slot_mapping = self.runner.slot_mapping_cpu[:num_actual_tokens].to(
322
self.runner.device, non_blocking=True).long()
323
324
+ if self.aot_sliding_window is None:
325
+ self.aot_sliding_window = (-1, -1)
326
+ # For the AOT scheduler we need the sliding window value to be
327
+ # constant for all layers to. We have to populate this on the first
328
+ # build() call so the layers are constructed (cannot populate)
329
+ # in __init__.
330
+ if self.aot_schedule:
331
+ sliding_window_configs = _get_sliding_window_configs(
332
+ self.runner.vllm_config)
333
+ if len(sliding_window_configs) == 1:
334
+ sliding_window_config = sliding_window_configs.pop()
335
+ if sliding_window_config is not None:
336
+ self.aot_sliding_window = sliding_window_config
337
+ elif len(sliding_window_configs) > 1:
338
+ self.aot_schedule = False
339
+
340
def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
341
max_seq_len, causal):
342
if self.aot_schedule:
@@ -318,6 +351,7 @@ class FlashAttentionMetadataBuilder:
351
page_size=self.page_size,
352
cu_seqlens_q=cu_query_lens,
353
causal=causal,
354
+ window_size=self.aot_sliding_window,
355
)
356
return None
357