iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[ROCm] Give turboquant boundary layers a layout-compatible backend (#54988)

Browse Source 
Signed-off-by: Stefan Koncarevic <stefan.koncarevic@amd.com>
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: Andreas Karatzas <akaratza@amd.com>
Co-authored-by: Codex <noreply@openai.com>
This commit is contained in:
stefankoncarevic
2026-09-23 22:00:07 -05:00
committed by GitHub
co-authored by Andreas Karatzas Codex
parent e6dc16cebd
commit e78e367c6e
3 changed files with 232 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/attention/test_rocm_attention_backends_selection.py
+143
View File
@@ -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
vllm/platforms/rocm.py
+82 -16
View File
@@ -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,
vllm/v1/attention/selector.py
+7
View File
@@ -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(