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] Allow V1 flash_attn to support cross-attention (#23297)
Signed-off-by: Russell Bryant <rbryant@redhat.com>
Russell Bryant committed
Aug 22, 2025 at 08:10 UTC
281710ef9a2a795d57bce997d89a3ed69287464e
1 file changed
+7
-10
vllm/v1/attention/backends/flash_attn.py
+7
-10
@@ -405,13 +405,6 @@ class FlashAttentionImpl(AttentionImpl):
405
406
FlashAttentionBackend.validate_head_size(head_size)
407
408
- if attn_type not in [
409
- AttentionType.DECODER, AttentionType.ENCODER_ONLY
410
- ]:
411
- raise NotImplementedError("Encoder/decoder cross-attention "
412
- "is not implemented for "
413
- "FlashAttentionImpl")
414
-
408
self.attn_type = attn_type
409
self.vllm_flash_attn_version = get_flash_attn_version()
410
if is_quantized_kv_cache(self.kv_cache_dtype) \
@@ -477,7 +470,7 @@ class FlashAttentionImpl(AttentionImpl):
470
num_actual_tokens = attn_metadata.num_actual_tokens
471
472
# Handle encoder attention differently - no KV cache needed
480
- if attn_type in (AttentionType.ENCODER_ONLY, ):
473
+ if attn_type in (AttentionType.ENCODER_ONLY, AttentionType.ENCODER):
474
# For encoder attention,
475
# we use direct Q, K, V tensors without caching
476
return self._forward_encoder_attention(query[:num_actual_tokens],
@@ -489,7 +482,11 @@ class FlashAttentionImpl(AttentionImpl):
482
# For decoder and cross-attention, use KV cache as before
483
key_cache, value_cache = kv_cache.unbind(0)
484
492
- if self.kv_sharing_target_layer_name is None:
485
+ # key and value may be None in the case of cross attention. They are
486
+ # calculated once based on the output from the encoder and then cached
487
+ # in KV cache.
488
+ if (self.kv_sharing_target_layer_name is None and key is not None
489
+ and value is not None):
490
# Reshape the input keys and values and store them in the cache.
491
# Skip this if sharing KV cache with an earlier attention layer.
492
# NOTE(woosuk): Here, key and value are padded while slot_mapping is
@@ -528,7 +525,7 @@ class FlashAttentionImpl(AttentionImpl):
525
block_table = attn_metadata.block_table
526
scheduler_metadata = attn_metadata.scheduler_metadata
527
531
- descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
528
+ descale_shape = (cu_seqlens_q.shape[0] - 1, self.num_kv_heads)
529
530
flash_attn_varlen_func(
531
q=query[:num_actual_tokens],