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 9, 14:45 UTC
Fix IndexError with encoder-decoder models when using Custom Paged Attention (#33112)
Signed-off-by: sstamenk <strahinja.stamenkovic@amd.com>
Strahinja Stamenkovic committed
Jan 27, 2026 at 03:33 UTC
c568581ff38af13924627fdd5c2cd9d9282fbe3d
2 files changed
+13
-5
vllm/v1/attention/backends/rocm_attn.py
+10
-3
@@ -330,7 +330,14 @@ class RocmAttentionImpl(AttentionImpl):
330
kv_cache, self.num_kv_heads, self.head_size
331
)
332
333
- if self.kv_sharing_target_layer_name is None:
333
+ # key and value may be None in the case of cross attention. They are
334
+ # calculated once based on the output from the encoder and then cached
335
+ # in KV cache.
336
+ if (
337
+ self.kv_sharing_target_layer_name is None
338
+ and key is not None
339
+ and value is not None
340
+ ):
341
# Reshape the input keys and values and store them in the cache.
342
# Skip this if sharing KV cache with an earlier attention layer.
343
@@ -382,8 +389,8 @@ class RocmAttentionImpl(AttentionImpl):
389
# Compute attention and update output up to `num_actual_tokens`.
390
chunked_prefill_paged_decode(
391
query=query[:num_actual_tokens],
385
- key=key[:num_actual_tokens],
386
- value=value[:num_actual_tokens],
392
+ key=key[:num_actual_tokens] if key is not None else None,
393
+ value=value[:num_actual_tokens] if value is not None else None,
394
output=output[:num_actual_tokens],
395
kv_cache_dtype=self.kv_cache_dtype,
396
key_cache=key_cache,
vllm/v1/attention/ops/chunked_prefill_paged_decode.py
+3
-2
@@ -302,8 +302,9 @@ def chunked_prefill_paged_decode(
302
block_size = value_cache.shape[3]
303
num_seqs = len(seq_lens)
304
num_query_heads = query.shape[1]
305
- num_kv_heads = key.shape[1]
306
- num_queries_per_kv = query.shape[1] // key.shape[1]
305
+ # key may be None in cross-attention decode (already cached from encoder)
306
+ num_kv_heads = key.shape[1] if key is not None else key_cache.shape[1]
307
+ num_queries_per_kv = num_query_heads // num_kv_heads
308
head_size = query.shape[2]
309
310
# Conversion of FP8 Tensor from uint8 storage to