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 2, 11:00 UTC
[ROCm][CI] Fix AITER unified attention for encoder-decoder cross-attention (#43945)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
May 29, 2026 at 03:43 UTC
ff990d0d322b0685a2e526168a99e980ef4eb856
1 file changed
+27
-3
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+27
-3
@@ -134,6 +134,30 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
134
self.unified_attention = unified_attention
135
self.supports_quant_query_input = True
136
137
+ def _split_kv_cache(
138
+ self, kv_cache: torch.Tensor
139
+ ) -> tuple[torch.Tensor, torch.Tensor]:
140
+ if self.attn_type != AttentionType.ENCODER_DECODER:
141
+ return kv_cache.unbind(1)
142
+
143
+ # NOTE: Encoder-decoder layers can share the same raw KV allocation with
144
+ # ROCM_ATTN decoder layers, whose physical layout is K/V first. Keep
145
+ # this cross-attention path on that physical layout so block IDs do not
146
+ # alias different bytes across the shared allocation.
147
+ num_blocks, _, block_size, num_kv_heads, head_size = kv_cache.shape
148
+ block_stride = block_size * num_kv_heads * head_size
149
+ kv_cache = kv_cache.as_strided(
150
+ (2, num_blocks, block_size, num_kv_heads, head_size),
151
+ (
152
+ num_blocks * block_stride,
153
+ block_stride,
154
+ num_kv_heads * head_size,
155
+ head_size,
156
+ 1,
157
+ ),
158
+ )
159
+ return kv_cache.unbind(0)
160
+
161
def forward(
162
self,
163
layer: torch.nn.Module,
@@ -194,7 +218,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
218
layer,
219
)
220
197
- key_cache, value_cache = kv_cache.unbind(1)
221
+ key_cache, value_cache = self._split_kv_cache(kv_cache)
222
223
softmax_scale = self.scale
224
if is_quantized_kv_cache(self.kv_cache_dtype):
@@ -243,7 +267,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
267
# For encoder attention,
268
# we use direct Q, K, V tensors without caching
269
return
246
- key_cache, value_cache = kv_cache.unbind(1)
270
+ key_cache, value_cache = self._split_kv_cache(kv_cache)
271
272
# Reshape the input keys and values and store them in the cache.
273
ops.reshape_and_cache_flash(
@@ -276,7 +300,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
300
# For encoder attention,
301
# we use direct Q, K, V tensors without caching
302
return
279
- key_cache, value_cache = kv_cache.unbind(1)
303
+ key_cache, value_cache = self._split_kv_cache(kv_cache)
304
flash_layout = True
305
306
is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)