[ROCm][Feature] Support KV connectors with ROCM_AITER_UNIFIED_ATTN (#53695)
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com> Signed-off-by: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Co-authored-by: coderabbitai[bot] <136622811+coderabbitai[bot]@users.noreply.github.com> Co-authored-by: Shanshan Shen <467638484@qq.com>
Simon Danielsson committed
Sep 10, 2026 at 19:48 UTC
8359e15aae32dee9dc1f259a9b2574fb72b5507e
4 files changed
+154
-3
tests/v1/attention/test_rocm_attention_backends_selection.py
+112
@@ -455,3 +455,115 @@ def test_sparse_not_supported(mock_vllm_config):
455
RocmPlatform.get_attn_backend_cls(
456
selected_backend=None, attn_selector_config=attn_selector_config
457
)
458
+
459
+
460
+def _kv_connector_selector_config() -> AttentionSelectorConfig:
461
+ return AttentionSelectorConfig(
462
+ head_size=128,
463
+ dtype=torch.float16,
464
+ kv_cache_dtype="auto",
465
+ block_size=16,
466
+ use_mla=False,
467
+ has_sink=False,
468
+ use_sparse=False,
469
+ use_kv_connector=True,
470
+ )
471
+
472
+
473
+def test_unified_attn_declares_kv_connector_support():
474
+ """ROCM_AITER_UNIFIED_ATTN opts into KV connectors and ROCM_ATTN does not."""
475
+ from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
476
+ RocmAiterUnifiedAttentionBackend,
477
+ )
478
+ from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
479
+
480
+ assert RocmAiterUnifiedAttentionBackend.supports_kv_connector() is True
481
+ assert RocmAttentionBackend.supports_kv_connector() is False
482
+
483
+
484
+def test_unified_attn_supports_kv_connector(mock_vllm_config, mock_get_cdna_version):
485
+ """ROCM_AITER_UNIFIED_ATTN can be selected with KV connectors."""
486
+ from vllm.platforms.rocm import RocmPlatform
487
+
488
+ backend_path = RocmPlatform.get_attn_backend_cls(
489
+ selected_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
490
+ attn_selector_config=_kv_connector_selector_config(),
491
+ )
492
+
493
+ assert backend_path == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path()
494
+
495
+
496
+def test_rocm_attn_rejects_kv_connector(mock_vllm_config, mock_get_cdna_version):
497
+ """Selecting ROCM_ATTN with a KV connector is illegal."""
498
+ from vllm.platforms.rocm import RocmPlatform
499
+
500
+ attn_selector_config = _kv_connector_selector_config()
501
+
502
+ with pytest.raises(ValueError, match="KV connector not supported"):
503
+ RocmPlatform.get_attn_backend_cls(
504
+ selected_backend=AttentionBackendEnum.ROCM_ATTN,
505
+ attn_selector_config=attn_selector_config,
506
+ )
507
+
508
+
509
+@pytest.mark.parametrize(
510
+ "aiter_found, expected_backend",
511
+ [
512
+ (True, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN),
513
+ (False, AttentionBackendEnum.TRITON_ATTN),
514
+ ],
515
+)
516
+def test_auto_selection_for_kv_connector(
517
+ aiter_found, expected_backend, mock_vllm_config, mock_get_cdna_version
518
+):
519
+ """Auto-selection with a KV connector and AITER enabled resolves to unified attn,
520
+ and to triton attn if AITER not enabled."""
521
+ from vllm.platforms.rocm import RocmPlatform
522
+
523
+ with patch(
524
+ "vllm._aiter_ops.is_aiter_found_and_supported", return_value=aiter_found
525
+ ):
526
+ backend_path = RocmPlatform.get_attn_backend_cls(
527
+ selected_backend=None,
528
+ attn_selector_config=_kv_connector_selector_config(),
529
+ )
530
+
531
+ assert backend_path == expected_backend.get_path()
532
+
533
+
534
+def test_unified_attn_prefers_block_contiguous_layout():
535
+ """Unified attn prefers a block-first KV layout, hence ok with kv connectors."""
536
+ from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
537
+ RocmAiterUnifiedAttentionBackend,
538
+ )
539
+ from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
540
+
541
+ unified_preferred = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()[0]
542
+ rocm_attn_preferred = RocmAttentionBackend.supported_kv_cache_layouts()[0]
543
+
544
+ assert unified_preferred.is_block_contiguous is True
545
+ assert rocm_attn_preferred.is_block_contiguous is False
546
+
547
+
548
+def test_unified_attn_drops_lhbnc_with_kv_connector():
549
+ """Connectors move a block as one contiguous byte range, which LHBNC breaks."""
550
+ from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
551
+ from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
552
+ RocmAiterUnifiedAttentionBackend,
553
+ )
554
+ from vllm.v1.kv_cache_interface import KVCacheLayout
555
+
556
+ assert KVCacheLayout.LHBNC in (
557
+ RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()
558
+ )
559
+
560
+ config = VllmConfig(
561
+ kv_transfer_config=KVTransferConfig(
562
+ kv_connector="ExampleConnector", kv_role="kv_both"
563
+ )
564
+ )
565
+ with set_current_vllm_config(config):
566
+ layouts = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()
567
+
568
+ assert layouts
569
+ assert all(layout.is_block_compact for layout in layouts)
tests/v1/kv_connector/unit/test_offloading_connector.py
+15
-1
@@ -15,12 +15,26 @@ from vllm.config import KVEventsConfig, KVTransferConfig
15
from vllm.distributed.kv_events import BlockStored, KVEventBatch
16
from vllm.platforms import current_platform
17
18
-CPU_BLOCK_SIZES: int = 64 if current_platform.is_xpu() else 48
18
+# The offloaded block size must stay a whole number of GPU blocks.
19
+CPU_BLOCK_SIZES: int
20
+if current_platform.is_xpu():
21
+ CPU_BLOCK_SIZES = 64
22
+elif current_platform.is_rocm():
23
+ # ROCM_AITER_UNIFIED_ATTN prefers 64-token GPU blocks, so 192 covers both that
24
+ # and the 16-token default.
25
+ CPU_BLOCK_SIZES = 192
26
+else:
27
+ CPU_BLOCK_SIZES = 48
28
+
29
_ATTN_BACKENDS: list[str] = []
30
if current_platform.is_cuda():
31
_ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"]
32
elif current_platform.is_rocm():
33
+ from vllm._aiter_ops import is_aiter_found_and_supported
34
+
35
_ATTN_BACKENDS = ["TRITON_ATTN"]
36
+ if is_aiter_found_and_supported():
37
+ _ATTN_BACKENDS.append("ROCM_AITER_UNIFIED_ATTN")
38
elif current_platform.is_xpu():
39
_ATTN_BACKENDS = ["FLASH_ATTN", "TRITON_ATTN"]
40
vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py
+11
-2
@@ -4,6 +4,7 @@
4
5
from typing import TYPE_CHECKING
6
7
+from vllm.utils.math_utils import round_up
8
from vllm.v1.core.kv_cache_utils import (
9
resolve_dcp_kv_block_size,
10
resolve_kv_cache_block_sizes,
@@ -90,8 +91,16 @@ def build_offloading_config(
91
)
92
93
tokens_per_block = unique_tokens_per_block.pop()
93
- assert tokens_per_chunk_int % tokens_per_block == 0
94
- blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
94
+ if tokens_per_chunk_int % tokens_per_block == 0:
95
+ blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
96
+ else:
97
+ raise ValueError(
98
+ f"'block_size'={tokens_per_chunk_int} in kv_connector_extra_config "
99
+ f"must be a multiple of the GPU KV cache block size "
100
+ f"({tokens_per_block} tokens). Use "
101
+ f"{round_up(tokens_per_chunk_int, tokens_per_block)} instead, or set "
102
+ f"'blocks_per_chunk' to express the chunk size in blocks."
103
+ )
104
105
worker_kv_bytes_per_block = 0
106
if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors:
vllm/v1/attention/backends/rocm_aiter_unified_attn.py
+16
@@ -8,6 +8,7 @@ import torch
8
9
from vllm import _custom_ops as ops
10
from vllm._aiter_ops import rocm_aiter_ops
11
+from vllm.config import get_current_vllm_config_or_none
12
from vllm.config.cache import CacheDType
13
from vllm.logger import init_logger
14
from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -32,6 +33,14 @@ from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
33
logger = init_logger(__name__)
34
35
36
+def _kv_connector_enabled() -> bool:
37
+ vllm_config = get_current_vllm_config_or_none()
38
+ if vllm_config is None:
39
+ return False
40
+ kv_transfer_config = vllm_config.kv_transfer_config
41
+ return kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance
42
+
43
+
44
class RocmAiterUnifiedAttentionMetadataBuilder(RocmAttentionMetadataBuilder):
45
def build_for_cudagraph_capture(
46
self, common_attn_metadata: CommonAttentionMetadata
@@ -84,6 +93,10 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
93
def supports_non_causal(cls) -> bool:
94
return False
95
96
+ @classmethod
97
+ def supports_kv_connector(cls) -> bool:
98
+ return True
99
+
100
forward_includes_kv_cache_update: bool = False
101
102
@staticmethod
@@ -110,6 +123,9 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
123
# K and V come out of the content dim as transposed views rather than
124
# copies, so the head dim may sit on either side of the block dim, but
125
# the layer must stay outermost.
126
+ if _kv_connector_enabled():
127
+ # Connectors like MoRI assume contiguous blocks
128
+ return (KVCacheLayout.LBHNC,)
129
return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC)
130
131
@staticmethod