xinyun/vllm
1
0
Fork 0
mirror of https://git.datalinker.icu/vllm-project/vllm.git synced 2026-07-18 04:17:10 +08:00
Code Issues Packages Projects Releases Wiki Activity

[Optimization] Truncate kv page indices for sliding window attention

Browse Source 
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
This commit is contained in:
Woosuk Kwon 2025-08-15 19:10:01 +00:00
parent df5afa82e5
commit 90d43db442
1 changed files with 20 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

25
vllm/v1/attention/backends/flashinfer.py
View File

@ -248,7 +248,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
self.block_table_arange = torch.arange(max_num_pages_per_req, self.block_table_arange = torch.arange(max_num_pages_per_req,
dtype=torch.int32, dtype=torch.int32,
device=self.device) device="cpu")
self.sliding_window_size = getattr(kv_cache_spec, "sliding_window", None)
def _get_workspace_buffer(self): def _get_workspace_buffer(self):
if self._workspace_buffer is None: if self._workspace_buffer is None:
@ -488,10 +490,23 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
shared_kv_last_page_len_cpu = None shared_kv_last_page_len_cpu = None
max_num_blocks = block_table_bounds_cpu.max() max_num_blocks = block_table_bounds_cpu.max()
block_table_bounds = block_table_bounds_cpu.to(self.device, arange = self.block_table_arange[:max_num_blocks].unsqueeze(0)
non_blocking=True) mask = arange < block_table_bounds_cpu.unsqueeze(1)
mask = (self.block_table_arange[:max_num_blocks].unsqueeze(0) if (self.sliding_window_size is not None and not use_cascade
< block_table_bounds.unsqueeze(1)) and num_decodes > 0):
# NOTE(woosuk): Since FlashInfer's decode kernel doesn't skip the kv
# outside the sliding window and only do masking, we manually
# manipulate the seq_lens and block table for skipping.
# NOTE: Don't apply this optimization to prefill requests.
decode_seq_lens_cpu = seq_lens_cpu[:num_decodes]
num_skipped_pages = (
torch.relu(decode_seq_lens_cpu - self.sliding_window_size) //
page_size)
block_table_bounds_cpu[:num_decodes] -= num_skipped_pages
mask[:num_decodes] &= (arange[:num_decodes]
>= num_skipped_pages.unsqueeze(1))
# write self.paged_kv_indices inplace # write self.paged_kv_indices inplace
num_actual_pages = torch.sum(mask) num_actual_pages = torch.sum(mask)
paged_kv_indices = self.paged_kv_indices[:num_actual_pages] paged_kv_indices = self.paged_kv_indices[:num_actual_pages]