diff --git a/tests/v1/attention/test_rocm_attention_backends_selection.py b/tests/v1/attention/test_rocm_attention_backends_selection.py index 2342a279a3..b99a33f9f7 100644 --- a/tests/v1/attention/test_rocm_attention_backends_selection.py +++ b/tests/v1/attention/test_rocm_attention_backends_selection.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for attention backend selectors.""" +from types import SimpleNamespace from unittest.mock import MagicMock, patch import pytest @@ -34,6 +35,15 @@ def mock_get_cdna_version(): yield +@pytest.fixture +def cleared_attention_selector_cache(): + from vllm.v1.attention.selector import _cached_get_attn_backend + + _cached_get_attn_backend.cache_clear() + yield + _cached_get_attn_backend.cache_clear() + + def test_aiter_unified_attention_uses_dedicated_metadata_builder(): from vllm.v1.attention.backends.rocm_aiter_unified_attn import ( RocmAiterUnifiedAttentionBackend, @@ -434,6 +444,139 @@ def test_aiter_fa_requires_mi3xx(mock_vllm_config): ) +@pytest.fixture +def turboquant_run_config(): + """Current config of a run whose KV cache dtype is a turboquant_* preset.""" + config = SimpleNamespace( + cache_config=SimpleNamespace(cache_dtype="turboquant_k8v4") + ) + with patch( + "vllm.config.get_current_vllm_config_or_none", + return_value=config, + ): + yield config + + +def test_turboquant_boundary_selection_is_not_cached_from_ordinary_run( + cleared_attention_selector_cache, +): + from vllm.config import CacheConfig, VllmConfig, set_current_vllm_config + from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts + from vllm.v1.attention.selector import get_attn_backend + from vllm.v1.kv_cache_layout import KVCacheLayout + + ordinary_config = VllmConfig(cache_config=CacheConfig(cache_dtype="auto")) + turboquant_config = VllmConfig( + cache_config=CacheConfig(cache_dtype="turboquant_k8v4") + ) + + backend_priorities = [ + AttentionBackendEnum.ROCM_ATTN, + AttentionBackendEnum.TRITON_ATTN, + AttentionBackendEnum.TURBOQUANT, + ] + with patch( + "vllm.platforms.rocm._get_backend_priorities", + return_value=backend_priorities, + ): + with set_current_vllm_config(ordinary_config): + ordinary_backend = get_attn_backend(128, torch.float16, "auto") + + with set_current_vllm_config(turboquant_config): + boundary_backend = get_attn_backend(128, torch.float16, "auto") + quantized_backend = get_attn_backend(128, torch.float16, "turboquant_k8v4") + + assert ordinary_backend is AttentionBackendEnum.ROCM_ATTN.get_class() + assert quantized_backend is AttentionBackendEnum.TURBOQUANT.get_class() + assert get_supported_kv_cache_layouts([boundary_backend, quantized_backend]) == [ + KVCacheLayout.LBNHC + ] + + +def test_turboquant_run_does_not_mask_unsupported_selected_backend( + turboquant_run_config, +): + from vllm.platforms.rocm import RocmPlatform + + attn_selector_config = AttentionSelectorConfig( + head_size=128, + dtype=torch.float16, + kv_cache_dtype="auto", + block_size=16, + ) + + with ( + patch("vllm.platforms.rocm.get_cdna_version", return_value=1), + pytest.raises(ValueError, match="compute capability not supported"), + ): + RocmPlatform.get_attn_backend_cls( + selected_backend=AttentionBackendEnum.ROCM_AITER_FA, + attn_selector_config=attn_selector_config, + ) + + +def test_turboquant_layout_check_respects_backend_override(): + from vllm.platforms.rocm import _shares_layout_with_turboquant + from vllm.v1.kv_cache_layout import KVCacheLayout + + boundary_backend = MagicMock() + boundary_backend.supported_kv_cache_layouts.return_value = (KVCacheLayout.LBHNC,) + turboquant_override = MagicMock() + turboquant_override.supported_kv_cache_layouts.return_value = (KVCacheLayout.LBHNC,) + + with patch.object( + AttentionBackendEnum.TURBOQUANT, + "get_class", + return_value=turboquant_override, + ) as get_turboquant_class: + assert _shares_layout_with_turboquant(boundary_backend) + get_turboquant_class.assert_called_once_with() + + +@pytest.mark.parametrize( + "selected_backend", + [None, "ROCM_ATTN", "ROCM_AITER_FA", "TURBOQUANT"], +) +def test_turboquant_boundary_layers_share_a_layout( + selected_backend, + turboquant_run_config, + mock_get_cdna_version, +): + """A turboquant_* run must resolve to backends with a common KV layout. + + The boundary layers keep the native dtype and pick their own backend while + every other layer picks TURBOQUANT. Engine startup hard-errors when the two + share no layout, so the boundary layers must not land on a native ROCm or + AITER backend, whichever backend the run asked for. + """ + from vllm.platforms.rocm import RocmPlatform + from vllm.utils.import_utils import resolve_obj_by_qualname + from vllm.v1.attention.backends.utils import get_supported_kv_cache_layouts + + backend_enum = ( + getattr(AttentionBackendEnum, selected_backend) if selected_backend else None + ) + + def resolve(kv_cache_dtype): + path = RocmPlatform.get_attn_backend_cls( + selected_backend=backend_enum, + attn_selector_config=AttentionSelectorConfig( + head_size=128, + dtype=torch.float16, + kv_cache_dtype=kv_cache_dtype, + block_size=16, + ), + ) + return resolve_obj_by_qualname(path) + + boundary_backend = resolve("auto") + quantized_backend = resolve("turboquant_k8v4") + + assert quantized_backend is AttentionBackendEnum.TURBOQUANT.get_class() + # Raises when the intersection is empty, which is the startup failure. + assert get_supported_kv_cache_layouts([boundary_backend, quantized_backend]) + + def test_sparse_not_supported(mock_vllm_config): """Test that sparse MLA without use_mla flag raises an error.""" from vllm.platforms.rocm import RocmPlatform diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index e724ee2cdf..726673b29b 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -23,10 +23,14 @@ if TYPE_CHECKING: from vllm.config import VllmConfig from vllm.config.kernel import IrOpPriorityConfig from vllm.utils.argparse_utils import FlexibleArgumentParser + from vllm.v1.attention.backend import AttentionBackend from vllm.v1.attention.selector import AttentionSelectorConfig logger = init_logger(__name__) +_KV_CACHE_DTYPE_REASON = "kv_cache_dtype not supported" +_TURBOQUANT_LAYOUT_REASON = "no KV cache layout in common with TURBOQUANT" + try: from amdsmi import ( AmdSmiException, @@ -505,6 +509,53 @@ def _get_backend_priorities( return backends +def _uses_turboquant(vllm_config: "VllmConfig | None") -> bool: + """Whether the run's KV cache dtype is one of the turboquant_* presets.""" + cache_config = getattr(vllm_config, "cache_config", None) + return cache_config is not None and str(cache_config.cache_dtype).startswith( + "turboquant_" + ) + + +def _shares_layout_with_turboquant(backend_class: type["AttentionBackend"]) -> bool: + """Whether this backend reads a KV cache layout TURBOQUANT reads too. + + A turboquant_* run keeps its boundary layers at the native dtype (see + TurboQuantConfig.get_boundary_skip_layers), so those layers pick a backend of + their own while every other layer picks TURBOQUANT. One layout has to serve + the whole worker, and get_supported_kv_cache_layouts() turns an empty + intersection into a hard error at engine startup. + """ + layouts = backend_class.supported_kv_cache_layouts() + turboquant_layouts = ( + AttentionBackendEnum.TURBOQUANT.get_class().supported_kv_cache_layouts() + ) + if layouts is None or turboquant_layouts is None: + return True + return not set(layouts).isdisjoint(turboquant_layouts) + + +def _get_invalid_reasons( + backend_class: type["AttentionBackend"], + device_capability: DeviceCapability, + attn_selector_config: "AttentionSelectorConfig", + *, + is_turboquant_run: bool, +) -> list[str]: + """Why this backend cannot serve the layer, empty when it can.""" + invalid_reasons = backend_class.validate_configuration( + device_capability=device_capability, + **attn_selector_config._asdict(), + ) + if ( + not invalid_reasons + and is_turboquant_run + and not _shares_layout_with_turboquant(backend_class) + ): + invalid_reasons = [_TURBOQUANT_LAYOUT_REASON] + return invalid_reasons + + class RocmPlatform(Platform): _enum = PlatformEnum.ROCM device_name: str = "rocm" @@ -611,12 +662,14 @@ class RocmPlatform(Platform): # TODO: Make this explicit in the selector in a future PR. if is_encoder_decoder and AttentionBackendEnum.ROCM_ATTN in backend_priorities: backend_priorities.remove(AttentionBackendEnum.ROCM_ATTN) + is_turboquant_run = _uses_turboquant(vllm_config) for priority, backend in enumerate(backend_priorities): try: - backend_class = backend.get_class() - invalid_reasons_i = backend_class.validate_configuration( - device_capability=device_capability, - **attn_selector_config._asdict(), + invalid_reasons_i = _get_invalid_reasons( + backend.get_class(), + device_capability, + attn_selector_config, + is_turboquant_run=is_turboquant_run, ) except ImportError: invalid_reasons_i = ["ImportError"] @@ -639,11 +692,16 @@ class RocmPlatform(Platform): # First try checking just the selected backend, if there is one. if selected_backend is not None: + # Keep lazy: vllm.config imports current_platform during initialization. + from vllm.config import get_current_vllm_config_or_none + + is_turboquant_run = _uses_turboquant(get_current_vllm_config_or_none()) try: - backend_class = selected_backend.get_class() - sel_invalid_reasons = backend_class.validate_configuration( - device_capability=device_capability, - **attn_selector_config._asdict(), + sel_invalid_reasons = _get_invalid_reasons( + selected_backend.get_class(), + device_capability, + attn_selector_config, + is_turboquant_run=is_turboquant_run, ) except ImportError: sel_invalid_reasons = ["ImportError"] @@ -653,21 +711,29 @@ class RocmPlatform(Platform): selected_backend.name, ) return selected_backend.get_path() - # Only tolerate the mismatch for turboquant_* KV-cache layers: - # boundary layers keep the native dtype (served by the selected - # backend) while turboquant_* layers need TURBOQUANT, so no single - # --attention-backend can serve every layer. For any other dtype - # the explicit selection is genuinely invalid -> fail loud. + # Only tolerate the mismatch when turboquant is in play: boundary + # layers keep the native dtype while every other layer needs + # TURBOQUANT, so no single --attention-backend can serve every layer. + # For any other dtype the selection is genuinely invalid -> fail loud. kv_dtype = attn_selector_config.kv_cache_dtype - if not (kv_dtype is not None and str(kv_dtype).startswith("turboquant")): + layer_is_turboquant = kv_dtype is not None and str(kv_dtype).startswith( + "turboquant" + ) + is_turboquant_fallback = ( + is_turboquant_run or layer_is_turboquant + ) and sel_invalid_reasons in ( + [_KV_CACHE_DTYPE_REASON], + [_TURBOQUANT_LAYOUT_REASON], + ) + if not is_turboquant_fallback: raise ValueError( f"Selected backend {selected_backend} is not valid for " f"this configuration. Reason: {sel_invalid_reasons}" ) # NOTE: pass a str (not the list) -- info_once hashes its args. logger.info_once( - "Selected backend %s is incompatible with this turboquant " - "layer (%s); using the auto-selected per-layer backend. " + "Selected backend %s is incompatible with this layer (%s) of " + "the turboquant run; using the auto-selected per-layer backend. " "Reason: %s", selected_backend.name, attn_selector_config.attn_type, diff --git a/vllm/v1/attention/selector.py b/vllm/v1/attention/selector.py index 695916d819..a7e8cd0fd2 100644 --- a/vllm/v1/attention/selector.py +++ b/vllm/v1/attention/selector.py @@ -194,6 +194,9 @@ def get_attn_backend( backend=backend, attn_selector_config=attn_selector_config, num_heads=num_heads, + _run_kv_cache_dtype=( + cache_config.cache_dtype if cache_config is not None else None + ), ) @@ -202,7 +205,11 @@ def _cached_get_attn_backend( backend, attn_selector_config: AttentionSelectorConfig, num_heads: int | None = None, + *, + _run_kv_cache_dtype: CacheDType | None = None, ) -> type[AttentionBackend]: + # Some platform selectors inspect the run-wide cache dtype. Keep it in the + # cache key even though they read the value from the current config. from vllm.platforms import current_platform attention_cls = current_platform.get_attn_backend_cls(