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 3, 22:45 UTC
[ROCm] Clean up a bit the AITER FA backend (#41942)
Signed-off-by: Patrick Schlangen <pschlan@amd.com>
pschlan-amd committed
May 11, 2026 at 16:45 UTC
5f1b313900fd5375857bdb7f709e1c5aad54443c
1 file changed
+15
-50
vllm/v1/attention/backends/rocm_aiter_fa.py
+15
-50
@@ -324,22 +324,17 @@ logger = init_logger(__name__)
324
@dataclass
325
class AiterFlashAttentionDecodeMetadata:
326
max_query_len: int
327
- min_query_len: int
328
- max_seq_len: int
329
- query_start_loc: torch.Tensor
327
328
329
@dataclass
330
class AiterFlashAttentionPrefillMetadata:
331
max_query_len: int
335
- min_query_len: int
332
max_seq_len: int
333
query_start_loc: torch.Tensor
334
335
336
@dataclass
337
class AiterChunkSlidingWindowMetadata:
342
- swa_seqlens: torch.Tensor
338
swa_cu_seqlens: torch.Tensor
339
swa_seq_starts: torch.Tensor
340
swa_token_to_batch: torch.Tensor
@@ -354,9 +349,7 @@ class AiterChunkContextMetadata:
349
cu_seq_lens_chunk: torch.Tensor
350
chunk_starts: torch.Tensor
351
token_to_batch: torch.Tensor
357
- seq_tot: list[int]
352
max_seq_lens: list[int]
359
- seq_lens: torch.Tensor
353
num_chunks: int
354
total_token_per_batch: list[int]
355
swa_metadata: AiterChunkSlidingWindowMetadata | None
@@ -365,7 +358,6 @@ class AiterChunkContextMetadata:
358
@dataclass
359
class AiterFlashAttentionChunkPrefillMetadata:
360
max_query_len: int
368
- min_query_len: int
361
max_seq_len: int
362
query_start_loc: torch.Tensor
363
chunk_context_metadata: AiterChunkContextMetadata
@@ -382,8 +374,6 @@ class AiterFlashAttentionMetadata:
374
# |-- query_len ---|
375
376
num_actual_tokens: int # Number of tokens excluding padding.
385
- num_actual_kv_tokens: int
386
- max_query_len: int
377
query_start_loc: torch.Tensor
378
max_seq_len: int
379
seq_lens: torch.Tensor
@@ -395,7 +385,6 @@ class AiterFlashAttentionMetadata:
385
num_decodes: int
386
num_decode_tokens: int
387
num_prefills: int
398
- num_prefill_tokens: int
388
num_extends: int
389
num_extend_tokens: int
390
@@ -405,8 +394,6 @@ class AiterFlashAttentionMetadata:
394
395
# For cascade attention.
396
use_cascade: bool
408
- common_prefix_len: int
409
- total_tokens: int
397
398
# Only for fp8 shuffle layout kv cache, we allocate kv_scale for each layer
399
# since we might integrate per token quant for kv cache in the future.
@@ -441,7 +428,6 @@ class AiterFlashAttentionMetadataBuilder(
428
# Sliding window size to be used with the AOT scheduler will be
429
# populated on first build() call.
430
self.aot_sliding_window: tuple[int, int] | None = None
444
- self.total_tokens: int = 0
431
self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
432
433
sliding_window_configs: set[tuple[int, int] | None] = set()
@@ -473,13 +459,9 @@ class AiterFlashAttentionMetadataBuilder(
459
def build_for_cudagraph_capture(
460
self, common_attn_metadata: CommonAttentionMetadata
461
):
476
- self.total_tokens = (
477
- self.model_config.max_model_len
478
- * self.vllm_config.scheduler_config.max_num_partial_prefills
462
+ return self.build(
463
+ common_prefix_len=0, common_attn_metadata=common_attn_metadata
464
)
480
- res = self.build(common_prefix_len=0, common_attn_metadata=common_attn_metadata)
481
- self.total_tokens = 0
482
- return res
465
466
def build(
467
self,
@@ -515,12 +497,18 @@ class AiterFlashAttentionMetadataBuilder(
497
num_prefills,
498
num_decode_tokens,
499
num_extend_tokens,
518
- num_prefill_tokens,
500
+ _,
501
) = split_ret
502
503
query_start_loc_cpu = common_attn_metadata.query_start_loc_cpu
504
523
- seq_lens = common_attn_metadata.seq_lens.cpu()
505
+ # Only copy seq_lens to CPU when prefill or extend is present to avoid a
506
+ # blocking device→host transfer.
507
+ seq_lens = (
508
+ common_attn_metadata.seq_lens.cpu()
509
+ if num_prefills > 0 or num_extends > 0
510
+ else None
511
+ )
512
513
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
514
@@ -528,26 +516,24 @@ class AiterFlashAttentionMetadataBuilder(
516
if num_decodes > 0:
517
decode_metadata = AiterFlashAttentionDecodeMetadata(
518
max_query_len=query_lens_cpu[:num_decodes].max().item(),
531
- min_query_len=query_lens_cpu[:num_decodes].min().item(),
532
- max_seq_len=seq_lens[:num_decodes].max().item(),
533
- query_start_loc=common_attn_metadata.query_start_loc[: num_decodes + 1],
519
)
520
521
prefill_metadata = None
522
if num_prefills > 0:
523
+ assert seq_lens is not None
524
query_lens_for_prefill = query_lens_cpu[num_decodes + num_extends :]
525
query_start_loc_device = common_attn_metadata.query_start_loc[
526
num_decodes + num_extends :
527
]
528
prefill_metadata = AiterFlashAttentionPrefillMetadata(
529
max_query_len=query_lens_for_prefill.max().item(),
544
- min_query_len=query_lens_for_prefill.min().item(),
530
max_seq_len=seq_lens[num_decodes + num_extends :].max().item(),
531
query_start_loc=query_start_loc_device - query_start_loc_device[0],
532
)
533
534
extend_metadata = None
535
if num_extends > 0:
536
+ assert seq_lens is not None
537
num_extends_slice = slice(num_decodes, num_decodes + num_extends)
538
query_lens_for_extend = query_lens_cpu[num_extends_slice]
539
seq_lens_for_extend = seq_lens[num_extends_slice]
@@ -591,9 +577,6 @@ class AiterFlashAttentionMetadataBuilder(
577
total_tokens = cu_seq_lens[-1].item()
578
579
swa_metadata = AiterChunkSlidingWindowMetadata(
594
- swa_seqlens=swa_seqlen_for_extend.to(
595
- self.device, non_blocking=True
596
- ),
580
swa_cu_seqlens=cu_seq_lens.to(self.device, non_blocking=True),
581
swa_seq_starts=seq_starts.to(self.device, non_blocking=True),
582
swa_token_to_batch=token_to_seq.to(self.device, non_blocking=True),
@@ -638,10 +621,8 @@ class AiterFlashAttentionMetadataBuilder(
621
workspace=self.extend_workspace,
622
cu_seq_lens_chunk=cu_seq_lens_cpu.to(self.device, non_blocking=True),
623
chunk_starts=chunk_starts.to(self.device, non_blocking=True),
641
- seq_tot=chunk_seq_lens.sum(dim=1).tolist(),
642
- max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
643
- seq_lens=chunk_seq_lens,
624
token_to_batch=token_to_batch_tensor.to(self.device, non_blocking=True),
625
+ max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),
626
num_chunks=num_chunks,
627
total_token_per_batch=cu_seq_lens_cpu[:, -1].tolist(),
628
swa_metadata=swa_metadata,
@@ -659,20 +640,15 @@ class AiterFlashAttentionMetadataBuilder(
640
)
641
extend_metadata = AiterFlashAttentionChunkPrefillMetadata(
642
max_query_len=query_lens_for_extend.max().item(),
662
- min_query_len=query_lens_for_extend.min().item(),
643
max_seq_len=seq_lens[num_extends_slice].max().item(),
644
query_start_loc=query_start_loc_device - query_start_loc_device[0],
645
chunk_context_metadata=chunk_context_metadata,
646
)
647
668
- num_actual_kv_tokens = torch.sum(seq_lens).item()
669
-
648
use_cascade = common_prefix_len > 0
649
650
attn_metadata = AiterFlashAttentionMetadata(
651
num_actual_tokens=common_attn_metadata.num_actual_tokens,
674
- num_actual_kv_tokens=num_actual_kv_tokens,
675
- max_query_len=common_attn_metadata.max_query_len,
652
query_start_loc=common_attn_metadata.query_start_loc,
653
max_seq_len=common_attn_metadata.max_seq_len,
654
seq_lens=common_attn_metadata.seq_lens,
@@ -682,15 +658,12 @@ class AiterFlashAttentionMetadataBuilder(
658
num_decodes=num_decodes,
659
num_decode_tokens=num_decode_tokens,
660
num_prefills=num_prefills,
685
- num_prefill_tokens=num_prefill_tokens,
661
num_extends=num_extends,
662
num_extend_tokens=num_extend_tokens,
663
decode_metadata=decode_metadata,
664
prefill_metadata=prefill_metadata,
665
extend_metadata=extend_metadata,
666
use_cascade=use_cascade,
692
- common_prefix_len=common_prefix_len,
693
- total_tokens=self.total_tokens,
667
k_scale=self.scale,
668
v_scale=self.scale,
669
)
@@ -713,15 +686,10 @@ class AiterFlashAttentionMetadataBuilder(
686
687
decode_metadata = AiterFlashAttentionDecodeMetadata(
688
max_query_len=common_attn_metadata.max_query_len,
716
- min_query_len=common_attn_metadata.max_query_len, # uniform batch
717
- max_seq_len=common_attn_metadata.max_seq_len,
718
- query_start_loc=common_attn_metadata.query_start_loc,
689
)
690
691
return AiterFlashAttentionMetadata(
692
num_actual_tokens=num_tokens,
723
- num_actual_kv_tokens=0, # not used in unified_attention path
724
- max_query_len=common_attn_metadata.max_query_len,
693
query_start_loc=common_attn_metadata.query_start_loc,
694
max_seq_len=common_attn_metadata.max_seq_len,
695
seq_lens=common_attn_metadata.seq_lens,
@@ -731,15 +699,12 @@ class AiterFlashAttentionMetadataBuilder(
699
num_decodes=num_reqs,
700
num_decode_tokens=num_tokens,
701
num_prefills=0,
734
- num_prefill_tokens=0,
702
num_extends=0,
703
num_extend_tokens=0,
704
decode_metadata=decode_metadata,
705
prefill_metadata=None,
706
extend_metadata=None,
707
use_cascade=False,
741
- common_prefix_len=0,
742
- total_tokens=self.total_tokens,
708
k_scale=self.scale,
709
v_scale=self.scale,
710
)
@@ -932,8 +897,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
897
output: torch.Tensor,
898
cu_seqlens_q: torch.Tensor,
899
max_seqlen_q: int,
935
- max_seqlen_k: int,
900
min_seqlen_q: int,
901
+ max_seqlen_k: int,
902
block_table: torch.Tensor,
903
slot_mapping: torch.Tensor,
904
k_scale: torch.Tensor,
@@ -1160,8 +1125,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1125
output=extend_outputs,
1126
cu_seqlens_q=attn_metadata.extend_metadata.query_start_loc,
1127
max_seqlen_q=attn_metadata.extend_metadata.max_query_len,
1163
- max_seqlen_k=attn_metadata.extend_metadata.max_seq_len,
1128
min_seqlen_q=1,
1129
+ max_seqlen_k=attn_metadata.extend_metadata.max_seq_len,
1130
block_table=attn_metadata.block_table[
1131
num_decodes : num_decodes + num_extends
1132
],