@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9d4dc4ca2f

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

[Kernel] Support GLM-5 dimensions for TRT-LLM ragged MLA prefill (#43525)

Signed-off-by: Mohammad Miadh Angkad <176301910+mmangkad@users.noreply.github.com> Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Matthew Bonanni <mbonanni@redhat.com>

Mohammad Miadh Angkad committed Jun 17, 2026 at 04:49 UTC 9d4dc4ca2fefae1a19648d7989f1dffb351e1b29

10 files changed +200 -77

docs/design/attention_backends.md

+3 -3

@@ -214,9 +214,9 @@ hardware and configuration.

215 | ------- | ----------- | ------ | ------------ | ----- |

220

221 > **‡** Automatic selection tries FlashAttention first. On Blackwell

222 > (SM100), the fallback order is TRT-LLM Ragged, FlashInfer, then

tests/v1/attention/test_mla_backends.py

+10 -2

@@ -765,7 +765,8 @@ def test_backend_correctness(

765 if not backends_to_test:

766 pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}")

767

768 - # Skip prefill backends that can't satisfy capability/deps/R1 constraints.

768 + # Skip prefill backends that can't satisfy capability/deps/dimension constraints.

769 + from vllm.v1.attention.backends.mla.prefill.base import MLADimensions

770 from vllm.v1.attention.backends.mla.prefill.selector import (

771 MLAPrefillSelectorConfig,

772 )

@@ -773,7 +774,14 @@ def test_backend_correctness(

774 try:

775 prefill_invalid_reasons = prefill_backend.get_class().validate_configuration(

776 current_platform.get_device_capability(),

776 - MLAPrefillSelectorConfig(dtype=torch.bfloat16, is_r1_compatible=True),

777 + MLAPrefillSelectorConfig(

778 + dtype=torch.bfloat16,

779 + mla_dimensions=MLADimensions(

780 + qk_nope_head_dim=128,

781 + qk_rope_head_dim=64,

782 + v_head_dim=128,

783 + ),

784 + ),

785 )

786 except ImportError:

787 prefill_invalid_reasons = ["ImportError"]

tests/v1/attention/test_mla_prefill_registry.py

-2

@@ -16,7 +16,6 @@ class CustomMLAPrefillBackend(MLAPrefillBackend):

16 """Mock custom MLA prefill backend for testing."""

18 supported_dtypes = [torch.bfloat16, torch.float16]

19 - requires_r1_mla_dimensions = False

20 @staticmethod

21 def get_name() -> str:

@@ -83,7 +82,6 @@ def test_register_custom_backend_as_decorator():

82 @register_mla_prefill_backend(MLAPrefillBackendEnum.CUSTOM)

83 class DecoratedPrefillBackend(MLAPrefillBackend):

84 supported_dtypes = [torch.bfloat16]

86 - requires_r1_mla_dimensions = False

86 @staticmethod

87 def get_name() -> str:

tests/v1/attention/test_mla_prefill_selector.py

+44 -25

@@ -9,12 +9,12 @@ import torch

10 from vllm.config import AttentionConfig, ModelConfig, VllmConfig

11 from vllm.platforms.interface import DeviceCapability

12 +from vllm.v1.attention.backends.mla.prefill.base import MLADimensions

13 from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

14 from vllm.v1.attention.backends.mla.prefill.selector import (

15 MLAPrefillSelectorConfig,

16 _auto_select_mla_prefill_backend,

17 get_mla_prefill_backend,

17 - is_deepseek_r1_mla_compatible,

18 )

@@ -149,11 +149,14 @@ class TestAutoSelectMLAPrefillBackend:

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

150

151 def test_blackwell_falls_back_to_trtllm(self):

152 - vllm_config = _make_vllm_config()

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

153 selector_config = MLAPrefillSelectorConfig(

154 dtype=torch.bfloat16,

156 - is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),

155 + mla_dimensions=MLADimensions(

156 + qk_nope_head_dim=128,

157 + qk_rope_head_dim=64,

158 + v_head_dim=128,

159 + ),

160 )

161

162 try:

@@ -177,11 +180,14 @@ class TestAutoSelectMLAPrefillBackend:

180 assert backend.get_name() == "TRTLLM_RAGGED"

181

182 def test_all_fail_raises_error(self):

180 - vllm_config = _make_vllm_config()

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

184 selector_config = MLAPrefillSelectorConfig(

185 dtype=torch.bfloat16,

184 - is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),

186 + mla_dimensions=MLADimensions(

187 + qk_nope_head_dim=128,

188 + qk_rope_head_dim=64,

189 + v_head_dim=128,

190 + ),

191 )

192

193 def mock_get_class(backend_enum): # noqa: ARG001

@@ -201,28 +207,26 @@ class TestAutoSelectMLAPrefillBackend:

207 class TestBackendValidation:

208 """Tests for backend validation logic."""

209

204 - def test_r1_dimension_requirement(self):

210 + def test_backend_supported_dimension_validation(self):

211 try:

212 from vllm.v1.attention.backends.mla.prefill.flashinfer import (

213 FlashInferPrefillBackend,

214 )

215 + from vllm.v1.attention.backends.mla.prefill.trtllm_ragged import (

216 + TrtllmRaggedPrefillBackend,

217 + )

218 except ImportError:

210 - pytest.skip("FlashInfer prefill backend not available")

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

220 return

221

213 - assert FlashInferPrefillBackend.requires_r1_mla_dimensions is True

214 -

215 - vllm_config = _make_vllm_config(

216 - model_config=_make_mock_model_config(

217 - qk_nope_head_dim=128,

218 - qk_rope_head_dim=64,

219 - v_head_dim=128,

220 - )

221 - )

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

223 selector_config = MLAPrefillSelectorConfig(

224 dtype=torch.bfloat16,

225 - is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),

225 + mla_dimensions=MLADimensions(

226 + qk_nope_head_dim=128,

227 + qk_rope_head_dim=64,

228 + v_head_dim=128,

229 + ),

230 )

231

232 with patch.object(FlashInferPrefillBackend, "is_available", return_value=True):

@@ -232,16 +236,13 @@ class TestBackendValidation:

236 )

237 assert len(invalid_reasons) == 0

238

235 - vllm_config_invalid = _make_vllm_config(

236 - model_config=_make_mock_model_config(

239 + selector_config_invalid = MLAPrefillSelectorConfig(

240 + dtype=torch.bfloat16,

241 + mla_dimensions=MLADimensions(

242 qk_nope_head_dim=64,

243 qk_rope_head_dim=64,

244 v_head_dim=128,

240 - )

241 - )

242 - selector_config_invalid = MLAPrefillSelectorConfig(

243 - dtype=torch.bfloat16,

244 - is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config_invalid),

245 + ),

246 )

247

248 with patch.object(FlashInferPrefillBackend, "is_available", return_value=True):

@@ -250,7 +251,25 @@ class TestBackendValidation:

251 selector_config_invalid,

252 )

253 assert len(invalid_reasons) == 1

253 - assert "DeepSeek R1 MLA dimensions" in invalid_reasons[0]

254 + assert "supported MLA dimensions" in invalid_reasons[0]

255 +

256 + selector_config_glm5 = MLAPrefillSelectorConfig(

257 + dtype=torch.bfloat16,

258 + mla_dimensions=MLADimensions(

259 + qk_nope_head_dim=192,

260 + qk_rope_head_dim=64,

261 + v_head_dim=256,

262 + ),

263 + )

264 +

265 + with patch.object(

266 + TrtllmRaggedPrefillBackend, "is_available", return_value=True

267 + ):

268 + invalid_reasons = TrtllmRaggedPrefillBackend.validate_configuration(

269 + capability,

270 + selector_config_glm5,

271 + )

272 + assert invalid_reasons == []

273

274

275 class TestMLAPrefillBackendParsing:

tools/pre_commit/generate_attention_backend_docs.py

+55 -11

@@ -384,6 +384,49 @@ def parse_mla_prefill_priorities() -> dict[str, list[str]]:

384 return priorities

385

386

387 +def parse_mla_dimensions_call(node: ast.AST) -> str | None:

388 + """Parse an MLADimensions(...) call into a compact display string."""

389 + if not isinstance(node, ast.Call):

390 + return None

391 +

392 + func = node.func

393 + if not isinstance(func, ast.Name) or func.id != "MLADimensions":

394 + return None

395 +

396 + dimensions: dict[str, int] = {}

397 + for keyword in node.keywords:

398 + if (

399 + keyword.arg is not None

400 + and isinstance(keyword.value, ast.Constant)

401 + and isinstance(keyword.value.value, int)

402 + ):

403 + dimensions[keyword.arg] = keyword.value.value

404 +

405 + qk_nope_head_dim = dimensions.get("qk_nope_head_dim")

406 + qk_rope_head_dim = dimensions.get("qk_rope_head_dim")

407 + v_head_dim = dimensions.get("v_head_dim")

408 + if qk_nope_head_dim is None or qk_rope_head_dim is None or v_head_dim is None:

409 + return None

410 +

411 + return (

412 + f"(qk_nope_head_dim={qk_nope_head_dim}, "

413 + f"qk_rope_head_dim={qk_rope_head_dim}, v_head_dim={v_head_dim})"

414 + )

415 +

416 +

417 +def parse_supported_mla_dimensions(node: ast.AST | None) -> list[str]:

418 + """Parse a supported_mla_dimensions class variable."""

419 + if not isinstance(node, ast.List):

420 + return []

421 +

422 + supported_dimensions = []

423 + for element in node.elts:

424 + dimensions = parse_mla_dimensions_call(element)

425 + if dimensions is not None:

426 + supported_dimensions.append(dimensions)

427 + return supported_dimensions

428 +

429 +

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

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

432

@@ -409,20 +452,20 @@ def parse_mla_prefill_backend_file(class_path: str) -> dict[str, Any] | None:

452

453 info: dict[str, Any] = {

454 "compute_capability": "Any",

412 - "requires_r1_dims": False,

455 + "supported_mla_dimensions": [],

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

457 }

458

459 # Parse class variables

460 for item in class_node.body:

418 - if isinstance(item, ast.Assign):

419 - for target in item.targets:

420 - if (

421 - isinstance(target, ast.Name)

422 - and target.id == "requires_r1_mla_dimensions"

423 - and isinstance(item.value, ast.Constant)

424 - ):

425 - info["requires_r1_dims"] = item.value.value

461 + if (

462 + isinstance(item, ast.AnnAssign)

463 + and isinstance(item.target, ast.Name)

464 + and item.target.id == "supported_mla_dimensions"

465 + ):

466 + info["supported_mla_dimensions"] = parse_supported_mla_dimensions(

467 + item.value

468 + )

469

470 # Parse supported_dtypes class variable

471 if (

@@ -515,8 +558,9 @@ def parse_mla_prefill_backends() -> list[dict[str, Any]]:

558 marker = "‡"

559

560 notes = ""

518 - if backend_info.get("requires_r1_dims"):

519 - notes = "DeepSeek R1 dims only"

561 + supported_mla_dimensions = backend_info.get("supported_mla_dimensions", [])

562 + if supported_mla_dimensions:

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

564 elif backend_name == "FLASH_ATTN":

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

566

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

+23 -4

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

3 """Abstract base class for MLA prefill backends."""

5 from abc import ABC, abstractmethod

6 +from dataclasses import dataclass

7 from typing import TYPE_CHECKING, ClassVar

9 import torch

@@ -19,6 +20,20 @@ if TYPE_CHECKING:

20 )

23 +@dataclass(frozen=True, kw_only=True)

24 +class MLADimensions:

25 + qk_nope_head_dim: int

26 + qk_rope_head_dim: int

27 + v_head_dim: int

28 +

29 + def __str__(self) -> str:

30 + return (

31 + f"(qk_nope_head_dim={self.qk_nope_head_dim}, "

32 + f"qk_rope_head_dim={self.qk_rope_head_dim}, "

33 + f"v_head_dim={self.v_head_dim})"

34 + )

35 +

36 +

37 class MLAPrefillBackend(ABC):

38 """Abstract base class for MLA prefill backends."""

@@ -26,7 +41,7 @@ class MLAPrefillBackend(ABC):

41 torch.float16,

42 torch.bfloat16,

43 ]

29 - requires_r1_mla_dimensions: ClassVar[bool] = False

44 + supported_mla_dimensions: ClassVar[list[MLADimensions]] = []

46 @staticmethod

47 @abstractmethod

@@ -71,10 +86,14 @@ class MLAPrefillBackend(ABC):

86 if not cls.is_available():

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

74 - if cls.requires_r1_mla_dimensions and not selector_config.is_r1_compatible:

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(

76 - "model does not have DeepSeek R1 MLA dimensions "

77 - "(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128)"

95 + "Model does not have supported MLA dimensions "

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

97 )

99 return invalid_reasons

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

+12 -3

@@ -2,12 +2,15 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """FlashInfer backend for MLA prefill."""

5 -from typing import TYPE_CHECKING

5 +from typing import TYPE_CHECKING, ClassVar

7 import torch

9 import vllm.envs as envs

10 -from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend

10 +from vllm.v1.attention.backends.mla.prefill.base import (

11 + MLADimensions,

12 + MLAPrefillBackend,

13 +)

14 from vllm.v1.attention.backends.utils import (

15 PerLayerParameters,

16 get_per_layer_parameters,

@@ -33,7 +36,13 @@ _DEFAULT_NUM_CHUNKS = 32

36 class FlashInferPrefillBackend(MLAPrefillBackend):

37 """FlashInfer backend for MLA prefill."""

36 - requires_r1_mla_dimensions = True

39 + supported_mla_dimensions: ClassVar[list[MLADimensions]] = [

40 + MLADimensions(

41 + qk_nope_head_dim=128,

42 + qk_rope_head_dim=64,

43 + v_head_dim=128,

44 + ),

45 + ]

47 @staticmethod

48 def get_name() -> str:

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

+24 -21

@@ -13,6 +13,7 @@ import torch

14 from vllm.logger import init_logger

15 from vllm.platforms.interface import DeviceCapability

16 +from vllm.v1.attention.backends.mla.prefill.base import MLADimensions

17 from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

19 if TYPE_CHECKING:

@@ -31,24 +32,17 @@ class MLAPrefillSelectorConfig(NamedTuple):

32 """

34 dtype: torch.dtype

34 - is_r1_compatible: bool

35 -

36 -

37 -def is_deepseek_r1_mla_compatible(vllm_config: "VllmConfig") -> bool:

38 - """Check if model has DeepSeek R1 compatible MLA dimensions.

35 + mla_dimensions: MLADimensions = MLADimensions(

36 + qk_nope_head_dim=0,

37 + qk_rope_head_dim=0,

38 + v_head_dim=0,

39 + )

40 - DeepSeek R1 MLA dimensions are:

41 - - qk_nope_head_dim = 128

42 - - qk_rope_head_dim = 64

43 - - v_head_dim = 128

44 - """

45 - if vllm_config.model_config is None:

46 - return False

47 - hf_text_config = vllm_config.model_config.hf_text_config

48 - qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)

49 - qk_rope_head_dim = getattr(hf_text_config, "qk_rope_head_dim", 1)

50 - v_head_dim = getattr(hf_text_config, "v_head_dim", 1)

51 - return qk_nope_head_dim == 128 and qk_rope_head_dim == 64 and v_head_dim == 128

41 + def __repr__(self):

42 + return (

43 + f"MLAPrefillSelectorConfig(dtype={self.dtype}, "

44 + f"mla_dimensions={self.mla_dimensions})"

45 + )

48 def _get_mla_prefill_backend_priorities(

@@ -101,10 +95,19 @@ def get_mla_prefill_backend(

96 attention_config = vllm_config.attention_config

104 - selector_config = MLAPrefillSelectorConfig(

105 - dtype=vllm_config.model_config.dtype,

106 - is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),

107 - )

98 + model_config = vllm_config.model_config

99 + if model_config is None:

100 + selector_config = MLAPrefillSelectorConfig(dtype=torch.get_default_dtype())

101 + else:

102 + hf_text_config = model_config.hf_text_config

103 + selector_config = MLAPrefillSelectorConfig(

104 + dtype=model_config.dtype,

105 + mla_dimensions=MLADimensions(

106 + qk_nope_head_dim=getattr(hf_text_config, "qk_nope_head_dim", 0),

107 + qk_rope_head_dim=getattr(hf_text_config, "qk_rope_head_dim", 0),

108 + v_head_dim=getattr(hf_text_config, "v_head_dim", 0),

109 + ),

110 + )

111

112 if attention_config.mla_prefill_backend is not None:

113 selected_backend = attention_config.mla_prefill_backend

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

+12 -3

@@ -2,11 +2,14 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """TokenSpeed CuTe DSL backend for MLA prefill."""

5 -from typing import TYPE_CHECKING

5 +from typing import TYPE_CHECKING, ClassVar

7 import torch

9 -from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend

9 +from vllm.v1.attention.backends.mla.prefill.base import (

10 + MLADimensions,

11 + MLAPrefillBackend,

12 +)

14 if TYPE_CHECKING:

15 from vllm.config import VllmConfig

@@ -19,7 +22,13 @@ if TYPE_CHECKING:

22 class TokenspeedMLAPrefillBackend(MLAPrefillBackend):

23 """TokenSpeed CuTe DSL backend for MLA prefill."""

22 - requires_r1_mla_dimensions = True

25 + supported_mla_dimensions: ClassVar[list[MLADimensions]] = [

26 + MLADimensions(

27 + qk_nope_head_dim=128,

28 + qk_rope_head_dim=64,

29 + v_head_dim=128,

30 + ),

31 + ]

33 @staticmethod

34 def get_name() -> str:

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

+17 -3

@@ -2,12 +2,15 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 """TRT-LLM Ragged backend for MLA prefill."""

5 -from typing import TYPE_CHECKING

5 +from typing import TYPE_CHECKING, ClassVar

7 import torch

9 import vllm.envs as envs

10 -from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend

10 +from vllm.v1.attention.backends.mla.prefill.base import (

11 + MLADimensions,

12 + MLAPrefillBackend,

13 +)

14 from vllm.v1.worker.workspace import current_workspace_manager

16 if TYPE_CHECKING:

@@ -21,7 +24,18 @@ if TYPE_CHECKING:

24 class TrtllmRaggedPrefillBackend(MLAPrefillBackend):

25 """TRT-LLM Ragged backend for MLA prefill."""

24 - requires_r1_mla_dimensions = True

27 + supported_mla_dimensions: ClassVar[list[MLADimensions]] = [

28 + MLADimensions(

29 + qk_nope_head_dim=128,

30 + qk_rope_head_dim=64,

31 + v_head_dim=128,

32 + ),

33 + MLADimensions(

34 + qk_nope_head_dim=192,

35 + qk_rope_head_dim=64,

36 + v_head_dim=256,

37 + ),

38 + ]

40 @staticmethod

41 def get_name() -> str: