[Bugfix][ROCm] Fix ROCM_AITER_FA & ROCM_AITER_UNIFIED_ATTN QK-Norm+RoPE+KVCache fusion for the packed KV-cache [BLOCKS, HEADS, BLOCK_SIZE, 2*HEAD_DIM] layout (#49373)
Signed-off-by: Jack Hu <Jack.Hu@amd.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Jack Hu committed
Aug 7, 2026 at 12:54 UTC
34c1cd20a58e96f197d4b5927336a2d9facafa58
3 files changed
+27
-14
tests/compile/passes/test_rocm_aiter_qk_norm_rope_kvcache_fusion.py
+11
-4
@@ -420,10 +420,9 @@ _FUSION_CONFIGS = [
420
@pytest.mark.parametrize("use_shuffle_kv_layout", ["1", "0"])
421
@pytest.mark.parametrize(
422
"kv_stride_order",
423
- [
424
- pytest.param((0, 1, 2, 3, 4), id="block_first"),
425
- pytest.param((1, 0, 2, 3, 4), id="kv_first"),
426
- ],
423
+ # FA/unified use the 4D packed cache (num_blocks, num_kv_heads, block_size,
424
+ # 2*head_size); no separate K/V dim to permute.
425
+ [pytest.param((0, 1, 2, 3), id="packed_4d")],
426
)
427
@pytest.mark.parametrize("enable_aiter_triton_rope", [True, False])
428
@pytest.mark.parametrize("block_size", [16, 32, 64])
@@ -458,6 +457,14 @@ def test_qk_norm_rope_kvcache_fusion(
457
and use_shuffle_kv_layout == "1"
458
):
459
pytest.skip("ROCM_AITER_UNIFIED_ATTN is NHD-only; shuffle env is ignored")
460
+ if (
461
+ attn_backend == AttentionBackendEnum.ROCM_AITER_FA
462
+ and use_shuffle_kv_layout == "1"
463
+ ):
464
+ pytest.skip(
465
+ "ROCM_AITER_FA gates the qk_norm+rope+kvcache fusion off under shuffle "
466
+ "layout (defers to reshape_and_cache_shuffle_triton), so nothing fuses"
467
+ )
468
_run_qk_norm_rope_kvcache_fusion_test(
469
attn_backend=attn_backend,
470
enable_aiter_triton_rope=enable_aiter_triton_rope,
vllm/v1/attention/backends/rocm_aiter_fa.py
+16
-7
@@ -1065,9 +1065,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1065
# Whenever making a change in this method, please benchmark the
1066
# performance to make sure it does not introduce any overhead.
1067
num_actual_tokens = attn_metadata.num_actual_tokens
1068
- # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
1069
- key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1068
1069
+ key_cache, value_cache = self._split_kv_cache(kv_cache)
1070
if is_quantized_kv_cache(self.kv_cache_dtype):
1071
key_cache = key_cache.view(current_platform.fp8_dtype())
1072
value_cache = value_cache.view(current_platform.fp8_dtype())
@@ -1385,6 +1384,12 @@ class AiterFlashAttentionImpl(AttentionImpl):
1384
1385
return output
1386
1387
+ def _split_kv_cache(
1388
+ self, kv_cache: torch.Tensor
1389
+ ) -> tuple[torch.Tensor, torch.Tensor]:
1390
+ # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
1391
+ return kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1392
+
1393
def do_kv_cache_update(
1394
self,
1395
layer: AttentionLayer,
@@ -1393,8 +1398,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1398
kv_cache: torch.Tensor,
1399
slot_mapping: torch.Tensor,
1400
):
1396
- # (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
1397
- key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1401
+ key_cache, value_cache = self._split_kv_cache(kv_cache)
1402
1403
# key and value may be None in the case of cross attention. They are
1404
# calculated once based on the output from the encoder and then cached
@@ -1449,7 +1453,12 @@ class AiterFlashAttentionImpl(AttentionImpl):
1453
)
1454
1455
def fused_qk_norm_rope_kvcache_supported(self):
1452
- return rocm_aiter_ops.is_enabled()
1456
+ # Only fuse when shuffle layout is off; the shuffle write path uses a
1457
+ # dedicated cache update, mirroring fused_rope_kvcache_supported.
1458
+ return (
1459
+ rocm_aiter_ops.is_enabled()
1460
+ and not rocm_aiter_ops.is_shuffle_kv_cache_enabled()
1461
+ )
1462
1463
def do_qk_norm_rope_kvcache_update(
1464
self,
@@ -1466,7 +1475,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1475
kv_cache: torch.Tensor,
1476
layer_slot_mapping: torch.Tensor,
1477
):
1469
- key_cache, value_cache = kv_cache.unbind(1)
1478
+ key_cache, value_cache = self._split_kv_cache(kv_cache)
1479
rocm_aiter_ops.do_qk_norm_rope_kvcache_update(
1480
qkv=qkv,
1481
q_weight=q_weight,
@@ -1502,7 +1511,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1511
layer_slot_mapping: torch.Tensor,
1512
):
1513
# (B, H, N, 2*hs) -> ((B, N, H, hs), (B, N, H, hs))
1505
- key_cache, value_cache = kv_cache.transpose(1, 2).split(self.head_size, dim=-1)
1514
+ key_cache, value_cache = self._split_kv_cache(kv_cache)
1515
flash_layout = True
1516
1517
is_fp8_kv_cache = is_quantized_kv_cache(self.kv_cache_dtype)
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
-3
@@ -329,9 +329,6 @@ class RocmAiterUnifiedAttentionImpl(RocmAttentionImpl):
329
kv_cache: torch.Tensor,
330
layer_slot_mapping: torch.Tensor,
331
):
332
- # _split_kv_cache picks the unbind dim per layout (incl. the K/V-first
333
- # encoder-decoder path). unified reads NHD, so never write the shuffle
334
- # layout here.
332
key_cache, value_cache = self._split_kv_cache(kv_cache)
333
rocm_aiter_ops.do_qk_norm_rope_kvcache_update(
334
qkv=qkv,