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 12, 06:00 UTC
[v1] Add real sliding window calculation to FlexAttention direct BlockMask building (#26015)
Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Signed-off-by: baonudesifeizhai <baonudesifeizhai@gmail.com> Co-authored-by: baonudesifeizhai <baonudesifeizhai@gmail.com>
Isotr0py committed
Dec 1, 2025 at 21:12 UTC
b95db244ee2d37390709aae27a2457a18984540a
2 files changed
+38
-7
tests/v1/attention/test_attention_backends.py
+11
-1
@@ -74,6 +74,9 @@ BATCH_SPECS = {
74
),
75
"large_decode": BatchSpec(seq_lens=[2048] * 32, query_lens=[1] * 32),
76
"large_prefill": BatchSpec(seq_lens=[4096] * 8, query_lens=[32] * 8),
77
+ "mixed_large": BatchSpec(
78
+ seq_lens=[1024, 2048, 4096, 1024, 2048, 4096], query_lens=[1, 1, 1, 32, 32, 32]
79
+ ),
80
"single_decode": BatchSpec(seq_lens=[1024], query_lens=[1]),
81
"single_prefill": BatchSpec(seq_lens=[1024], query_lens=[64]),
82
}
@@ -587,7 +590,14 @@ SLIDING_WINDOW_BACKENDS_TO_TEST = [
590
591
@pytest.mark.parametrize(
592
"batch_spec_name",
590
- ["small_decode", "small_prefill", "mixed_medium", "large_decode", "large_prefill"],
593
+ [
594
+ "small_decode",
595
+ "small_prefill",
596
+ "mixed_medium",
597
+ "large_decode",
598
+ "large_prefill",
599
+ "mixed_large",
600
+ ],
601
)
602
@pytest.mark.parametrize("model", ["microsoft/Phi-tiny-MoE-instruct"])
603
@pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])
vllm/v1/attention/backends/flex_attention.py
+27
-6
@@ -4,6 +4,7 @@
4
5
import math
6
from dataclasses import dataclass
7
+from functools import cached_property
8
from typing import ClassVar
9
10
import torch
@@ -315,6 +316,14 @@ class FlexAttentionMetadata:
316
transformed_score_mod: _score_mod_signature | None = None
317
sliding_window: int | None = None
318
319
+ @cached_property
320
+ def logical_block_ids(self):
321
+ return torch.arange(
322
+ cdiv(self.max_seq_len, self.block_size),
323
+ device=self.block_table.device,
324
+ dtype=torch.long,
325
+ )
326
+
327
def _convert_physical_to_logical(
328
self,
329
request_lookup: torch.Tensor,
@@ -493,6 +502,7 @@ class FlexAttentionMetadata:
502
503
The direct path works as follows:
504
1. For each query token, fetch blocks from block_table using max_seq_len
505
+ and exclude out of sliding window blocks if needed.
506
(this fetches more blocks than needed for shorter sequences)
507
2. Group query tokens into chunks of q_block_size
508
3. For each group, deduplicate the blocks using unique_static_unsorted
@@ -517,6 +527,23 @@ class FlexAttentionMetadata:
527
used_pages = self.block_table[
528
self.doc_ids, : cdiv(self.max_seq_len, self.block_size)
529
]
530
+
531
+ if self.sliding_window and self.causal:
532
+ device = used_pages.device
533
+ assert self.doc_ids is not None
534
+ token_indices = torch.arange(
535
+ self.doc_ids.shape[0], device=device, dtype=torch.long
536
+ )
537
+ logical_q_idx = (
538
+ token_indices
539
+ - self.query_start_loc[self.doc_ids]
540
+ + self.decode_offset[self.doc_ids]
541
+ )
542
+ min_kv_idx = torch.clamp(logical_q_idx - (self.sliding_window - 1), min=0)
543
+ min_block_idx = min_kv_idx // self.block_size
544
+ sliding_mask = self.logical_block_ids >= min_block_idx[:, None]
545
+ used_pages.masked_fill_(~sliding_mask, 0)
546
+
547
used_pages_padded = pad_to_multiple(
548
used_pages, multiple=self.q_block_size, dim=0
549
)
@@ -785,12 +812,6 @@ class FlexAttentionImpl(AttentionImpl):
812
if attn_metadata.sliding_window != self.sliding_window:
813
attn_metadata.sliding_window = self.sliding_window
814
if attn_metadata.direct_build:
788
- # TODO: Support skipping the computation of sliding window
789
- # in direct block mask building code path.
790
- logger.warning_once(
791
- "Using direct block mask building with sliding window, "
792
- "which is suboptimal now. Performance may be degraded."
793
- )
815
# update mask mod in attention metadata
816
attn_metadata.mask_mod = attn_metadata.get_mask_mod()
817
attn_metadata.block_mask = attn_metadata._build_block_mask_direct()