From 8359e15aae32dee9dc1f259a9b2574fb72b5507e Mon Sep 17 00:00:00 2001 From: Simon Danielsson <70206058+simondanielsson@users.noreply.github.com> Date: Thu, 10 Sep 2026 19:48:42 +0200 Subject: [PATCH] [ROCm][Feature] Support KV connectors with ROCM_AITER_UNIFIED_ATTN (#53695) Signed-off-by: simondanielsson 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> --- .../test_rocm_attention_backends_selection.py | 112 ++++++++++++++++++ .../unit/test_offloading_connector.py | 16 ++- .../kv_connector/v1/offloading/config.py | 13 +- .../backends/rocm_aiter_unified_attn.py | 16 +++ 4 files changed, 154 insertions(+), 3 deletions(-) diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 8668fa26bc..2342a279a3 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/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) diff --git a/tests/v1/kv_connector/unit/test_offloading_connector.py b/tests/v1/kv_connector/unit/test_offloading_connector.py index 08ef476992..cde672bcad 100644 --- a/tests/v1/kv_connector/unit/test_offloading_connector.py +++ b/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"] diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py index e5ed896078..bc78b72988 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading/config.py +++ b/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: diff --git a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py b/vllm/v1/attention/backends/rocm_aiter_unified_attn.py index cbd33e7bc6..e0e4b38550 100644 --- a/vllm/v1/attention/backends/rocm_aiter_unified_attn.py +++ b/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