From 072d49bb6eb017a1e33bf3d016a979972438a90f Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Thu, 8 Oct 2026 23:25:21 +0800 Subject: [PATCH] [Platform] Support platform-specific MLA prefill backend selection (#56959) Signed-off-by: wangxiyuan Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../gen_files/generate_attention_backends.py | 96 +++--- .../v1/attention/test_mla_prefill_selector.py | 311 +++++++++++------- vllm/platforms/cpu.py | 20 ++ vllm/platforms/interface.py | 84 +++++ vllm/platforms/rocm.py | 55 ++++ .../backends/mla/prefill/aiter_flash_attn.py | 5 +- .../v1/attention/backends/mla/prefill/base.py | 18 +- .../backends/mla/prefill/flashinfer.py | 7 +- .../backends/mla/prefill/selector.py | 142 ++------ .../backends/mla/prefill/tokenspeed_mla.py | 7 +- .../backends/mla/prefill/trtllm_ragged.py | 7 +- 11 files changed, 469 insertions(+), 283 deletions(-) diff --git a/docs/mkdocs/gen_files/generate_attention_backends.py b/docs/mkdocs/gen_files/generate_attention_backends.py index 725d7072ca..135dcf3a35 100644 --- a/docs/mkdocs/gen_files/generate_attention_backends.py +++ b/docs/mkdocs/gen_files/generate_attention_backends.py @@ -51,7 +51,7 @@ BACKEND_KV_DTYPE_EXCLUDES: dict[str, set[str]] = { MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill" MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py" -MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py" +PLATFORM_INTERFACE_FILE = REPO_ROOT / "vllm" / "platforms" / "interface.py" # --------------------------------------------------------------------------- @@ -320,67 +320,83 @@ def parse_mla_prefill_registry() -> dict[str, str]: def parse_mla_prefill_priorities() -> dict[str, list[str]]: - """Parse MLA prefill backend priorities from selector.py. + """Parse MLA prefill backend priorities from platforms/interface.py. Returns: A dict with keys like 'blackwell' and 'default' containing lists of backend enum names in priority order. """ - if not MLA_PREFILL_SELECTOR_FILE.exists(): + if not PLATFORM_INTERFACE_FILE.exists(): return {} try: - tree = ast.parse(MLA_PREFILL_SELECTOR_FILE.read_text()) + tree = ast.parse(PLATFORM_INTERFACE_FILE.read_text()) except Exception: return {} + def extract_backends(stmt: ast.AST) -> list[str] | None: + """Read a backend enum list from a return/assignment statement.""" + if isinstance(stmt, ast.Return) or ( + isinstance(stmt, ast.Assign) + and len(stmt.targets) == 1 + and isinstance(stmt.targets[0], ast.Name) + ): + value = stmt.value + else: + return None + if not isinstance(value, ast.List): + return None + return [elt.attr for elt in value.elts if isinstance(elt, ast.Attribute)] + + def extract_last_list(nodes: list[ast.stmt], recurse: bool) -> list[str] | None: + """Return the last backend list found among the given statements.""" + result = None + for stmt in nodes: + backends = extract_backends(stmt) + if backends is not None: + result = backends + elif recurse and isinstance(stmt, ast.If): + nested = extract_last_list([*stmt.body, *stmt.orelse], True) + if nested is not None: + result = nested + return result + + def is_blackwell_check(stmt: ast.AST) -> bool: + return ( + isinstance(stmt, ast.If) + and isinstance(stmt.test, ast.Compare) + and isinstance(stmt.test.left, ast.Attribute) + and stmt.test.left.attr == "major" + and stmt.test.comparators + and isinstance(stmt.test.comparators[0], ast.Constant) + and stmt.test.comparators[0].value == 10 + ) + priorities: dict[str, list[str]] = {} for node in ast.walk(tree): - if not isinstance(node, ast.FunctionDef): + if not isinstance(node, ast.ClassDef) or node.name != "Platform": continue - if node.name != "_get_mla_prefill_backend_priorities": + + method = find_method(node, "get_mla_prefill_backend_cls") + if method is None: continue # Look for if statements checking device_capability.major - for stmt in ast.walk(node): + for stmt in ast.walk(method): if not isinstance(stmt, ast.If): continue - # Check if it's a capability.major == 10 check (Blackwell) - is_blackwell = ( - isinstance(stmt.test, ast.Compare) - and isinstance(stmt.test.left, ast.Attribute) - and stmt.test.left.attr == "major" - and stmt.test.comparators - and isinstance(stmt.test.comparators[0], ast.Constant) - and stmt.test.comparators[0].value == 10 - ) - - # Extract backends from return statements - for body_stmt in stmt.body: - if isinstance(body_stmt, ast.Return) and isinstance( - body_stmt.value, ast.List - ): - backends = [] - for elt in body_stmt.value.elts: - if isinstance(elt, ast.Attribute): - backends.append(elt.attr) - if is_blackwell: - priorities["blackwell"] = backends - else: - priorities["default"] = backends - - # Extract from else branch - for else_stmt in stmt.orelse: - if isinstance(else_stmt, ast.Return) and isinstance( - else_stmt.value, ast.List - ): - backends = [] - for elt in else_stmt.value.elts: - if isinstance(elt, ast.Attribute): - backends.append(elt.attr) + if is_blackwell_check(stmt): + # The Blackwell lists live in the nested if/else inside the + # major == 10 branch. + backends = extract_last_list(stmt.body, True) + if backends is not None: + priorities["blackwell"] = backends + else: + backends = extract_last_list([*stmt.body, *stmt.orelse], False) + if backends is not None: priorities["default"] = backends return priorities diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index d5b81695aa..3c6e1a9295 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Tests for MLA prefill backend selector.""" +from contextlib import suppress from unittest.mock import MagicMock, patch import pytest @@ -9,13 +10,12 @@ import torch from vllm.config import AttentionConfig, ModelConfig, VllmConfig from vllm.platforms import current_platform -from vllm.platforms.interface import DeviceCapability +from vllm.platforms.interface import DeviceCapability, Platform from vllm.v1.attention.backends.mla.prefill.base import MLADimensions from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.mla.prefill.selector import ( MLAPrefillSelectorConfig, _auto_select_mla_prefill_backend, - _get_mla_prefill_backend_priorities, get_mla_prefill_backend, ) @@ -55,51 +55,170 @@ def _make_vllm_config( return mock_vllm_config +class TestPlatformMLAPrefillHooks: + """Tests for the Platform.get_mla_prefill_backend_cls hook itself.""" + + def _get_tried_backends(self, capability, mla_dimensions): + """Run the default hook, returning the backend order tried.""" + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=mla_dimensions, + ) + tried: list[MLAPrefillBackendEnum] = [] + + def mock_get_class(backend_enum): + tried.append(backend_enum) + cls = MagicMock() + cls.validate_configuration.return_value = ["not available"] + return cls + + with ( + patch.object(Platform, "get_device_capability", return_value=capability), + patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class), + suppress(ValueError), + ): + Platform.get_mla_prefill_backend_cls(selector_config) + return tried + + def test_priorities_no_capability(self): + tried = self._get_tried_backends( + None, + MLADimensions(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), + ) + assert tried == [MLAPrefillBackendEnum.FLASH_ATTN] + + def test_priorities_blackwell_glm_dimensions(self): + tried = self._get_tried_backends( + DeviceCapability(major=10, minor=0), + MLADimensions(qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256), + ) + assert tried == [ + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] + + def test_priorities_blackwell_deepseek_dimensions(self): + tried = self._get_tried_backends( + DeviceCapability(major=10, minor=0), + MLADimensions(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), + ) + assert tried == [ + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] + + def test_priorities_hopper(self): + tried = self._get_tried_backends( + DeviceCapability(major=9, minor=0), + MLADimensions(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128), + ) + assert tried == [MLAPrefillBackendEnum.FLASH_ATTN] + + def test_skips_import_error_backends(self): + selector_config = MLAPrefillSelectorConfig( + dtype=torch.bfloat16, + mla_dimensions=MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ), + ) + valid = MagicMock() + valid.validate_configuration.return_value = [] + + def mock_get_class(backend_enum): + if backend_enum is MLAPrefillBackendEnum.FLASH_ATTN: + return valid + raise ImportError("not installed") + + with ( + patch.object( + Platform, + "get_device_capability", + return_value=DeviceCapability(major=10, minor=0), + ), + patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class), + ): + selected = Platform.get_mla_prefill_backend_cls(selector_config) + assert selected is valid + + def test_all_invalid_raises_error(self): + selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) + + invalid = MagicMock() + invalid.validate_configuration.return_value = ["not available"] + + with ( + patch.object( + Platform, + "get_device_capability", + return_value=DeviceCapability(major=10, minor=0), + ), + patch.object(MLAPrefillBackendEnum, "get_class", return_value=invalid), + pytest.raises(ValueError, match="No valid MLA"), + ): + Platform.get_mla_prefill_backend_cls(selector_config) + + def test_cpu_platform_returns_fixed_backend(self): + """CPU backends are returned via the fixed-backend hook.""" + try: + cpu_cls = MLAPrefillBackendEnum.CPU.get_class() + zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + except ImportError: + pytest.skip("CPU MLA prefill backend not available") + return + + from vllm.platforms.cpu import CpuPlatform + + expected = zen_cls if zen_cls.is_available() else cpu_cls + assert ( + CpuPlatform.get_mla_prefill_backend_cls( + MLAPrefillSelectorConfig(dtype=torch.bfloat16) + ) + is expected + ) + + class TestGetMLAPrefillBackend: """Tests for get_mla_prefill_backend (public API).""" def test_cpu_uses_sdpa_prefill(self): vllm_config = _make_vllm_config() - zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() - - # A Zen CPU with zentorch selects the Zen backend, so expect whichever - # one this host offers. - expected = ( - zen_cls if zen_cls.is_available() else MLAPrefillBackendEnum.CPU.get_class() - ) with patch("vllm.platforms.current_platform") as mock_platform: - mock_platform.is_cpu.return_value = True - - backend = get_mla_prefill_backend(vllm_config) - assert backend is expected - - def test_zen_cpu_prefers_zentorch_prefill(self): - vllm_config = _make_vllm_config() - zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() - - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object(zen_cls, "is_available", return_value=True), - ): - mock_platform.is_cpu.return_value = True - - backend = get_mla_prefill_backend(vllm_config) - assert backend is zen_cls - - def test_zen_cpu_falls_back_when_zentorch_unavailable(self): - vllm_config = _make_vllm_config() - zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() - - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object(zen_cls, "is_available", return_value=False), - ): - mock_platform.is_cpu.return_value = True + mock_platform.get_mla_prefill_backend_cls.return_value = ( + MLAPrefillBackendEnum.CPU.get_class() + ) backend = get_mla_prefill_backend(vllm_config) assert backend is MLAPrefillBackendEnum.CPU.get_class() + def test_zen_cpu_prefers_zentorch_prefill(self): + from vllm.platforms.cpu import CpuPlatform + + zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + + with patch.object(zen_cls, "is_available", return_value=True): + backend = CpuPlatform.get_mla_prefill_backend_cls( + MLAPrefillSelectorConfig(dtype=torch.bfloat16) + ) + assert backend is zen_cls + + def test_zen_cpu_falls_back_when_zentorch_unavailable(self): + from vllm.platforms.cpu import CpuPlatform + + zen_cls = MLAPrefillBackendEnum.ZEN_CPU.get_class() + + with patch.object(zen_cls, "is_available", return_value=False): + backend = CpuPlatform.get_mla_prefill_backend_cls( + MLAPrefillSelectorConfig(dtype=torch.bfloat16) + ) + assert backend is MLAPrefillBackendEnum.CPU.get_class() + def test_no_device_capability_returns_flash_attn(self): vllm_config = _make_vllm_config() @@ -108,16 +227,8 @@ class TestGetMLAPrefillBackend: def get_name(): return "FLASH_ATTN" - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object( - MLAPrefillBackendEnum.FLASH_ATTN, - "get_class", - return_value=FlashAttnBackend, - ), - ): - mock_platform.get_device_capability.return_value = None - mock_platform.is_cpu.return_value = False + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = FlashAttnBackend backend = get_mla_prefill_backend(vllm_config) assert backend.get_name() == "FLASH_ATTN" @@ -214,7 +325,7 @@ class TestGetMLAPrefillBackend: mock_platform.get_device_capability.return_value = DeviceCapability( major=9, minor=0 ) - mock_platform.is_rocm.return_value = False + mock_platform.get_mla_prefill_backend_cls.return_value = flash_attn_cls backend = get_mla_prefill_backend(vllm_config) assert backend.get_name() == "FLASH_ATTN" @@ -224,7 +335,6 @@ class TestAutoSelectMLAPrefillBackend: """Tests for fallback and error paths in auto-selection.""" def test_blackwell_glm_dimensions_use_trtllm(self): - capability = DeviceCapability(major=10, minor=0) selector_config = MLAPrefillSelectorConfig( dtype=torch.bfloat16, mla_dimensions=MLADimensions( @@ -235,7 +345,6 @@ class TestAutoSelectMLAPrefillBackend: ) try: - flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() trtllm_cls = MLAPrefillBackendEnum.TRTLLM_RAGGED.get_class() except ImportError: pytest.skip("MLA prefill backend not available") @@ -243,19 +352,13 @@ class TestAutoSelectMLAPrefillBackend: with ( patch("vllm.platforms.current_platform") as mock_platform, - patch.object(flash_attn_cls, "is_available", return_value=True), patch.object(trtllm_cls, "validate_configuration", return_value=[]), ): - # Force the non-ROCm priority on the Blackwell. - mock_platform.is_rocm.return_value = False - backend = _auto_select_mla_prefill_backend( - capability, - selector_config, - ) + mock_platform.get_mla_prefill_backend_cls.return_value = trtllm_cls + backend = _auto_select_mla_prefill_backend(selector_config) assert backend.get_name() == "TRTLLM_RAGGED" - def test_all_fail_raises_error(self): - capability = DeviceCapability(major=10, minor=0) + def test_platform_hook_raises_propagates(self): selector_config = MLAPrefillSelectorConfig( dtype=torch.bfloat16, mla_dimensions=MLADimensions( @@ -265,18 +368,14 @@ class TestAutoSelectMLAPrefillBackend: ), ) - def mock_get_class(backend_enum): # noqa: ARG001 - cls = MagicMock() - cls.validate_configuration.return_value = ["not available"] - return cls - - with patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class): - _auto_select_mla_prefill_backend.cache_clear() - with pytest.raises(ValueError, match="No valid MLA"): - _auto_select_mla_prefill_backend( - capability, - selector_config, - ) + with ( + patch("vllm.platforms.current_platform") as mock_platform, + pytest.raises(ValueError, match="No valid MLA"), + ): + mock_platform.get_mla_prefill_backend_cls.side_effect = ValueError( + "No valid MLA prefill backend found" + ) + _auto_select_mla_prefill_backend(selector_config) class TestBackendValidation: @@ -391,18 +490,29 @@ class TestROCmAiterFAPrefillSelection: def test_rocm_priorities_prefer_aiter_fa(self): """On ROCm, ROCM_AITER_FA is tried first, FLASH_ATTN as fallback.""" - with patch("vllm.platforms.current_platform") as mock_platform: - mock_platform.is_rocm.return_value = True - priorities = _get_mla_prefill_backend_priorities( - DeviceCapability(major=9, minor=5), - MLADimensions( - qk_nope_head_dim=128, - qk_rope_head_dim=64, - v_head_dim=128, - ), - ) + from vllm.platforms.rocm import RocmPlatform - assert priorities == [ + selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) + tried: list[MLAPrefillBackendEnum] = [] + + def mock_get_class(backend_enum): + tried.append(backend_enum) + cls = MagicMock() + cls.validate_configuration.return_value = ["not available"] + return cls + + with ( + patch.object( + RocmPlatform, + "get_device_capability", + return_value=MagicMock(), + ), + patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class), + suppress(ValueError), + ): + RocmPlatform.get_mla_prefill_backend_cls(selector_config) + + assert tried == [ MLAPrefillBackendEnum.ROCM_AITER_FA, MLAPrefillBackendEnum.FLASH_ATTN, ] @@ -459,50 +569,29 @@ class TestROCmAiterFAPrefillSelection: AiterFlashAttnPrefillBackend, ) - # gfx gating is simulated via the mocked validate_configuration, - # not the capability. - capability = MagicMock() selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object( - AiterFlashAttnPrefillBackend, - "validate_configuration", - return_value=[], - ), - ): - mock_platform.is_rocm.return_value = True - backend = _auto_select_mla_prefill_backend(capability, selector_config) + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = ( + AiterFlashAttnPrefillBackend + ) + backend = _auto_select_mla_prefill_backend(selector_config) assert backend.get_name() == "ROCM_AITER_FA" def test_auto_select_falls_back_to_flash_attn_when_aiter_invalid(self): - from vllm.v1.attention.backends.mla.prefill.aiter_flash_attn import ( - AiterFlashAttnPrefillBackend, - ) - try: flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class() except ImportError: pytest.skip("FLASH_ATTN backend not available") return - # the fallback is forced by the mocked validate_configuration, + # the fallback is forced by the mocked get_mla_prefill_backend_cls, # not the capability. - capability = MagicMock() selector_config = MLAPrefillSelectorConfig(dtype=torch.bfloat16) - with ( - patch("vllm.platforms.current_platform") as mock_platform, - patch.object( - AiterFlashAttnPrefillBackend, - "validate_configuration", - return_value=["compute capability not supported"], - ), - patch.object(flash_attn_cls, "validate_configuration", return_value=[]), - ): - mock_platform.is_rocm.return_value = True - backend = _auto_select_mla_prefill_backend(capability, selector_config) + with patch("vllm.platforms.current_platform") as mock_platform: + mock_platform.get_mla_prefill_backend_cls.return_value = flash_attn_cls + backend = _auto_select_mla_prefill_backend(selector_config) assert backend.get_name() == "FLASH_ATTN" diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 25d28daf93..e0257acb80 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -18,6 +18,7 @@ from vllm.utils.cpu_resource_utils import ( get_visible_memory_node, ) from vllm.utils.mem_constants import GiB_bytes +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum from vllm.v1.attention.backends.registry import AttentionBackendEnum from .interface import CpuArchEnum, Platform, PlatformEnum @@ -26,6 +27,10 @@ logger = init_logger(__name__) if TYPE_CHECKING: from vllm.config import VllmConfig + from vllm.v1.attention.backends.mla.prefill.base import ( + MLAPrefillBackend, + ) + from vllm.v1.attention.backends.mla.prefill.selector import MLAPrefillSelectorConfig from vllm.v1.attention.selector import AttentionSelectorConfig else: VllmConfig = None @@ -173,6 +178,21 @@ class CpuPlatform(Platform): logger.info("Cannot use %s backend on CPU.", selected_backend) return AttentionBackendEnum.CPU_ATTN.get_path() + @classmethod + def get_mla_prefill_backend_cls( + cls, + mla_selector_config: "MLAPrefillSelectorConfig", + ) -> "type[MLAPrefillBackend]": + """Get the MLA prefill backend class of a device.""" + for backend_enum in (MLAPrefillBackendEnum.ZEN_CPU, MLAPrefillBackendEnum.CPU): + try: + cpu_backend_cls = backend_enum.get_class() + except ImportError: + continue + if cpu_backend_cls.is_available(): + return cpu_backend_cls + raise ValueError("No valid CPU MLA prefill backend found.") + @classmethod def get_device_total_memory(cls, device_id: int = 0) -> int: meminfo = get_memory_node_info(device_id) diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index f4583daba2..dc3200759c 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -26,6 +26,8 @@ if TYPE_CHECKING: from vllm.sampling_params import SamplingParams from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.attention.backend import AttentionBackend + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + from vllm.v1.attention.backends.mla.prefill.selector import MLAPrefillSelectorConfig from vllm.v1.attention.selector import AttentionSelectorConfig else: FlexibleArgumentParser = object @@ -429,6 +431,88 @@ class Platform: ) return AttentionBackendEnum.TORCH_SDPA + @classmethod + def get_mla_prefill_backend_cls( + cls, + mla_selector_config: "MLAPrefillSelectorConfig", + ) -> "type[MLAPrefillBackend]": + """Get the MLA prefill backend class of a device. + + Platforms that have a single, fixed MLA prefill backend can override + this to return it directly. The default implementation selects the + best backend from a device-capability-based priority list. + + Returns: + The selected backend class. + + Raises: + ValueError: If no MLA prefill backend is valid for the + configuration. + + """ + from vllm.v1.attention.backends.mla.prefill.base import MLADimensions + from vllm.v1.attention.backends.mla.prefill.registry import ( + MLAPrefillBackendEnum, + ) + + device_capability = cls.get_device_capability() + if device_capability is None: + return MLAPrefillBackendEnum.FLASH_ATTN.get_class() + + if device_capability.major == 10: # Blackwell + if mla_selector_config.mla_dimensions == MLADimensions( + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=256, + ): + priorities = [ + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] + else: + priorities = [ + MLAPrefillBackendEnum.FLASH_ATTN, + MLAPrefillBackendEnum.TRTLLM_RAGGED, + MLAPrefillBackendEnum.FLASHINFER, + MLAPrefillBackendEnum.TOKENSPEED_MLA, + ] + else: # Hopper (SM90) and older + priorities = [MLAPrefillBackendEnum.FLASH_ATTN] + + all_invalid_reasons: dict[str, list[str]] = {} + + for backend_enum in priorities: + try: + backend_cls = backend_enum.get_class() + invalid_reasons = backend_cls.validate_configuration( + device_capability, mla_selector_config + ) + except ImportError: + invalid_reasons = ["ImportError"] + if not invalid_reasons: + return backend_cls + all_invalid_reasons[backend_enum.name] = invalid_reasons + + reasons_str = ( + "{" + + ", ".join( + f"{name}: [{', '.join(reasons)}]" + for name, reasons in all_invalid_reasons.items() + ) + + "}" + ) + logger.debug_once( + "Some MLA prefill backends are not valid with %s. Reasons: %s.", + repr(mla_selector_config), + reasons_str, + ) + raise ValueError( + f"No valid MLA prefill backend found with {mla_selector_config!r}. " + f"Reasons: {reasons_str}." + ) + @classmethod def get_device_capability( cls, diff --git a/vllm/platforms/rocm.py b/vllm/platforms/rocm.py index f394e56bf6..5840ffe3b7 100644 --- a/vllm/platforms/rocm.py +++ b/vllm/platforms/rocm.py @@ -24,6 +24,10 @@ if TYPE_CHECKING: from vllm.config.kernel import IrOpPriorityConfig from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm.v1.attention.backend import AttentionBackend + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend + from vllm.v1.attention.backends.mla.prefill.selector import ( + MLAPrefillSelectorConfig, + ) from vllm.v1.attention.selector import AttentionSelectorConfig logger = init_logger(__name__) @@ -904,6 +908,57 @@ class RocmPlatform(Platform): logger.info_once("Using Torch SDPA backend for ViT model.") return AttentionBackendEnum.TORCH_SDPA + @classmethod + def get_mla_prefill_backend_cls( + cls, + mla_selector_config: "MLAPrefillSelectorConfig", + ) -> "type[MLAPrefillBackend]": + """On ROCm, prefer the AITER FlashAttention backend, falling back to + FlashAttention. + + Raises: + ValueError: If neither backend is valid for the configuration. + + """ + from vllm.v1.attention.backends.mla.prefill.registry import ( + MLAPrefillBackendEnum, + ) + + device_capability = cls.get_device_capability() + all_invalid_reasons: dict[str, list[str]] = {} + for backend_enum in ( + MLAPrefillBackendEnum.ROCM_AITER_FA, + MLAPrefillBackendEnum.FLASH_ATTN, + ): + try: + backend_cls = backend_enum.get_class() + invalid_reasons = backend_cls.validate_configuration( + device_capability, mla_selector_config + ) + except ImportError: + invalid_reasons = ["ImportError"] + if not invalid_reasons: + return backend_cls + all_invalid_reasons[backend_enum.name] = invalid_reasons + + reasons_str = ( + "{" + + ", ".join( + f"{name}: [{', '.join(reasons)}]" + for name, reasons in all_invalid_reasons.items() + ) + + "}" + ) + logger.debug_once( + "Some MLA prefill backends are not valid with %s. Reasons: %s.", + repr(mla_selector_config), + reasons_str, + ) + raise ValueError( + f"No valid MLA prefill backend found with {mla_selector_config!r}. " + f"Reasons: {reasons_str}." + ) + @classmethod def set_device(cls, device: torch.device) -> None: """Set the device for the current platform.""" diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py index 1d03b9e773..8b0e337bb2 100644 --- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -31,7 +31,10 @@ class AiterFlashAttnPrefillBackend(MLAPrefillBackend): return "ROCM_AITER_FA" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: if not current_platform.is_rocm(): return False from vllm.platforms.rocm import on_mi3xx diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py index 1b589cb295..4b4361214f 100644 --- a/vllm/v1/attention/backends/mla/prefill/base.py +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -49,7 +49,10 @@ class MLAPrefillBackend(ABC): raise NotImplementedError @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: return True @classmethod @@ -90,16 +93,19 @@ class MLAPrefillBackend(ABC): @classmethod def validate_configuration( cls, - device_capability: "DeviceCapability", + device_capability: "DeviceCapability | None", selector_config: "MLAPrefillSelectorConfig", ) -> list[str]: invalid_reasons: list[str] = [] if not cls.supports_compute_capability(device_capability): - invalid_reasons.append( - f"compute capability {device_capability.major}." - f"{device_capability.minor} not supported" - ) + if device_capability is None: + invalid_reasons.append("device capability not available") + else: + invalid_reasons.append( + f"compute capability {device_capability.major}." + f"{device_capability.minor} not supported" + ) if not cls.supports_dtype(selector_config.dtype): invalid_reasons.append(f"dtype {selector_config.dtype} not supported") diff --git a/vllm/v1/attention/backends/mla/prefill/flashinfer.py b/vllm/v1/attention/backends/mla/prefill/flashinfer.py index 777dd33756..5cf47ea2ca 100644 --- a/vllm/v1/attention/backends/mla/prefill/flashinfer.py +++ b/vllm/v1/attention/backends/mla/prefill/flashinfer.py @@ -51,8 +51,11 @@ class FlashInferPrefillBackend(MLAPrefillBackend): return "FLASHINFER" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: - return device_capability.major == 10 + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: + return device_capability is not None and device_capability.major == 10 @classmethod def is_available(cls) -> bool: diff --git a/vllm/v1/attention/backends/mla/prefill/selector.py b/vllm/v1/attention/backends/mla/prefill/selector.py index 0cb29f51b4..f0f0015809 100644 --- a/vllm/v1/attention/backends/mla/prefill/selector.py +++ b/vllm/v1/attention/backends/mla/prefill/selector.py @@ -12,9 +12,7 @@ from typing import TYPE_CHECKING, NamedTuple import torch from vllm.logger import init_logger -from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.prefill.base import MLADimensions -from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum if TYPE_CHECKING: from vllm.config import VllmConfig @@ -45,54 +43,8 @@ class MLAPrefillSelectorConfig(NamedTuple): ) -def _get_mla_prefill_backend_priorities( - device_capability: DeviceCapability, - mla_dimensions: MLADimensions, -) -> list[MLAPrefillBackendEnum]: - """Get MLA prefill backend priorities based on device capability. - - Args: - device_capability: The device's compute capability. - mla_dimensions: The model's MLA head dimensions. - - Returns: - List of backends in priority order (highest priority first). - - """ - from vllm.platforms import current_platform - - if current_platform.is_rocm(): - return [ - MLAPrefillBackendEnum.ROCM_AITER_FA, - MLAPrefillBackendEnum.FLASH_ATTN, - ] - - if device_capability.major == 10: # Blackwell - if mla_dimensions == MLADimensions( - qk_nope_head_dim=192, - qk_rope_head_dim=64, - v_head_dim=256, - ): - return [ - MLAPrefillBackendEnum.TRTLLM_RAGGED, - MLAPrefillBackendEnum.FLASH_ATTN, - MLAPrefillBackendEnum.FLASHINFER, - MLAPrefillBackendEnum.TOKENSPEED_MLA, - ] - return [ - MLAPrefillBackendEnum.FLASH_ATTN, - MLAPrefillBackendEnum.TRTLLM_RAGGED, - MLAPrefillBackendEnum.FLASHINFER, - MLAPrefillBackendEnum.TOKENSPEED_MLA, - ] - else: # Hopper (SM90) and older - return [ - MLAPrefillBackendEnum.FLASH_ATTN, - ] - - def get_mla_prefill_backend( - vllm_config: "VllmConfig", + vllm_config: "VllmConfig | None", ) -> "type[MLAPrefillBackend]": """Select the MLA prefill backend based on configuration and device. @@ -101,7 +53,8 @@ def get_mla_prefill_backend( priority-based selection. Args: - vllm_config: The vLLM configuration. + vllm_config: The vLLM configuration. May be None before model + resolution, in which case defaults are used. Returns: The selected prefill backend class. @@ -109,30 +62,7 @@ def get_mla_prefill_backend( """ from vllm.platforms import current_platform - if current_platform.is_cpu(): - # CPUs have no compute capability, so the capability-driven priority - # path below does not apply. Prefer an accelerator-specific CPU backend - # when its kernels are present, else the generic SDPA one. - for backend_enum in (MLAPrefillBackendEnum.ZEN_CPU, MLAPrefillBackendEnum.CPU): - try: - cpu_backend_cls = backend_enum.get_class() - except ImportError: - continue - if cpu_backend_cls.is_available(): - logger.info_once("Using %s MLA prefill backend.", backend_enum.name) - return cpu_backend_cls - raise ValueError("No valid CPU MLA prefill backend found.") - - device_capability = current_platform.get_device_capability() - if device_capability is None: - logger.info_once( - "Device capability not available, using FlashAttention MLA prefill backend." - ) - return MLAPrefillBackendEnum.FLASH_ATTN.get_class() - - attention_config = vllm_config.attention_config - - model_config = vllm_config.model_config + model_config = vllm_config.model_config if vllm_config is not None else None if model_config is None: selector_config = MLAPrefillSelectorConfig(dtype=torch.get_default_dtype()) else: @@ -146,11 +76,21 @@ def get_mla_prefill_backend( ), ) - if attention_config.mla_prefill_backend is not None: + if current_platform.is_cpu(): + # CPUs have a single fixed MLA prefill backend, so explicit backend + # configuration does not apply; always auto-select via the platform. + return _auto_select_mla_prefill_backend(selector_config) + + attention_config = vllm_config.attention_config if vllm_config is not None else None + if ( + attention_config is not None + and attention_config.mla_prefill_backend is not None + ): selected_backend = attention_config.mla_prefill_backend backend_cls: type[MLAPrefillBackend] | None = None try: backend_cls = selected_backend.get_class() + device_capability = current_platform.get_device_capability() invalid_reasons = backend_cls.validate_configuration( device_capability, selector_config ) @@ -166,63 +106,27 @@ def get_mla_prefill_backend( logger.info_once("Using %s MLA prefill backend.", selected_backend.name) return backend_cls - return _auto_select_mla_prefill_backend( - device_capability, - selector_config, - ) + return _auto_select_mla_prefill_backend(selector_config) @cache def _auto_select_mla_prefill_backend( - device_capability: DeviceCapability, selector_config: MLAPrefillSelectorConfig, ) -> "type[MLAPrefillBackend]": """Auto-select the best available MLA prefill backend. Args: - device_capability: The device's compute capability. selector_config: Hashable configuration for backend selection. Returns: The selected prefill backend class. + Raises: + ValueError: If the platform has no valid MLA prefill backend. + """ - priorities = _get_mla_prefill_backend_priorities( - device_capability, - selector_config.mla_dimensions, - ) - all_invalid_reasons: dict[str, list[str]] = {} + from vllm.platforms import current_platform - for backend_enum in priorities: - backend_cls: type[MLAPrefillBackend] | None = None - try: - backend_cls = backend_enum.get_class() - invalid_reasons = backend_cls.validate_configuration( - device_capability, selector_config - ) - except ImportError: - invalid_reasons = ["ImportError"] - if not invalid_reasons: - assert backend_cls is not None - logger.info_once("Using %s MLA prefill backend.", backend_enum.name) - return backend_cls - all_invalid_reasons[backend_enum.name] = invalid_reasons - - reasons_str = ( - "{" - + ", ".join( - f"{name}: [{', '.join(reasons)}]" - for name, reasons in all_invalid_reasons.items() - ) - + "}" - ) - config_str = repr(selector_config) - logger.debug_once( - "Some MLA prefill backends are not valid with %s. Reasons: %s.", - config_str, - reasons_str, - ) - - raise ValueError( - f"No valid MLA prefill backend found with {config_str}. Reasons: {reasons_str}." - ) + selected_backend = current_platform.get_mla_prefill_backend_cls(selector_config) + logger.info_once("Using %s MLA prefill backend.", selected_backend.get_name()) + return selected_backend diff --git a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py index 6985b0a70c..d65721b81b 100644 --- a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py @@ -35,8 +35,11 @@ class TokenspeedMLAPrefillBackend(MLAPrefillBackend): return "TOKENSPEED_MLA" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: - return device_capability.major == 10 + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: + return device_capability is not None and device_capability.major == 10 _INSTALL_HINT = ( "tokenspeed_mla package is not installed. " diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py index 167738b07a..3cb498be2b 100644 --- a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -43,8 +43,11 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): return "TRTLLM_RAGGED" @classmethod - def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool: - return device_capability.major == 10 + def supports_compute_capability( + cls, + device_capability: "DeviceCapability | None", + ) -> bool: + return device_capability is not None and device_capability.major == 10 @classmethod def is_available(cls) -> bool: