[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>

tests/v1/attention/test_rocm_attention_backends_selection.py

@@ -455,3 +455,115 @@ def test_sparse_not_supported(mock_vllm_config):
         RocmPlatform.get_attn_backend_cls(
             selected_backend=None, attn_selector_config=attn_selector_config
         )
+
+
+def _kv_connector_selector_config() -> AttentionSelectorConfig:
+    return AttentionSelectorConfig(
+        head_size=128,
+        dtype=torch.float16,
+        kv_cache_dtype="auto",
+        block_size=16,
+        use_mla=False,
+        has_sink=False,
+        use_sparse=False,
+        use_kv_connector=True,
+    )
+
+
+def test_unified_attn_declares_kv_connector_support():
+    """ROCM_AITER_UNIFIED_ATTN opts into KV connectors and ROCM_ATTN does not."""
+    from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
+        RocmAiterUnifiedAttentionBackend,
+    )
+    from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
+
+    assert RocmAiterUnifiedAttentionBackend.supports_kv_connector() is True
+    assert RocmAttentionBackend.supports_kv_connector() is False
+
+
+def test_unified_attn_supports_kv_connector(mock_vllm_config, mock_get_cdna_version):
+    """ROCM_AITER_UNIFIED_ATTN can be selected with KV connectors."""
+    from vllm.platforms.rocm import RocmPlatform
+
+    backend_path = RocmPlatform.get_attn_backend_cls(
+        selected_backend=AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN,
+        attn_selector_config=_kv_connector_selector_config(),
+    )
+
+    assert backend_path == AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN.get_path()
+
+
+def test_rocm_attn_rejects_kv_connector(mock_vllm_config, mock_get_cdna_version):
+    """Selecting ROCM_ATTN with a KV connector is illegal."""
+    from vllm.platforms.rocm import RocmPlatform
+
+    attn_selector_config = _kv_connector_selector_config()
+
+    with pytest.raises(ValueError, match="KV connector not supported"):
+        RocmPlatform.get_attn_backend_cls(
+            selected_backend=AttentionBackendEnum.ROCM_ATTN,
+            attn_selector_config=attn_selector_config,
+        )
+
+
+@pytest.mark.parametrize(
+    "aiter_found, expected_backend",
+    [
+        (True, AttentionBackendEnum.ROCM_AITER_UNIFIED_ATTN),
+        (False, AttentionBackendEnum.TRITON_ATTN),
+    ],
+)
+def test_auto_selection_for_kv_connector(
+    aiter_found, expected_backend, mock_vllm_config, mock_get_cdna_version
+):
+    """Auto-selection with a KV connector and AITER enabled resolves to unified attn,
+    and to triton attn if AITER not enabled."""
+    from vllm.platforms.rocm import RocmPlatform
+
+    with patch(
+        "vllm._aiter_ops.is_aiter_found_and_supported", return_value=aiter_found
+    ):
+        backend_path = RocmPlatform.get_attn_backend_cls(
+            selected_backend=None,
+            attn_selector_config=_kv_connector_selector_config(),
+        )
+
+    assert backend_path == expected_backend.get_path()
+
+
+def test_unified_attn_prefers_block_contiguous_layout():
+    """Unified attn prefers a block-first KV layout, hence ok with kv connectors."""
+    from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
+        RocmAiterUnifiedAttentionBackend,
+    )
+    from vllm.v1.attention.backends.rocm_attn import RocmAttentionBackend
+
+    unified_preferred = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()[0]
+    rocm_attn_preferred = RocmAttentionBackend.supported_kv_cache_layouts()[0]
+
+    assert unified_preferred.is_block_contiguous is True
+    assert rocm_attn_preferred.is_block_contiguous is False
+
+
+def test_unified_attn_drops_lhbnc_with_kv_connector():
+    """Connectors move a block as one contiguous byte range, which LHBNC breaks."""
+    from vllm.config import KVTransferConfig, VllmConfig, set_current_vllm_config
+    from vllm.v1.attention.backends.rocm_aiter_unified_attn import (
+        RocmAiterUnifiedAttentionBackend,
+    )
+    from vllm.v1.kv_cache_interface import KVCacheLayout
+
+    assert KVCacheLayout.LHBNC in (
+        RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()
+    )
+
+    config = VllmConfig(
+        kv_transfer_config=KVTransferConfig(
+            kv_connector="ExampleConnector", kv_role="kv_both"
+        )
+    )
+    with set_current_vllm_config(config):
+        layouts = RocmAiterUnifiedAttentionBackend.supported_kv_cache_layouts()
+
+    assert layouts
+    assert all(layout.is_block_compact for layout in layouts)

tests/v1/kv_connector/unit/test_offloading_connector.py

@@ -15,12 +15,26 @@ from vllm.config import KVEventsConfig, KVTransferConfig
 from vllm.distributed.kv_events import BlockStored, KVEventBatch
 from vllm.platforms import current_platform
 
-CPU_BLOCK_SIZES: int = 64 if current_platform.is_xpu() else 48
+# The offloaded block size must stay a whole number of GPU blocks.
+CPU_BLOCK_SIZES: int
+if current_platform.is_xpu():
+    CPU_BLOCK_SIZES = 64
+elif current_platform.is_rocm():
+    # ROCM_AITER_UNIFIED_ATTN prefers 64-token GPU blocks, so 192 covers both that
+    # and the 16-token default.
+    CPU_BLOCK_SIZES = 192
+else:
+    CPU_BLOCK_SIZES = 48
+
 _ATTN_BACKENDS: list[str] = []
 if current_platform.is_cuda():
     _ATTN_BACKENDS = ["FLASH_ATTN", "FLASHINFER", "TRITON_ATTN"]
 elif current_platform.is_rocm():
+    from vllm._aiter_ops import is_aiter_found_and_supported
+
     _ATTN_BACKENDS = ["TRITON_ATTN"]
+    if is_aiter_found_and_supported():
+        _ATTN_BACKENDS.append("ROCM_AITER_UNIFIED_ATTN")
 elif current_platform.is_xpu():
     _ATTN_BACKENDS = ["FLASH_ATTN", "TRITON_ATTN"]
 

vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py

@@ -4,6 +4,7 @@
 
 from typing import TYPE_CHECKING
 
+from vllm.utils.math_utils import round_up
 from vllm.v1.core.kv_cache_utils import (
     resolve_dcp_kv_block_size,
     resolve_kv_cache_block_sizes,
@@ -90,8 +91,16 @@ def build_offloading_config(
         )
 
         tokens_per_block = unique_tokens_per_block.pop()
-        assert tokens_per_chunk_int % tokens_per_block == 0
-        blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
+        if tokens_per_chunk_int % tokens_per_block == 0:
+            blocks_per_chunk = tokens_per_chunk_int // tokens_per_block
+        else:
+            raise ValueError(
+                f"'block_size'={tokens_per_chunk_int} in kv_connector_extra_config "
+                f"must be a multiple of the GPU KV cache block size "
+                f"({tokens_per_block} tokens). Use "
+                f"{round_up(tokens_per_chunk_int, tokens_per_block)} instead, or set "
+                f"'blocks_per_chunk' to express the chunk size in blocks."
+            )
 
     worker_kv_bytes_per_block = 0
     if kv_cache_config.num_blocks > 0 and kv_cache_config.kv_cache_tensors:

vllm/v1/attention/backends/rocm_aiter_unified_attn.py

@@ -8,6 +8,7 @@ import torch
 
 from vllm import _custom_ops as ops
 from vllm._aiter_ops import rocm_aiter_ops
+from vllm.config import get_current_vllm_config_or_none
 from vllm.config.cache import CacheDType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.quantization.utils.quant_utils import (
@@ -32,6 +33,14 @@ from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheLayout
 logger = init_logger(__name__)
 
 
+def _kv_connector_enabled() -> bool:
+    vllm_config = get_current_vllm_config_or_none()
+    if vllm_config is None:
+        return False
+    kv_transfer_config = vllm_config.kv_transfer_config
+    return kv_transfer_config is not None and kv_transfer_config.is_kv_transfer_instance
+
+
 class RocmAiterUnifiedAttentionMetadataBuilder(RocmAttentionMetadataBuilder):
     def build_for_cudagraph_capture(
         self, common_attn_metadata: CommonAttentionMetadata
@@ -84,6 +93,10 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
     def supports_non_causal(cls) -> bool:
         return False
 
+    @classmethod
+    def supports_kv_connector(cls) -> bool:
+        return True
+
     forward_includes_kv_cache_update: bool = False
 
     @staticmethod
@@ -110,6 +123,9 @@ class RocmAiterUnifiedAttentionBackend(RocmAttentionBackend):
         # K and V come out of the content dim as transposed views rather than
         # copies, so the head dim may sit on either side of the block dim, but
         # the layer must stay outermost.
+        if _kv_connector_enabled():
+            # Connectors like MoRI assume contiguous blocks
+            return (KVCacheLayout.LBHNC,)
         return (KVCacheLayout.LBHNC, KVCacheLayout.LHBNC)
 
     @staticmethod