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, 16:00 UTC
[ROCm] Support for Whisper v1 with Aiter Unified Attention and Aiter Flash Attention (#28376)
Signed-off-by: apinge <Tong.Qiu2@amd.com>
tongqiu committed
Nov 24, 2025 at 11:26 UTC
5253f4276f333474f43d7f1cdaad6104d8f88f1f
3 files changed
+28
-15
vllm/v1/attention/backends/rocm_aiter_fa.py
+14
-8
@@ -517,12 +517,9 @@ class AiterFlashAttentionImpl(AttentionImpl):
517
assert self.num_heads % self.num_kv_heads == 0
518
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
519
520
- if attn_type != AttentionType.DECODER:
520
+ if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
521
raise NotImplementedError(
522
- "Encoder self-attention and "
523
- "encoder/decoder cross-attention "
524
- "are not implemented for "
525
- "FlashAttentionImpl"
522
+ "Encoder self-attention is not implemented for FlashAttentionImpl"
523
)
524
525
def extend_forward(
@@ -678,7 +675,14 @@ class AiterFlashAttentionImpl(AttentionImpl):
675
# performance to make sure it does not introduce any overhead.
676
num_actual_tokens = attn_metadata.num_actual_tokens
677
key_cache, value_cache = kv_cache.unbind(0)
681
- if self.kv_sharing_target_layer_name is None:
678
+ # key and value may be None in the case of cross attention. They are
679
+ # calculated once based on the output from the encoder and then cached
680
+ # in KV cache.
681
+ if (
682
+ self.kv_sharing_target_layer_name is None
683
+ and key is not None
684
+ and value is not None
685
+ ):
686
# Reshape the input keys and values and store them in the cache.
687
# Skip this if sharing KV cache with an earlier attention layer.
688
# NOTE(woosuk): Here, key and value are padded while slot_mapping
@@ -704,8 +708,10 @@ class AiterFlashAttentionImpl(AttentionImpl):
708
709
# decode:extend:prefill
710
query = query[:num_actual_tokens]
707
- key = key[:num_actual_tokens]
708
- value = value[:num_actual_tokens]
711
+ if key is not None:
712
+ key = key[:num_actual_tokens]
713
+ if value is not None:
714
+ value = value[:num_actual_tokens]
715
716
output_actual_tokens = output[:num_actual_tokens]
717
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+12
-2
@@ -142,7 +142,14 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
142
143
key_cache, value_cache = kv_cache.unbind(0)
144
145
- if self.kv_sharing_target_layer_name is None:
145
+ # key and value may be None in the case of cross attention. They are
146
+ # calculated once based on the output from the encoder and then cached
147
+ # in KV cache.
148
+ if (
149
+ self.kv_sharing_target_layer_name is None
150
+ and key is not None
151
+ and value is not None
152
+ ):
153
# Reshape the input keys and values and store them in the cache.
154
# Skip this if sharing KV cache with an earlier attention layer.
155
ops.reshape_and_cache_flash(
@@ -169,7 +176,10 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
176
max_seqlen_k = attn_metadata.max_seq_len
177
block_table = attn_metadata.block_table
178
172
- descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
179
+ descale_shape = (
180
+ cu_seqlens_q.shape[0] - 1,
181
+ key.shape[1] if key is not None else self.num_kv_heads,
182
+ )
183
184
self.unified_attention(
185
q=query[:num_actual_tokens],
vllm/v1/attention/backends/rocm_attn.py
+2
-5
@@ -238,12 +238,9 @@ class RocmAttentionImpl(AttentionImpl):
238
239
RocmAttentionBackend.validate_head_size(head_size)
240
241
- if attn_type != AttentionType.DECODER:
241
+ if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
242
raise NotImplementedError(
243
- "Encoder self-attention and "
244
- "encoder/decoder cross-attention "
245
- "are not implemented for "
246
- "RocmAttentionImpl"
243
+ "Encoder self-attention is not implemented for RocmAttentionImpl"
244
)
245
246
self.fp8_dtype = current_platform.fp8_dtype()