[ROCm] Fix encoder-decoder cross-attention KV layout aliasing (#47035)
Signed-off-by: Djordje Ramic <djoramic@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
djramic committed
Jul 3, 2026 at 20:53 UTC
f63dca68385c41148f386c86eac796c83d5ddd3b
3 files changed
+91
-43
tests/entrypoints/speech_to_text/transcription/test_transcription_validation_whisper.py
+2
-1
@@ -68,7 +68,8 @@ def server(request):
68
if request.param is not None:
69
args += ["--attention-backend", request.param]
70
if "AITER" in request.param:
71
- env_dict = _AITER_ENV
71
+ # TODO: re-enable once AITER reenables fp16 unified attention.
72
+ pytest.skip("ROCM_AITER_UNIFIED_ATTN does not support fp16")
73
with RemoteOpenAIServer(MODEL_NAME, args, env_dict=env_dict) as remote_server:
74
yield remote_server
75
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+56
-41
@@ -148,26 +148,10 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
148
def _split_kv_cache(
149
self, kv_cache: torch.Tensor
150
) -> tuple[torch.Tensor, torch.Tensor]:
151
- if self.attn_type != AttentionType.ENCODER_DECODER:
152
- return kv_cache.unbind(1)
153
-
154
- # NOTE: Encoder-decoder layers can share the same raw KV allocation with
155
- # ROCM_ATTN decoder layers, whose physical layout is K/V first. Keep
156
- # this cross-attention path on that physical layout so block IDs do not
157
- # alias different bytes across the shared allocation.
158
- num_blocks, _, block_size, num_kv_heads, head_size = kv_cache.shape
159
- block_stride = block_size * num_kv_heads * head_size
160
- kv_cache = kv_cache.as_strided(
161
- (2, num_blocks, block_size, num_kv_heads, head_size),
162
- (
163
- num_blocks * block_stride,
164
- block_stride,
165
- num_kv_heads * head_size,
166
- head_size,
167
- 1,
168
- ),
169
- )
170
- return kv_cache.unbind(0)
151
+ # Blocks-first ``(num_blocks, 2, ...)``. The model runner normalizes any
152
+ # shared decoder/cross-attention allocation to this layout, so no
153
+ # per-backend restriding is needed here.
154
+ return kv_cache.unbind(1)
155
156
def forward(
157
self,
@@ -242,27 +226,58 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
226
max_seqlen_k = attn_metadata.max_seq_len
227
block_table = attn_metadata.block_table
228
245
- self.unified_attention(
246
- q=query[:num_actual_tokens],
247
- k=key_cache,
248
- v=value_cache,
249
- out=output[:num_actual_tokens],
250
- cu_seqlens_q=cu_seqlens_q,
251
- max_seqlen_q=max_seqlen_q,
252
- seqused_k=seqused_k,
253
- max_seqlen_k=max_seqlen_k,
254
- softmax_scale=softmax_scale,
255
- causal=True,
256
- alibi_slopes=self.alibi_slopes,
257
- window_size=self.sliding_window,
258
- block_table=block_table,
259
- softcap=self.logits_soft_cap,
260
- q_descale=layer._q_scale if query.dtype == self.fp8_dtype else None,
261
- k_descale=layer._k_scale,
262
- v_descale=layer._v_scale,
263
- sinks=self.sinks,
264
- output_scale=output_scale,
265
- )
229
+ if attn_metadata.causal:
230
+ self.unified_attention(
231
+ q=query[:num_actual_tokens],
232
+ k=key_cache,
233
+ v=value_cache,
234
+ out=output[:num_actual_tokens],
235
+ cu_seqlens_q=cu_seqlens_q,
236
+ max_seqlen_q=max_seqlen_q,
237
+ seqused_k=seqused_k,
238
+ max_seqlen_k=max_seqlen_k,
239
+ softmax_scale=softmax_scale,
240
+ causal=True,
241
+ alibi_slopes=self.alibi_slopes,
242
+ window_size=self.sliding_window,
243
+ block_table=block_table,
244
+ softcap=self.logits_soft_cap,
245
+ q_descale=layer._q_scale if query.dtype == self.fp8_dtype else None,
246
+ k_descale=layer._k_scale,
247
+ v_descale=layer._v_scale,
248
+ sinks=self.sinks,
249
+ output_scale=output_scale,
250
+ )
251
+ else:
252
+ # The aiter kernel is causal-only. Non-causal cross-attention
253
+ # (ENCODER_DECODER, e.g. Whisper) falls back to the vLLM Triton
254
+ # unified kernel, which shares this layout and honors the flag.
255
+ from vllm.v1.attention.ops.triton_unified_attention import (
256
+ unified_attention as triton_unified_attention,
257
+ )
258
+
259
+ descale_shape = (cu_seqlens_q.shape[0] - 1, key_cache.shape[2])
260
+ triton_unified_attention(
261
+ q=query[:num_actual_tokens],
262
+ k=key_cache,
263
+ v=value_cache,
264
+ out=output[:num_actual_tokens],
265
+ cu_seqlens_q=cu_seqlens_q,
266
+ max_seqlen_q=max_seqlen_q,
267
+ seqused_k=seqused_k,
268
+ max_seqlen_k=max_seqlen_k,
269
+ softmax_scale=softmax_scale,
270
+ causal=attn_metadata.causal,
271
+ alibi_slopes=self.alibi_slopes,
272
+ window_size=self.sliding_window,
273
+ block_table=block_table,
274
+ softcap=self.logits_soft_cap,
275
+ q_descale=layer._q_scale if query.dtype == self.fp8_dtype else None,
276
+ k_descale=layer._k_scale.expand(descale_shape),
277
+ v_descale=layer._v_scale.expand(descale_shape),
278
+ sinks=self.sinks,
279
+ output_scale=output_scale,
280
+ )
281
282
return output
283
vllm/v1/worker/gpu_model_runner.py
+33
-1
@@ -7196,11 +7196,43 @@ class GPUModelRunner(
7196
else:
7197
raise NotImplementedError
7198
7199
- if has_attn and has_mamba:
7199
+ # Reconcile divergent KV layouts to blocks-first. Triggered by hybrid
7200
+ # attention/mamba models, and by encoder-decoder models whose shared
7201
+ # decoder/cross-attention allocation mixes K/V-first and blocks-first
7202
+ # backends (see _has_mixed_attention_kv_layout).
7203
+ if has_attn and (
7204
+ has_mamba or self._has_mixed_attention_kv_layout(kernel_block_sizes)
7205
+ ):
7206
self._update_hybrid_attention_mamba_layout(kv_caches, kernel_block_sizes)
7207
7208
return kv_caches
7209
7210
+ def _has_mixed_attention_kv_layout(self, kernel_block_sizes: list[int]) -> bool:
7211
+ """Whether attention groups disagree on the physical KV cache layout.
7212
+
7213
+ Encoder-decoder models (e.g. Whisper) share one raw KV allocation
7214
+ between a decoder self-attention layer (K/V-first ROCM_ATTN, block dim
7215
+ 1) and a cross-attention layer (blocks-first, block dim 0). Mixed block
7216
+ dims mean a block ID maps to different bytes per layer, so the shared
7217
+ buffer must be normalized to a single (blocks-first) layout.
7218
+ """
7219
+ block_dims: set[int] = set()
7220
+ for group in self._kv_cache_spec_attn_group_iterator():
7221
+ kv_cache_spec = group.kv_cache_spec
7222
+ if not isinstance(kv_cache_spec, AttentionSpec):
7223
+ continue
7224
+ if group.kv_cache_group_id == len(kernel_block_sizes):
7225
+ continue
7226
+ block_dims.add(
7227
+ group.backend.get_kv_cache_block_dim(
7228
+ kernel_block_sizes[group.kv_cache_group_id],
7229
+ kv_cache_spec.num_kv_heads,
7230
+ kv_cache_spec.head_size,
7231
+ cache_dtype_str=self.cache_config.cache_dtype,
7232
+ )
7233
+ )
7234
+ return len(block_dims) > 1
7235
+
7236
def _update_hybrid_attention_mamba_layout(
7237
self, kv_caches: dict[str, torch.Tensor], kernel_block_sizes: list[int]
7238
) -> None: