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 14, 05:30 UTC
Optimize input preparation for FlashInfer [2/N] (#23174)
Signed-off-by: Woosuk Kwon <woosuk.kwon@berkeley.edu>
Woosuk Kwon committed
Aug 27, 2025 at 02:52 UTC
6578e873655859462758c5c51e51f876f2aa24a3
1 file changed
+54
-26
vllm/v1/attention/backends/flashinfer.py
+54
-26
@@ -6,6 +6,7 @@ from __future__ import annotations
6
from dataclasses import dataclass
7
from typing import ClassVar, Optional, Union
8
9
+import numpy as np
10
import torch
11
from flashinfer import (BatchDecodeWithPagedKVCacheWrapper,
12
BatchPrefillWithPagedKVCacheWrapper,
@@ -22,6 +23,7 @@ from vllm.logger import init_logger
23
from vllm.model_executor.layers.quantization.utils.quant_utils import (
24
QuantKey, kFp8StaticTensorSym, kNvfp4Quant)
25
from vllm.platforms import current_platform
26
+from vllm.triton_utils import tl, triton
27
from vllm.utils import cdiv, is_pin_memory_available
28
from vllm.utils.flashinfer import (supports_trtllm_attention,
29
use_trtllm_attention)
@@ -230,6 +232,7 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
232
dtype=torch.int32,
233
device="cpu",
234
pin_memory=pin_memory)
235
+ self.paged_kv_indptr_np = self.paged_kv_indptr_cpu.numpy()
236
self.paged_kv_indices_cpu = torch.zeros(max_num_pages,
237
dtype=torch.int32,
238
device="cpu",
@@ -238,10 +241,8 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
241
dtype=torch.int32,
242
device="cpu",
243
pin_memory=pin_memory)
241
-
242
- self.block_table_arange = torch.arange(max_num_pages_per_req,
243
- dtype=torch.int32,
244
- device=self.device)
244
+ self.paged_kv_last_page_len_np = (
245
+ self.paged_kv_last_page_len_cpu.numpy())
246
247
def _get_workspace_buffer(self):
248
if self._workspace_buffer is None:
@@ -317,9 +318,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
318
max_seq_len = common_attn_metadata.max_seq_len
319
seq_lens = common_attn_metadata.seq_lens
320
seq_lens_cpu = common_attn_metadata.seq_lens_cpu
321
+ seq_lens_np = seq_lens_cpu.numpy()
322
block_table_tensor = common_attn_metadata.block_table_tensor
323
322
- block_table_bounds_cpu = (seq_lens_cpu + page_size - 1) // page_size
324
+ num_blocks_np = (seq_lens_np + (page_size - 1)) // page_size
325
326
use_cascade = common_prefix_len > 0
327
if use_cascade:
@@ -342,37 +344,41 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
344
345
# Remove the blocks of the shared prefix from all requests.
346
block_table_tensor = block_table_tensor[:, num_common_kv_blocks:]
345
- block_table_bounds_cpu -= num_common_kv_blocks
347
+ num_blocks_np -= num_common_kv_blocks
348
else:
349
shared_qo_indptr_cpu = None
350
shared_kv_page_indptr_cpu = None
351
shared_kv_page_indices_cpu = None
352
shared_kv_last_page_len_cpu = None
353
352
- max_num_blocks = block_table_bounds_cpu.max().item()
353
- block_table_bounds = block_table_bounds_cpu.to(self.device,
354
- non_blocking=True)
355
- mask = (self.block_table_arange[:max_num_blocks].unsqueeze(0)
356
- < block_table_bounds.unsqueeze(1))
354
+ # write self.paged_kv_indptr_cpu inplace (0-index is always 0)
355
+ np.cumsum(
356
+ num_blocks_np,
357
+ dtype=np.int32,
358
+ out=self.paged_kv_indptr_np[1:num_reqs + 1],
359
+ )
360
+ paged_kv_indptr = self.paged_kv_indptr[:num_reqs + 1]
361
+ paged_kv_indptr.copy_(self.paged_kv_indptr_cpu[:num_reqs + 1],
362
+ non_blocking=True)
363
+
364
# write self.paged_kv_indices inplace
358
- num_actual_pages = torch.sum(mask)
365
+ num_actual_pages = num_blocks_np.sum().item()
366
paged_kv_indices = self.paged_kv_indices[:num_actual_pages]
360
- torch.masked_select(block_table_tensor[:, :max_num_blocks],
361
- mask,
362
- out=paged_kv_indices)
363
-
364
- # write self.paged_kv_indptr_cpu inplace (0-index is always 0)
365
- torch.cumsum(block_table_bounds_cpu,
366
- dim=0,
367
- dtype=torch.int32,
368
- out=self.paged_kv_indptr_cpu[1:1 + num_reqs])
367
+ _copy_page_indices_kernel[(num_reqs, )](
368
+ paged_kv_indices,
369
+ block_table_tensor,
370
+ block_table_tensor.stride(0),
371
+ paged_kv_indptr,
372
+ BLOCK_SIZE=1024,
373
+ )
374
370
- paged_kv_last_page_len_cpu = seq_lens_cpu % page_size
375
# write self.paged_kv_last_page_len_cpu inplace
372
- torch.where(paged_kv_last_page_len_cpu == 0,
373
- torch.tensor(page_size),
374
- paged_kv_last_page_len_cpu,
375
- out=self.paged_kv_last_page_len_cpu[:num_reqs])
376
+ paged_kv_last_page_len_np = seq_lens_np % page_size
377
+ self.paged_kv_last_page_len_np[:num_reqs] = np.where(
378
+ paged_kv_last_page_len_np == 0,
379
+ page_size,
380
+ paged_kv_last_page_len_np,
381
+ )
382
383
# Check if any layer uses sinks (requires TRTLLM attention)
384
has_sinks = self.global_hyperparameters.has_sinks
@@ -1002,3 +1008,25 @@ def fast_plan_decode(
1008
self._sm_scale = sm_scale
1009
self._rope_scale = rope_scale
1010
self._rope_theta = rope_theta
1011
+
1012
+
1013
+@triton.jit
1014
+def _copy_page_indices_kernel(
1015
+ page_indices,
1016
+ block_table,
1017
+ block_table_stride,
1018
+ cu_num_blocks,
1019
+ BLOCK_SIZE: tl.constexpr,
1020
+):
1021
+ req_idx = tl.program_id(0)
1022
+ row_ptr = block_table + req_idx * block_table_stride
1023
+ start_idx = tl.load(cu_num_blocks + req_idx)
1024
+ end_idx = tl.load(cu_num_blocks + req_idx + 1)
1025
+ num_blocks = end_idx - start_idx
1026
+
1027
+ offset = tl.arange(0, BLOCK_SIZE)
1028
+ for i in tl.range(0, num_blocks, BLOCK_SIZE):
1029
+ block_ids = tl.load(row_ptr + i + offset, mask=i + offset < num_blocks)
1030
+ tl.store(page_indices + start_idx + i + offset,
1031
+ block_ids,
1032
+ mask=i + offset < num_blocks)