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
[Attention] Optimize make_local_attention_virtual_batches for Flash Attention (#23185)
Signed-off-by: linzebing <linzebing1995@gmail.com>
Zebing Lin committed
Aug 19, 2025 at 22:57 UTC
a634733f67b39fd9c1da1a861ba39f75efb576f3
1 file changed
+7
-7
vllm/v1/attention/backends/utils.py
+7
-7
@@ -464,8 +464,9 @@ def make_local_attention_virtual_batches(
464
attn_chunk_size)[arange > 0]
465
466
# convert from q_seqlens to cu_seqlens_q
467
- cu_seqlens_q_local = np.pad(np.cumsum(seqlens_q_local), (1, 0))\
468
- .astype(np.int32)
467
+ cu_seqlens_q_local = np.empty(virtual_batches + 1, dtype=np.int32)
468
+ np.cumsum(seqlens_q_local, out=cu_seqlens_q_local[1:])
469
+ cu_seqlens_q_local[0] = 0
470
471
# compute the seqlens_k_local,
472
# basically a full local attention block for all but the last block in each
@@ -508,11 +509,10 @@ def make_local_attention_virtual_batches(
509
# [ 22, 23 ], < local-batch 6, (batch 2, starting from k[4])
510
# [ 24, 25 ], < local-batch 7, (batch 2, starting from k[8])
511
# ]
511
- block_indices= np.broadcast_to(
512
- np.arange(pages_per_local_batch, dtype=np.int32),
513
- (virtual_batches, pages_per_local_batch)) \
514
- + np.expand_dims(block_starts, axis=1)
515
- block_indices = block_indices.flatten().clip(max=block_table.shape[1] - 1)
512
+ block_indices = (block_starts[:, None] +
513
+ np.arange(pages_per_local_batch, dtype=np.int32))
514
+ block_indices = block_indices.reshape(-1).clip(max=block_table.shape[1] -
515
+ 1)
516
batch_indices = np.repeat(np.arange(actual_batch_size, dtype=np.int32),
517
local_blocks * pages_per_local_batch)
518
block_table_local = block_table[batch_indices, block_indices]\