[ROCm][CI] Fix AITER unified attention for encoder-decoder cross-attention (#43945)

Signed-off-by: Andreas Karatzas <akaratza@amd.com>

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

@@ -134,6 +134,30 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
         self.unified_attention = unified_attention
         self.supports_quant_query_input = True
 
+    def _split_kv_cache(
+        self, kv_cache: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if self.attn_type != AttentionType.ENCODER_DECODER:
+            return kv_cache.unbind(1)
+
+        # NOTE: Encoder-decoder layers can share the same raw KV allocation with
+        # ROCM_ATTN decoder layers, whose physical layout is K/V first. Keep
+        # this cross-attention path on that physical layout so block IDs do not
+        # alias different bytes across the shared allocation.
+        num_blocks, _, block_size, num_kv_heads, head_size = kv_cache.shape
+        block_stride = block_size * num_kv_heads * head_size
+        kv_cache = kv_cache.as_strided(
+            (2, num_blocks, block_size, num_kv_heads, head_size),
+            (
+                num_blocks * block_stride,
+                block_stride,
+                num_kv_heads * head_size,
+                head_size,
+                1,
+            ),
+        )
+        return kv_cache.unbind(0)
+
     def forward(
         self,
         layer: torch.nn.Module,
@@ -194,7 +218,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
                 layer,
             )
 
-        key_cache, value_cache = kv_cache.unbind(1)
+        key_cache, value_cache = self._split_kv_cache(kv_cache)
 
         softmax_scale = self.scale
         if is_quantized_kv_cache(self.kv_cache_dtype):
@@ -243,7 +267,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
             # For encoder attention,
             # we use direct Q, K, V tensors without caching
             return
-        key_cache, value_cache = kv_cache.unbind(1)
+        key_cache, value_cache = self._split_kv_cache(kv_cache)
 
         # Reshape the input keys and values and store them in the cache.
         ops.reshape_and_cache_flash(
@@ -276,7 +300,7 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
             # For encoder attention,
             # we use direct Q, K, V tensors without caching
             return
-        key_cache, value_cache = kv_cache.unbind(1)
+        key_cache, value_cache = self._split_kv_cache(kv_cache)
         flash_layout = True
 
         is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)