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 15, 13:30 UTC
[ROCm] [Attention] Cleanup ROCm output passing (#16431)
Signed-off-by: Luka Govedič <lgovedic@redhat.com>
Luka Govedič committed
Apr 18, 2025 at 01:46 UTC
aaec845f8ed7f445c66ba0d28c84bec9d184f5ed
1 file changed
+18
-23
vllm/attention/backends/rocm_flash_attn.py
+18
-23
@@ -27,6 +27,7 @@ _PARTITION_SIZE_ROCM = 256
27
28
29
class ROCmFlashAttentionBackend(AttentionBackend):
30
+ accept_output_buffer: bool = True
31
32
@staticmethod
33
def get_name() -> str:
@@ -515,7 +516,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
516
517
from vllm.attention.ops.triton_flash_attention import ( # noqa: F401
518
triton_attention)
518
- self.attn_func = triton_attention
519
+ self.triton_attn_func = triton_attention
520
logger.debug("Using Triton FA in ROCmBackend")
521
if self.sliding_window != (-1, -1):
522
logger.warning("ROCm Triton FA does not currently support "
@@ -531,7 +532,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
532
else:
533
try:
534
from flash_attn import flash_attn_varlen_func # noqa: F401
534
- self.attn_func = flash_attn_varlen_func
535
+ self.fa_attn_func = flash_attn_varlen_func
536
logger.debug("Using CK FA in ROCmBackend")
537
except ModuleNotFoundError:
538
self.use_naive_attn = True
@@ -542,7 +543,7 @@ class ROCmFlashAttentionImpl(AttentionImpl):
543
"ROCm Naive FlashAttention does not support "
544
"attention logits soft capping.")
545
545
- self.attn_func = _sdpa_attention
546
+ self.sdpa_attn_func = _sdpa_attention
547
logger.debug("Using naive (SDPA) attention in ROCmBackend")
548
549
def repeat_kv(self, x: torch.Tensor, n_rep: int) -> torch.Tensor:
@@ -613,6 +614,8 @@ class ROCmFlashAttentionImpl(AttentionImpl):
614
Returns:
615
shape = [num_tokens, num_heads * head_size]
616
"""
617
+ assert output is not None, "Output tensor must be provided."
618
+
619
query = query.view(-1, self.num_heads, self.head_size)
620
if key is not None:
621
assert value is not None
@@ -656,7 +659,6 @@ class ROCmFlashAttentionImpl(AttentionImpl):
659
assert attn_metadata.num_encoder_tokens is not None
660
num_prefill_tokens = attn_metadata.num_encoder_tokens
661
659
- output = torch.empty_like(query)
662
# Query for decode. KV is not needed because it is already cached.
663
decode_query = query[num_prefill_tokens:]
664
# QKV for prefill.
@@ -704,11 +706,11 @@ class ROCmFlashAttentionImpl(AttentionImpl):
706
query.dtype,
707
seq_lens,
708
make_attn_mask=causal_mask) # type: ignore
707
- out, _ = self.attn_func(
709
+ self.triton_attn_func(
710
query,
711
key,
712
value,
711
- None,
713
+ output[:num_prefill_tokens],
714
query_seq_start_loc,
715
key_seq_start_loc,
716
query_max_seq_len,
@@ -733,10 +735,11 @@ class ROCmFlashAttentionImpl(AttentionImpl):
735
key = key.movedim(0, key.dim() - 2)
736
value = value.movedim(0, value.dim() - 2)
737
# sdpa math backend attention
736
- out = self.attn_func(
738
+ self.sdpa_attn_func(
739
query,
740
key,
741
value,
742
+ output[:num_prefill_tokens],
743
query_seq_start_loc,
744
num_prefill_tokens,
745
self.num_heads,
@@ -745,7 +748,8 @@ class ROCmFlashAttentionImpl(AttentionImpl):
748
attn_masks,
749
)
750
else:
748
- out = self.attn_func(
751
+ # upstream FA does not support an output arg, copy
752
+ output[:num_prefill_tokens] = self.fa_attn_func(
753
q=query,
754
k=key,
755
v=value,
@@ -760,12 +764,6 @@ class ROCmFlashAttentionImpl(AttentionImpl):
764
softcap=self.logits_soft_cap,
765
)
766
763
- # common code for prefill
764
- assert output[:num_prefill_tokens].shape == out.shape
765
- if output.shape[0] > num_prefill_tokens:
766
- output[:num_prefill_tokens] = out
767
- else:
768
- output = out
767
else:
768
# prefix-enabled attention -
769
# not applicable for encoder-only models
@@ -818,14 +816,10 @@ class ROCmFlashAttentionImpl(AttentionImpl):
816
device=output.device,
817
)
818
max_logits = torch.empty_like(exp_sums)
821
- if num_prefill_tokens > 0:
822
- out = output[num_prefill_tokens:]
823
- else:
824
- out = output
819
820
query_start_loc = None
821
ops.paged_attention_rocm(
828
- out,
822
+ output[num_prefill_tokens:],
823
exp_sums,
824
max_logits,
825
tmp_output,
@@ -878,7 +872,8 @@ def _sdpa_attention(
872
query: torch.Tensor,
873
key: torch.Tensor,
874
value: torch.Tensor,
881
- seq_lens: List[int],
875
+ output: torch.Tensor,
876
+ seq_lens: torch.Tensor,
877
num_tokens: int,
878
num_heads: int,
879
head_size: int,
@@ -886,9 +881,9 @@ def _sdpa_attention(
881
attn_masks: Optional[List[torch.Tensor]] = None,
882
) -> torch.Tensor:
883
start = 0
889
- output = torch.empty((num_tokens, num_heads, head_size),
890
- dtype=query.dtype,
891
- device=query.device)
884
+ assert output.shape == (num_tokens, num_heads, head_size)
885
+ assert output.dtype == query.dtype
886
+ assert output.device == query.device
887
888
for i, seq_len in enumerate(seq_lens):
889
end = start + seq_len