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

[CI][Bugfix] Fix FlashAttention reported MLA dimension support (#48631)

Browse Source 
Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>
This commit is contained in:
Matthew Bonanni
2026-07-14 21:33:02 +00:00
committed by GitHub
parent 05d4f8bba3
commit 313d01f507
5 changed files with 213 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/attention_backends.md
+1 -1
View File
@@ -205,7 +205,7 @@ hardware and configuration.
| Backend | Description | Dtypes | Compute Cap. | Notes |
| ------- | ----------- | ------ | ------------ | ----- |
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
| `FLASH_ATTN`‡ | FlashAttention varlen (FA2/FA3/FA4) | fp16, bf16 | Any | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) (FA2/FA3/FA4) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) (FA2/FA3 only) |
| `TRTLLM_RAGGED` | TensorRT-LLM ragged attention | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) or (qk_nope_head_dim=192, qk_rope_head_dim=64, v_head_dim=256) only |
| `FLASHINFER` | FlashInfer CUTLASS backend | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
| `TOKENSPEED_MLA` | | fp16, bf16 | 10.x | (qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) only |
tests/v1/attention/test_mla_prefill_selector.py
+34 -18
View File
@@ -123,55 +123,71 @@ class TestGetMLAPrefillBackend:
):
get_mla_prefill_backend(vllm_config)
def test_auto_selection_on_hopper(self):
@pytest.mark.parametrize(
("qk_nope_head_dim", "v_head_dim"),
[(128, 128), (192, 256)],
ids=["deepseek", "glm"],
)
def test_auto_selection_on_hopper(self, qk_nope_head_dim: int, v_head_dim: int):
try:
flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()
except ImportError:
pytest.skip("FLASH_ATTN backend not available")
return
vllm_config = _make_vllm_config()
vllm_config = _make_vllm_config(
_make_mock_model_config(
qk_nope_head_dim=qk_nope_head_dim,
v_head_dim=v_head_dim,
)
)
with patch("vllm.platforms.current_platform") as mock_platform:
with (
patch("vllm.platforms.current_platform") as mock_platform,
patch.object(flash_attn_cls, "is_available", return_value=True),
patch(
"vllm.v1.attention.backends.mla.prefill.flash_attn."
"get_flash_attn_version",
return_value=3,
),
):
mock_platform.get_device_capability.return_value = DeviceCapability(
major=9, minor=0
)
mock_platform.is_rocm.return_value = False
with patch.object(
flash_attn_cls,
"validate_configuration",
return_value=[],
):
backend = get_mla_prefill_backend(vllm_config)
assert backend.get_name() == "FLASH_ATTN"
backend = get_mla_prefill_backend(vllm_config)
assert backend.get_name() == "FLASH_ATTN"
class TestAutoSelectMLAPrefillBackend:
"""Tests for fallback and error paths in auto-selection."""
def test_blackwell_falls_back_to_trtllm(self):
def test_blackwell_glm_dimensions_fall_back_to_trtllm(self):
capability = DeviceCapability(major=10, minor=0)
selector_config = MLAPrefillSelectorConfig(
dtype=torch.bfloat16,
mla_dimensions=MLADimensions(
qk_nope_head_dim=128,
qk_nope_head_dim=192,
qk_rope_head_dim=64,
v_head_dim=128,
v_head_dim=256,
),
)
try:
flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()
trtllm_cls = MLAPrefillBackendEnum.TRTLLM_RAGGED.get_class()
except ImportError:
pytest.skip("TRTLLM_RAGGED backend not available")
pytest.skip("MLA prefill backend not available")
return
with (
patch("vllm.platforms.current_platform") as mock_platform,
patch.object(
MLAPrefillBackendEnum.FLASH_ATTN,
"get_class",
side_effect=ImportError("FLASH_ATTN not available"),
patch.object(flash_attn_cls, "is_available", return_value=True),
patch(
"vllm.v1.attention.backends.mla.prefill.flash_attn."
"get_flash_attn_version",
return_value=4,
),
patch.object(trtllm_cls, "validate_configuration", return_value=[]),
):
tools/pre_commit/generate_attention_backend_docs.py
+139
View File
@@ -446,6 +446,119 @@ def parse_supported_mla_dimensions(node: ast.AST | None) -> list[str]:
return supported_dimensions
def _parse_mla_dimension_reference(
node: ast.AST,
named_dimensions: dict[str, str],
) -> str | None:
if isinstance(node, ast.Name):
return named_dimensions.get(node.id)
return parse_mla_dimensions_call(node)
def _parse_mla_dimensions_return(
statements: list[ast.stmt],
named_dimensions: dict[str, str],
) -> list[str]:
for statement in statements:
if not isinstance(statement, ast.Return):
continue
value = statement.value
if not (
isinstance(value, ast.Compare)
and isinstance(value.left, ast.Name)
and value.left.id == "mla_dimensions"
and len(value.ops) == 1
and len(value.comparators) == 1
):
continue
comparator = value.comparators[0]
if isinstance(value.ops[0], ast.Eq):
dimension = _parse_mla_dimension_reference(
comparator,
named_dimensions,
)
return [dimension] if dimension is not None else []
if isinstance(value.ops[0], ast.In) and isinstance(
comparator, ast.List | ast.Tuple | ast.Set
):
dimensions = [
_parse_mla_dimension_reference(element, named_dimensions)
for element in comparator.elts
]
return [dimension for dimension in dimensions if dimension is not None]
return []
def _parse_fa_version_condition(node: ast.AST) -> int | None:
if not (
isinstance(node, ast.Compare)
and len(node.ops) == 1
and isinstance(node.ops[0], ast.Eq)
and len(node.comparators) == 1
and isinstance(node.comparators[0], ast.Constant)
and isinstance(node.comparators[0].value, int)
):
return None
left = node.left
is_fa_version = isinstance(left, ast.Name) and left.id == "fa_version"
is_get_fa_version = (
isinstance(left, ast.Call)
and isinstance(left.func, ast.Name)
and left.func.id == "get_flash_attn_version"
)
if is_fa_version or is_get_fa_version:
return node.comparators[0].value
return None
def parse_supports_mla_dimensions(
method: ast.FunctionDef | None,
) -> tuple[list[str], dict[int | None, list[str]]]:
"""Parse dimensions and FA-version branches from a support method."""
if method is None:
return [], {}
named_dimensions: dict[str, str] = {}
for statement in method.body:
if not (
isinstance(statement, ast.Assign)
and len(statement.targets) == 1
and isinstance(statement.targets[0], ast.Name)
):
continue
dimension = parse_mla_dimensions_call(statement.value)
if dimension is not None:
named_dimensions[statement.targets[0].id] = dimension
support_by_version: dict[int | None, list[str]] = {}
for statement in method.body:
if isinstance(statement, ast.Return):
dimensions = _parse_mla_dimensions_return([statement], named_dimensions)
if dimensions:
support_by_version[None] = dimensions
elif isinstance(statement, ast.If):
fa_version = _parse_fa_version_condition(statement.test)
if fa_version is None:
continue
dimensions = _parse_mla_dimensions_return(statement.body, named_dimensions)
if dimensions:
support_by_version[fa_version] = dimensions
default_dimensions = _parse_mla_dimensions_return(
statement.orelse, named_dimensions
)
if default_dimensions:
support_by_version[None] = default_dimensions
supported_dimensions: list[str] = []
for dimensions in support_by_version.values():
for dimension in dimensions:
if dimension not in supported_dimensions:
supported_dimensions.append(dimension)
return supported_dimensions, support_by_version
def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:
"""Parse a single MLA prefill backend file to extract its properties.
@@ -472,6 +585,7 @@ def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:
info: dict[str, Any] = {
"compute_capability": "Any",
"supported_mla_dimensions": [],
"mla_dimension_support_by_fa_version": {},
"dtypes": "fp16, bf16", # Default from base class
}
@@ -501,6 +615,13 @@ def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:
if dtypes:
info["dtypes"] = ", ".join(dtypes)
dimensions, support_by_version = parse_supports_mla_dimensions(
find_method(class_node, "supports_mla_dimensions")
)
if dimensions:
info["supported_mla_dimensions"] = dimensions
info["mla_dimension_support_by_fa_version"] = support_by_version
# Parse get_name static method
get_name_method = find_method(class_node, "get_name")
if get_name_method:
@@ -580,6 +701,24 @@ def parse_mla_prefill_backends() -> list[dict[str, Any]]:
supported_mla_dimensions = backend_info.get("supported_mla_dimensions", [])
if supported_mla_dimensions:
notes = " or ".join(supported_mla_dimensions) + " only"
support_by_version = backend_info.get(
"mla_dimension_support_by_fa_version", {}
)
if any(version is not None for version in support_by_version):
default_dimensions = support_by_version.get(None, [])
dimension_notes = []
for dimension in supported_mla_dimensions:
versions = [
f"FA{version}"
for version in (2, 3, 4)
if dimension
in support_by_version.get(version, default_dimensions)
]
version_limit = "" if len(versions) == 3 else " only"
dimension_notes.append(
f"{dimension} ({'/'.join(versions)}{version_limit})"
)
notes = " or ".join(dimension_notes)
elif backend_name == "FLASH_ATTN":
notes = "FA4 on SM100+, FA3 on SM90, FA2 otherwise"
vllm/v1/attention/backends/mla/prefill/base.py
+20 -8
View File
@@ -56,6 +56,13 @@ class MLAPrefillBackend(ABC):
def supports_dtype(cls, dtype: torch.dtype) -> bool:
return dtype in cls.supported_dtypes
@classmethod
def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool:
return (
not cls.supported_mla_dimensions
or mla_dimensions in cls.supported_mla_dimensions
)
@classmethod
def is_available(cls) -> bool:
return True
@@ -86,15 +93,20 @@ class MLAPrefillBackend(ABC):
if not cls.is_available():
invalid_reasons.append("required dependencies not available")
if (
cls.supported_mla_dimensions
and selector_config.mla_dimensions not in cls.supported_mla_dimensions
):
supported = ", ".join(str(dims) for dims in cls.supported_mla_dimensions)
invalid_reasons.append(
"Model does not have supported MLA dimensions "
f"(got {selector_config.mla_dimensions}; supported: {supported})"
mla_dimensions = selector_config.mla_dimensions
if not cls.supports_mla_dimensions(mla_dimensions):
reason = (
f"Model does not have supported MLA dimensions (got {mla_dimensions}"
)
if (
cls.supported_mla_dimensions
and mla_dimensions not in cls.supported_mla_dimensions
):
supported = ", ".join(
str(dims) for dims in cls.supported_mla_dimensions
)
reason += f"; supported: {supported}"
invalid_reasons.append(reason + ")")
return invalid_reasons
vllm/v1/attention/backends/mla/prefill/flash_attn.py
+19 -9
View File
@@ -3,7 +3,7 @@
"""FlashAttention backend for MLA prefill."""
import functools
from typing import TYPE_CHECKING, ClassVar
from typing import TYPE_CHECKING
import torch
@@ -43,14 +43,6 @@ else:
class FlashAttnPrefillBackend(MLAPrefillBackend):
"""FlashAttention backend for MLA prefill."""
supported_mla_dimensions: ClassVar[list[MLADimensions]] = [
MLADimensions(
qk_nope_head_dim=128,
qk_rope_head_dim=64,
v_head_dim=128,
),
]
@staticmethod
def get_name() -> str:
return "FLASH_ATTN"
@@ -59,6 +51,24 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):
def is_available(cls) -> bool:
return is_flash_attn_varlen_func_available()
@classmethod
def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool:
dims_deepseek = MLADimensions(
qk_nope_head_dim=128,
qk_rope_head_dim=64,
v_head_dim=128,
)
dims_glm = MLADimensions(
qk_nope_head_dim=192,
qk_rope_head_dim=64,
v_head_dim=256,
)
fa_version = get_flash_attn_version()
if fa_version == 4:
return mla_dimensions == dims_deepseek
else:
return mla_dimensions in [dims_deepseek, dims_glm]
def __init__(
self,
num_heads: int,