From 313d01f507fc86306960aadf7939d2b7376b3a90 Mon Sep 17 00:00:00 2001 From: Matthew Bonanni Date: Tue, 14 Jul 2026 17:33:02 -0400 Subject: [PATCH] [CI][Bugfix] Fix FlashAttention reported MLA dimension support (#48631) Signed-off-by: Matthew Bonanni --- docs/design/attention_backends.md | 2 +- .../v1/attention/test_mla_prefill_selector.py | 52 ++++--- .../generate_attention_backend_docs.py | 139 ++++++++++++++++++ .../v1/attention/backends/mla/prefill/base.py | 28 +++- .../backends/mla/prefill/flash_attn.py | 28 ++-- 5 files changed, 213 insertions(+), 36 deletions(-) diff --git a/docs/design/attention_backends.md b/docs/design/attention_backends.md index c8d02fb640..9d524a489b 100644 --- a/docs/design/attention_backends.md +++ b/docs/design/attention_backends.md @@ -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 | diff --git a/tests/v1/attention/test_mla_prefill_selector.py b/tests/v1/attention/test_mla_prefill_selector.py index c893203246..8677197b7a 100644 --- a/tests/v1/attention/test_mla_prefill_selector.py +++ b/tests/v1/attention/test_mla_prefill_selector.py @@ -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=[]), ): diff --git a/tools/pre_commit/generate_attention_backend_docs.py b/tools/pre_commit/generate_attention_backend_docs.py index 6f6e1341ac..39d173f761 100644 --- a/tools/pre_commit/generate_attention_backend_docs.py +++ b/tools/pre_commit/generate_attention_backend_docs.py @@ -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" diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py index c56f3d46c7..dfcc037c8b 100644 --- a/vllm/v1/attention/backends/mla/prefill/base.py +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -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 diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 4e35072acd..afc4efa1ae 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -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,