@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 313d01f507

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

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Matthew Bonanni committed Jul 14, 2026 at 17:33 UTC 313d01f507fc86306960aadf7939d2b7376b3a90

5 files changed +213 -36

docs/design/attention_backends.md

+1 -1

@@ -205,7 +205,7 @@ hardware and configuration.

205

207 | ------- | ----------- | ------ | ------------ | ----- |

tests/v1/attention/test_mla_prefill_selector.py

+34 -18

@@ -123,55 +123,71 @@ class TestGetMLAPrefillBackend:

123 ):

124 get_mla_prefill_backend(vllm_config)

125

126 - def test_auto_selection_on_hopper(self):

126 + @pytest.mark.parametrize(

127 + ("qk_nope_head_dim", "v_head_dim"),

128 + [(128, 128), (192, 256)],

129 + ids=["deepseek", "glm"],

130 + )

131 + def test_auto_selection_on_hopper(self, qk_nope_head_dim: int, v_head_dim: int):

132 try:

133 flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()

134 except ImportError:

135 pytest.skip("FLASH_ATTN backend not available")

136 return

137

133 - vllm_config = _make_vllm_config()

138 + vllm_config = _make_vllm_config(

139 + _make_mock_model_config(

140 + qk_nope_head_dim=qk_nope_head_dim,

141 + v_head_dim=v_head_dim,

142 + )

143 + )

144

135 - with patch("vllm.platforms.current_platform") as mock_platform:

145 + with (

146 + patch("vllm.platforms.current_platform") as mock_platform,

147 + patch.object(flash_attn_cls, "is_available", return_value=True),

148 + patch(

149 + "vllm.v1.attention.backends.mla.prefill.flash_attn."

150 + "get_flash_attn_version",

151 + return_value=3,

152 + ),

153 + ):

154 mock_platform.get_device_capability.return_value = DeviceCapability(

155 major=9, minor=0

156 )

157 + mock_platform.is_rocm.return_value = False

158

140 - with patch.object(

141 - flash_attn_cls,

142 - "validate_configuration",

143 - return_value=[],

144 - ):

145 - backend = get_mla_prefill_backend(vllm_config)

146 - assert backend.get_name() == "FLASH_ATTN"

159 + backend = get_mla_prefill_backend(vllm_config)

160 + assert backend.get_name() == "FLASH_ATTN"

161

162

163 class TestAutoSelectMLAPrefillBackend:

164 """Tests for fallback and error paths in auto-selection."""

165

152 - def test_blackwell_falls_back_to_trtllm(self):

166 + def test_blackwell_glm_dimensions_fall_back_to_trtllm(self):

167 capability = DeviceCapability(major=10, minor=0)

168 selector_config = MLAPrefillSelectorConfig(

169 dtype=torch.bfloat16,

170 mla_dimensions=MLADimensions(

157 - qk_nope_head_dim=128,

171 + qk_nope_head_dim=192,

172 qk_rope_head_dim=64,

159 - v_head_dim=128,

173 + v_head_dim=256,

174 ),

175 )

176

177 try:

178 + flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()

179 trtllm_cls = MLAPrefillBackendEnum.TRTLLM_RAGGED.get_class()

180 except ImportError:

166 - pytest.skip("TRTLLM_RAGGED backend not available")

181 + pytest.skip("MLA prefill backend not available")

182 return

183

184 with (

185 patch("vllm.platforms.current_platform") as mock_platform,

171 - patch.object(

172 - MLAPrefillBackendEnum.FLASH_ATTN,

173 - "get_class",

174 - side_effect=ImportError("FLASH_ATTN not available"),

186 + patch.object(flash_attn_cls, "is_available", return_value=True),

187 + patch(

188 + "vllm.v1.attention.backends.mla.prefill.flash_attn."

189 + "get_flash_attn_version",

190 + return_value=4,

191 ),

192 patch.object(trtllm_cls, "validate_configuration", return_value=[]),

193 ):

tools/pre_commit/generate_attention_backend_docs.py

+139

@@ -446,6 +446,119 @@ def parse_supported_mla_dimensions(node: ast.AST | None) -> list[str]:

446 return supported_dimensions

447

448

449 +def _parse_mla_dimension_reference(

450 + node: ast.AST,

451 + named_dimensions: dict[str, str],

452 +) -> str | None:

453 + if isinstance(node, ast.Name):

454 + return named_dimensions.get(node.id)

455 + return parse_mla_dimensions_call(node)

456 +

457 +

458 +def _parse_mla_dimensions_return(

459 + statements: list[ast.stmt],

460 + named_dimensions: dict[str, str],

461 +) -> list[str]:

462 + for statement in statements:

463 + if not isinstance(statement, ast.Return):

464 + continue

465 + value = statement.value

466 + if not (

467 + isinstance(value, ast.Compare)

468 + and isinstance(value.left, ast.Name)

469 + and value.left.id == "mla_dimensions"

470 + and len(value.ops) == 1

471 + and len(value.comparators) == 1

472 + ):

473 + continue

474 +

475 + comparator = value.comparators[0]

476 + if isinstance(value.ops[0], ast.Eq):

477 + dimension = _parse_mla_dimension_reference(

478 + comparator,

479 + named_dimensions,

480 + )

481 + return [dimension] if dimension is not None else []

482 + if isinstance(value.ops[0], ast.In) and isinstance(

483 + comparator, ast.List | ast.Tuple | ast.Set

484 + ):

485 + dimensions = [

486 + _parse_mla_dimension_reference(element, named_dimensions)

487 + for element in comparator.elts

488 + ]

489 + return [dimension for dimension in dimensions if dimension is not None]

490 + return []

491 +

492 +

493 +def _parse_fa_version_condition(node: ast.AST) -> int | None:

494 + if not (

495 + isinstance(node, ast.Compare)

496 + and len(node.ops) == 1

497 + and isinstance(node.ops[0], ast.Eq)

498 + and len(node.comparators) == 1

499 + and isinstance(node.comparators[0], ast.Constant)

500 + and isinstance(node.comparators[0].value, int)

501 + ):

502 + return None

503 +

504 + left = node.left

505 + is_fa_version = isinstance(left, ast.Name) and left.id == "fa_version"

506 + is_get_fa_version = (

507 + isinstance(left, ast.Call)

508 + and isinstance(left.func, ast.Name)

509 + and left.func.id == "get_flash_attn_version"

510 + )

511 + if is_fa_version or is_get_fa_version:

512 + return node.comparators[0].value

513 + return None

514 +

515 +

516 +def parse_supports_mla_dimensions(

517 + method: ast.FunctionDef | None,

518 +) -> tuple[list[str], dict[int | None, list[str]]]:

519 + """Parse dimensions and FA-version branches from a support method."""

520 + if method is None:

521 + return [], {}

522 +

523 + named_dimensions: dict[str, str] = {}

524 + for statement in method.body:

525 + if not (

526 + isinstance(statement, ast.Assign)

527 + and len(statement.targets) == 1

528 + and isinstance(statement.targets[0], ast.Name)

529 + ):

530 + continue

531 + dimension = parse_mla_dimensions_call(statement.value)

532 + if dimension is not None:

533 + named_dimensions[statement.targets[0].id] = dimension

534 +

535 + support_by_version: dict[int | None, list[str]] = {}

536 + for statement in method.body:

537 + if isinstance(statement, ast.Return):

538 + dimensions = _parse_mla_dimensions_return([statement], named_dimensions)

539 + if dimensions:

540 + support_by_version[None] = dimensions

541 + elif isinstance(statement, ast.If):

542 + fa_version = _parse_fa_version_condition(statement.test)

543 + if fa_version is None:

544 + continue

545 + dimensions = _parse_mla_dimensions_return(statement.body, named_dimensions)

546 + if dimensions:

547 + support_by_version[fa_version] = dimensions

548 + default_dimensions = _parse_mla_dimensions_return(

549 + statement.orelse, named_dimensions

550 + )

551 + if default_dimensions:

552 + support_by_version[None] = default_dimensions

553 +

554 + supported_dimensions: list[str] = []

555 + for dimensions in support_by_version.values():

556 + for dimension in dimensions:

557 + if dimension not in supported_dimensions:

558 + supported_dimensions.append(dimension)

559 + return supported_dimensions, support_by_version

560 +

561 +

562 def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:

563 """Parse a single MLA prefill backend file to extract its properties.

564

@@ -472,6 +585,7 @@ def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:

585 info: dict[str, Any] = {

586 "compute_capability": "Any",

587 "supported_mla_dimensions": [],

588 + "mla_dimension_support_by_fa_version": {},

589 "dtypes": "fp16, bf16", # Default from base class

590 }

591

@@ -501,6 +615,13 @@ def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:

615 if dtypes:

616 info["dtypes"] = ", ".join(dtypes)

617

618 + dimensions, support_by_version = parse_supports_mla_dimensions(

619 + find_method(class_node, "supports_mla_dimensions")

620 + )

621 + if dimensions:

622 + info["supported_mla_dimensions"] = dimensions

623 + info["mla_dimension_support_by_fa_version"] = support_by_version

624 +

625 # Parse get_name static method

626 get_name_method = find_method(class_node, "get_name")

627 if get_name_method:

@@ -580,6 +701,24 @@ def parse_mla_prefill_backends() -> list[dict[str, Any]]:

701 supported_mla_dimensions = backend_info.get("supported_mla_dimensions", [])

702 if supported_mla_dimensions:

703 notes = " or ".join(supported_mla_dimensions) + " only"

704 + support_by_version = backend_info.get(

705 + "mla_dimension_support_by_fa_version", {}

706 + )

707 + if any(version is not None for version in support_by_version):

708 + default_dimensions = support_by_version.get(None, [])

709 + dimension_notes = []

710 + for dimension in supported_mla_dimensions:

711 + versions = [

712 + f"FA{version}"

713 + for version in (2, 3, 4)

714 + if dimension

715 + in support_by_version.get(version, default_dimensions)

716 + ]

717 + version_limit = "" if len(versions) == 3 else " only"

718 + dimension_notes.append(

719 + f"{dimension} ({'/'.join(versions)}{version_limit})"

720 + )

721 + notes = " or ".join(dimension_notes)

722 elif backend_name == "FLASH_ATTN":

723 notes = "FA4 on SM100+, FA3 on SM90, FA2 otherwise"

724

vllm/v1/attention/backends/mla/prefill/base.py

+20 -8

@@ -56,6 +56,13 @@ class MLAPrefillBackend(ABC):

56 def supports_dtype(cls, dtype: torch.dtype) -> bool:

57 return dtype in cls.supported_dtypes

59 + @classmethod

60 + def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool:

61 + return (

62 + not cls.supported_mla_dimensions

63 + or mla_dimensions in cls.supported_mla_dimensions

64 + )

65 +

66 @classmethod

67 def is_available(cls) -> bool:

68 return True

@@ -86,15 +93,20 @@ class MLAPrefillBackend(ABC):

93 if not cls.is_available():

94 invalid_reasons.append("required dependencies not available")

89 - if (

90 - cls.supported_mla_dimensions

91 - and selector_config.mla_dimensions not in cls.supported_mla_dimensions

92 - ):

93 - supported = ", ".join(str(dims) for dims in cls.supported_mla_dimensions)

94 - invalid_reasons.append(

95 - "Model does not have supported MLA dimensions "

96 - f"(got {selector_config.mla_dimensions}; supported: {supported})"

96 + mla_dimensions = selector_config.mla_dimensions

97 + if not cls.supports_mla_dimensions(mla_dimensions):

98 + reason = (

99 + f"Model does not have supported MLA dimensions (got {mla_dimensions}"

100 )

101 + if (

102 + cls.supported_mla_dimensions

103 + and mla_dimensions not in cls.supported_mla_dimensions

104 + ):

105 + supported = ", ".join(

106 + str(dims) for dims in cls.supported_mla_dimensions

107 + )

108 + reason += f"; supported: {supported}"

109 + invalid_reasons.append(reason + ")")

110

111 return invalid_reasons

112

vllm/v1/attention/backends/mla/prefill/flash_attn.py

+19 -9

@@ -3,7 +3,7 @@

3 """FlashAttention backend for MLA prefill."""

5 import functools

6 -from typing import TYPE_CHECKING, ClassVar

6 +from typing import TYPE_CHECKING

8 import torch

@@ -43,14 +43,6 @@ else:

43 class FlashAttnPrefillBackend(MLAPrefillBackend):

44 """FlashAttention backend for MLA prefill."""

46 - supported_mla_dimensions: ClassVar[list[MLADimensions]] = [

47 - MLADimensions(

48 - qk_nope_head_dim=128,

49 - qk_rope_head_dim=64,

50 - v_head_dim=128,

51 - ),

52 - ]

53 -

46 @staticmethod

47 def get_name() -> str:

48 return "FLASH_ATTN"

@@ -59,6 +51,24 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):

51 def is_available(cls) -> bool:

52 return is_flash_attn_varlen_func_available()

54 + @classmethod

55 + def supports_mla_dimensions(cls, mla_dimensions: MLADimensions) -> bool:

56 + dims_deepseek = MLADimensions(

57 + qk_nope_head_dim=128,

58 + qk_rope_head_dim=64,

59 + v_head_dim=128,

60 + )

61 + dims_glm = MLADimensions(

62 + qk_nope_head_dim=192,

63 + qk_rope_head_dim=64,

64 + v_head_dim=256,

65 + )

66 + fa_version = get_flash_attn_version()

67 + if fa_version == 4:

68 + return mla_dimensions == dims_deepseek

69 + else:

70 + return mla_dimensions in [dims_deepseek, dims_glm]

71 +

72 def __init__(

73 self,

74 num_heads: int,