[ROCm][Bugfix] Support shared KV prefill in AITER attention (#55887)
Signed-off-by: Andreas Karatzas <Andreas.Karatzas@amd.com>
Andreas Karatzas committed
Sep 9, 2026 at 16:51 UTC
dcd544486b7f4672b3c2e60ea289f19d86cf355c
2 files changed
+220
-2
tests/kernels/attention/test_rocm_aiter_fa.py
+148
@@ -687,6 +687,154 @@ def test_aiter_fa_sliding_window_matches_reference():
687
torch.testing.assert_close(output, ref, atol=atol, rtol=rtol)
688
689
690
+@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
691
+@pytest.mark.parametrize(
692
+ "shuffle,sliding_window", [(False, None), (False, 16), (True, None)]
693
+)
694
+@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"])
695
+@pytest.mark.parametrize("batch", ["prefill", "extend", "mixed"])
696
+def test_aiter_fa_shared_kv_matches_reference(
697
+ monkeypatch, shuffle, sliding_window, kv_cache_dtype, batch
698
+):
699
+ """Shared layers read cached K/V for prefill and extend without modifying it."""
700
+ from tests.v1.attention.utils import (
701
+ BatchSpec,
702
+ create_common_attn_metadata,
703
+ create_standard_kv_cache_spec,
704
+ create_vllm_config,
705
+ )
706
+ from vllm._aiter_ops import rocm_aiter_ops
707
+ from vllm.config import set_current_vllm_config
708
+ from vllm.model_executor.layers.attention import Attention
709
+ from vllm.v1.attention.backends.rocm_aiter_fa import (
710
+ AiterFlashAttentionBackend,
711
+ AiterFlashAttentionMetadataBuilder,
712
+ )
713
+
714
+ _assert_aiter_supported()
715
+ monkeypatch.setattr(rocm_aiter_ops, "_SHUFFLE_KV_CACHE_ENABLED", shuffle)
716
+ set_random_seed(42)
717
+ batch_spec = {
718
+ "prefill": BatchSpec(seq_lens=[32, 24], query_lens=[32, 24]),
719
+ "extend": BatchSpec(seq_lens=[80, 56], query_lens=[32, 24]),
720
+ "mixed": BatchSpec(seq_lens=[48, 80, 24], query_lens=[1, 32, 24]),
721
+ }[batch]
722
+ config = create_vllm_config(
723
+ model_name="Qwen/Qwen3-0.6B",
724
+ dtype="bfloat16",
725
+ max_num_batched_tokens=1024,
726
+ )
727
+ config.cache_config.cache_dtype = kv_cache_dtype
728
+ num_heads = config.model_config.get_num_attention_heads(config.parallel_config)
729
+ num_kv_heads = config.model_config.get_num_kv_heads(config.parallel_config)
730
+ head_size = config.model_config.get_head_size()
731
+ scale = head_size**-0.5
732
+ with set_current_vllm_config(config):
733
+ target = Attention(
734
+ num_heads,
735
+ head_size,
736
+ scale,
737
+ num_kv_heads,
738
+ cache_config=config.cache_config,
739
+ per_layer_sliding_window=sliding_window,
740
+ prefix="target",
741
+ attn_backend=AiterFlashAttentionBackend,
742
+ )
743
+ layer = Attention(
744
+ num_heads,
745
+ head_size,
746
+ scale,
747
+ num_kv_heads,
748
+ cache_config=config.cache_config,
749
+ per_layer_sliding_window=sliding_window,
750
+ kv_sharing_target_layer_name="target",
751
+ prefix="shared",
752
+ attn_backend=AiterFlashAttentionBackend,
753
+ )
754
+ builder = AiterFlashAttentionMetadataBuilder(
755
+ create_standard_kv_cache_spec(config),
756
+ ["target", "shared"],
757
+ config,
758
+ torch.device("cuda"),
759
+ )
760
+ common = create_common_attn_metadata(
761
+ batch_spec, BLOCK_SIZE, torch.device("cuda"), max_block_idx=32
762
+ )
763
+ cache_dtype = (
764
+ current_platform.fp8_dtype() if kv_cache_dtype == "fp8" else torch.bfloat16
765
+ )
766
+ kv_cache = (
767
+ torch.empty(
768
+ 2, 32, BLOCK_SIZE, num_kv_heads * head_size, dtype=cache_dtype
769
+ ).transpose(0, 1)
770
+ if shuffle
771
+ else torch.empty(
772
+ 32, num_kv_heads, BLOCK_SIZE, 2 * head_size, dtype=cache_dtype
773
+ )
774
+ )
775
+ target.kv_cache = kv_cache
776
+ layer.kv_cache = kv_cache
777
+ k_scale = 0.5 if kv_cache_dtype == "fp8" and not shuffle else 1.0
778
+ v_scale = 0.25 if kv_cache_dtype == "fp8" and not shuffle else 1.0
779
+ for attn_layer in (target, layer):
780
+ attn_layer._k_scale.fill_(k_scale)
781
+ attn_layer._v_scale.fill_(v_scale)
782
+ key = torch.randn(
783
+ 32 * BLOCK_SIZE, num_kv_heads, head_size, dtype=torch.bfloat16
784
+ )
785
+ value = torch.randn_like(key)
786
+ target.impl.do_kv_cache_update(
787
+ target,
788
+ key,
789
+ value,
790
+ kv_cache,
791
+ torch.arange(32 * BLOCK_SIZE, dtype=torch.int64),
792
+ )
793
+ key_cache = ((key / k_scale).to(cache_dtype).to(key.dtype) * k_scale).reshape(
794
+ 32, BLOCK_SIZE, num_kv_heads, head_size
795
+ )
796
+ value_cache = (
797
+ (value / v_scale).to(cache_dtype).to(value.dtype) * v_scale
798
+ ).reshape(32, BLOCK_SIZE, num_kv_heads, head_size)
799
+ with torch.device("cpu"):
800
+ metadata = builder.build(0, common)
801
+ cache_before = kv_cache.view(torch.uint8).clone()
802
+ query = torch.randn(
803
+ batch_spec.compute_num_tokens(),
804
+ num_heads,
805
+ head_size,
806
+ dtype=torch.bfloat16,
807
+ )
808
+ output = torch.empty_like(query)
809
+ layer.impl.forward(layer, query, None, None, kv_cache, metadata, output)
810
+ expected = ref_paged_attn(
811
+ query,
812
+ key_cache.contiguous(),
813
+ value_cache.contiguous(),
814
+ batch_spec.query_lens,
815
+ batch_spec.seq_lens,
816
+ common.block_table_tensor,
817
+ scale,
818
+ sliding_window,
819
+ )
820
+ num_decode_tokens = metadata.num_decode_tokens
821
+ if kv_cache_dtype == "fp8" and num_decode_tokens:
822
+ # Use the established FP8 decode tolerance; gathered prefill and
823
+ # extend tokens must still satisfy the tighter BF16 tolerance.
824
+ torch.testing.assert_close(
825
+ output[:num_decode_tokens],
826
+ expected[:num_decode_tokens],
827
+ atol=6e-2,
828
+ rtol=1e-1,
829
+ )
830
+ output = output[num_decode_tokens:]
831
+ expected = expected[num_decode_tokens:]
832
+ torch.testing.assert_close(output, expected, atol=2e-2, rtol=2e-2)
833
+ torch.testing.assert_close(
834
+ kv_cache.view(torch.uint8), cache_before, atol=0, rtol=0
835
+ )
836
+
837
+
838
# Decode path test --------------------------------------------------------
839
@pytest.mark.skipif(not on_mi3xx(), reason="MI300/MI350 ROCm only")
840
@pytest.mark.parametrize(
vllm/v1/attention/backends/rocm_aiter_fa.py
+72
-2
@@ -403,6 +403,14 @@ class AiterFlashAttentionChunkPrefillMetadata:
403
chunk_context_metadata: AiterChunkContextMetadata
404
405
406
+@dataclass
407
+class AiterKVSharingMetadata:
408
+ workspace: torch.Tensor
409
+ query_start_loc: torch.Tensor
410
+ token_to_batch: torch.Tensor
411
+ seq_starts: torch.Tensor
412
+
413
+
414
@dataclass
415
class AiterFlashAttentionMetadata:
416
# NOTE(sang): Definition of context_len, query_len, and seq_len.
@@ -439,6 +447,7 @@ class AiterFlashAttentionMetadata:
447
# since we might integrate per token quant for kv cache in the future.
448
k_scale: dict[str, torch.Tensor] | None
449
v_scale: dict[str, torch.Tensor] | None
450
+ kv_sharing_metadata: AiterKVSharingMetadata | None = None
451
452
453
class AiterFlashAttentionMetadataBuilder(
@@ -471,6 +480,7 @@ class AiterFlashAttentionMetadataBuilder(
480
self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)
481
482
sliding_window_configs: set[tuple[int, int] | None] = set()
483
+ kv_sharing_shape = None
484
layers = get_layers_from_vllm_config(self.vllm_config, Attention)
485
for name, layer in layers.items():
486
if name not in layer_names:
@@ -480,6 +490,8 @@ class AiterFlashAttentionMetadataBuilder(
490
"with Aiter Flash Attention Impl."
491
)
492
sliding_window_configs.add(layer.impl.sliding_window)
493
+ if layer.kv_sharing_target_layer_name is not None:
494
+ kv_sharing_shape = (layer.impl.num_kv_heads, layer.impl.head_size)
495
496
while len(sliding_window_configs) > 0:
497
sliding_window_config = sliding_window_configs.pop()
@@ -495,6 +507,19 @@ class AiterFlashAttentionMetadataBuilder(
507
device=device,
508
)
509
self.scale = torch.tensor([1.0], dtype=torch.float, device=self.device)
510
+ self.kv_sharing_workspace = (
511
+ torch.empty(
512
+ (
513
+ 2,
514
+ vllm_config.scheduler_config.max_num_batched_tokens,
515
+ *kv_sharing_shape,
516
+ ),
517
+ dtype=self.model_config.dtype,
518
+ device=device,
519
+ )
520
+ if kv_sharing_shape is not None
521
+ else None
522
+ )
523
524
def build_for_cudagraph_capture(
525
self, common_attn_metadata: CommonAttentionMetadata
@@ -554,6 +579,23 @@ class AiterFlashAttentionMetadataBuilder(
579
580
query_lens_cpu = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1]
581
582
+ kv_sharing_metadata = None
583
+ if self.kv_sharing_workspace is not None and seq_lens is not None:
584
+ query_lens = query_lens_cpu[num_decodes:]
585
+ token_to_batch = torch.repeat_interleave(
586
+ torch.arange(len(query_lens), dtype=torch.int32, device="cpu"),
587
+ query_lens,
588
+ )
589
+ query_start_loc = common_attn_metadata.query_start_loc[num_decodes:]
590
+ kv_sharing_metadata = AiterKVSharingMetadata(
591
+ workspace=self.kv_sharing_workspace,
592
+ query_start_loc=query_start_loc - query_start_loc[0],
593
+ token_to_batch=async_tensor_h2d(token_to_batch, self.device),
594
+ seq_starts=async_tensor_h2d(
595
+ seq_lens[num_decodes:] - query_lens, self.device
596
+ ),
597
+ )
598
+
599
decode_metadata = None
600
if num_decodes > 0:
601
decode_max_query_len = query_lens_cpu[:num_decodes].max().item()
@@ -719,6 +761,7 @@ class AiterFlashAttentionMetadataBuilder(
761
use_cascade=use_cascade,
762
k_scale=self.scale,
763
v_scale=self.scale,
764
+ kv_sharing_metadata=kv_sharing_metadata,
765
)
766
return attn_metadata
767
@@ -1135,8 +1178,8 @@ class AiterFlashAttentionImpl(AttentionImpl):
1178
self,
1179
layer: torch.nn.Module,
1180
query: torch.Tensor,
1138
- key: torch.Tensor,
1139
- value: torch.Tensor,
1181
+ key: torch.Tensor | None,
1182
+ value: torch.Tensor | None,
1183
kv_cache: torch.Tensor,
1184
attn_metadata: AiterFlashAttentionMetadata,
1185
output: torch.Tensor,
@@ -1199,10 +1242,36 @@ class AiterFlashAttentionImpl(AttentionImpl):
1242
1243
num_decode_tokens = attn_metadata.num_decode_tokens
1244
num_extend_tokens = attn_metadata.num_extend_tokens
1245
+ if self.kv_sharing_target_layer_name is not None and (
1246
+ num_prefills > 0 or num_extends > 0
1247
+ ):
1248
+ # Shared layers project only Q. Fetch the new K/V tokens from the
1249
+ # target cache for the direct prefill and extend suffix kernels.
1250
+ shared = attn_metadata.kv_sharing_metadata
1251
+ assert shared is not None
1252
+ key, value = shared.workspace[:, :num_actual_tokens].unbind(0)
1253
+ cp_mha_gather_cache(
1254
+ key_cache=key_cache,
1255
+ value_cache=value_cache,
1256
+ key=key[num_decode_tokens:],
1257
+ value=value[num_decode_tokens:],
1258
+ block_tables=attn_metadata.block_table[num_decodes:],
1259
+ k_scales=layer._k_scale,
1260
+ v_scales=layer._v_scale,
1261
+ cu_seqlens_kv=shared.query_start_loc,
1262
+ token_to_batch=shared.token_to_batch,
1263
+ seq_starts=shared.seq_starts,
1264
+ dequant=is_quantized_kv_cache(self.kv_cache_dtype),
1265
+ kv_cache_layout="SHUFFLE"
1266
+ if rocm_aiter_ops.is_shuffle_kv_cache_enabled()
1267
+ else "NHD",
1268
+ total_tokens=num_actual_tokens - num_decode_tokens,
1269
+ )
1270
if not attn_metadata.use_cascade:
1271
# calculate for pure prefills
1272
if num_prefills > 0:
1273
assert attn_metadata.prefill_metadata is not None
1274
+ assert key is not None and value is not None
1275
1276
prefill_query = query[num_decode_tokens + num_extend_tokens :]
1277
prefill_key = key[num_decode_tokens + num_extend_tokens :]
@@ -1229,6 +1298,7 @@ class AiterFlashAttentionImpl(AttentionImpl):
1298
# calculate for extends
1299
if num_extends > 0:
1300
assert attn_metadata.extend_metadata is not None
1301
+ assert key is not None and value is not None
1302
extend_tokens_slice = slice(
1303
num_decode_tokens, num_decode_tokens + num_extend_tokens
1304
)