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, 23:45 UTC
[Attention] Optimize FlashInfer MetadataBuilder Build call (#21137)
Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com>
Lucas Wilkinson committed
Jul 24, 2025 at 06:21 UTC
61b8cea3b42feab021d506e9143551de18f9165c
3 files changed
+94
-78
tests/v1/attention/test_attention_backends.py
+10
-3
@@ -11,7 +11,8 @@ from tests.v1.attention.utils import (BatchSpec, _Backend,
11
create_vllm_config,
12
get_attention_backend)
13
from vllm.utils import STR_DTYPE_TO_TORCH_DTYPE, cdiv
14
-from vllm.v1.attention.backends.utils import CommonAttentionMetadata
14
+from vllm.v1.attention.backends.utils import (CommonAttentionMetadata,
15
+ set_kv_cache_layout)
16
from vllm.v1.kv_cache_interface import FullAttentionSpec
17
18
BACKENDS_TO_TEST = [
@@ -212,7 +213,7 @@ def run_attention_backend(backend: _Backend, kv_cache_spec: FullAttentionSpec,
213
214
from vllm.v1.attention.backends.flashinfer import PerLayerParameters
215
215
- def mock_get_per_layer_parameters(vllm_config):
216
+ def mock_get_per_layer_parameters(vllm_config, impl_cls):
217
# Return mock parameters for a single layer
218
head_size = vllm_config.model_config.get_head_size()
219
return {
@@ -297,7 +298,8 @@ def test_backend_correctness(batch_spec_name: str, model: str):
298
5. Comparing the vLLM backend's output to the ground-truth SDPA output.
299
"""
300
batch_spec = BATCH_SPECS[batch_spec_name]
300
- vllm_config = create_vllm_config(model_name=model)
301
+ vllm_config = create_vllm_config(model_name=model,
302
+ max_model_len=max(batch_spec.seq_lens))
303
device = torch.device("cuda:0")
304
305
kv_cache_spec = create_standard_kv_cache_spec(vllm_config)
@@ -419,6 +421,11 @@ def test_backend_correctness(batch_spec_name: str, model: str):
421
if backend_name == _Backend.FLASHINFER_VLLM_V1:
422
kv_cache_for_backend = kv_cache.transpose(0, 1)
423
424
+ # For FlashInfer default to HND layout and
425
+ kv_cache_for_backend = kv_cache_for_backend.transpose(
426
+ 2, 3).contiguous().transpose(2, 3)
427
+ set_kv_cache_layout("HND")
428
+
429
backend_output = run_attention_backend(backend_name, kv_cache_spec,
430
vllm_config, device,
431
common_attn_metadata,
tests/v1/attention/utils.py
+1
-1
@@ -66,7 +66,7 @@ def create_common_attn_metadata(
66
num_computed_tokens_cpu = torch.tensor(context_lens, dtype=torch.int32)
67
68
# Create block table (random for testing)
69
- max_blocks = max(batch_spec.seq_lens) // block_size + 1
69
+ max_blocks = (max(batch_spec.seq_lens) + block_size - 1) // block_size
70
block_table_tensor = torch.randint(0,
71
max_block_idx,
72
(batch_spec.batch_size, max_blocks),
vllm/v1/attention/backends/flashinfer.py
+83
-74
@@ -18,6 +18,7 @@ from vllm.attention.backends.abstract import (AttentionBackend, AttentionImpl,
18
from vllm.config import VllmConfig
19
from vllm.logger import init_logger
20
from vllm.platforms import current_platform
21
+from vllm.utils import cdiv
22
from vllm.v1.attention.backends.flash_attn import use_cascade_attention
23
from vllm.v1.attention.backends.utils import (
24
AttentionMetadataBuilder, CommonAttentionMetadata, PerLayerParameters,
@@ -158,7 +159,7 @@ class FlashInferMetadata:
159
# (batch_size + 1,). The cumulative subquery lengths of the sequences in
160
# the batch, used to index into subquery. E.g., if the subquery length
161
# is [4, 6], it is [0, 4, 10].
161
- qo_indptr: torch.Tensor
162
+ qo_indptr_cpu: torch.Tensor
163
# An example for paged_kv_indices, paged_kv_indptr:
164
# request 1, page indices [0, 5, 8]
165
# request 2, page indices [1, 6, 7]
@@ -167,13 +168,13 @@ class FlashInferMetadata:
168
# [0, 5, 8, 1, 6, 7, 3, 4]
169
# paged_kv_indptr is used to index into paged_kv_indices:
170
# [0, 3, 6, 8]
170
- # The indptr of the paged kv cache, shape: [batch_size + 1]
171
- paged_kv_indptr: torch.Tensor
172
- # The page indices of the paged kv cache
171
+ # The indptr of the paged kv cache, shape: [batch_size + 1] (CPU for plan)
172
+ paged_kv_indptr_cpu: torch.Tensor
173
+ # The page indices of the paged kv cache (on device for plan)
174
paged_kv_indices: torch.Tensor
175
# The number of entries in the last page of each request in
175
- # the paged kv cache, shape: [batch_size]
176
- paged_kv_last_page_len: torch.Tensor
176
+ # the paged kv cache, shape: [batch_size] (CPU for plan)
177
+ paged_kv_last_page_len_cpu: torch.Tensor
178
# The number of query/output heads
179
num_qo_heads: int
180
# The number of key/value heads
@@ -201,22 +202,17 @@ class FlashInferMetadata:
202
num_prefills: int
203
num_prefill_tokens: int
204
204
- # For cascade attention.
205
+ # For cascade attention (CPU for planning).
206
use_cascade: bool
206
- shared_qo_indptr: Optional[torch.Tensor] = None
207
- shared_kv_page_indptr: Optional[torch.Tensor] = None
208
- shared_kv_page_indices: Optional[torch.Tensor] = None
209
- shared_kv_last_page_len: Optional[torch.Tensor] = None
207
+ shared_qo_indptr_cpu: Optional[torch.Tensor] = None
208
+ shared_kv_page_indptr_cpu: Optional[torch.Tensor] = None
209
+ shared_kv_page_indices_cpu: Optional[torch.Tensor] = None
210
+ shared_kv_last_page_len_cpu: Optional[torch.Tensor] = None
211
212
prefill_wrapper: Optional[BatchPrefillWithPagedKVCacheWrapper] = None
213
decode_wrapper: Optional[BatchDecodeWithPagedKVCacheWrapper] = None
214
cascade_wrapper: Optional[MultiLevelCascadeAttentionWrapper] = None
215
215
- @property
216
- def query_start_loc(self):
217
- # The GPUModelRunner expects to be able to access this property.
218
- return self.qo_indptr
219
-
216
def __post_init__(self):
217
if self.head_dim is not None:
218
FlashInferBackend.validate_head_size(self.head_dim)
@@ -238,6 +234,12 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
234
self.vllm_config = vllm_config
235
self.cache_config = vllm_config.cache_config
236
self.kv_cache_spec = kv_cache_spec
237
+ max_num_blocks_per_request = cdiv(
238
+ vllm_config.model_config.max_model_len,
239
+ self.kv_cache_spec.block_size)
240
+ self.block_table_arange = torch.arange(max_num_blocks_per_request,
241
+ dtype=torch.int32,
242
+ device=self.device)
243
244
def reorder_batch(self, input_batch: InputBatch,
245
scheduler_output: SchedulerOutput) -> bool:
@@ -285,21 +287,25 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
287
if self.global_hyperparameters is None:
288
self.global_hyperparameters = infer_global_hyperparameters(
289
get_per_layer_parameters(self.vllm_config, FlashInferImpl))
290
+
291
if attn_metadata.use_cascade:
292
attn_metadata.cascade_wrapper = self._get_cascade_wrapper()
293
attn_metadata.cascade_wrapper.plan(
291
- [attn_metadata.shared_qo_indptr, attn_metadata.qo_indptr],
294
[
293
- attn_metadata.shared_kv_page_indptr,
294
- attn_metadata.paged_kv_indptr
295
+ attn_metadata.shared_qo_indptr_cpu,
296
+ attn_metadata.qo_indptr_cpu
297
+ ],
298
+ [
299
+ attn_metadata.shared_kv_page_indptr_cpu,
300
+ attn_metadata.paged_kv_indptr_cpu
301
],
302
[
297
- attn_metadata.shared_kv_page_indices,
303
+ attn_metadata.shared_kv_page_indices_cpu,
304
attn_metadata.paged_kv_indices
305
],
306
[
301
- attn_metadata.shared_kv_last_page_len,
302
- attn_metadata.paged_kv_last_page_len
307
+ attn_metadata.shared_kv_last_page_len_cpu,
308
+ attn_metadata.paged_kv_last_page_len_cpu
309
],
310
attn_metadata.num_qo_heads,
311
attn_metadata.num_kv_heads,
@@ -320,22 +326,22 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
326
# Decodes are first so prefills start after the last decode
327
prefill_start = num_decodes
328
attn_metadata.prefill_wrapper = self._get_prefill_wrapper()
323
- assert attn_metadata.qo_indptr[prefill_start:].shape[
329
+ assert attn_metadata.qo_indptr_cpu[prefill_start:].shape[
330
0] == num_prefills + 1
325
- assert attn_metadata.paged_kv_indptr[prefill_start:].shape[
331
+ assert attn_metadata.paged_kv_indptr_cpu[prefill_start:].shape[
332
0] == num_prefills + 1
327
- assert attn_metadata.paged_kv_last_page_len[
333
+ assert attn_metadata.paged_kv_last_page_len_cpu[
334
prefill_start:].shape[0] == num_prefills
335
# Since prefill_wrapper.run() will be called with
336
# query[num_decode_tokens:] we need to adjust the qo_indptr
337
# to be relative to the start of the prefill queries.
332
- qo_indptr = attn_metadata.qo_indptr[
333
- prefill_start:] - attn_metadata.qo_indptr[prefill_start]
338
+ qo_indptr_cpu = attn_metadata.qo_indptr_cpu[
339
+ prefill_start:] - attn_metadata.qo_indptr_cpu[prefill_start]
340
attn_metadata.prefill_wrapper.plan(
335
- qo_indptr,
336
- attn_metadata.paged_kv_indptr[prefill_start:],
341
+ qo_indptr_cpu,
342
+ attn_metadata.paged_kv_indptr_cpu[prefill_start:],
343
attn_metadata.paged_kv_indices,
338
- attn_metadata.paged_kv_last_page_len[prefill_start:],
344
+ attn_metadata.paged_kv_last_page_len_cpu[prefill_start:],
345
attn_metadata.num_qo_heads,
346
attn_metadata.num_kv_heads,
347
attn_metadata.head_dim,
@@ -357,9 +363,9 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
363
attn_metadata.num_qo_heads, attn_metadata.num_kv_heads,
364
attn_metadata.head_dim):
365
attn_metadata.decode_wrapper.plan(
360
- attn_metadata.paged_kv_indptr[:num_decodes + 1],
366
+ attn_metadata.paged_kv_indptr_cpu[:num_decodes + 1],
367
attn_metadata.paged_kv_indices,
362
- attn_metadata.paged_kv_last_page_len[:num_decodes],
368
+ attn_metadata.paged_kv_last_page_len_cpu[:num_decodes],
369
attn_metadata.num_qo_heads,
370
attn_metadata.num_kv_heads,
371
attn_metadata.head_dim,
@@ -383,55 +389,58 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
389
split_decodes_and_prefills(common_attn_metadata)
390
391
page_size = self.kv_cache_spec.block_size
386
- device = self.device
387
- qo_indptr = common_attn_metadata.query_start_loc
392
max_seq_len = common_attn_metadata.seq_lens_cpu.max()
393
seq_lens = common_attn_metadata.seq_lens
394
+ seq_lens_cpu = common_attn_metadata.seq_lens_cpu
395
block_table_tensor = common_attn_metadata.block_table_tensor
396
392
- block_table_bounds = (seq_lens + page_size - 1) // page_size
397
+ block_table_bounds_cpu = (seq_lens_cpu + page_size - 1) // page_size
398
399
use_cascade = common_prefix_len > 0
400
if use_cascade:
401
# Grab the blocks of the shared prefix from the first request.
402
assert common_prefix_len % page_size == 0
403
num_common_kv_blocks = common_prefix_len // page_size
399
- shared_qo_indptr = torch.tensor([0, num_actual_tokens],
400
- dtype=torch.int32,
401
- device=device)
402
- shared_kv_page_indptr = torch.tensor([0, num_common_kv_blocks],
403
- dtype=torch.int32,
404
- device=device)
405
- shared_kv_page_indices = block_table_tensor[
404
+
405
+ # Create CPU versions directly for cascade (no GPU versions needed)
406
+ shared_qo_indptr_cpu = torch.tensor([0, num_actual_tokens],
407
+ dtype=torch.int32,
408
+ device='cpu')
409
+ shared_kv_page_indptr_cpu = torch.tensor([0, num_common_kv_blocks],
410
+ dtype=torch.int32,
411
+ device='cpu')
412
+ shared_kv_page_indices_cpu = block_table_tensor[
413
0, :num_common_kv_blocks]
407
- shared_kv_last_page_len = torch.tensor([page_size],
408
- dtype=torch.int32,
409
- device=device)
414
+ shared_kv_last_page_len_cpu = torch.tensor([page_size],
415
+ dtype=torch.int32,
416
+ device='cpu')
417
+
418
# Remove the blocks of the shared prefix from all requests.
419
block_table_tensor = block_table_tensor[:, num_common_kv_blocks:]
412
- block_table_bounds -= num_common_kv_blocks
420
+ block_table_bounds_cpu -= num_common_kv_blocks
421
else:
414
- shared_qo_indptr = None
415
- shared_kv_page_indptr = None
416
- shared_kv_page_indices = None
417
- shared_kv_last_page_len = None
418
-
419
- mask = (torch.arange(block_table_tensor.size(1),
420
- dtype=block_table_tensor.dtype,
421
- device=block_table_tensor.device).unsqueeze(0)
422
+ shared_qo_indptr_cpu = None
423
+ shared_kv_page_indptr_cpu = None
424
+ shared_kv_page_indices_cpu = None
425
+ shared_kv_last_page_len_cpu = None
426
+
427
+ max_num_blocks = block_table_bounds_cpu.max()
428
+ block_table_bounds = block_table_bounds_cpu.to(self.device,
429
+ non_blocking=True)
430
+ mask = (self.block_table_arange[:max_num_blocks].unsqueeze(0)
431
< block_table_bounds.unsqueeze(1))
423
- paged_kv_indices = block_table_tensor[mask]
424
-
425
- paged_kv_indptr = torch.cat([
426
- torch.zeros(1,
427
- dtype=block_table_bounds.dtype,
428
- device=block_table_bounds.device),
429
- block_table_bounds.cumsum(dim=0, dtype=torch.int32)
430
- ])
431
-
432
- paged_kv_last_page_len = seq_lens % page_size
433
- paged_kv_last_page_len = torch.where(paged_kv_last_page_len == 0,
434
- page_size, paged_kv_last_page_len)
432
+ paged_kv_indices = block_table_tensor[:, :max_num_blocks][mask]
433
+
434
+ paged_kv_indptr_cpu = torch.zeros(len(block_table_bounds_cpu) + 1,
435
+ dtype=torch.int32,
436
+ device='cpu')
437
+ paged_kv_indptr_cpu[1:] = block_table_bounds_cpu.cumsum(
438
+ dim=0, dtype=torch.int32)
439
+
440
+ paged_kv_last_page_len_cpu = seq_lens_cpu % page_size
441
+ paged_kv_last_page_len_cpu = torch.where(
442
+ paged_kv_last_page_len_cpu == 0, page_size,
443
+ paged_kv_last_page_len_cpu)
444
cache_dtype = self.cache_config.cache_dtype
445
if cache_dtype.startswith("fp8"):
446
kv_cache_dtype = FlashInferBackend.get_fp8_dtype_for_flashinfer(
@@ -440,10 +449,10 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
449
kv_cache_dtype = self.kv_cache_spec.dtype
450
attn_metadata = FlashInferMetadata(
451
num_actual_tokens=num_actual_tokens,
443
- qo_indptr=qo_indptr,
444
- paged_kv_indptr=paged_kv_indptr,
452
+ qo_indptr_cpu=common_attn_metadata.query_start_loc_cpu,
453
+ paged_kv_indptr_cpu=paged_kv_indptr_cpu,
454
paged_kv_indices=paged_kv_indices,
446
- paged_kv_last_page_len=paged_kv_last_page_len,
455
+ paged_kv_last_page_len_cpu=paged_kv_last_page_len_cpu,
456
num_qo_heads=self.vllm_config.model_config.get_num_attention_heads(
457
self.vllm_config.parallel_config),
458
num_kv_heads=self.kv_cache_spec.num_kv_heads,
@@ -457,14 +466,14 @@ class FlashInferMetadataBuilder(AttentionMetadataBuilder[FlashInferMetadata]):
466
num_prefills=num_prefills,
467
num_prefill_tokens=num_prefill_tokens,
468
use_cascade=use_cascade,
460
- shared_qo_indptr=shared_qo_indptr,
461
- shared_kv_page_indptr=shared_kv_page_indptr,
462
- shared_kv_page_indices=shared_kv_page_indices,
463
- shared_kv_last_page_len=shared_kv_last_page_len,
469
+ shared_qo_indptr_cpu=shared_qo_indptr_cpu,
470
+ shared_kv_page_indptr_cpu=shared_kv_page_indptr_cpu,
471
+ shared_kv_page_indices_cpu=shared_kv_page_indices_cpu,
472
+ shared_kv_last_page_len_cpu=shared_kv_last_page_len_cpu,
473
max_seq_len=max_seq_len,
474
seq_lens=seq_lens,
475
block_table_tensor=block_table_tensor,
467
- workspace_buffer=self._workspace_buffer,
476
+ workspace_buffer=self._get_workspace_buffer(),
477
)
478
479
self._plan(num_prefills, num_decodes, attn_metadata)