diff --git a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py index b5e354ddfd..900f90eb42 100644 --- a/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py +++ b/tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py @@ -420,10 +420,9 @@ _FUSION_CONFIGS = [ @pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"]) @pytest.mark.parametrize( "kv_stride_order", - [ - pytest.param((0, 1, 2, 3, 4), id="block_first"), - pytest.param((1, 0, 2, 3, 4), id="kv_first"), - ], + # FA/unified use the 4D packed cache (num_blocks, num_kv_heads, block_size, + # 2*head_size); no separate K/V dim to permute. + [pytest.param((0, 1, 2, 3), id="packed_4d")], ) @pytest.mark.parametrize("enable_aiter_triton_rope", [True, False]) @pytest.mark.parametrize("block_size", [16, 32, 64]) @@ -458,6 +457,14 @@ def test_qk_norm_rope_kvcache_fusion( and use_shuffle_kv_layout == "1" ): pytest.skip("ROCM_AITER_UNIFIED_ATTN is NHD-only; shuffle env is ignored") + if ( + attn_backend == AttentionBackendEnum.ROCM_AITER_FA + and use_shuffle_kv_layout == "1" + ): + pytest.skip( + "ROCM_AITER_FA gates the qk_norm+rope+kvcache fusion off under shuffle " + "layout (defers to reshape_and_cache_shuffle_triton), so nothing fuses" + ) _run_qk_norm_rope_kvcache_fusion_test( attn_backend=attn_backend, enable_aiter_triton_rope=enable_aiter_triton_rope, diff --git a/vllm/v1/attention/backends/rocm_aiter_fa.py b/vllm/v1/attention/backends/rocm_aiter_fa.py index 5a204c7ed6..96dae0c266 100644 --- a/vllm/v1/attention/backends/rocm_aiter_fa.py +++ b/vllm/v1/attention/backends/rocm_aiter_fa.py @@ -1065,9 +1065,8 @@ class AiterFlashAttentionImpl(AttentionImpl): # Whenever making a change in this method, please benchmark the # performance to make sure it does not introduce any overhead. num_actual_tokens = attn_metadata.num_actual_tokens - # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) - key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + key_cache, value_cache = self._split_kv_cache(kv_cache) if is_quantized_kv_cache(self.kv_cache_dtype): key_cache = key_cache.view(current_platform.fp8_dtype()) value_cache = value_cache.view(current_platform.fp8_dtype()) @@ -1385,6 +1384,12 @@ class AiterFlashAttentionImpl(AttentionImpl): return output + def _split_kv_cache( + self, kv_cache: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) + return kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + def do_kv_cache_update( self, layer: AttentionLayer, @@ -1393,8 +1398,7 @@ class AiterFlashAttentionImpl(AttentionImpl): kv_cache: torch.Tensor, slot_mapping: torch.Tensor, ): - # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) - key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1) + key_cache, value_cache = self._split_kv_cache(kv_cache) # key and value may be None in the case of cross attention. They are # calculated once based on the output from the encoder and then cached @@ -1449,7 +1453,12 @@ class AiterFlashAttentionImpl(AttentionImpl): ) def fused_qk_norm_rope_kvcache_supported(self): - return rocm_aiter_ops.is_enabled() + # Only fuse when shuffle layout is off; the shuffle write path uses a + # dedicated cache update, mirroring fused_rope_kvcache_supported. + return ( + rocm_aiter_ops.is_enabled() + and not rocm_aiter_ops.is_shuffle_kv_cache_enabled() + ) def do_qk_norm_rope_kvcache_update( self, @@ -1466,7 +1475,7 @@ class AiterFlashAttentionImpl(AttentionImpl): kv_cache: torch.Tensor, layer_slot_mapping: torch.Tensor, ): - key_cache, value_cache = kv_cache.unbind(1) + key_cache, value_cache = self._split_kv_cache(kv_cache) rocm_aiter_ops.do_qk_norm_rope_kvcache_update( qkv=qkv, q_weight=q_weight, @@ -1502,7 +1511,7 @@ class AiterFlashAttentionImpl(AttentionImpl): layer_slot_mapping: torch.Tensor, ): # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs)) - key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-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) diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index 08cd869686..2ce0cf4585 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py @@ -329,9 +329,6 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl): kv_cache: torch.Tensor, layer_slot_mapping: torch.Tensor, ): - # _split_kv_cache picks the unbind dim per layout (incl. the K/V-first - # encoder-decoder path). unified reads NHD, so never write the shuffle - # layout here. key_cache, value_cache = self._split_kv_cache(kv_cache) rocm_aiter_ops.do_qk_norm_rope_kvcache_update( qkv=qkv,