diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index f43a42f75d..984fc20eca 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/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)