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, 18:30 UTC
fix cross attention (#28346)
Signed-off-by: fsx950223 <fsx950223@outlook.com>
who who who committed
Nov 21, 2025 at 20:55 UTC
fc9f821d2062d412474ced64b9087c881651eb30
1 file changed
+9
-8
vllm/v1/attention/backends/triton_attn.py
+9
-8
@@ -244,14 +244,11 @@ class TritonAttentionImpl(AttentionImpl):
244
245
self.num_queries_per_kv = self.num_heads // self.num_kv_heads
246
247
- if attn_type != AttentionType.DECODER:
247
+ if attn_type not in [AttentionType.DECODER, AttentionType.ENCODER_DECODER]:
248
raise NotImplementedError(
249
- "Encoder self-attention and "
250
- "encoder/decoder cross-attention "
251
- "are not implemented for "
252
- "TritonAttentionImpl"
249
+ "Encoder self-attention is not implemented for TritonAttentionImpl"
250
)
254
-
251
+ self.attn_type = attn_type
252
self.fp8_dtype = current_platform.fp8_dtype()
253
254
self.sinks = sinks
@@ -312,7 +309,11 @@ class TritonAttentionImpl(AttentionImpl):
309
num_actual_tokens = attn_metadata.num_actual_tokens
310
key_cache, value_cache = kv_cache.unbind(1)
311
315
- if self.kv_sharing_target_layer_name is None:
312
+ if (
313
+ self.kv_sharing_target_layer_name is None
314
+ and key is not None
315
+ and value is not None
316
+ ):
317
# Reshape the input keys and values and store them in the cache.
318
# Skip this if sharing KV cache with an earlier attention layer.
319
if self.kv_cache_dtype.startswith("fp8"):
@@ -346,7 +347,7 @@ class TritonAttentionImpl(AttentionImpl):
347
max_seqlen_k = attn_metadata.max_seq_len
348
block_table = attn_metadata.block_table
349
349
- descale_shape = (cu_seqlens_q.shape[0] - 1, key.shape[1])
350
+ descale_shape = (cu_seqlens_q.shape[0] - 1, key_cache.shape[2])
351
352
unified_attention(
353
q=query[:num_actual_tokens],