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
[Kernel][AMD] Avoid D2H copy and cumsum kernel (#22683)
Signed-off-by: Xiaozhu <mxz297@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
Xiaozhu Meng committed
Aug 12, 2025 at 12:53 UTC
6bd8ebf026600e9851026c8850f88c5e10acfab1
1 file changed
+20
-12
vllm/v1/attention/backends/rocm_aiter_fa.py
+20
-12
@@ -214,12 +214,14 @@ class AiterFlashAttentionMetadata:
214
# |-- query_len ---|
215
216
num_actual_tokens: int # Number of tokens excluding padding.
217
+ num_actual_kv_tokens: int
218
max_query_len: int
219
query_start_loc: torch.Tensor
220
max_seq_len: int
221
seq_lens: torch.Tensor
222
slot_mapping: torch.Tensor
223
block_table: torch.Tensor
224
+ cu_seq_lens: Optional[torch.Tensor]
225
226
# For cascade attention.
227
use_cascade: bool
@@ -272,6 +274,20 @@ class AiterFlashAttentionMetadataBuilder(
274
seq_lens = common_attn_metadata.seq_lens
275
block_table_tensor = common_attn_metadata.block_table_tensor
276
slot_mapping = common_attn_metadata.slot_mapping
277
+ if max_query_len > 1:
278
+ # We pre-compute cumulative seq len needed for prefill attention
279
+ # here to avoid recomputing it for every layer
280
+ cu_seq_lens = torch.zeros(seq_lens.shape[0] + 1,
281
+ dtype=torch.int32,
282
+ device=seq_lens.device)
283
+ torch.cumsum(seq_lens,
284
+ dim=0,
285
+ dtype=cu_seq_lens.dtype,
286
+ out=cu_seq_lens[1:])
287
+ num_actual_kv_tokens = int(cu_seq_lens[-1].item())
288
+ else:
289
+ cu_seq_lens = None
290
+ num_actual_kv_tokens = 0
291
292
def schedule(batch_size, cu_query_lens, max_query_len, seqlens,
293
max_seq_len, causal):
@@ -281,12 +297,14 @@ class AiterFlashAttentionMetadataBuilder(
297
298
attn_metadata = AiterFlashAttentionMetadata(
299
num_actual_tokens=num_actual_tokens,
300
+ num_actual_kv_tokens=num_actual_kv_tokens,
301
max_query_len=max_query_len,
302
query_start_loc=query_start_loc,
303
max_seq_len=max_seq_len,
304
seq_lens=seq_lens,
305
block_table=block_table_tensor,
306
slot_mapping=slot_mapping,
307
+ cu_seq_lens=cu_seq_lens,
308
use_cascade=use_cascade,
309
common_prefix_len=common_prefix_len,
310
total_tokens=self.total_tokens,
@@ -475,16 +493,6 @@ class AiterFlashAttentionImpl(AttentionImpl):
493
block_table = attn_metadata.block_table
494
495
if max_seqlen_q > 1:
478
-
479
- cu_seq_lens = torch.zeros(seqused_k.shape[0] + 1,
480
- dtype=torch.int32,
481
- device=query.device)
482
-
483
- torch.cumsum(seqused_k,
484
- dim=0,
485
- dtype=cu_seq_lens.dtype,
486
- out=cu_seq_lens[1:])
487
-
496
torch.ops.vllm.flash_attn_varlen_func(
497
query[:num_actual_tokens],
498
key_cache,
@@ -497,10 +505,10 @@ class AiterFlashAttentionImpl(AttentionImpl):
505
alibi_slopes=self.alibi_slopes,
506
window_size=self.sliding_window,
507
block_table=block_table,
500
- cu_seqlens_k=cu_seq_lens,
508
+ cu_seqlens_k=attn_metadata.cu_seq_lens,
509
k_scale=layer._k_scale,
510
v_scale=layer._v_scale,
503
- total_tokens=attn_metadata.total_tokens,
511
+ total_tokens=attn_metadata.num_actual_kv_tokens,
512
)
513
514
_, num_heads, head_size = query.shape