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

Fork 0

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

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 4, 08:45 UTC

[Attention] Abstract the MLA prefill backends and eliminate cuDNN (#32623)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Michael Goin <mgoin64@gmail.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Matthew Bonanni committed May 1, 2026 at 13:36 UTC f3fef123504db07b3ac83ad4ef677915b53e8386

16 files changed +1629 -708

docs/design/attention_backends.md

+12 -8

@@ -192,21 +192,25 @@ MLA uses separate backends for prefill and decode phases.

192

193 ### Prefill Backends

194

195 -The prefill backend is selected at runtime based on hardware and

196 -configuration.

195 +To explicitly select a prefill backend, use

196 +`-ac.mla_prefill_backend=<BACKEND>` (e.g., `FLASH_ATTN`, `FLASHINFER`).

197 +Otherwise, the prefill backend is selected automatically at runtime based on

198 +hardware and configuration.

199

199 -| ------- | ----------- | ------------ | ------ | ------- | ----- |

201 +| ------- | ----------- | ------ | ------------ | ----- |

205

206 > **‡** TRT-LLM Ragged is the default on Blackwell (SM100).

207 > On other GPUs, FlashAttention is used as the default.

208

209 ### Decode Backends

210

211 +MLA decode backends are selected using the standard

212 +`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).

213 +

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

216 | `CUTLASS_MLA` | fp16, bf16 | `auto`, `float16`, `bfloat16`, `fp8`, `fp8_e4m3` | 128 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 10.x |

tests/engine/test_arg_utils.py

-3

@@ -333,8 +333,6 @@ def test_attention_config():

333 "true",

334 "--attention-config.flash_attn_max_num_splits_for_cuda_graph",

335 "16",

336 - "--attention-config.use_cudnn_prefill",

337 - "true",

336 "--attention-config.use_trtllm_ragged_deepseek_prefill",

337 "true",

338 "--attention-config.use_trtllm_attention",

@@ -352,7 +350,6 @@ def test_attention_config():

350 assert engine_args.attention_config.flash_attn_version == 3

351 assert engine_args.attention_config.use_prefill_decode_attention is True

352 assert engine_args.attention_config.flash_attn_max_num_splits_for_cuda_graph == 16

355 - assert engine_args.attention_config.use_cudnn_prefill is True

353 assert engine_args.attention_config.use_trtllm_ragged_deepseek_prefill is True

354 assert engine_args.attention_config.use_trtllm_attention is True

355 assert engine_args.attention_config.disable_flashinfer_prefill is True

tests/v1/attention/test_mla_backends.py

@@ -672,6 +672,7 @@ def run_attention_backend(

672 def test_backend_correctness(

673 default_vllm_config,

674 dist_init,

675 + workspace_init,

676 batch_spec_name: str,

677 model: str,

678 tensor_parallel_size: int,

tests/v1/attention/test_mla_prefill_selector.py new

+304

@@ -0,0 +1,304 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +"""Tests for MLA prefill backend selector."""

4 +

5 +from unittest.mock import MagicMock, patch

6 +

7 +import pytest

8 +import torch

9 +

10 +from vllm.config import AttentionConfig, ModelConfig, VllmConfig

11 +from vllm.platforms.interface import DeviceCapability

12 +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

13 +from vllm.v1.attention.backends.mla.prefill.selector import (

14 + MLAPrefillSelectorConfig,

15 + _auto_select_mla_prefill_backend,

16 + get_mla_prefill_backend,

17 + is_deepseek_r1_mla_compatible,

18 +)

19 +

20 +

21 +@pytest.fixture(autouse=True)

22 +def clear_cache():

23 + """Clear lru cache to ensure each test case runs without caching."""

24 + _auto_select_mla_prefill_backend.cache_clear()

25 +

26 +

27 +def _make_mock_model_config(

28 + qk_nope_head_dim: int = 128,

29 + qk_rope_head_dim: int = 64,

30 + v_head_dim: int = 128,

31 + dtype: torch.dtype = torch.bfloat16,

32 +) -> ModelConfig:

33 + mock_config = MagicMock(spec=ModelConfig)

34 + mock_config.dtype = dtype

35 + mock_config.hf_text_config = MagicMock()

36 + mock_config.hf_text_config.qk_nope_head_dim = qk_nope_head_dim

37 + mock_config.hf_text_config.qk_rope_head_dim = qk_rope_head_dim

38 + mock_config.hf_text_config.v_head_dim = v_head_dim

39 + return mock_config

40 +

41 +

42 +def _make_vllm_config(

43 + model_config: ModelConfig | None = None,

44 + mla_prefill_backend: MLAPrefillBackendEnum | None = None,

45 +) -> VllmConfig:

46 + if model_config is None:

47 + model_config = _make_mock_model_config()

48 +

49 + attention_config = AttentionConfig(mla_prefill_backend=mla_prefill_backend)

50 + mock_vllm_config = MagicMock(spec=VllmConfig)

51 + mock_vllm_config.model_config = model_config

52 + mock_vllm_config.attention_config = attention_config

53 + return mock_vllm_config

54 +

55 +

56 +class TestGetMLAPrefillBackend:

57 + """Tests for get_mla_prefill_backend (public API)."""

58 +

59 + def test_no_device_capability_returns_flash_attn(self):

60 + vllm_config = _make_vllm_config()

61 +

62 + with patch("vllm.platforms.current_platform") as mock_platform:

63 + mock_platform.get_device_capability.return_value = None

64 +

65 + backend = get_mla_prefill_backend(vllm_config)

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

67 +

68 + def test_explicit_flash_attn_selection(self):

69 + try:

70 + flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()

71 + except ImportError:

72 + pytest.skip("FLASH_ATTN backend not available")

73 + return

74 +

75 + vllm_config = _make_vllm_config(

76 + mla_prefill_backend=MLAPrefillBackendEnum.FLASH_ATTN,

77 + )

78 +

79 + with patch("vllm.platforms.current_platform") as mock_platform:

80 + mock_platform.get_device_capability.return_value = DeviceCapability(

81 + major=9, minor=0

82 + )

83 +

84 + with patch.object(

85 + flash_attn_cls,

86 + "validate_configuration",

87 + return_value=[],

88 + ):

89 + backend = get_mla_prefill_backend(vllm_config)

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

91 +

92 + def test_explicit_backend_invalid_raises_error(self):

93 + vllm_config = _make_vllm_config(

94 + mla_prefill_backend=MLAPrefillBackendEnum.FLASHINFER,

95 + )

96 +

97 + with patch("vllm.platforms.current_platform") as mock_platform:

98 + mock_platform.get_device_capability.return_value = DeviceCapability(

99 + major=9, minor=0

100 + )

101 +

102 + with pytest.raises(ValueError, match="is not valid"):

103 + get_mla_prefill_backend(vllm_config)

104 +

105 + def test_explicit_backend_import_error_raises(self):

106 + vllm_config = _make_vllm_config(

107 + mla_prefill_backend=MLAPrefillBackendEnum.TRTLLM_RAGGED,

108 + )

109 +

110 + with patch("vllm.platforms.current_platform") as mock_platform:

111 + mock_platform.get_device_capability.return_value = DeviceCapability(

112 + major=10, minor=0

113 + )

114 +

115 + with (

116 + patch.object(

117 + MLAPrefillBackendEnum.TRTLLM_RAGGED,

118 + "get_class",

119 + side_effect=ImportError("trtllm not installed"),

120 + ),

121 + pytest.raises(ValueError, match="is not valid"),

122 + ):

123 + get_mla_prefill_backend(vllm_config)

124 +

125 + def test_auto_selection_on_hopper(self):

126 + try:

127 + flash_attn_cls = MLAPrefillBackendEnum.FLASH_ATTN.get_class()

128 + except ImportError:

129 + pytest.skip("FLASH_ATTN backend not available")

130 + return

131 +

132 + vllm_config = _make_vllm_config()

133 +

134 + with patch("vllm.platforms.current_platform") as mock_platform:

135 + mock_platform.get_device_capability.return_value = DeviceCapability(

136 + major=9, minor=0

137 + )

138 +

139 + with patch.object(

140 + flash_attn_cls,

141 + "validate_configuration",

142 + return_value=[],

143 + ):

144 + backend = get_mla_prefill_backend(vllm_config)

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

146 +

147 +

148 +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()

153 + capability = DeviceCapability(major=10, minor=0)

154 + selector_config = MLAPrefillSelectorConfig(

155 + dtype=torch.bfloat16,

156 + is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),

157 + )

158 +

159 + try:

160 + trtllm_cls = MLAPrefillBackendEnum.TRTLLM_RAGGED.get_class()

161 + except ImportError:

162 + pytest.skip("TRTLLM_RAGGED backend not available")

163 + return

164 +

165 + with (

166 + patch.object(

167 + MLAPrefillBackendEnum.FLASH_ATTN,

168 + "get_class",

169 + side_effect=ImportError("FLASH_ATTN not available"),

170 + ),

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

172 + ):

173 + backend = _auto_select_mla_prefill_backend(

174 + capability,

175 + selector_config,

176 + )

177 + assert backend.get_name() == "TRTLLM_RAGGED"

178 +

179 + def test_all_fail_raises_error(self):

180 + vllm_config = _make_vllm_config()

181 + capability = DeviceCapability(major=10, minor=0)

182 + selector_config = MLAPrefillSelectorConfig(

183 + dtype=torch.bfloat16,

184 + is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),

185 + )

186 +

187 + def mock_get_class(backend_enum): # noqa: ARG001

188 + cls = MagicMock()

189 + cls.validate_configuration.return_value = ["not available"]

190 + return cls

191 +

192 + with patch.object(MLAPrefillBackendEnum, "get_class", mock_get_class):

193 + _auto_select_mla_prefill_backend.cache_clear()

194 + with pytest.raises(ValueError, match="No valid MLA"):

195 + _auto_select_mla_prefill_backend(

196 + capability,

197 + selector_config,

198 + )

199 +

200 +

201 +class TestBackendValidation:

202 + """Tests for backend validation logic."""

203 +

204 + def test_r1_dimension_requirement(self):

205 + try:

206 + from vllm.v1.attention.backends.mla.prefill.flashinfer import (

207 + FlashInferPrefillBackend,

208 + )

209 + except ImportError:

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

211 + return

212 +

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),

226 + )

227 +

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

229 + invalid_reasons = FlashInferPrefillBackend.validate_configuration(

230 + capability,

231 + selector_config,

232 + )

233 + assert len(invalid_reasons) == 0

234 +

235 + vllm_config_invalid = _make_vllm_config(

236 + model_config=_make_mock_model_config(

237 + qk_nope_head_dim=64,

238 + qk_rope_head_dim=64,

239 + 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 + with patch.object(FlashInferPrefillBackend, "is_available", return_value=True):

248 + invalid_reasons = FlashInferPrefillBackend.validate_configuration(

249 + capability,

250 + selector_config_invalid,

251 + )

252 + assert len(invalid_reasons) == 1

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

254 +

255 +

256 +class TestMLAPrefillBackendParsing:

257 + """Tests for string-based mla_prefill_backend parsing from CLI args."""

258 +

259 + def test_valid_string_parses_to_enum(self):

260 + config = AttentionConfig(

261 + mla_prefill_backend="FLASH_ATTN", # type: ignore[arg-type]

262 + )

263 + assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN

264 +

265 + def test_invalid_string_raises_error(self):

266 + with pytest.raises(ValueError, match="Unknown MLA prefill backend"):

267 + AttentionConfig(

268 + mla_prefill_backend="NONEXISTENT", # type: ignore[arg-type]

269 + )

270 +

271 +

272 +class TestDeprecatedFlagMigration:

273 + """Tests for _migrate_deprecated_mla_prefill_flags in AttentionConfig."""

274 +

275 + def test_no_deprecated_flags_leaves_backend_none(self):

276 + config = AttentionConfig()

277 + assert config.mla_prefill_backend is None

278 +

279 + def test_use_trtllm_ragged_migrates_to_trtllm_ragged(self):

280 + config = AttentionConfig(use_trtllm_ragged_deepseek_prefill=True)

281 + assert config.mla_prefill_backend == MLAPrefillBackendEnum.TRTLLM_RAGGED

282 +

283 + def test_disable_flashinfer_prefill_migrates_to_flash_attn(self):

284 + config = AttentionConfig(disable_flashinfer_prefill=True)

285 + assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN

286 +

287 + def test_explicit_backend_ignores_deprecated_flags(self):

288 + config = AttentionConfig(

289 + mla_prefill_backend=MLAPrefillBackendEnum.FLASH_ATTN,

290 + use_cudnn_prefill=True,

291 + )

292 + assert config.mla_prefill_backend == MLAPrefillBackendEnum.FLASH_ATTN

293 +

294 + def test_cudnn_raises_error(self):

295 + match = "cuDNN MLA prefill backend has been removed"

296 + with pytest.raises(ValueError, match=match):

297 + AttentionConfig(use_cudnn_prefill=True)

298 +

299 + def test_trtllm_takes_priority_over_disable_flashinfer(self):

300 + config = AttentionConfig(

301 + use_trtllm_ragged_deepseek_prefill=True,

302 + disable_flashinfer_prefill=True,

303 + )

304 + assert config.mla_prefill_backend == MLAPrefillBackendEnum.TRTLLM_RAGGED

tools/pre_commit/generate_attention_backend_docs.py

+253 -88

@@ -30,7 +30,6 @@ REPO_ROOT = Path(__file__).parent.parent.parent

30 RELEVANT_PATTERNS = [

31 "vllm/v1/attention/backends/*.py",

32 "vllm/v1/attention/backends/**/*.py",

33 - "vllm/v1/attention/backends/fa_utils.py",

33 "vllm/model_executor/layers/attention/mla_attention.py",

34 "vllm/platforms/cuda.py",

35 "tools/pre_commit/generate_attention_backend_docs.py",

@@ -68,6 +67,11 @@ def is_relevant_file(filepath: str) -> bool:

67 return any(fnmatch.fnmatch(path_str, pattern) for pattern in RELEVANT_PATTERNS)

70 +MLA_PREFILL_DIR = BACKENDS_DIR / "mla" / "prefill"

71 +MLA_PREFILL_REGISTRY_FILE = MLA_PREFILL_DIR / "registry.py"

72 +MLA_PREFILL_SELECTOR_FILE = MLA_PREFILL_DIR / "selector.py"

73 +

74 +

75 # ---------------------------------------------------------------------------

76 # AST utility helpers

77 # ---------------------------------------------------------------------------

@@ -293,6 +297,242 @@ def get_file_from_class_path(class_path: str) -> Path | None:

297 return py_file if py_file.exists() else None

298

299

300 +def parse_mla_prefill_registry() -> dict[str, str]:

301 + """Parse MLAPrefillBackendEnum from the prefill registry.

302 +

303 + Returns:

304 + A dict mapping backend names to their class paths.

305 + """

306 + if not MLA_PREFILL_REGISTRY_FILE.exists():

307 + return {}

308 +

309 + try:

310 + tree = ast.parse(MLA_PREFILL_REGISTRY_FILE.read_text())

311 + except Exception:

312 + return {}

313 +

314 + for node in ast.walk(tree):

315 + if isinstance(node, ast.ClassDef) and node.name == "MLAPrefillBackendEnum":

316 + return _extract_enum_values(node)

317 + return {}

318 +

319 +

320 +def parse_mla_prefill_priorities() -> dict[str, list[str]]:

321 + """Parse MLA prefill backend priorities from selector.py.

322 +

323 + Returns:

324 + A dict with keys like 'blackwell' and 'default' containing

325 + lists of backend enum names in priority order.

326 + """

327 + if not MLA_PREFILL_SELECTOR_FILE.exists():

328 + return {}

329 +

330 + try:

331 + tree = ast.parse(MLA_PREFILL_SELECTOR_FILE.read_text())

332 + except Exception:

333 + return {}

334 +

335 + priorities: dict[str, list[str]] = {}

336 +

337 + for node in ast.walk(tree):

338 + if not isinstance(node, ast.FunctionDef):

339 + continue

340 + if node.name != "_get_mla_prefill_backend_priorities":

341 + continue

342 +

343 + # Look for if statements checking device_capability.major

344 + for stmt in ast.walk(node):

345 + if not isinstance(stmt, ast.If):

346 + continue

347 +

348 + # Check if it's a capability.major == 10 check (Blackwell)

349 + is_blackwell = (

350 + isinstance(stmt.test, ast.Compare)

351 + and isinstance(stmt.test.left, ast.Attribute)

352 + and stmt.test.left.attr == "major"

353 + and stmt.test.comparators

354 + and isinstance(stmt.test.comparators[0], ast.Constant)

355 + and stmt.test.comparators[0].value == 10

356 + )

357 +

358 + # Extract backends from return statements

359 + for body_stmt in stmt.body:

360 + if isinstance(body_stmt, ast.Return) and isinstance(

361 + body_stmt.value, ast.List

362 + ):

363 + backends = []

364 + for elt in body_stmt.value.elts:

365 + if isinstance(elt, ast.Attribute):

366 + backends.append(elt.attr)

367 + if is_blackwell:

368 + priorities["blackwell"] = backends

369 + else:

370 + priorities["default"] = backends

371 +

372 + # Extract from else branch

373 + for else_stmt in stmt.orelse:

374 + if isinstance(else_stmt, ast.Return) and isinstance(

375 + else_stmt.value, ast.List

376 + ):

377 + backends = []

378 + for elt in else_stmt.value.elts:

379 + if isinstance(elt, ast.Attribute):

380 + backends.append(elt.attr)

381 + priorities["default"] = backends

382 +

383 + return priorities

384 +

385 +

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

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

388 +

389 + Args:

390 + class_path: The fully qualified class path.

391 +

392 + Returns:

393 + A dict with backend properties, or None if parsing fails.

394 + """

395 + file_path = get_file_from_class_path(class_path)

396 + if file_path is None:

397 + return None

398 +

399 + try:

400 + tree = ast.parse(file_path.read_text())

401 + except Exception:

402 + return None

403 +

404 + class_name = class_path.rsplit(".", 1)[1]

405 + class_node = find_class_in_ast(tree, class_name)

406 + if class_node is None:

407 + return None

408 +

409 + info: dict[str, Any] = {

410 + "compute_capability": "Any",

411 + "requires_r1_dims": False,

412 + "dtypes": "fp16, bf16", # Default from base class

413 + }

414 +

415 + # Parse class variables

416 + for item in class_node.body:

417 + if isinstance(item, ast.Assign):

418 + for target in item.targets:

419 + if (

420 + isinstance(target, ast.Name)

421 + and target.id == "requires_r1_mla_dimensions"

422 + and isinstance(item.value, ast.Constant)

423 + ):

424 + info["requires_r1_dims"] = item.value.value

425 +

426 + # Parse supported_dtypes class variable

427 + if (

428 + isinstance(item, ast.AnnAssign)

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

430 + and item.target.id == "supported_dtypes"

431 + and isinstance(item.value, ast.List)

432 + ):

433 + dtype_map = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32"}

434 + dtypes = []

435 + for elt in item.value.elts:

436 + if isinstance(elt, ast.Attribute):

437 + dtypes.append(dtype_map.get(elt.attr, elt.attr))

438 + if dtypes:

439 + info["dtypes"] = ", ".join(dtypes)

440 +

441 + # Parse get_name static method

442 + get_name_method = find_method(class_node, "get_name")

443 + if get_name_method:

444 + for n in ast.walk(get_name_method):

445 + if isinstance(n, ast.Return) and isinstance(n.value, ast.Constant):

446 + info["name"] = n.value.value

447 +

448 + # Parse supports_compute_capability classmethod

449 + cc_method = find_method(class_node, "supports_compute_capability")

450 + if cc_method:

451 + for n in ast.walk(cc_method):

452 + # Look for capability.major == 10 style checks

453 + if (

454 + isinstance(n, ast.Compare)

455 + and isinstance(n.left, ast.Attribute)

456 + and n.left.attr == "major"

457 + and n.comparators

458 + and isinstance(n.comparators[0], ast.Constant)

459 + ):

460 + major = n.comparators[0].value

461 + info["compute_capability"] = f"{major}.x"

462 +

463 + return info

464 +

465 +

466 +def parse_mla_prefill_backends() -> list[dict[str, Any]]:

467 + """Parse MLA prefill backend options from the prefill registry.

468 +

469 + MLA uses different backends for prefill vs decode. The decode backends are

470 + registered in the main registry, but prefill backends have their own

471 + registry at vllm/v1/attention/backends/mla/prefill/registry.py.

472 +

473 + Returns a list of prefill backend info dicts with their requirements.

474 + """

475 + registry = parse_mla_prefill_registry()

476 + priorities = parse_mla_prefill_priorities()

477 +

478 + if not registry:

479 + return []

480 +

481 + # Get the priority order (Blackwell order shows all backends)

482 + priority_order = priorities.get("blackwell", list(registry.keys()))

483 +

484 + prefill_backends: list[dict[str, Any]] = []

485 +

486 + # Backend-specific metadata that can't be easily parsed from code

487 + backend_metadata = {

488 + "TRTLLM_RAGGED": {

489 + "description": "TensorRT-LLM ragged attention",

490 + },

491 + "FLASHINFER": {

492 + "description": "FlashInfer CUTLASS backend",

493 + },

494 + "FLASH_ATTN": {

495 + "description": "FlashAttention varlen (FA2/FA3/FA4)",

496 + },

497 + }

498 +

499 + for backend_name in priority_order:

500 + if backend_name not in registry:

501 + continue

502 +

503 + class_path = registry[backend_name]

504 + backend_info = parse_mla_prefill_backend_file(class_path)

505 + if backend_info is None:

506 + continue

507 +

508 + metadata = backend_metadata.get(backend_name, {})

509 + display_name = backend_info.get("name", backend_name)

510 +

511 + # Add marker for default Blackwell backend

512 + marker = ""

513 + if backend_name == priority_order[0] and priorities.get("blackwell"):

514 + marker = "‡"

515 +

516 + notes = ""

517 + if backend_info.get("requires_r1_dims"):

518 + notes = "DeepSeek R1 dims only"

519 + elif backend_name == "FLASH_ATTN":

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

521 +

522 + prefill_backends.append(

523 + {

524 + "name": display_name,

525 + "marker": marker,

526 + "description": metadata.get("description", ""),

527 + "dtypes": backend_info.get("dtypes", "fp16, bf16"),

528 + "compute_capability": backend_info.get("compute_capability", "Any"),

529 + "notes": notes,

530 + }

531 + )

532 +

533 + return prefill_backends

534 +

535 +

536 # ---------------------------------------------------------------------------

537 # Backend feature extraction from AST

538 # ---------------------------------------------------------------------------

@@ -807,86 +1047,6 @@ def parse_flashinfer_trtllm_features() -> dict[str, dict[str, Any]]:

1047 }

1048

1049

810 -def parse_mla_prefill_backends() -> list[dict[str, Any]]:

811 - """Parse MLA prefill backend options from mla_attention.py.

812 -

813 - MLA uses different backends for prefill vs decode. The decode backends are

814 - registered in the registry, but prefill backends are selected at runtime

815 - based on conditions in MLACommonImpl.__init__.

816 -

817 - Returns a list of prefill backend info dicts with their requirements.

818 - """

819 - if not MLA_ATTENTION_FILE.exists():

820 - return []

821 -

822 - try:

823 - tree = ast.parse(MLA_ATTENTION_FILE.read_text())

824 - except Exception:

825 - return []

826 -

827 - # Find compute capability requirements by parsing use_* functions

828 - trtllm_cc = _find_cc_in_function(tree, "use_trtllm_ragged_deepseek_prefill")

829 - flashinfer_cc = _find_cc_in_function(tree, "use_flashinfer_prefill")

830 - cudnn_cc = _find_cc_in_function(tree, "use_cudnn_prefill")

831 -

832 - # Build prefill backend list based on what we found

833 - # Order matches the priority in MLACommonImpl.__init__

834 - prefill_backends: list[dict[str, Any]] = []

835 -

836 - # TRT-LLM Ragged (highest priority if available)

837 - if trtllm_cc:

838 - prefill_backends.append(

839 - {

840 - "name": "TRT-LLM Ragged‡",

841 - "description": "TensorRT-LLM ragged attention",

842 - "compute_capability": trtllm_cc,

843 - "enable": "Default on SM100",

844 - "disable": "`-ac.use_trtllm_ragged_deepseek_prefill=0`",

845 - "notes": "DeepSeek R1 dims only",

846 - }

847 - )

848 -

849 - # FlashInfer prefill

850 - if flashinfer_cc:

851 - prefill_backends.append(

852 - {

853 - "name": "FlashInfer",

854 - "description": "FlashInfer CUTLASS backend",

855 - "compute_capability": flashinfer_cc,

856 - "enable": "`-ac.disable_flashinfer_prefill=0`",

857 - "disable": "`-ac.disable_flashinfer_prefill=1`",

858 - "notes": "DeepSeek R1 dims only",

859 - }

860 - )

861 -

862 - # cuDNN prefill

863 - if cudnn_cc:

864 - prefill_backends.append(

865 - {

866 - "name": "cuDNN",

867 - "description": "cuDNN-based attention",

868 - "compute_capability": cudnn_cc,

869 - "enable": "`-ac.use_cudnn_prefill=1`",

870 - "disable": "`-ac.use_cudnn_prefill=0`",

871 - "notes": "",

872 - }

873 - )

874 -

875 - # FlashAttention is always available as fallback

876 - prefill_backends.append(

877 - {

878 - "name": "FlashAttention",

879 - "description": "FlashAttention varlen (FA2/FA3)",

880 - "compute_capability": "Any",

881 - "enable": "Default fallback",

882 - "disable": "Use other backends",

883 - "notes": "FA3 on SM90, FA2 otherwise",

884 - }

885 - )

886 -

887 - return prefill_backends

888 -

889 -

1050 # ---------------------------------------------------------------------------

1051 # Backend variant expansion (FA2/FA3/FA4, FlashInfer native/TRTLLM)

1052 # ---------------------------------------------------------------------------

@@ -1415,20 +1575,22 @@ def generate_mla_section(

1575 "",

1576 "### Prefill Backends",

1577 "",

1418 - "The prefill backend is selected at runtime based on hardware and",

1419 - "configuration.",

1578 + "To explicitly select a prefill backend, use",

1579 + "`-ac.mla_prefill_backend=<BACKEND>` (e.g., `FLASH_ATTN`, `FLASHINFER`).",

1580 + "Otherwise, the prefill backend is selected automatically at runtime based on",

1581 + "hardware and configuration.",

1582 "",

1422 - "| ------- | ----------- | ------------ | ------ | ------- | ----- |",

1584 + "| ------- | ----------- | ------ | ------------ | ----- |",

1585 ]

1586

1587 for backend in prefill_backends:

1426 - row = "| {} | {} | {} | {} | {} | {} |".format(

1588 + row = "| `{}`{} | {} | {} | {} | {} |".format(

1589 backend["name"],

1590 + backend.get("marker", ""),

1591 backend["description"],

1592 + backend.get("dtypes", "fp16, bf16"),

1593 backend["compute_capability"],

1430 - backend["enable"],

1431 - backend["disable"],

1594 backend.get("notes", ""),

1595 )

1596 lines.append(row.replace(" ", " "))

@@ -1441,6 +1603,9 @@ def generate_mla_section(

1603 "",

1604 "### Decode Backends",

1605 "",

1606 + "MLA decode backends are selected using the standard",

1607 + "`-ac.backend=<BACKEND>` argument (e.g., `FLASHMLA`, `TRITON_MLA`).",

1608 + "",

1609 ]

1610 )

1611

vllm/config/attention.py

+57 -2

@@ -6,8 +6,12 @@ from typing import Any, Literal

6 from pydantic import field_validator

8 from vllm.config.utils import config

9 +from vllm.logger import init_logger

10 +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

11 from vllm.v1.attention.backends.registry import AttentionBackendEnum

13 +logger = init_logger(__name__)

14 +

16 @config

17 class AttentionConfig:

@@ -33,7 +37,7 @@ class AttentionConfig:

37 and buffers can be pre-allocated to avoid inflating the memory estimate."""

39 use_cudnn_prefill: bool = False

36 - """Whether to use cudnn prefill."""

40 + """Deprecated: cuDNN prefill backend has been removed."""

42 use_trtllm_ragged_deepseek_prefill: bool = False

43 """Whether to use TRTLLM ragged deepseek prefill."""

@@ -42,12 +46,18 @@ class AttentionConfig:

46 """If set to True/False, use or don't use the TRTLLM attention backend

47 in flashinfer. If None, auto-detect the attention backend in flashinfer."""

45 - disable_flashinfer_prefill: bool = True

49 + disable_flashinfer_prefill: bool | None = None

50 """Whether to disable flashinfer prefill."""

52 disable_flashinfer_q_quantization: bool = False

53 """If set, when using fp8 kv, do not quantize Q to fp8."""

55 + mla_prefill_backend: MLAPrefillBackendEnum | None = None

56 + """MLA prefill backend to use. If None, will be selected automatically.

57 + Valid options: FLASH_ATTN (FA3/FA4), FLASHINFER, TRTLLM_RAGGED.

58 + This option supersedes use_trtllm_ragged_deepseek_prefill

59 + and disable_flashinfer_prefill which are deprecated."""

60 +

61 use_prefill_query_quantization: bool = False

62 """If set, quantize query for attention in prefill."""

@@ -84,3 +94,48 @@ class AttentionConfig:

94 return None

95 return AttentionBackendEnum[value.upper()]

96 return value

97 +

98 + @field_validator("mla_prefill_backend", mode="before")

99 + @classmethod

100 + def validate_mla_prefill_backend_before(cls, value: Any) -> Any:

101 + """Enable parsing of the `mla_prefill_backend` enum type from string."""

102 + if isinstance(value, str):

103 + return MLAPrefillBackendEnum[value.upper()]

104 + return value

105 +

106 + def __post_init__(self) -> None:

107 + self._migrate_deprecated_mla_prefill_flags()

108 +

109 + def _migrate_deprecated_mla_prefill_flags(self) -> None:

110 + """Migrate deprecated MLA prefill flags to mla_prefill_backend."""

111 + # If the new option is already set, it takes precedence

112 + if self.mla_prefill_backend is not None:

113 + return

114 +

115 + # Check for deprecated flags and migrate them.

116 + # Only the first flag encountered sets the backend.

117 + if self.use_cudnn_prefill:

118 + raise ValueError(

119 + "The cuDNN MLA prefill backend has been removed. "

120 + "Use --attention-config.mla_prefill_backend=FLASH_ATTN or "

121 + "FLASHINFER or TRTLLM_RAGGED instead."

122 + )

123 +

124 + if self.use_trtllm_ragged_deepseek_prefill:

125 + if self.mla_prefill_backend is None:

126 + self.mla_prefill_backend = MLAPrefillBackendEnum.TRTLLM_RAGGED

127 + logger.warning_once(

128 + "use_trtllm_ragged_deepseek_prefill is deprecated and "

129 + "will be removed in v0.22. Use "

130 + "--attention-config.mla_prefill_backend=TRTLLM_RAGGED "

131 + "instead."

132 + )

133 +

134 + if self.disable_flashinfer_prefill:

135 + if self.mla_prefill_backend is None:

136 + self.mla_prefill_backend = MLAPrefillBackendEnum.FLASH_ATTN

137 + logger.warning_once(

138 + "disable_flashinfer_prefill is deprecated and will be removed "

139 + "in v0.22. Use --attention-config.mla_prefill_backend="

140 + "FLASH_ATTN instead."

141 + )

vllm/model_executor/layers/attention/mla_attention.py

+58 -595

@@ -189,12 +189,9 @@ return curr_o @ W_O

189

190 import functools

191 from abc import abstractmethod

192 -from dataclasses import dataclass, field

192 +from dataclasses import dataclass

193 from enum import Enum

194 -from typing import TYPE_CHECKING, ClassVar, Generic, TypeVar, cast

195 -

196 -if TYPE_CHECKING:

197 - from flashinfer import BatchPrefillWithRaggedKVCacheWrapper

194 +from typing import ClassVar, Generic, TypeVar, cast

195

196 import torch

197 import torch.nn as nn

@@ -242,7 +239,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (

239 kNvfp4Dynamic,

240 )

241 from vllm.platforms import current_platform

245 -from vllm.utils.flashinfer import has_flashinfer, has_nvidia_artifactory

242 +from vllm.utils.flashinfer import has_flashinfer

243 from vllm.utils.math_utils import cdiv, round_down

244 from vllm.utils.torch_utils import (

245 LayerNameType,

@@ -262,11 +259,9 @@ from vllm.v1.attention.backend import (

259 MLAAttentionImpl,

260 SparseMLAAttentionImpl,

261 )

265 -from vllm.v1.attention.backends.fa_utils import get_flash_attn_version

262 +from vllm.v1.attention.backends.mla.prefill import MLAPrefillBackend

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

264 get_dcp_local_seq_lens,

268 - get_per_layer_parameters,

269 - infer_global_hyperparameters,

265 split_decodes_and_prefills,

266 )

267 from vllm.v1.attention.ops.common import cp_lse_ag_out_rs

@@ -1123,33 +1118,6 @@ class QueryLenSupport(Enum):

1118 VARLEN = "varlen"

1119

1120

1126 -try:

1127 - from vllm.vllm_flash_attn import ( # type: ignore[attr-defined]

1128 - flash_attn_varlen_func,

1129 - )

1130 -

1131 - is_vllm_fa = True

1132 -except ImportError:

1133 - is_vllm_fa = False

1134 - flash_attn_varlen_func = None # type: ignore[assignment]

1135 - # On ROCm, vllm_flash_attn is not available, try upstream flash_attn instead.

1136 - # On CUDA, vllm_flash_attn should always be available (built with vLLM),

1137 - # so we don't attempt the fallback there.

1138 - if current_platform.is_rocm():

1139 - try:

1140 - from flash_attn import flash_attn_varlen_func # type: ignore[no-redef]

1141 - except ImportError:

1142 - logger.debug(

1143 - "flash_attn not available on ROCm; "

1144 - "MLA models using TRITON_MLA will require flash_attn. "

1145 - "AITER_MLA backends use aiter kernels instead."

1146 - )

1147 - elif current_platform.is_xpu():

1148 - from vllm._xpu_ops import xpu_ops

1149 -

1150 - flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func # type: ignore[no-redef,attr-defined,assignment]

1151 -

1152 -

1121 def dynamic_per_batched_tensor_quant(

1122 x: torch.Tensor, dtype: torch.dtype = torch.float8_e4m3fn

1123 ):

@@ -1161,9 +1129,6 @@ def dynamic_per_batched_tensor_quant(

1129 return x_scl_sat.to(dtype).contiguous(), scale.float().reciprocal()

1130

1131

1164 -logger = init_logger(__name__)

1165 -

1166 -

1132 @CustomOp.register(

1133 "mla_decode_concat_quant_fp8",

1134 dynamic_arg_dims={"decode_ql_nope": 0, "decode_q_pe": 0},

@@ -1197,9 +1162,6 @@ class _DecodeConcatQuantFP8(QuantFP8):

1162 forward_hip = _make_forward(QuantFP8.forward_hip) # type: ignore[arg-type]

1163

1164

1200 -CUDNN_WORKSPACE_SIZE = 12800

1201 -

1202 -

1165 class MLACommonBackend(AttentionBackend):

1166 @staticmethod

1167 def get_name() -> str:

@@ -1268,26 +1230,9 @@ class MLACommonPrefillMetadata:

1230 query_start_loc: torch.Tensor

1231 max_query_len: int

1232 chunked_context: ChunkedContextMetadata | None = None

1271 - query_seq_lens: torch.Tensor | None = None

1272 - workspace_buffer: torch.Tensor | None = None

1233 q_data_type: torch.dtype | None = None

1234 output_dtype: torch.dtype | None = None

1275 -

1276 -

1277 -@dataclass

1278 -class FlashInferPrefillMetadata(MLACommonPrefillMetadata):

1279 - prefill_main: "BatchPrefillWithRaggedKVCacheWrapper | None" = None

1280 - prefill_chunks: "list[BatchPrefillWithRaggedKVCacheWrapper]" = field(

1281 - default_factory=list

1282 - )

1283 -

1284 -

1285 -@dataclass

1286 -class CudnnPrefillMetadata(MLACommonPrefillMetadata):

1287 - class ChunkedContextMetadata(MLACommonPrefillMetadata.ChunkedContextMetadata):

1288 - seq_lens: torch.Tensor

1289 -

1290 - cudnn_workspace: torch.Tensor | None = None

1235 + prefill_backend: MLAPrefillBackend | None = None

1236

1237

1238 @dataclass

@@ -1333,13 +1278,8 @@ class MLACommonMetadata(AttentionMetadata, Generic[D]):

1278 # The dimension of the attention heads

1279 head_dim: int | None = None

1280

1281 + prefill: MLACommonPrefillMetadata | None = None

1282 decode: D | None = None

1337 - prefill: (

1338 - MLACommonPrefillMetadata

1339 - | FlashInferPrefillMetadata

1340 - | CudnnPrefillMetadata

1341 - | None

1342 - ) = None

1283

1284 def __post_init__(self):

1285 if self.head_dim is not None and not MLACommonBackend.supports_head_size(

@@ -1352,64 +1292,6 @@ M = TypeVar("M", bound=MLACommonMetadata)

1292 A = TypeVar("A", bound=AttentionMetadata)

1293

1294

1355 -def is_deepseek_r1_mla_compatible(vllm_config: VllmConfig) -> bool:

1356 - # Check if model has DeepSeek R1 compatible MLA dimensions:

1357 - # qk_nope_head_dim = 128, qk_rope_head_dim = 64, v_head_dim = 128

1358 - # which results in query/key head dim = 192.

1359 - if vllm_config.model_config is None:

1360 - return False

1361 - hf_text_config = vllm_config.model_config.hf_text_config

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

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

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

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

1366 -

1367 -

1368 -@functools.cache

1369 -def use_flashinfer_prefill() -> bool:

1370 - from vllm.config import get_current_vllm_config

1371 -

1372 - vllm_config = get_current_vllm_config()

1373 - if not (

1374 - not vllm_config.attention_config.disable_flashinfer_prefill

1375 - and has_flashinfer()

1376 - and not vllm_config.attention_config.use_cudnn_prefill

1377 - and current_platform.is_device_capability_family(100)

1378 - ):

1379 - return False

1380 -

1381 - return is_deepseek_r1_mla_compatible(vllm_config)

1382 -

1383 -

1384 -@functools.cache

1385 -def use_cudnn_prefill() -> bool:

1386 - from vllm.config import get_current_vllm_config

1387 -

1388 - vllm_config = get_current_vllm_config()

1389 - return (

1390 - has_flashinfer()

1391 - and vllm_config.attention_config.use_cudnn_prefill

1392 - and current_platform.is_device_capability_family(100)

1393 - and has_nvidia_artifactory()

1394 - )

1395 -

1396 -

1397 -@functools.cache

1398 -def use_trtllm_ragged_deepseek_prefill() -> bool:

1399 - """Check if TRT-LLM ragged DeepSeek prefill should be used."""

1400 - from vllm.config import get_current_vllm_config

1401 -

1402 - vllm_config = get_current_vllm_config()

1403 - if not (

1404 - has_flashinfer()

1405 - and vllm_config.attention_config.use_trtllm_ragged_deepseek_prefill

1406 - and current_platform.is_device_capability_family(100)

1407 - ):

1408 - return False

1409 -

1410 - return is_deepseek_r1_mla_compatible(vllm_config)

1411 -

1412 -

1295 @dataclass

1296 class MLADims:

1297 q_lora_rank: int | None

@@ -1447,15 +1329,14 @@ def get_mla_dims(model_config: ModelConfig) -> MLADims:

1329

1330 @functools.cache

1331 def backend_supports_prefill_query_quantization() -> bool:

1450 - """Check if the selected MLA backend supports prefill query quantization.

1332 + """Check if the selected MLA prefill backend supports query quantization.

1333

1334 Currently supported backends:

1453 - - FlashInfer prefill

1454 - - TRT-LLM ragged DeepSeek prefill

1335 + - FlashInfer

1336 + - TRT-LLM Ragged

1337

1338 Not supported:

1457 - - cuDNN Prefill

1458 - - FlashAttention

1339 + - FlashAttention (FA3/FA4)

1340 - Non-GB200 devices (FP8 prefill requires device capability 100)

1341 """

1342 # FP8 prefill query quantization requires GB200 (device capability 100)

@@ -1463,7 +1344,15 @@ def backend_supports_prefill_query_quantization() -> bool:

1344 if not current_platform.is_device_capability_family(100):

1345 return False

1346

1466 - return use_flashinfer_prefill() or use_trtllm_ragged_deepseek_prefill()

1347 + from vllm.config import get_current_vllm_config

1348 + from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend

1349 +

1350 + vllm_config = get_current_vllm_config()

1351 + backend_cls = get_mla_prefill_backend(vllm_config)

1352 + return backend_cls.get_name() in (

1353 + "FLASHINFER",

1354 + "TRTLLM_RAGGED",

1355 + )

1356

1357

1358 class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

@@ -1574,7 +1463,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1463 metadata_cls if metadata_cls is not None else MLACommonMetadata

1464 )

1465 self.kv_cache_spec = kv_cache_spec

1577 - scheduler_config = vllm_config.scheduler_config

1466 self.model_config = vllm_config.model_config

1467 parallel_config = vllm_config.parallel_config

1468 self.compilation_config = vllm_config.compilation_config

@@ -1634,139 +1522,32 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1522 device=device,

1523 )

1524

1637 - self._use_cudnn_prefill = use_cudnn_prefill()

1638 - self._use_fi_prefill = use_flashinfer_prefill()

1639 - self._use_trtllm_ragged_prefill = use_trtllm_ragged_deepseek_prefill()

1640 - self.prefill_metadata_cls = (

1641 - FlashInferPrefillMetadata

1642 - if self._use_fi_prefill

1643 - else CudnnPrefillMetadata

1644 - if self._use_cudnn_prefill

1645 - else MLACommonPrefillMetadata

1646 - )

1647 -

1648 - if self._use_fi_prefill:

1649 - self._workspace_buffer = torch.empty(

1650 - envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,

1651 - dtype=torch.uint8,

1652 - device=device,

1653 - )

1654 -

1655 - self._fi_prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None

1656 - self._fi_prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []

1657 -

1658 - self._global_hyperparameters = infer_global_hyperparameters(

1659 - get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl) # type: ignore[type-abstract]

1660 - )

1661 -

1662 - if self._use_trtllm_ragged_prefill:

1663 - self._workspace_buffer = torch.empty(

1664 - envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,

1665 - dtype=torch.uint8,

1666 - device=device,

1667 - )

1525 + from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend

1526

1669 - if self._use_cudnn_prefill:

1670 - self.cudnn_workspace = torch.empty(

1671 - CUDNN_WORKSPACE_SIZE * scheduler_config.max_num_seqs,

1672 - dtype=torch.int8,

1673 - device=device,

1674 - )

1527 + prefill_backend_cls = get_mla_prefill_backend(vllm_config)

1528 + self._prefill_backend = prefill_backend_cls(

1529 + num_heads=self.num_heads,

1530 + scale=self.model_config.get_head_size() ** -0.5,

1531 + kv_lora_rank=self.mla_dims.kv_lora_rank,

1532 + qk_nope_head_dim=self.mla_dims.qk_nope_head_dim,

1533 + qk_rope_head_dim=self.mla_dims.qk_rope_head_dim,

1534 + v_head_dim=self.mla_dims.v_head_dim,

1535 + vllm_config=vllm_config,

1536 + device=device,

1537 + layer_names=layer_names,

1538 + )

1539

1540 supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY

1541 self._init_reorder_batch_threshold(

1542 self.reorder_batch_threshold, supports_spec_decode, supports_dcp_with_varlen

1543 )

1544

1681 - # Validate consistency between query_len_support and reorder_batch_threshold

1545 if self.query_len_support == QueryLenSupport.SINGLE_ONLY:

1546 assert self.reorder_batch_threshold == 1, (

1547 f"reorder_batch_threshold must be 1 when query_len_support is "

1548 f"SINGLE_ONLY, got {self.reorder_batch_threshold}"

1549 )

1550

1688 - def _build_fi_prefill_wrappers(self, prefill: FlashInferPrefillMetadata):

1689 - qo_indptr = prefill.query_start_loc

1690 -

1691 - has_context = False

1692 - if prefill.chunked_context is not None:

1693 - chunked_context = prefill.chunked_context

1694 - has_context = True

1695 -

1696 - if self._fi_prefill_main is None:

1697 - from flashinfer import BatchPrefillWithRaggedKVCacheWrapper

1698 -

1699 - self._fi_prefill_main = BatchPrefillWithRaggedKVCacheWrapper(

1700 - self._workspace_buffer, "NHD", backend="cutlass"

1701 - )

1702 -

1703 - if has_context:

1704 - num_chunks = chunked_context.cu_seq_lens.shape[0]

1705 - # Allocate more prefill chunk wrappers if needed

1706 - if len(self._fi_prefill_chunks) < num_chunks:

1707 - from flashinfer import BatchPrefillWithRaggedKVCacheWrapper

1708 -

1709 - for _ in range(len(self._fi_prefill_chunks), num_chunks):

1710 - self._fi_prefill_chunks.append(

1711 - BatchPrefillWithRaggedKVCacheWrapper(

1712 - self._workspace_buffer, "NHD", backend="cutlass"

1713 - )

1714 - )

1715 - assert num_chunks <= len(self._fi_prefill_chunks)

1716 -

1717 - # In MLA, the non-latent num_qo_heads == num_kv_heads

1718 - num_qo_heads = self.num_heads

1719 - num_kv_heads = num_qo_heads

1720 -

1721 - # Sanity: Verify that num_kv_heads == 1 since it is latent space

1722 - assert self.kv_cache_spec.num_kv_heads == 1

1723 -

1724 - # Get non-latent head_dim_qk and head_dim_vo

1725 - head_dim_qk = self.mla_dims.qk_nope_head_dim + self.mla_dims.qk_rope_head_dim

1726 - head_dim_vo = self.mla_dims.v_head_dim

1727 -

1728 - # For main run, qo_indptr == kv_indptr

1729 - kv_indptr = qo_indptr.clone()

1730 -

1731 - # Prepare main prefill

1732 - self._fi_prefill_main.plan(

1733 - qo_indptr=qo_indptr,

1734 - kv_indptr=kv_indptr,

1735 - num_qo_heads=num_qo_heads,

1736 - num_kv_heads=num_kv_heads,

1737 - head_dim_qk=head_dim_qk,

1738 - head_dim_vo=head_dim_vo,

1739 - causal=True, # This is main run

1740 - sm_scale=self._global_hyperparameters.sm_scale,

1741 - window_left=self._global_hyperparameters.window_left,

1742 - logits_soft_cap=self._global_hyperparameters.logits_soft_cap,

1743 - q_data_type=self.q_data_type,

1744 - o_data_type=prefill.output_dtype,

1745 - )

1746 -

1747 - # Prepare context prefills

1748 - if has_context:

1749 - for i in range(num_chunks):

1750 - kv_indptr_chunk = chunked_context.cu_seq_lens[i]

1751 -

1752 - self._fi_prefill_chunks[i].plan(

1753 - qo_indptr=qo_indptr,

1754 - kv_indptr=kv_indptr_chunk,

1755 - num_qo_heads=num_qo_heads,

1756 - num_kv_heads=num_kv_heads,

1757 - head_dim_qk=head_dim_qk,

1758 - head_dim_vo=head_dim_vo,

1759 - causal=False, # This is context run

1760 - sm_scale=self._global_hyperparameters.sm_scale,

1761 - window_left=self._global_hyperparameters.window_left,

1762 - logits_soft_cap=self._global_hyperparameters.logits_soft_cap,

1763 - q_data_type=self.q_data_type,

1764 - o_data_type=prefill.output_dtype,

1765 - )

1766 -

1767 - prefill.prefill_main = self._fi_prefill_main

1768 - prefill.prefill_chunks = self._fi_prefill_chunks

1769 -

1551 def _build_decode(

1552 self,

1553 block_table_tensor: torch.Tensor,

@@ -1972,18 +1753,14 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1753 dtype=torch.int32,

1754 )

1755

1975 - chunked_context_metadata_cls = (

1976 - CudnnPrefillMetadata.ChunkedContextMetadata

1977 - if self._use_cudnn_prefill

1978 - else MLACommonPrefillMetadata.ChunkedContextMetadata

1979 - )

1756 prefill_tokens_with_context = None

1757 if num_prefills_with_context_cpu > 0:

1758 prefill_tokens_with_context = prefill_query_start_loc_cpu[

1759 num_prefills_with_context_cpu

1760 ].item()

1761 + _ChunkedMetadata = MLACommonPrefillMetadata.ChunkedContextMetadata

1762 if self.dcp_world_size > 1:

1986 - chunked_context_metadata = chunked_context_metadata_cls(

1763 + chunked_context_metadata = _ChunkedMetadata(

1764 cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),

1765 starts=local_chunk_starts.to(device, non_blocking=True),

1766 seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),

@@ -2004,7 +1781,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1781 prefill_tokens_with_context=prefill_tokens_with_context,

1782 )

1783 else:

2007 - chunked_context_metadata = chunked_context_metadata_cls(

1784 + chunked_context_metadata = _ChunkedMetadata(

1785 cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),

1786 starts=chunk_starts.to(device, non_blocking=True),

1787 seq_tot=chunk_seq_lens.sum(dim=1).tolist(),

@@ -2018,35 +1795,22 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1795 prefill_tokens_with_context=prefill_tokens_with_context,

1796 )

1797

2021 - if self._use_cudnn_prefill:

2022 - chunked_context_metadata.seq_lens = chunk_seq_lens

2023 -

1798 assert (

1799 max(chunked_context_metadata.max_seq_lens)

1800 <= self.chunked_prefill_workspace_size

1801 )

1802

2029 - prefill_metadata = self.prefill_metadata_cls(

1803 + prefill_metadata = MLACommonPrefillMetadata(

1804 block_table=block_table_tensor[reqs_start:, ...],

1805 query_start_loc=prefill_query_start_loc,

1806 max_query_len=max_query_len,

1807 chunked_context=chunked_context_metadata,

1808 output_dtype=self.model_config.dtype,

1809 q_data_type=self.q_data_type,

1810 + prefill_backend=self._prefill_backend,

1811 )

1812

2038 - if self._use_cudnn_prefill:

2039 - assert isinstance(prefill_metadata, CudnnPrefillMetadata)

2040 - prefill_metadata.query_seq_lens = (

2041 - prefill_query_start_loc[1:] - prefill_query_start_loc[:-1]

2042 - )

2043 - prefill_metadata.cudnn_workspace = self.cudnn_workspace

2044 -

2045 - if self._use_trtllm_ragged_prefill:

2046 - prefill_metadata.query_seq_lens = (

2047 - prefill_query_start_loc[1:] - prefill_query_start_loc[:-1]

2048 - )

2049 - prefill_metadata.workspace_buffer = self._workspace_buffer

1813 + self._prefill_backend.prepare_metadata(prefill_metadata)

1814

1815 decode_metadata = None

1816 if num_decodes > 0:

@@ -2091,10 +1855,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1855 decode=decode_metadata,

1856 )

1857

2094 - if self._use_fi_prefill and num_prefills > 0:

2095 - assert isinstance(attn_metadata.prefill, FlashInferPrefillMetadata)

2096 - self._build_fi_prefill_wrappers(attn_metadata.prefill)

2097 -

1858 return attn_metadata # type: ignore[return-value]

1859

1860

@@ -2240,308 +2000,12 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2000 and (self.qk_rope_head_dim == 64)

2001 )

2002

2243 - if use_trtllm_ragged_deepseek_prefill():

2244 - logger.info_once("Using TRT-LLM ragged DeepSeek prefill for MLA")

2245 - self._run_prefill_context_chunk = (

2246 - self._run_prefill_context_chunk_trtllm_ragged

2247 - )

2248 - self._run_prefill_new_tokens = self._run_prefill_new_tokens_trtllm_ragged

2249 - self._pad_v = False

2250 - elif use_flashinfer_prefill():

2251 - logger.info_once("Using FlashInfer prefill for MLA")

2252 - self._run_prefill_context_chunk = self._run_prefill_context_chunk_fi

2253 - self._run_prefill_new_tokens = self._run_prefill_new_tokens_fi

2254 - self._pad_v = False

2255 - elif use_cudnn_prefill():

2256 - logger.info_once("Using CUDNN prefill for MLA")

2257 - self._run_prefill_context_chunk = self._run_prefill_context_chunk_cudnn

2258 - self._run_prefill_new_tokens = self._run_prefill_new_tokens_cudnn

2259 - self._pad_v = False

2260 - else: # Use FlashAttention

2261 - if flash_attn_varlen_func is None:

2262 - raise RuntimeError(

2263 - "MLA attention requires FlashAttention but it is not "

2264 - "available. Please install flash_attn or use "

2265 - "--attention-backend ROCM_AITER_MLA."

2266 - )

2267 - logger.info_once("Using FlashAttention prefill for MLA")

2268 - self._run_prefill_context_chunk = self._run_prefill_context_chunk_fa

2269 - self._run_prefill_new_tokens = self._run_prefill_new_tokens_fa

2270 -

2271 - # Handle the differences between the flash_attn_varlen from

2272 - # flash_attn and the one from vllm_flash_attn. The former is used on

2273 - # RoCM and the latter has an additional parameter to control

2274 - # FA2 vs FA3

2275 - self.flash_attn_varlen_func = flash_attn_varlen_func

2276 - self.vllm_flash_attn_version = get_flash_attn_version(

2277 - head_size=self.qk_head_dim

2278 - )

2279 - if self.vllm_flash_attn_version is not None:

2280 - self.flash_attn_varlen_func = functools.partial(

2281 - flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version

2282 - )

2283 -

2284 - # For MLA the v head dim is smaller than qk head dim so we pad out

2285 - # v with 0s to match the qk head dim for attention backends that do

2286 - # not support different headdims.

2287 - # FA3 on Hopper (SM90) and FA4 natively handle diff headdims.

2288 - device_capability = current_platform.get_device_capability()

2289 - self._pad_v = self.vllm_flash_attn_version is None or not (

2290 - (

2291 - self.vllm_flash_attn_version == 3

2292 - and device_capability is not None

2293 - and device_capability[0] == 9

2294 - )

2295 - or self.vllm_flash_attn_version == 4

2296 - )

2297 -

2003 self.dcp_world_size: int = -1

2004

2005 self.cp_kv_cache_interleave_size: int = (

2006 get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size

2007 )

2008

2304 - def _flash_attn_varlen_diff_headdims(

2305 - self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs

2306 - ):

2307 - maybe_padded_v = v

2308 - if self._pad_v:

2309 - maybe_padded_v = torch.nn.functional.pad(

2310 - v, [0, q.shape[-1] - v.shape[-1]], value=0

2311 - )

2312 -

2313 - if is_vllm_fa:

2314 - kwargs["return_softmax_lse"] = return_softmax_lse

2315 - else:

2316 - # ROCm leverages the upstream flash_attn, which takes a parameter

2317 - # called "return_attn_probs" instead of return_softmax_lse

2318 - kwargs["return_attn_probs"] = return_softmax_lse

2319 - if envs.VLLM_BATCH_INVARIANT:

2320 - kwargs["num_splits"] = 1

2321 -

2322 - attn_out = self.flash_attn_varlen_func(

2323 - q=q,

2324 - k=k,

2325 - v=maybe_padded_v,

2326 - softmax_scale=softmax_scale,

2327 - **kwargs,

2328 - )

2329 -

2330 - # Unpack the output if there is multiple results

2331 - lse = None

2332 - if isinstance(attn_out, tuple):

2333 - attn_out, lse = attn_out[0], attn_out[1]

2334 -

2335 - # Remain consistent with old `flash_attn_varlen_func` where there

2336 - # is only one output tensor if `return_softmax_lse` is False.

2337 - if return_softmax_lse:

2338 - return attn_out, lse

2339 - return attn_out

2340 -

2341 - def _run_prefill_new_tokens_fa(

2342 - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse

2343 - ):

2344 - return self._flash_attn_varlen_diff_headdims(

2345 - q=q,

2346 - k=k,

2347 - v=v,

2348 - cu_seqlens_q=prefill.query_start_loc,

2349 - cu_seqlens_k=prefill.query_start_loc,

2350 - max_seqlen_q=prefill.max_query_len,

2351 - max_seqlen_k=prefill.max_query_len,

2352 - softmax_scale=self.scale,

2353 - causal=True,

2354 - return_softmax_lse=return_softmax_lse,

2355 - )

2356 -

2357 - def _run_prefill_new_tokens_fi(

2358 - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse

2359 - ):

2360 - assert isinstance(prefill, FlashInferPrefillMetadata)

2361 - assert prefill.prefill_main is not None

2362 -

2363 - ret = prefill.prefill_main.run(

2364 - q=q,

2365 - k=k,

2366 - v=v,

2367 - return_lse=return_softmax_lse,

2368 - )

2369 -

2370 - if isinstance(ret, tuple):

2371 - return ret[0], ret[1].transpose(0, 1).contiguous()

2372 - return ret

2373 -

2374 - def _run_prefill_new_tokens_cudnn(

2375 - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse

2376 - ):

2377 - assert isinstance(prefill, CudnnPrefillMetadata)

2378 - assert prefill.query_seq_lens is not None

2379 - from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache

2380 -

2381 - output, lse = cudnn_batch_prefill_with_kv_cache(

2382 - q=q,

2383 - k_cache=k,

2384 - v_cache=v,

2385 - scale=self.scale,

2386 - workspace_buffer=prefill.cudnn_workspace,

2387 - max_token_per_sequence=prefill.max_query_len,

2388 - max_sequence_kv=prefill.max_query_len,

2389 - actual_seq_lens_q=prefill.query_seq_lens.view(-1, 1, 1, 1),

2390 - actual_seq_lens_kv=prefill.query_seq_lens.view(-1, 1, 1, 1),

2391 - causal=True,

2392 - # Do not support False for now

2393 - return_lse=True,

2394 - # Indicates actual_seq_lens are on GPU or CPU.

2395 - is_cuda_graph_compatible=True,

2396 - )

2397 - if return_softmax_lse:

2398 - return output, lse

2399 - return output

2400 -

2401 - def _run_prefill_context_chunk_fa(

2402 - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

2403 - ):

2404 - assert prefill.chunked_context is not None

2405 - return self._flash_attn_varlen_diff_headdims(

2406 - q=q,

2407 - k=k,

2408 - v=v,

2409 - cu_seqlens_q=prefill.query_start_loc,

2410 - cu_seqlens_k=prefill.chunked_context.cu_seq_lens[chunk_idx],

2411 - max_seqlen_q=prefill.max_query_len,

2412 - max_seqlen_k=prefill.chunked_context.max_seq_lens[chunk_idx],

2413 - softmax_scale=self.scale,

2414 - causal=False, # Context is unmasked

2415 - return_softmax_lse=True,

2416 - )

2417 -

2418 - def _run_prefill_context_chunk_fi(

2419 - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

2420 - ):

2421 - assert isinstance(prefill, FlashInferPrefillMetadata)

2422 -

2423 - attn_out, lse = prefill.prefill_chunks[chunk_idx].run(

2424 - q=q,

2425 - k=k,

2426 - v=v,

2427 - return_lse=True,

2428 - )

2429 -

2430 - # Convert from (q_len, num_heads) to (num_heads, q_len)

2431 - return attn_out, lse.transpose(0, 1).contiguous()

2432 -

2433 - def _run_prefill_context_chunk_cudnn(

2434 - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

2435 - ):

2436 - assert isinstance(prefill, CudnnPrefillMetadata)

2437 - assert prefill.chunked_context is not None

2438 - assert prefill.chunked_context.seq_lens[chunk_idx] is not None

2439 - assert prefill.query_seq_lens is not None

2440 - from flashinfer.prefill import cudnn_batch_prefill_with_kv_cache

2441 -

2442 - return cudnn_batch_prefill_with_kv_cache(

2443 - q=q,

2444 - k_cache=k,

2445 - v_cache=v,

2446 - scale=self.scale,

2447 - workspace_buffer=prefill.cudnn_workspace,

2448 - max_token_per_sequence=prefill.max_query_len,

2449 - max_sequence_kv=prefill.chunked_context.max_seq_lens[chunk_idx],

2450 - actual_seq_lens_q=prefill.query_seq_lens.view(-1, 1, 1, 1),

2451 - actual_seq_lens_kv=prefill.chunked_context.seq_lens[chunk_idx].view(

2452 - -1, 1, 1, 1

2453 - ),

2454 - causal=False,

2455 - return_lse=True,

2456 - # Indicates actual_seq_lens are on GPU or CPU.

2457 - is_cuda_graph_compatible=True,

2458 - )

2459 -

2460 - def _run_prefill_new_tokens_trtllm_ragged(

2461 - self, prefill: MLACommonPrefillMetadata, q, k, v, return_softmax_lse

2462 - ):

2463 - """TRT-LLM ragged attention for new tokens (causal)."""

2464 - from flashinfer.prefill import trtllm_ragged_attention_deepseek

2465 -

2466 - assert prefill.query_seq_lens is not None

2467 - assert prefill.workspace_buffer is not None

2468 - # allocate BF16 / FP16 output tensor for TRT-LLM ragged attention

2469 - out = torch.empty(

2470 - q.shape[0],

2471 - q.shape[1],

2472 - v.shape[2],

2473 - device=q.device,

2474 - dtype=prefill.output_dtype,

2475 - )

2476 -

2477 - ret = trtllm_ragged_attention_deepseek(

2478 - query=q,

2479 - key=k,

2480 - value=v,

2481 - workspace_buffer=prefill.workspace_buffer,

2482 - seq_lens=prefill.query_seq_lens,

2483 - max_q_len=prefill.max_query_len,

2484 - max_kv_len=prefill.max_query_len,

2485 - bmm1_scale=self.scale,

2486 - bmm2_scale=1.0,

2487 - o_sf_scale=1.0,

2488 - batch_size=prefill.query_seq_lens.shape[0],

2489 - window_left=-1,

2490 - cum_seq_lens_q=prefill.query_start_loc,

2491 - cum_seq_lens_kv=prefill.query_start_loc,

2492 - enable_pdl=False,

2493 - is_causal=True,

2494 - return_lse=return_softmax_lse,

2495 - out=out,

2496 - )

2497 -

2498 - if isinstance(ret, tuple):

2499 - # Convert from (q_len, num_heads) to (num_heads, q_len)

2500 - return ret[0], ret[1].transpose(0, 1).contiguous()

2501 - return ret

2502 -

2503 - def _run_prefill_context_chunk_trtllm_ragged(

2504 - self, prefill: MLACommonPrefillMetadata, chunk_idx: int, q, k, v

2505 - ):

2506 - """TRT-LLM ragged attention for context chunks (non-causal)."""

2507 - from flashinfer.prefill import trtllm_ragged_attention_deepseek

2508 -

2509 - assert prefill.chunked_context is not None

2510 - assert prefill.chunked_context.seq_lens[chunk_idx] is not None

2511 - assert prefill.workspace_buffer is not None

2512 -

2513 - out = torch.empty(

2514 - q.shape[0],

2515 - q.shape[1],

2516 - v.shape[2],

2517 - device=q.device,

2518 - dtype=prefill.output_dtype,

2519 - )

2520 -

2521 - attn_out, lse = trtllm_ragged_attention_deepseek(

2522 - query=q,

2523 - key=k,

2524 - value=v,

2525 - workspace_buffer=prefill.workspace_buffer,

2526 - seq_lens=prefill.chunked_context.seq_lens[chunk_idx],

2527 - max_q_len=prefill.max_query_len,

2528 - max_kv_len=prefill.chunked_context.max_seq_lens[chunk_idx],

2529 - bmm1_scale=self.scale,

2530 - bmm2_scale=1.0,

2531 - o_sf_scale=1.0,

2532 - batch_size=prefill.chunked_context.seq_lens[chunk_idx].shape[0],

2533 - window_left=-1,

2534 - cum_seq_lens_q=prefill.query_start_loc,

2535 - cum_seq_lens_kv=prefill.chunked_context.cu_seq_lens[chunk_idx],

2536 - enable_pdl=False,

2537 - is_causal=False,

2538 - return_lse=True,

2539 - out=out,

2540 - )

2541 -

2542 - # Convert from (q_len, num_heads) to (num_heads, q_len)

2543 - return attn_out, lse.transpose(0, 1).contiguous()

2544 -

2009 def _concat_k_nope_k_pe(

2010 self, k_nope: torch.Tensor, k_pe: torch.Tensor

2011 ) -> torch.Tensor:

@@ -2582,6 +2046,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2046 ):

2047 assert attn_metadata.prefill is not None

2048 prefill_metadata = attn_metadata.prefill

2049 + assert prefill_metadata.prefill_backend is not None

2050 assert prefill_metadata.chunked_context is not None

2051

2052 use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()

@@ -2649,12 +2114,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2114

2115 k = self._concat_k_nope_k_pe(k_nope, k_pe)

2116

2652 - attn_output, attn_softmax_lse = self._run_prefill_context_chunk(

2653 - prefill=prefill_metadata,

2654 - chunk_idx=i,

2655 - q=q,

2656 - k=k,

2657 - v=v,

2117 + attn_output, attn_softmax_lse = (

2118 + prefill_metadata.prefill_backend.run_prefill_context_chunk(

2119 + chunk_idx=i,

2120 + q=q,

2121 + k=k,

2122 + v=v,

2123 + )

2124 )

2125

2126 if output is None:

@@ -2687,6 +2153,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2153 assert k_scale is None, "DCP not support scaled kvcache now."

2154 assert attn_metadata.prefill is not None

2155 prefill_metadata = attn_metadata.prefill

2156 + assert prefill_metadata.prefill_backend is not None

2157 assert prefill_metadata.chunked_context is not None

2158 assert prefill_metadata.chunked_context.padded_local_chunk_seq_lens is not None

2159 assert prefill_metadata.chunked_context.local_context_lens_allranks is not None

@@ -2753,12 +2220,13 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2220 k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

2221 k = self._concat_k_nope_k_pe(k_nope, k_pe)

2222

2756 - attn_output, attn_softmax_lse = self._run_prefill_context_chunk(

2757 - prefill=prefill_metadata,

2758 - chunk_idx=i,

2759 - q=q,

2760 - k=k,

2761 - v=v,

2223 + attn_output, attn_softmax_lse = (

2224 + prefill_metadata.prefill_backend.run_prefill_context_chunk(

2225 + chunk_idx=i,

2226 + q=q,

2227 + k=k,

2228 + v=v,

2229 + )

2230 )

2231

2232 if output is None:

@@ -2790,11 +2258,11 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2258 k_scale: torch.Tensor,

2259 output: torch.Tensor,

2260 ) -> None:

2793 - # TODO (zyongye): Prefill function here

2261 assert attn_metadata.prefill is not None

2262 assert self.dcp_world_size != -1

2263

2264 prefill_metadata = attn_metadata.prefill

2265 + assert prefill_metadata.prefill_backend is not None

2266 use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype()

2267

2268 # Convert q to FP8 if FP8 prefill attention is enabled

@@ -2813,8 +2281,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2281 k = k.to(prefill_metadata.q_data_type)

2282 v = v.to(prefill_metadata.q_data_type)

2283

2816 - output_prefill = self._run_prefill_new_tokens(

2817 - prefill=prefill_metadata,

2284 + output_prefill = prefill_metadata.prefill_backend.run_prefill_new_tokens(

2285 q=q,

2286 k=k,

2287 v=v,

@@ -2839,11 +2306,6 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2306 q, kv_c_and_k_pe_cache, attn_metadata, k_scale

2307 )

2308

2842 - # unpad if necessary

2843 - if self._pad_v:

2844 - context_output = context_output[..., : v.shape[-1]]

2845 - suffix_output = suffix_output[..., : v.shape[-1]]

2846 -

2309 output = output.view(-1, self.num_heads, self.v_head_dim)

2310 merge_attn_states(

2311 output=output,

@@ -2854,7 +2316,8 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2316 prefill_tokens_with_context=prefill_metadata.chunked_context.prefill_tokens_with_context,

2317 )

2318 else:

2857 - output_prefill = output_prefill[..., : v.shape[-1]].flatten(start_dim=-2)

2319 + assert isinstance(output_prefill, torch.Tensor)

2320 + output_prefill = output_prefill.flatten(start_dim=-2)

2321 output.copy_(output_prefill)

2322

2323 @abstractmethod

vllm/platforms/interface.py

@@ -86,6 +86,9 @@ class DeviceCapability(NamedTuple):

86 return NotImplemented

87 return (self.major, self.minor) > (other.major, other.minor)

89 + def __hash__(self) -> int:

90 + return hash((self.major, self.minor))

91 +

92 def as_version_str(self) -> str:

93 return f"{self.major}.{self.minor}"

vllm/v1/attention/backends/mla/prefill/__init__.py new

+11

@@ -0,0 +1,11 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend

4 +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

5 +from vllm.v1.attention.backends.mla.prefill.selector import get_mla_prefill_backend

6 +

7 +__all__ = [

8 + "MLAPrefillBackend",

9 + "MLAPrefillBackendEnum",

10 + "get_mla_prefill_backend",

11 +]

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

+125

@@ -0,0 +1,125 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

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

4 +

5 +from abc import ABC, abstractmethod

6 +from typing import TYPE_CHECKING, ClassVar

7 +

8 +import torch

9 +

10 +if TYPE_CHECKING:

11 + from vllm.config import VllmConfig

12 + from vllm.model_executor.layers.attention.mla_attention import (

13 + MLACommonPrefillMetadata,

14 + )

15 + from vllm.platforms.interface import DeviceCapability

16 + from vllm.v1.attention.backends.mla.prefill.selector import (

17 + MLAPrefillSelectorConfig,

18 + )

19 +

20 +

21 +class MLAPrefillBackend(ABC):

22 + """Abstract base class for MLA prefill backends."""

23 +

24 + supported_dtypes: ClassVar[list[torch.dtype]] = [

25 + torch.float16,

26 + torch.bfloat16,

27 + ]

28 + requires_r1_mla_dimensions: ClassVar[bool] = False

29 +

30 + @staticmethod

31 + @abstractmethod

32 + def get_name() -> str:

33 + raise NotImplementedError

34 +

35 + @classmethod

36 + def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:

37 + return True

38 +

39 + @classmethod

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

41 + return dtype in cls.supported_dtypes

42 +

43 + @classmethod

44 + def is_available(cls) -> bool:

45 + return True

46 +

47 + @classmethod

48 + def validate_configuration(

49 + cls,

50 + device_capability: "DeviceCapability",

51 + selector_config: "MLAPrefillSelectorConfig",

52 + ) -> list[str]:

53 + invalid_reasons: list[str] = []

54 +

55 + if not cls.supports_compute_capability(device_capability):

56 + invalid_reasons.append(

57 + f"compute capability {device_capability.major}."

58 + f"{device_capability.minor} not supported"

59 + )

60 +

61 + if not cls.supports_dtype(selector_config.dtype):

62 + invalid_reasons.append(f"dtype {selector_config.dtype} not supported")

63 +

64 + if not cls.is_available():

65 + invalid_reasons.append("required dependencies not available")

66 +

67 + if cls.requires_r1_mla_dimensions and not selector_config.is_r1_compatible:

68 + invalid_reasons.append(

69 + "model does not have DeepSeek R1 MLA dimensions "

70 + "(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128)"

71 + )

72 +

73 + return invalid_reasons

74 +

75 + def __init__(

76 + self,

77 + num_heads: int,

78 + scale: float,

79 + kv_lora_rank: int,

80 + qk_nope_head_dim: int,

81 + qk_rope_head_dim: int,

82 + v_head_dim: int,

83 + vllm_config: "VllmConfig",

84 + device: torch.device,

85 + layer_names: list[str] | None = None,

86 + ) -> None:

87 + self.num_heads = num_heads

88 + self.scale = scale

89 + self.kv_lora_rank = kv_lora_rank

90 + self.qk_nope_head_dim = qk_nope_head_dim

91 + self.qk_rope_head_dim = qk_rope_head_dim

92 + self.v_head_dim = v_head_dim

93 + self.vllm_config = vllm_config

94 + self.device = device

95 + self.layer_names = layer_names

96 +

97 + def prepare_metadata( # noqa: B027

98 + self,

99 + prefill_metadata: "MLACommonPrefillMetadata",

100 + ) -> None:

101 + """Prepare backend-specific metadata before the forward pass.

102 +

103 + Called by the metadata builder after constructing the prefill metadata.

104 + """

105 + self._prefill_metadata = prefill_metadata

106 +

107 + @abstractmethod

108 + def run_prefill_new_tokens(

109 + self,

110 + q: torch.Tensor,

111 + k: torch.Tensor,

112 + v: torch.Tensor,

113 + return_softmax_lse: bool,

114 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

115 + raise NotImplementedError

116 +

117 + @abstractmethod

118 + def run_prefill_context_chunk(

119 + self,

120 + chunk_idx: int,

121 + q: torch.Tensor,

122 + k: torch.Tensor,

123 + v: torch.Tensor,

124 + ) -> tuple[torch.Tensor, torch.Tensor]:

125 + raise NotImplementedError

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

+180

@@ -0,0 +1,180 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

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

4 +

5 +import functools

6 +from typing import TYPE_CHECKING

7 +

8 +import torch

9 +

10 +import vllm.envs as envs

11 +from vllm.platforms import current_platform

12 +from vllm.v1.attention.backends.fa_utils import (

13 + get_flash_attn_version,

14 + is_flash_attn_varlen_func_available,

15 +)

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

17 +

18 +if TYPE_CHECKING:

19 + from vllm.config import VllmConfig

20 +

21 +if is_flash_attn_varlen_func_available():

22 + from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func

23 +else:

24 + flash_attn_varlen_func = None # type: ignore[assignment]

25 +

26 +

27 +class FlashAttnPrefillBackend(MLAPrefillBackend):

28 + """FlashAttention backend for MLA prefill."""

29 +

30 + @staticmethod

31 + def get_name() -> str:

32 + return "FLASH_ATTN"

33 +

34 + @classmethod

35 + def is_available(cls) -> bool:

36 + return is_flash_attn_varlen_func_available()

37 +

38 + def __init__(

39 + self,

40 + num_heads: int,

41 + scale: float,

42 + kv_lora_rank: int,

43 + qk_nope_head_dim: int,

44 + qk_rope_head_dim: int,

45 + v_head_dim: int,

46 + vllm_config: "VllmConfig",

47 + device: torch.device,

48 + layer_names: list[str] | None = None,

49 + ) -> None:

50 + super().__init__(

51 + num_heads=num_heads,

52 + scale=scale,

53 + kv_lora_rank=kv_lora_rank,

54 + qk_nope_head_dim=qk_nope_head_dim,

55 + qk_rope_head_dim=qk_rope_head_dim,

56 + v_head_dim=v_head_dim,

57 + vllm_config=vllm_config,

58 + device=device,

59 + layer_names=layer_names,

60 + )

61 +

62 + # Handle the differences between the flash_attn_varlen from

63 + # flash_attn and the one from vllm_flash_attn

64 + assert flash_attn_varlen_func is not None, (

65 + "FlashAttnPrefillBackend requires flash_attn_varlen_func. "

66 + "Ensure FlashAttnPrefillBackend.is_available() is checked first."

67 + )

68 + qk_head_dim = qk_nope_head_dim + qk_rope_head_dim

69 + self.flash_attn_varlen_func = flash_attn_varlen_func

70 + self.vllm_flash_attn_version = get_flash_attn_version(head_size=qk_head_dim)

71 + if self.vllm_flash_attn_version is not None:

72 + self.flash_attn_varlen_func = functools.partial(

73 + flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version

74 + )

75 +

76 + # Determine if we need to pad V

77 + # For MLA the v head dim is smaller than qk head dim so we pad out

78 + # v with 0s to match the qk head dim for attention backends that do

79 + # not support different headdims.

80 + # FA3 on Hopper (SM90) and FA4 natively handle diff headdims.

81 + device_capability = current_platform.get_device_capability()

82 + self.requires_v_padding = self.vllm_flash_attn_version is None or not (

83 + (

84 + self.vllm_flash_attn_version == 3

85 + and device_capability is not None

86 + and device_capability[0] == 9

87 + )

88 + or self.vllm_flash_attn_version == 4

89 + )

90 +

91 + # Track whether we're using vllm's FA or upstream (for ROCm)

92 + self._is_vllm_fa = current_platform.is_cuda() or current_platform.is_xpu()

93 +

94 + def _flash_attn_varlen_diff_headdims(

95 + self,

96 + q: torch.Tensor,

97 + k: torch.Tensor,

98 + v: torch.Tensor,

99 + return_softmax_lse: bool = False,

100 + softmax_scale: float | None = None,

101 + **kwargs,

102 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

103 + maybe_padded_v = v

104 + if self.requires_v_padding:

105 + maybe_padded_v = torch.nn.functional.pad(

106 + v, [0, q.shape[-1] - v.shape[-1]], value=0

107 + )

108 +

109 + if self._is_vllm_fa:

110 + kwargs["return_softmax_lse"] = return_softmax_lse

111 + else:

112 + # ROCm leverages the upstream flash_attn, which takes a parameter

113 + # called "return_attn_probs" instead of return_softmax_lse

114 + kwargs["return_attn_probs"] = return_softmax_lse

115 + if envs.VLLM_BATCH_INVARIANT:

116 + kwargs["num_splits"] = 1

117 +

118 + attn_out = self.flash_attn_varlen_func(

119 + q=q,

120 + k=k,

121 + v=maybe_padded_v,

122 + softmax_scale=softmax_scale,

123 + **kwargs,

124 + )

125 +

126 + # Unpack the output if there are multiple results

127 + lse = None

128 + if isinstance(attn_out, tuple):

129 + attn_out, lse = attn_out[0], attn_out[1]

130 +

131 + # Unpad output back to v_head_dim if we padded V

132 + if self.requires_v_padding:

133 + attn_out = attn_out[..., : v.shape[-1]]

134 +

135 + # Remain consistent with old `flash_attn_varlen_func` where there

136 + # is only one output tensor if `return_softmax_lse` is False.

137 + if return_softmax_lse:

138 + return attn_out, lse

139 + return attn_out

140 +

141 + def run_prefill_new_tokens(

142 + self,

143 + q: torch.Tensor,

144 + k: torch.Tensor,

145 + v: torch.Tensor,

146 + return_softmax_lse: bool,

147 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

148 + return self._flash_attn_varlen_diff_headdims(

149 + q=q,

150 + k=k,

151 + v=v,

152 + cu_seqlens_q=self._prefill_metadata.query_start_loc,

153 + cu_seqlens_k=self._prefill_metadata.query_start_loc,

154 + max_seqlen_q=self._prefill_metadata.max_query_len,

155 + max_seqlen_k=self._prefill_metadata.max_query_len,

156 + softmax_scale=self.scale,

157 + causal=True,

158 + return_softmax_lse=return_softmax_lse,

159 + )

160 +

161 + def run_prefill_context_chunk(

162 + self,

163 + chunk_idx: int,

164 + q: torch.Tensor,

165 + k: torch.Tensor,

166 + v: torch.Tensor,

167 + ) -> tuple[torch.Tensor, torch.Tensor]:

168 + assert self._prefill_metadata.chunked_context is not None

169 + return self._flash_attn_varlen_diff_headdims(

170 + q=q,

171 + k=k,

172 + v=v,

173 + cu_seqlens_q=self._prefill_metadata.query_start_loc,

174 + cu_seqlens_k=self._prefill_metadata.chunked_context.cu_seq_lens[chunk_idx],

175 + max_seqlen_q=self._prefill_metadata.max_query_len,

176 + max_seqlen_k=self._prefill_metadata.chunked_context.max_seq_lens[chunk_idx],

177 + softmax_scale=self.scale,

178 + causal=False, # Context is unmasked

179 + return_softmax_lse=True,

180 + )

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

+211

@@ -0,0 +1,211 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

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

4 +

5 +from typing import TYPE_CHECKING

6 +

7 +import torch

8 +

9 +import vllm.envs as envs

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

11 +from vllm.v1.attention.backends.utils import (

12 + get_per_layer_parameters,

13 + infer_global_hyperparameters,

14 +)

15 +from vllm.v1.worker.workspace import current_workspace_manager

16 +

17 +if TYPE_CHECKING:

18 + from vllm.config import VllmConfig

19 + from vllm.model_executor.layers.attention.mla_attention import (

20 + MLACommonPrefillMetadata,

21 + )

22 + from vllm.platforms.interface import DeviceCapability

23 +

24 +try:

25 + from flashinfer import BatchPrefillWithRaggedKVCacheWrapper

26 +except ImportError:

27 + BatchPrefillWithRaggedKVCacheWrapper = object # type: ignore[misc,assignment]

28 +

29 +_DEFAULT_NUM_CHUNKS = 32

30 +

31 +

32 +class FlashInferPrefillBackend(MLAPrefillBackend):

33 + """FlashInfer backend for MLA prefill."""

34 +

35 + requires_r1_mla_dimensions = True

36 +

37 + @staticmethod

38 + def get_name() -> str:

39 + return "FLASHINFER"

40 +

41 + @classmethod

42 + def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:

43 + return device_capability.major == 10

44 +

45 + @classmethod

46 + def is_available(cls) -> bool:

47 + try:

48 + from flashinfer import (

49 + BatchPrefillWithRaggedKVCacheWrapper, # noqa: F401

50 + )

51 +

52 + return True

53 + except ImportError:

54 + return False

55 +

56 + def __init__(

57 + self,

58 + num_heads: int,

59 + scale: float,

60 + kv_lora_rank: int,

61 + qk_nope_head_dim: int,

62 + qk_rope_head_dim: int,

63 + v_head_dim: int,

64 + vllm_config: "VllmConfig",

65 + device: torch.device,

66 + layer_names: list[str] | None = None,

67 + ) -> None:

68 + super().__init__(

69 + num_heads=num_heads,

70 + scale=scale,

71 + kv_lora_rank=kv_lora_rank,

72 + qk_nope_head_dim=qk_nope_head_dim,

73 + qk_rope_head_dim=qk_rope_head_dim,

74 + v_head_dim=v_head_dim,

75 + vllm_config=vllm_config,

76 + device=device,

77 + layer_names=layer_names,

78 + )

79 +

80 + self._prefill_main: BatchPrefillWithRaggedKVCacheWrapper | None = None

81 + self._prefill_chunks: list[BatchPrefillWithRaggedKVCacheWrapper] = []

82 + if layer_names is None:

83 + raise ValueError(

84 + "FlashInferPrefillBackend requires layer_names to "

85 + "initialize global hyperparameters."

86 + )

87 +

88 + from vllm.model_executor.layers.attention.mla_attention import (

89 + MLACommonImpl,

90 + )

91 +

92 + self._global_hyperparameters = infer_global_hyperparameters(

93 + get_per_layer_parameters(vllm_config, layer_names, MLACommonImpl) # type: ignore[type-abstract]

94 + )

95 +

96 + def _ensure_chunks(

97 + self,

98 + num_chunks: int,

99 + workspace_buffer: torch.Tensor,

100 + ) -> None:

101 + if len(self._prefill_chunks) < num_chunks:

102 + for _ in range(len(self._prefill_chunks), num_chunks):

103 + self._prefill_chunks.append(

104 + BatchPrefillWithRaggedKVCacheWrapper(

105 + workspace_buffer, "NHD", backend="cutlass"

106 + )

107 + )

108 +

109 + def prepare_metadata(

110 + self,

111 + prefill_metadata: "MLACommonPrefillMetadata",

112 + ) -> None:

113 + qo_indptr = prefill_metadata.query_start_loc

114 + has_context = prefill_metadata.chunked_context is not None

115 + (workspace_buffer,) = current_workspace_manager().get_simultaneous(

116 + ((envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), torch.uint8),

117 + )

118 +

119 + if self._prefill_main is None:

120 + self._prefill_main = BatchPrefillWithRaggedKVCacheWrapper(

121 + workspace_buffer, "NHD", backend="cutlass"

122 + )

123 + self._ensure_chunks(_DEFAULT_NUM_CHUNKS, workspace_buffer)

124 +

125 + if has_context:

126 + chunked_context = prefill_metadata.chunked_context

127 + assert chunked_context is not None

128 + num_chunks = chunked_context.cu_seq_lens.shape[0]

129 + self._ensure_chunks(num_chunks, workspace_buffer)

130 +

131 + num_qo_heads = self.num_heads

132 + num_kv_heads = num_qo_heads

133 +

134 + head_dim_qk = self.qk_nope_head_dim + self.qk_rope_head_dim

135 + head_dim_vo = self.v_head_dim

136 + kv_indptr = qo_indptr.clone()

137 +

138 + assert self._prefill_main is not None

139 + self._prefill_main.plan(

140 + qo_indptr=qo_indptr,

141 + kv_indptr=kv_indptr,

142 + num_qo_heads=num_qo_heads,

143 + num_kv_heads=num_kv_heads,

144 + head_dim_qk=head_dim_qk,

145 + head_dim_vo=head_dim_vo,

146 + causal=True,

147 + sm_scale=self._global_hyperparameters.sm_scale,

148 + window_left=self._global_hyperparameters.window_left,

149 + logits_soft_cap=self._global_hyperparameters.logits_soft_cap,

150 + q_data_type=prefill_metadata.q_data_type,

151 + o_data_type=prefill_metadata.output_dtype,

152 + )

153 +

154 + if has_context:

155 + chunked_context = prefill_metadata.chunked_context

156 + assert chunked_context is not None

157 + for i in range(num_chunks):

158 + kv_indptr_chunk = chunked_context.cu_seq_lens[i]

159 +

160 + self._prefill_chunks[i].plan(

161 + qo_indptr=qo_indptr,

162 + kv_indptr=kv_indptr_chunk,

163 + num_qo_heads=num_qo_heads,

164 + num_kv_heads=num_kv_heads,

165 + head_dim_qk=head_dim_qk,

166 + head_dim_vo=head_dim_vo,

167 + causal=False,

168 + sm_scale=self._global_hyperparameters.sm_scale,

169 + window_left=self._global_hyperparameters.window_left,

170 + logits_soft_cap=self._global_hyperparameters.logits_soft_cap,

171 + q_data_type=prefill_metadata.q_data_type,

172 + o_data_type=prefill_metadata.output_dtype,

173 + )

174 +

175 + def run_prefill_new_tokens(

176 + self,

177 + q: torch.Tensor,

178 + k: torch.Tensor,

179 + v: torch.Tensor,

180 + return_softmax_lse: bool,

181 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

182 + assert self._prefill_main is not None

183 +

184 + ret = self._prefill_main.run(

185 + q=q,

186 + k=k,

187 + v=v,

188 + return_lse=return_softmax_lse,

189 + )

190 +

191 + if isinstance(ret, tuple):

192 + # Convert from (q_len, num_heads) to (num_heads, q_len)

193 + return ret[0], ret[1].transpose(0, 1).contiguous()

194 + return ret

195 +

196 + def run_prefill_context_chunk(

197 + self,

198 + chunk_idx: int,

199 + q: torch.Tensor,

200 + k: torch.Tensor,

201 + v: torch.Tensor,

202 + ) -> tuple[torch.Tensor, torch.Tensor]:

203 + attn_out, lse = self._prefill_chunks[chunk_idx].run(

204 + q=q,

205 + k=k,

206 + v=v,

207 + return_lse=True,

208 + )

209 +

210 + # Convert from (q_len, num_heads) to (num_heads, q_len)

211 + return attn_out, lse.transpose(0, 1).contiguous()

vllm/v1/attention/backends/mla/prefill/registry.py new

+53

@@ -0,0 +1,53 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +"""Registry for MLA prefill backends.

4 +

5 +This module provides an enumeration of all available MLA prefill backends

6 +and utilities for loading them.

7 +"""

8 +

9 +from enum import Enum, EnumMeta

10 +from typing import TYPE_CHECKING

11 +

12 +from vllm.utils.import_utils import resolve_obj_by_qualname

13 +

14 +if TYPE_CHECKING:

15 + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend

16 +

17 +

18 +class _MLAPrefillBackendEnumMeta(EnumMeta):

19 + """Metaclass for MLAPrefillBackendEnum to provide better error messages."""

20 +

21 + def __getitem__(cls, name: str):

22 + try:

23 + return super().__getitem__(name)

24 + except KeyError:

25 + members = cls.__members__.keys()

26 + valid_backends = ", ".join(members)

27 + raise ValueError(

28 + f"Unknown MLA prefill backend: '{name}'. "

29 + f"Valid options are: {valid_backends}"

30 + ) from None

31 +

32 +

33 +class MLAPrefillBackendEnum(Enum, metaclass=_MLAPrefillBackendEnumMeta):

34 + """Enumeration of all supported MLA prefill backends."""

35 +

36 + FLASH_ATTN = (

37 + "vllm.v1.attention.backends.mla.prefill.flash_attn.FlashAttnPrefillBackend"

38 + )

39 + FLASHINFER = (

40 + "vllm.v1.attention.backends.mla.prefill.flashinfer.FlashInferPrefillBackend"

41 + )

42 + TRTLLM_RAGGED = (

43 + "vllm.v1.attention.backends.mla.prefill.trtllm_ragged."

44 + "TrtllmRaggedPrefillBackend"

45 + )

46 +

47 + def get_path(self) -> str:

48 + """Get the fully qualified class path for this backend."""

49 + return self.value

50 +

51 + def get_class(self) -> "type[MLAPrefillBackend]":

52 + """Lazy load and return the backend class."""

53 + return resolve_obj_by_qualname(self.get_path())

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

+183

@@ -0,0 +1,183 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

3 +"""Selector for MLA prefill backends.

4 +

5 +This module provides functions for selecting the appropriate MLA prefill

6 +backend based on device capabilities and configuration.

7 +"""

8 +

9 +from functools import cache

10 +from typing import TYPE_CHECKING, NamedTuple

11 +

12 +import torch

13 +

14 +from vllm.logger import init_logger

15 +from vllm.platforms.interface import DeviceCapability

16 +from vllm.v1.attention.backends.mla.prefill.registry import MLAPrefillBackendEnum

17 +

18 +if TYPE_CHECKING:

19 + from vllm.config import VllmConfig

20 + from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend

21 +

22 +logger = init_logger(__name__)

23 +

24 +

25 +class MLAPrefillSelectorConfig(NamedTuple):

26 + """Hashable configuration for MLA prefill backend selection.

27 +

28 + This is analogous to AttentionSelectorConfig and contains model-specific

29 + configuration needed to select an MLA prefill backend, extracted from

30 + VllmConfig into a hashable form for caching.

31 + """

32 +

33 + 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.

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

52 +

53 +

54 +def _get_mla_prefill_backend_priorities(

55 + device_capability: DeviceCapability,

56 +) -> list[MLAPrefillBackendEnum]:

57 + """Get MLA prefill backend priorities based on device capability.

58 +

59 + Args:

60 + device_capability: The device's compute capability.

61 +

62 + Returns:

63 + List of backends in priority order (highest priority first).

64 + """

65 + if device_capability.major == 10: # Blackwell

66 + return [

67 + MLAPrefillBackendEnum.FLASH_ATTN,

68 + MLAPrefillBackendEnum.TRTLLM_RAGGED,

69 + MLAPrefillBackendEnum.FLASHINFER,

70 + ]

71 + else: # Hopper (SM90) and older

72 + return [

73 + MLAPrefillBackendEnum.FLASH_ATTN,

74 + ]

75 +

76 +

77 +def get_mla_prefill_backend(

78 + vllm_config: "VllmConfig",

79 +) -> "type[MLAPrefillBackend]":

80 + """Select the MLA prefill backend based on configuration and device.

81 +

82 + This function first checks for explicit user preferences via

83 + mla_prefill_backend in AttentionConfig, then falls back to automatic

84 + priority-based selection.

85 +

86 + Args:

87 + vllm_config: The vLLM configuration.

88 +

89 + Returns:

90 + The selected prefill backend class.

91 + """

92 + from vllm.platforms import current_platform

93 +

94 + device_capability = current_platform.get_device_capability()

95 + if device_capability is None:

96 + logger.info_once(

97 + "Device capability not available, using FlashAttention MLA prefill backend."

98 + )

99 + return MLAPrefillBackendEnum.FLASH_ATTN.get_class()

100 +

101 + attention_config = vllm_config.attention_config

102 +

103 + selector_config = MLAPrefillSelectorConfig(

104 + dtype=vllm_config.model_config.dtype,

105 + is_r1_compatible=is_deepseek_r1_mla_compatible(vllm_config),

106 + )

107 +

108 + if attention_config.mla_prefill_backend is not None:

109 + selected_backend = attention_config.mla_prefill_backend

110 + backend_cls: type[MLAPrefillBackend] | None = None

111 + try:

112 + backend_cls = selected_backend.get_class()

113 + invalid_reasons = backend_cls.validate_configuration(

114 + device_capability, selector_config

115 + )

116 + except ImportError:

117 + invalid_reasons = ["ImportError"]

118 + if invalid_reasons:

119 + raise ValueError(

120 + f"Selected MLA prefill backend {selected_backend.name} "

121 + f"is not valid for this configuration. "

122 + f"Reason: {invalid_reasons}"

123 + )

124 + assert backend_cls is not None

125 + logger.info("Using %s MLA prefill backend.", selected_backend.name)

126 + return backend_cls

127 +

128 + return _auto_select_mla_prefill_backend(

129 + device_capability,

130 + selector_config,

131 + )

132 +

133 +

134 +@cache

135 +def _auto_select_mla_prefill_backend(

136 + device_capability: DeviceCapability,

137 + selector_config: MLAPrefillSelectorConfig,

138 +) -> "type[MLAPrefillBackend]":

139 + """Auto-select the best available MLA prefill backend.

140 +

141 + Args:

142 + device_capability: The device's compute capability.

143 + selector_config: Hashable configuration for backend selection.

144 +

145 + Returns:

146 + The selected prefill backend class.

147 + """

148 + priorities = _get_mla_prefill_backend_priorities(device_capability)

149 + all_invalid_reasons: dict[str, list[str]] = {}

150 +

151 + for backend_enum in priorities:

152 + backend_cls: type[MLAPrefillBackend] | None = None

153 + try:

154 + backend_cls = backend_enum.get_class()

155 + invalid_reasons = backend_cls.validate_configuration(

156 + device_capability, selector_config

157 + )

158 + except ImportError:

159 + invalid_reasons = ["ImportError"]

160 + if not invalid_reasons:

161 + assert backend_cls is not None

162 + logger.info_once("Using %s MLA prefill backend.", backend_enum.name)

163 + return backend_cls

164 + all_invalid_reasons[backend_enum.name] = invalid_reasons

165 +

166 + reasons_str = (

167 + "{"

168 + + ", ".join(

169 + f"{name}: [{', '.join(reasons)}]"

170 + for name, reasons in all_invalid_reasons.items()

171 + )

172 + + "}"

173 + )

174 + config_str = repr(selector_config)

175 + logger.debug_once(

176 + "Some MLA prefill backends are not valid with %s. Reasons: %s.",

177 + config_str,

178 + reasons_str,

179 + )

180 +

181 + raise ValueError(

182 + f"No valid MLA prefill backend found with {config_str}. Reasons: {reasons_str}."

183 + )

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

+178

@@ -0,0 +1,178 @@

1 +# SPDX-License-Identifier: Apache-2.0

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

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

4 +

5 +from typing import TYPE_CHECKING

6 +

7 +import torch

8 +

9 +import vllm.envs as envs

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

11 +from vllm.v1.worker.workspace import current_workspace_manager

12 +

13 +if TYPE_CHECKING:

14 + from vllm.config import VllmConfig

15 + from vllm.model_executor.layers.attention.mla_attention import (

16 + MLACommonPrefillMetadata,

17 + )

18 + from vllm.platforms.interface import DeviceCapability

19 +

20 +

21 +class TrtllmRaggedPrefillBackend(MLAPrefillBackend):

22 + """TRT-LLM Ragged backend for MLA prefill."""

23 +

24 + requires_r1_mla_dimensions = True

25 +

26 + @staticmethod

27 + def get_name() -> str:

28 + return "TRTLLM_RAGGED"

29 +

30 + @classmethod

31 + def supports_compute_capability(cls, device_capability: "DeviceCapability") -> bool:

32 + return device_capability.major == 10

33 +

34 + @classmethod

35 + def is_available(cls) -> bool:

36 + try:

37 + from flashinfer.prefill import (

38 + trtllm_ragged_attention_deepseek, # noqa: F401

39 + )

40 +

41 + return True

42 + except ImportError:

43 + return False

44 +

45 + def __init__(

46 + self,

47 + num_heads: int,

48 + scale: float,

49 + kv_lora_rank: int,

50 + qk_nope_head_dim: int,

51 + qk_rope_head_dim: int,

52 + v_head_dim: int,

53 + vllm_config: "VllmConfig",

54 + device: torch.device,

55 + layer_names: list[str] | None = None,

56 + ) -> None:

57 + super().__init__(

58 + num_heads=num_heads,

59 + scale=scale,

60 + kv_lora_rank=kv_lora_rank,

61 + qk_nope_head_dim=qk_nope_head_dim,

62 + qk_rope_head_dim=qk_rope_head_dim,

63 + v_head_dim=v_head_dim,

64 + vllm_config=vllm_config,

65 + device=device,

66 + layer_names=layer_names,

67 + )

68 +

69 + def _get_workspace_buffer(self) -> torch.Tensor:

70 + (workspace_buffer,) = current_workspace_manager().get_simultaneous(

71 + (

72 + (envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,),

73 + torch.uint8,

74 + ),

75 + )

76 + return workspace_buffer

77 +

78 + def prepare_metadata(

79 + self,

80 + prefill_metadata: "MLACommonPrefillMetadata",

81 + ) -> None:

82 + super().prepare_metadata(prefill_metadata)

83 + self._query_seq_lens = (

84 + prefill_metadata.query_start_loc[1:] - prefill_metadata.query_start_loc[:-1]

85 + )

86 +

87 + def run_prefill_new_tokens(

88 + self,

89 + q: torch.Tensor,

90 + k: torch.Tensor,

91 + v: torch.Tensor,

92 + return_softmax_lse: bool,

93 + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:

94 + from flashinfer.prefill import trtllm_ragged_attention_deepseek

95 +

96 + workspace_buffer = self._get_workspace_buffer()

97 + out = torch.empty(

98 + q.shape[0],

99 + q.shape[1],

100 + v.shape[2],

101 + device=q.device,

102 + dtype=self._prefill_metadata.output_dtype,

103 + )

104 +

105 + ret = trtllm_ragged_attention_deepseek(

106 + query=q,

107 + key=k,

108 + value=v,

109 + workspace_buffer=workspace_buffer,

110 + seq_lens=self._query_seq_lens,

111 + max_q_len=self._prefill_metadata.max_query_len,

112 + max_kv_len=self._prefill_metadata.max_query_len,

113 + bmm1_scale=self.scale,

114 + bmm2_scale=1.0,

115 + o_sf_scale=1.0,

116 + batch_size=self._query_seq_lens.shape[0],

117 + window_left=-1,

118 + cum_seq_lens_q=self._prefill_metadata.query_start_loc,

119 + cum_seq_lens_kv=self._prefill_metadata.query_start_loc,

120 + enable_pdl=False,

121 + is_causal=True,

122 + return_lse=return_softmax_lse,

123 + out=out,

124 + )

125 +

126 + if isinstance(ret, tuple):

127 + # Convert from (q_len, num_heads) to (num_heads, q_len)

128 + return ret[0], ret[1].transpose(0, 1).contiguous()

129 + return ret

130 +

131 + def run_prefill_context_chunk(

132 + self,

133 + chunk_idx: int,

134 + q: torch.Tensor,

135 + k: torch.Tensor,

136 + v: torch.Tensor,

137 + ) -> tuple[torch.Tensor, torch.Tensor]:

138 + from flashinfer.prefill import trtllm_ragged_attention_deepseek

139 +

140 + assert self._prefill_metadata.chunked_context is not None

141 + assert self._prefill_metadata.chunked_context.seq_lens[chunk_idx] is not None

142 + workspace_buffer = self._get_workspace_buffer()

143 +

144 + out = torch.empty(

145 + q.shape[0],

146 + q.shape[1],

147 + v.shape[2],

148 + device=q.device,

149 + dtype=self._prefill_metadata.output_dtype,

150 + )

151 +

152 + attn_out, lse = trtllm_ragged_attention_deepseek(

153 + query=q,

154 + key=k,

155 + value=v,

156 + workspace_buffer=workspace_buffer,

157 + seq_lens=self._prefill_metadata.chunked_context.seq_lens[chunk_idx],

158 + max_q_len=self._prefill_metadata.max_query_len,

159 + max_kv_len=self._prefill_metadata.chunked_context.max_seq_lens[chunk_idx],

160 + bmm1_scale=self.scale,

161 + bmm2_scale=1.0,

162 + o_sf_scale=1.0,

163 + batch_size=self._prefill_metadata.chunked_context.seq_lens[chunk_idx].shape[

164 + 0

165 + ],

166 + window_left=-1,

167 + cum_seq_lens_q=self._prefill_metadata.query_start_loc,

168 + cum_seq_lens_kv=self._prefill_metadata.chunked_context.cu_seq_lens[

169 + chunk_idx

170 + ],

171 + enable_pdl=False,

172 + is_causal=False,

173 + return_lse=True,

174 + out=out,

175 + )

176 +

177 + # Convert from (q_len, num_heads) to (num_heads, q_len)

178 + return attn_out, lse.transpose(0, 1).contiguous()

vllm/v1/attention/backends/mla/triton_mla.py

-12

@@ -123,18 +123,6 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):

123

124 self._sm_count = current_platform.num_compute_units()

125

126 - def _flash_attn_varlen_diff_headdims(

127 - self, q, k, v, return_softmax_lse=False, softmax_scale=None, **kwargs

128 - ):

129 - return super()._flash_attn_varlen_diff_headdims(

130 - q,

131 - k,

132 - v,

133 - return_softmax_lse=return_softmax_lse,

134 - softmax_scale=softmax_scale,

135 - **kwargs,

136 - )

137 -

126 def forward_mqa(

127 self,

128 q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],