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, 12:15 UTC
[Attention][Async] Eliminate `seq_lens_cpu` in FlashAttention metadata building with DCP > 1 (#29449)
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
Matthew Bonanni committed
Nov 26, 2025 at 21:48 UTC
77740191de965329e143e501321637a3e242e2f6
2 files changed
+18
-15
vllm/v1/attention/backends/flash_attn.py
+14
-13
@@ -328,7 +328,6 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
328
max_seq_len = common_attn_metadata.max_seq_len
329
query_start_loc = common_attn_metadata.query_start_loc
330
seq_lens = common_attn_metadata.seq_lens
331
- seq_lens_cpu = common_attn_metadata.seq_lens_cpu
331
block_table_tensor = common_attn_metadata.block_table_tensor
332
slot_mapping = common_attn_metadata.slot_mapping
333
causal = common_attn_metadata.causal
@@ -401,20 +400,23 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
400
prefix_scheduler_metadata = None
401
402
if self.dcp_world_size > 1:
404
- query_kv_lens_cpu = (
405
- common_attn_metadata.query_start_loc_cpu[1:]
406
- - common_attn_metadata.query_start_loc_cpu[:-1]
407
- )
408
- dcp_context_kv_lens_cpu = seq_lens_cpu - query_kv_lens_cpu
403
+ query_kv_lens = query_start_loc[1:] - query_start_loc[:-1]
404
+ dcp_context_kv_lens = seq_lens - query_kv_lens
405
410
- dcp_context_kv_lens_cpu = get_dcp_local_seq_lens(
411
- dcp_context_kv_lens_cpu,
406
+ dcp_context_kv_lens = get_dcp_local_seq_lens(
407
+ dcp_context_kv_lens,
408
self.dcp_world_size,
409
self.dcp_rank,
410
self.cp_kv_cache_interleave_size,
411
)
416
- dcp_context_kv_lens = dcp_context_kv_lens_cpu.to(self.device)
417
- max_dcp_context_kv_len = dcp_context_kv_lens.max().item()
412
+ # After DCP distribution, the maximum number of tokens for any rank is
413
+ # ceil(L / (N * I)) * I, where L is max_seq_len, N is dcp_world_size,
414
+ # and I is cp_kv_cache_interleave_size.
415
+ # This eliminates GPU->CPU sync while minimizing workspace over-allocation.
416
+ num_partitions = self.dcp_world_size * self.cp_kv_cache_interleave_size
417
+ max_dcp_context_kv_len = (
418
+ (max_seq_len + num_partitions - 1) // num_partitions
419
+ ) * self.cp_kv_cache_interleave_size
420
421
scheduler_metadata = schedule(
422
batch_size=num_reqs,
@@ -431,9 +433,8 @@ class FlashAttentionMetadataBuilder(AttentionMetadataBuilder[FlashAttentionMetad
433
prefix_kv_lens = torch.tensor(
434
[common_prefix_len], dtype=torch.int32, device=self.device
435
)
434
- suffix_kv_lens = (seq_lens_cpu[:num_reqs] - common_prefix_len).to(
435
- self.device, non_blocking=True
436
- )
436
+ # Use GPU tensor directly - no CPU sync needed
437
+ suffix_kv_lens = seq_lens[:num_reqs] - common_prefix_len
438
prefix_scheduler_metadata = schedule(
439
batch_size=1,
440
cu_query_lens=cu_prefix_query_lens,
vllm/v1/attention/backends/utils.py
+4
-2
@@ -1095,12 +1095,14 @@ def get_dcp_local_seq_lens(
1095
num_requests = seq_lens.size(0)
1096
if dcp_rank is None:
1097
rank_offsets = (
1098
- torch.arange(dcp_size, dtype=torch.int32)
1098
+ torch.arange(dcp_size, dtype=torch.int32, device=seq_lens.device)
1099
.unsqueeze(0)
1100
.repeat(num_requests, 1)
1101
)
1102
else:
1103
- rank_offsets = torch.Tensor([[dcp_rank]]).to(dtype=torch.int32)
1103
+ rank_offsets = torch.tensor(
1104
+ [[dcp_rank]], dtype=torch.int32, device=seq_lens.device
1105
+ )
1106
seq_lens_tiled = (
1107
seq_lens.to(torch.int32).unsqueeze(-1).repeat(1, rank_offsets.shape[1])
1108
)