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

[Platform] Support platform-specific MLA prefill backend selection (#56959)

Browse Source 
Signed-off-by: wangxiyuan <wangxiyuan1007@gmail.com>
Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
This commit is contained in:
wangxiyuan
2026-10-08 15:25:21 +00:00
committed by GitHub
co-authored by mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
parent fd928e5f3c
commit 072d49bb6e
11 changed files with 469 additions and 283 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/mkdocs/gen_files/generate_attention_backends.py
+56 -40
View File
@@ -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
tests/v1/attention/test_mla_prefill_selector.py
+200 -111
View File
@@ -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"
vllm/platforms/cpu.py
+20
View File
@@ -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)
vllm/platforms/interface.py
+84
View File
@@ -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,
vllm/platforms/rocm.py
+55
View File
@@ -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."""
vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py
+4 -1
View File
@@ -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
vllm/v1/attention/backends/mla/prefill/base.py
+12 -6
View File
@@ -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")
vllm/v1/attention/backends/mla/prefill/flashinfer.py
+5 -2
View File
@@ -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:
vllm/v1/attention/backends/mla/prefill/selector.py
+23 -119
View File
@@ -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
vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py
+5 -2
View File
@@ -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. "
vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py
+5 -2
View File
@@ -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: