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

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 8, 08:45 UTC

[Attention] Add FlashInfer Sparse MLA backend (#33451)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Signed-off-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <lwilkins@redhat.com> Co-authored-by: Lucas Wilkinson <LucasWilkinson@users.noreply.github.com>

Matthew Bonanni committed Feb 12, 2026 at 12:21 UTC f2c47886fdbabfeae7ddad871ee7889ee472d026

24 files changed +1180 -407

benchmarks/attention_benchmarks/benchmark.py

+28 -19

@@ -43,6 +43,7 @@ from common import (

43 ModelParameterSweep,

44 ParameterSweep,

45 ResultsFormatter,

46 + batch_spec_sort_key,

47 is_mla_backend,

48 )

@@ -218,10 +219,13 @@ def run_model_parameter_sweep(

219 by_param_and_spec[key].append(r)

220 break

221

221 - # Sort by param value then spec

222 + # Sort by param value then spec (batch_size, q_len, kv_len)

223 sorted_keys = sorted(

224 by_param_and_spec.keys(),

224 - key=lambda x: (int(x[0]) if x[0].isdigit() else x[0], x[1]),

225 + key=lambda x: (

226 + int(x[0]) if x[0].isdigit() else x[0],

227 + batch_spec_sort_key(x[1]),

228 + ),

229 )

230

231 current_param_value = None

@@ -330,7 +334,7 @@ def run_parameter_sweep(

334 by_spec[spec] = []

335 by_spec[spec].append(r)

336

333 - for spec in sorted(by_spec.keys()):

337 + for spec in sorted(by_spec.keys(), key=batch_spec_sort_key):

338 results = by_spec[spec]

339 best = min(results, key=lambda r: r.mean_time)

340 console.print(

@@ -496,15 +500,18 @@ def main():

500 if "description" in yaml_config:

501 console.print(f"[dim]{yaml_config['description']}[/]")

502

499 - # Override args with YAML values

500 - # (YAML takes precedence unless CLI arg was explicitly set)

501 - # Backend(s)

502 - if "backend" in yaml_config:

503 - args.backend = yaml_config["backend"]

504 - args.backends = None

505 - elif "backends" in yaml_config:

506 - args.backends = yaml_config["backends"]

507 - args.backend = None

503 + # Override args with YAML values, but CLI args take precedence

504 + # Check if CLI provided backends (they would be non-None and not default)

505 + cli_backends_provided = args.backends is not None or args.backend is not None

506 +

507 + # Backend(s) - only use YAML if CLI didn't specify

508 + if not cli_backends_provided:

509 + if "backend" in yaml_config:

510 + args.backend = yaml_config["backend"]

511 + args.backends = None

512 + elif "backends" in yaml_config:

513 + args.backends = yaml_config["backends"]

514 + args.backend = None

515

516 # Check for special modes

517 if "mode" in yaml_config:

@@ -544,13 +551,15 @@ def main():

551 args.num_kv_heads = model.get("num_kv_heads", args.num_kv_heads)

552 args.block_size = model.get("block_size", args.block_size)

553

547 - # Benchmark settings

548 - if "benchmark" in yaml_config:

549 - bench = yaml_config["benchmark"]

550 - args.device = bench.get("device", args.device)

551 - args.repeats = bench.get("repeats", args.repeats)

552 - args.warmup_iters = bench.get("warmup_iters", args.warmup_iters)

553 - args.profile_memory = bench.get("profile_memory", args.profile_memory)

554 + # Benchmark settings (top-level keys)

555 + if "device" in yaml_config:

556 + args.device = yaml_config["device"]

557 + if "repeats" in yaml_config:

558 + args.repeats = yaml_config["repeats"]

559 + if "warmup_iters" in yaml_config:

560 + args.warmup_iters = yaml_config["warmup_iters"]

561 + if "profile_memory" in yaml_config:

562 + args.profile_memory = yaml_config["profile_memory"]

563

564 # Parameter sweep configuration

565 if "parameter_sweep" in yaml_config:

benchmarks/attention_benchmarks/common.py

+63 -5

@@ -16,13 +16,32 @@ from batch_spec import get_batch_type, parse_batch_spec

16 from rich.console import Console

17 from rich.table import Table

19 +

20 +def batch_spec_sort_key(spec: str) -> tuple[int, int, int]:

21 + """

22 + Extract sorting key from batch spec: (batch_size, max_q_len, max_kv_len).

23 +

24 + This ensures results are sorted by batch size first, then query length,

25 + then sequence length, rather than alphabetically.

26 + """

27 + try:

28 + requests = parse_batch_spec(spec)

29 + batch_size = len(requests)

30 + max_q_len = max(r.q_len for r in requests) if requests else 0

31 + max_kv_len = max(r.kv_len for r in requests) if requests else 0

32 + return (batch_size, max_q_len, max_kv_len)

33 + except Exception:

34 + # Fallback for unparseable specs

35 + return (0, 0, 0)

36 +

37 +

38 # Mock classes for vLLM attention infrastructure

41 class MockHfConfig:

42 """Mock HuggingFace config that satisfies vLLM's requirements."""

25 - def __init__(self, mla_dims: dict):

44 + def __init__(self, mla_dims: dict, index_topk: int | None = None):

45 self.num_attention_heads = mla_dims["num_q_heads"]

46 self.num_key_value_heads = mla_dims["num_kv_heads"]

47 self.hidden_size = mla_dims["head_dim"] * mla_dims["num_q_heads"]

@@ -33,6 +52,8 @@ class MockHfConfig:

52 self.qk_rope_head_dim = mla_dims["qk_rope_head_dim"]

53 self.v_head_dim = mla_dims["v_head_dim"]

54 self.qk_head_dim = mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"]

55 + if index_topk is not None:

56 + self.index_topk = index_topk

58 def get_text_config(self):

59 return self

@@ -83,6 +104,38 @@ class MockKVBProj:

104 return (result,) # Return as tuple to match ColumnParallelLinear API

105

106

107 +class MockIndexer:

108 + """Mock Indexer for sparse MLA backends.

109 +

110 + Provides topk_indices_buffer that sparse MLA backends use to determine

111 + which KV cache slots to attend to for each token.

112 + """

113 +

114 + def __init__(

115 + self,

116 + max_num_tokens: int,

117 + topk_tokens: int,

118 + device: torch.device,

119 + ):

120 + self.topk_tokens = topk_tokens

121 + self.topk_indices_buffer = torch.zeros(

122 + (max_num_tokens, topk_tokens),

123 + dtype=torch.int32,

124 + device=device,

125 + )

126 +

127 + def fill_random_indices(self, num_tokens: int, max_kv_len: int):

128 + """Fill topk_indices_buffer with random valid indices for benchmarking."""

129 + indices = torch.randint(

130 + 0,

131 + max_kv_len,

132 + (num_tokens, self.topk_tokens),

133 + dtype=torch.int32,

134 + device=self.topk_indices_buffer.device,

135 + )

136 + self.topk_indices_buffer[:num_tokens] = indices

137 +

138 +

139 class MockLayer(AttentionLayerBase):

140 """Mock attention layer with scale parameters and impl.

141

@@ -327,6 +380,9 @@ class ResultsFormatter:

380 specs_order.append(spec)

381 by_spec[spec][r.config.backend] = r

382

383 + # Sort specs by (batch_size, q_len, kv_len) instead of alphabetically

384 + specs_order = sorted(by_spec.keys(), key=batch_spec_sort_key)

385 +

386 # Create shortened backend names for display

387 def shorten_backend_name(name: str) -> str:

388 """Shorten long backend names for table display."""

@@ -493,10 +549,11 @@ def get_attention_scale(head_dim: int) -> float:

549

550 def is_mla_backend(backend: str) -> bool:

551 """

496 - Check if backend is an MLA backend using the backend's is_mla() property.

552 + Check if backend is an MLA backend using the AttentionBackendEnum.

553

554 Args:

499 - backend: Backend name (e.g., "CUTLASS_MLA", "FLASHINFER_MLA")

555 + backend: Backend name matching AttentionBackendEnum exactly

556 + (e.g., "FLASHMLA_SPARSE")

557

558 Returns:

559 True if the backend is an MLA backend, False otherwise

@@ -504,7 +561,8 @@ def is_mla_backend(backend: str) -> bool:

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

562

563 try:

507 - backend_class = AttentionBackendEnum[backend.upper()].get_class()

564 + backend_enum = AttentionBackendEnum[backend]

565 + backend_class = backend_enum.get_class()

566 return backend_class.is_mla()

509 - except (KeyError, ValueError, ImportError):

567 + except (KeyError, ValueError, ImportError, AttributeError):

568 return False

benchmarks/attention_benchmarks/configs/mla_decode.yaml

+19 -10

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

3 model:

4 name: "deepseek-v3"

5 num_layers: 60

6 - num_q_heads: 128

6 + num_q_heads: 128 # Base value, can be swept for TP simulation

7 num_kv_heads: 1 # MLA uses single latent KV

8 head_dim: 576

9 kv_lora_rank: 512

@@ -12,6 +12,13 @@ model:

12 v_head_dim: 128

13 block_size: 128 # CUTLASS MLA and FlashAttn MLA use 128

15 +# Model parameter sweep: simulate tensor parallelism by varying num_q_heads

16 +# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads

17 +model_parameter_sweep:

18 + param_name: "num_q_heads"

19 + values: [128, 64, 32, 16]

20 + label_format: "{backend}_{value}h"

21 +

22 batch_specs:

23 # Small batches, varying sequence lengths

24 - "16q1s512" # 16 requests, 512 KV cache

@@ -34,28 +41,30 @@ batch_specs:

41 # Very large batches

42 - "128q1s1k" # 128 requests, 1k KV cache

43 - "128q1s2k" # 128 requests, 2k KV cache

44 + - "128q1s4k" # 128 requests, 4k KV cache

45 + - "128q1s8k" # 128 requests, 8k KV cache

47 # Long context

48 - "32q1s16k" # 32 requests, 16k KV cache

49 - "32q1s32k" # 32 requests, 32k KV cache

51 backends:

43 - - cutlass_mla

44 - - flashinfer_mla

45 - - flashattn_mla # Hopper only

46 - - flashmla # Hopper only

52 + - CUTLASS_MLA

53 + - FLASHINFER_MLA

54 + - FLASH_ATTN_MLA # Hopper only

55 + - FLASHMLA # Hopper only

57 device: "cuda:0"

49 -repeats: 5

50 -warmup_iters: 3

58 +repeats: 100

59 +warmup_iters: 10

60 profile_memory: true

62 # Backend-specific tuning

54 -cutlass_mla:

63 +CUTLASS_MLA:

64 num_kv_splits: auto # or specific value like 4, 8, 16

57 -flashattn_mla:

66 +FLASH_ATTN_MLA:

67 reorder_batch_threshold: 512

60 -flashmla:

69 +FLASHMLA:

70 reorder_batch_threshold: 1

benchmarks/attention_benchmarks/configs/mla_mixed_batch.yaml

+4 -4

@@ -45,10 +45,10 @@ batch_specs:

45 - "4q4k_60q1s4k" # 4 prefill + 60 decode

47 backends:

48 - - cutlass_mla

49 - - flashinfer_mla

50 - - flashattn_mla # Hopper only

51 - - flashmla # Hopper only

48 + - CUTLASS_MLA

49 + - FLASHINFER_MLA

50 + - FLASH_ATTN_MLA # Hopper only

51 + - FLASHMLA # Hopper only

53 device: "cuda:0"

54 repeats: 5

benchmarks/attention_benchmarks/configs/mla_prefill.yaml new

+62

@@ -0,0 +1,62 @@

1 +# MLA prefill-only benchmark configuration for sparse backends

2 +

3 +model:

4 + name: "deepseek-v3"

5 + num_layers: 60

6 + num_q_heads: 128

7 + num_kv_heads: 1

8 + head_dim: 576

9 + kv_lora_rank: 512

10 + qk_nope_head_dim: 128

11 + qk_rope_head_dim: 64

12 + v_head_dim: 128

13 + block_size: 128

14 +

15 +# Model parameter sweep: simulate tensor parallelism by varying num_q_heads

16 +# TP=1: 128 heads, TP=2: 64 heads, TP=4: 32 heads, TP=8: 16 heads

17 +model_parameter_sweep:

18 + param_name: "num_q_heads"

19 + values: [128, 64, 32, 16]

20 + label_format: "{backend}_{value}h"

21 +

22 +batch_specs:

23 + # Pure prefill

24 + - "1q512"

25 + - "1q1k"

26 + - "1q2k"

27 + - "1q4k"

28 + - "1q8k"

29 +

30 + # Batched pure prefill

31 + - "2q512"

32 + - "2q1k"

33 + - "2q2k"

34 + - "2q4k"

35 + - "2q8k"

36 + - "4q512"

37 + - "4q1k"

38 + - "4q2k"

39 + - "4q4k"

40 + - "4q8k"

41 + - "8q512"

42 + - "8q1k"

43 + - "8q2k"

44 + - "8q4k"

45 + - "8q8k"

46 +

47 + # Extend

48 + - "1q512s4k"

49 + - "1q512s8k"

50 + - "1q1ks8k"

51 + - "1q2ks8k"

52 + - "1q2ks16k"

53 + - "1q4ks16k"

54 +

55 +backends:

56 + - FLASHMLA_SPARSE

57 + - FLASHINFER_MLA_SPARSE

58 +

59 +device: "cuda:0"

60 +repeats: 10

61 +warmup_iters: 3

62 +profile_memory: true

benchmarks/attention_benchmarks/configs/reorder_threshold.yaml

+5 -6

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

6 description: "Decode vs Prefill pipeline crossover analysis"

8 # Test FlashAttn MLA

9 -backend: flashattn_mla

9 +backend: FLASH_ATTN_MLA

11 # Mode: decode_vs_prefill comparison (special sweep mode)

12 # For each batch spec, we'll test both decode and prefill pipelines

@@ -62,11 +62,10 @@ model:

62 block_size: 128

64 # Benchmark settings

65 -benchmark:

66 - device: "cuda:0"

67 - repeats: 15 # More repeats for spec decode variance

68 - warmup_iters: 5

69 - profile_memory: false

65 +device: "cuda:0"

66 +repeats: 15 # More repeats for spec decode variance

67 +warmup_iters: 5

68 +profile_memory: false

70 # Output

71 output:

benchmarks/attention_benchmarks/configs/speculative_decode.yaml

+7 -8

@@ -41,18 +41,17 @@ batch_specs:

42 # Backends that support query length > 1

43 backends:

44 - - flashattn_mla # reorder_batch_threshold = 512

45 - - flashmla # reorder_batch_threshold = 1 (tunable)

44 + - FLASH_ATTN_MLA # reorder_batch_threshold = 512

45 + - FLASHMLA # reorder_batch_threshold = 1 (tunable)

47 # FlashInfer-MLA also supports uniform spec-as-decode but with different mechanism

48 -# - flashinfer_mla

48 +# - FLASHINFER_MLA

50 # Benchmark settings

51 -benchmark:

52 - device: "cuda:0"

53 - repeats: 10 # More repeats for statistical significance

54 - warmup_iters: 5

55 - profile_memory: false

51 +device: "cuda:0"

52 +repeats: 10 # More repeats for statistical significance

53 +warmup_iters: 5

54 +profile_memory: false

56 # Test these threshold values for optimization

57 parameter_sweep:

benchmarks/attention_benchmarks/configs/standard_attention.yaml

+4 -4

@@ -36,11 +36,11 @@ batch_specs:

36 - "q1ks2k" # 1k query, 2k sequence

37 - "2q1ks4k" # 2 requests: 1k query, 4k sequence

39 -# Available backends: flash, triton, flashinfer

39 +# Available backends: FLASH_ATTN, TRITON_ATTN, FLASHINFER

40 backends:

41 - - flash

42 - - triton

43 - - flashinfer

41 + - FLASH_ATTN

42 + - TRITON_ATTN

43 + - FLASHINFER

45 device: "cuda:0"

46 repeats: 5

benchmarks/attention_benchmarks/mla_runner.py

+136 -81

@@ -8,14 +8,13 @@ This module provides helpers for running MLA backends without

8 needing full VllmConfig integration.

9 """

11 -import importlib

12 -

11 import numpy as np

12 import torch

13 from batch_spec import parse_batch_spec

14 from common import (

15 BenchmarkResult,

16 MockHfConfig,

17 + MockIndexer,

18 MockKVBProj,

19 MockLayer,

20 setup_mla_dims,

@@ -62,6 +61,7 @@ def create_minimal_vllm_config(

61 block_size: int = 128,

62 max_num_seqs: int = 256,

63 mla_dims: dict | None = None,

64 + index_topk: int | None = None,

65 ) -> VllmConfig:

66 """

67 Create minimal VllmConfig for MLA benchmarks.

@@ -73,6 +73,8 @@ def create_minimal_vllm_config(

73 max_num_seqs: Maximum number of sequences

74 mla_dims: Optional custom MLA dimensions dict. If not provided, uses

75 setup_mla_dims(model_name)

76 + index_topk: Optional topk value for sparse MLA backends. If provided,

77 + the config will include index_topk for sparse attention.

79 Returns:

80 VllmConfig for benchmarking

@@ -82,7 +84,7 @@ def create_minimal_vllm_config(

84 mla_dims = setup_mla_dims(model_name)

86 # Create mock HF config first (avoids downloading from HuggingFace)

85 - mock_hf_config = MockHfConfig(mla_dims)

87 + mock_hf_config = MockHfConfig(mla_dims, index_topk=index_topk)

89 # Create a temporary minimal config.json to avoid HF downloads

90 # This ensures consistent ModelConfig construction without network access

@@ -120,16 +122,12 @@ def create_minimal_vllm_config(

122 seed=0,

123 max_model_len=32768,

124 quantization=None,

123 - quantization_param_path=None,

125 enforce_eager=False,

125 - max_context_len_to_capture=None,

126 - max_seq_len_to_capture=8192,

126 max_logprobs=20,

127 disable_sliding_window=False,

128 skip_tokenizer_init=True,

129 served_model_name=None,

130 limit_mm_per_prompt=None,

132 - use_async_output_proc=True,

131 config_format="auto",

132 )

133 finally:

@@ -180,56 +178,65 @@ def create_minimal_vllm_config(

178 # ============================================================================

179

180

183 -# Backend name to class name prefix mapping

184 -_BACKEND_NAME_MAP = {

185 - "flashattn_mla": "FlashAttnMLA",

186 - "flashmla": "FlashMLA",

187 - "flashinfer_mla": "FlashInferMLA",

188 - "cutlass_mla": "CutlassMLA",

189 -}

190 -

191 -# Special properties that differ from defaults

181 +# Backend-specific properties that can't be inferred from the backend class

182 +# Keys are AttentionBackendEnum names (uppercase)

183 _BACKEND_PROPERTIES = {

193 - "flashmla": {

184 + "FLASHMLA": {

185 "query_format": "concat", # Single concatenated tensor (vs tuple)

195 - "block_size": 64, # FlashMLA uses fixed block size

186 },

197 - "flashinfer_mla": {

198 - "block_size": 64, # FlashInfer MLA only supports 32 or 64

187 + "FLASHMLA_SPARSE": {

188 + "query_format": "concat", # Single concatenated tensor (vs tuple)

189 },

190 }

191

192

193 def _get_backend_config(backend: str) -> dict:

194 """

205 - Get backend configuration using naming conventions.

206 -

207 - All MLA backends follow the pattern:

208 - - Module: vllm.v1.attention.backends.mla.{backend}

209 - - Impl: {Name}Impl

210 - - Metadata: {Name}Metadata (or MLACommonMetadata)

211 - - DecodeMetadata: {Name}DecodeMetadata (or MLACommonDecodeMetadata)

212 - - MetadataBuilder: {Name}MetadataBuilder

195 + Get backend configuration from AttentionBackendEnum.

196 +

197 + Uses the registry to get the backend class and extract configuration

198 + from its methods (get_impl_cls, get_builder_cls, is_sparse, etc.).

199 +

200 + Args:

201 + backend: Backend name matching AttentionBackendEnum exactly

202 + (e.g., "FLASHMLA_SPARSE")

203 +

204 + Returns:

205 + Dict with backend configuration

206 """

214 - if backend not in _BACKEND_NAME_MAP:

215 - raise ValueError(f"Unknown backend: {backend}")

207 + from vllm.v1.attention.backends.registry import AttentionBackendEnum

208

217 - name = _BACKEND_NAME_MAP[backend]

209 + try:

210 + backend_enum = AttentionBackendEnum[backend]

211 + backend_class = backend_enum.get_class()

212 + except (KeyError, ValueError) as e:

213 + valid_backends = [e.name for e in AttentionBackendEnum if e.name != "CUSTOM"]

214 + raise ValueError(

215 + f"Unknown backend: {backend}. "

216 + f"Valid MLA backends: {[b for b in valid_backends if 'MLA' in b]}"

217 + ) from e

218 +

219 + # Get block size from backend class

220 + block_sizes = backend_class.get_supported_kernel_block_sizes()

221 + # Use first supported block size (backends typically support one for MLA)

222 + block_size = block_sizes[0] if block_sizes else None

223 + if hasattr(block_size, "value"):

224 + # Handle MultipleOf enum

225 + block_size = None

226 +

227 + # Check if sparse via class method if available

228 + is_sparse = getattr(backend_class, "is_sparse", lambda: False)()

229 +

230 + # Get properties that can't be inferred

231 props = _BACKEND_PROPERTIES.get(backend, {})

232

220 - # Check if backend uses common metadata (FlashInfer, CUTLASS)

221 - uses_common = backend in ("flashinfer_mla", "cutlass_mla")

222 -

233 return {

224 - "module": f"vllm.v1.attention.backends.mla.{backend}",

225 - "impl_class": f"{name}Impl",

226 - "metadata_class": "MLACommonMetadata" if uses_common else f"{name}Metadata",

227 - "decode_metadata_class": "MLACommonDecodeMetadata"

228 - if uses_common

229 - else f"{name}DecodeMetadata",

230 - "builder_class": f"{name}MetadataBuilder",

234 + "backend_class": backend_class,

235 + "impl_class": backend_class.get_impl_cls(),

236 + "builder_class": backend_class.get_builder_cls(),

237 "query_format": props.get("query_format", "tuple"),

232 - "block_size": props.get("block_size", None),

238 + "block_size": block_size,

239 + "is_sparse": is_sparse,

240 }

241

242

@@ -447,22 +454,26 @@ def _create_backend_impl(

454 mla_dims: dict,

455 vllm_config: VllmConfig,

456 device: torch.device,

457 + max_num_tokens: int = 8192,

458 + index_topk: int | None = None,

459 ):

460 """

461 Create backend implementation instance.

462

463 Args:

455 - backend_cfg: Backend configuration dict

464 + backend_cfg: Backend configuration dict from _get_backend_config()

465 mla_dims: MLA dimension configuration

466 vllm_config: VllmConfig instance

467 device: Target device

468 + max_num_tokens: Maximum number of tokens for sparse indexer buffer

469 + index_topk: Topk value for sparse MLA backends

470

471 Returns:

461 - Tuple of (impl, layer, builder_instance)

472 + Tuple of (impl, layer, builder_instance, indexer)

473 """

463 - # Import backend classes

464 - backend_module = importlib.import_module(backend_cfg["module"])

465 - impl_class = getattr(backend_module, backend_cfg["impl_class"])

474 + # Get classes from backend config (already resolved by _get_backend_config)

475 + impl_class = backend_cfg["impl_class"]

476 + builder_class = backend_cfg["builder_class"]

477

478 # Calculate scale

479 scale = 1.0 / np.sqrt(mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"])

@@ -474,26 +485,44 @@ def _create_backend_impl(

485 v_head_dim=mla_dims["v_head_dim"],

486 )

487

488 + # Create indexer for sparse backends

489 + indexer = None

490 + if backend_cfg.get("is_sparse", False):

491 + if index_topk is None:

492 + index_topk = 2048 # Default topk for sparse MLA

493 + indexer = MockIndexer(

494 + max_num_tokens=max_num_tokens,

495 + topk_tokens=index_topk,

496 + device=device,

497 + )

498 +

499 + # Build impl kwargs

500 + impl_kwargs = {

501 + "num_heads": mla_dims["num_q_heads"],

502 + "head_size": mla_dims["head_dim"],

503 + "scale": scale,

504 + "num_kv_heads": mla_dims["num_kv_heads"],

505 + "alibi_slopes": None,

506 + "sliding_window": None,

507 + "kv_cache_dtype": "auto",

508 + "logits_soft_cap": None,

509 + "attn_type": "decoder",

510 + "kv_sharing_target_layer_name": None,

511 + "q_lora_rank": None,

512 + "kv_lora_rank": mla_dims["kv_lora_rank"],

513 + "qk_nope_head_dim": mla_dims["qk_nope_head_dim"],

514 + "qk_rope_head_dim": mla_dims["qk_rope_head_dim"],

515 + "qk_head_dim": mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],

516 + "v_head_dim": mla_dims["v_head_dim"],

517 + "kv_b_proj": mock_kv_b_proj,

518 + }

519 +

520 + # Add indexer for sparse backends

521 + if indexer is not None:

522 + impl_kwargs["indexer"] = indexer

523 +

524 # Create impl

478 - impl = impl_class(

479 - num_heads=mla_dims["num_q_heads"],

480 - head_size=mla_dims["head_dim"],

481 - scale=scale,

482 - num_kv_heads=mla_dims["num_kv_heads"],

483 - alibi_slopes=None,

484 - sliding_window=None,

485 - kv_cache_dtype="auto",

486 - logits_soft_cap=None,

487 - attn_type="decoder",

488 - kv_sharing_target_layer_name=None,

489 - q_lora_rank=None,

490 - kv_lora_rank=mla_dims["kv_lora_rank"],

491 - qk_nope_head_dim=mla_dims["qk_nope_head_dim"],

492 - qk_rope_head_dim=mla_dims["qk_rope_head_dim"],

493 - qk_head_dim=mla_dims["qk_nope_head_dim"] + mla_dims["qk_rope_head_dim"],

494 - v_head_dim=mla_dims["v_head_dim"],

495 - kv_b_proj=mock_kv_b_proj,

496 - )

525 + impl = impl_class(**impl_kwargs)

526

527 # Initialize DCP attributes

528 if not hasattr(impl, "dcp_world_size") or impl.dcp_world_size in (None, -1):

@@ -515,9 +544,7 @@ def _create_backend_impl(

544

545 # Create builder instance if needed

546 builder_instance = None

518 - if backend_cfg["builder_class"]:

519 - builder_class = getattr(backend_module, backend_cfg["builder_class"])

520 -

547 + if builder_class:

548 # Populate static_forward_context so builder can find the layer

549 # MockLayer inherits from AttentionLayerBase, so isinstance checks pass

550 vllm_config.compilation_config.static_forward_context = {"placeholder": layer}

@@ -529,7 +556,7 @@ def _create_backend_impl(

556 device=device,

557 )

558

532 - return impl, layer, builder_instance

559 + return impl, layer, builder_instance, indexer

560

561

562 # ============================================================================

@@ -594,6 +621,7 @@ def _run_single_benchmark(

621 backend_cfg: dict,

622 mla_dims: dict,

623 device: torch.device,

624 + indexer=None,

625 ) -> BenchmarkResult:

626 """

627 Run a single benchmark iteration.

@@ -606,6 +634,7 @@ def _run_single_benchmark(

634 backend_cfg: Backend configuration dict

635 mla_dims: MLA dimension configuration

636 device: Target device

637 + indexer: Optional MockIndexer for sparse backends

638

639 Returns:

640 BenchmarkResult with timing statistics

@@ -613,7 +642,9 @@ def _run_single_benchmark(

642 # Parse batch spec

643 requests = parse_batch_spec(config.batch_spec)

644 q_lens = [r.q_len for r in requests]

645 + kv_lens = [r.kv_len for r in requests]

646 total_q = sum(q_lens)

647 + max_kv_len = max(kv_lens)

648

649 # Determine block size

650 block_size = backend_cfg["block_size"] or config.block_size

@@ -641,8 +672,16 @@ def _run_single_benchmark(

672 torch.bfloat16,

673 )

674

644 - # Determine which forward method to use based on metadata

645 - if metadata.decode is not None:

675 + # Fill indexer with random indices for sparse backends

676 + is_sparse = backend_cfg.get("is_sparse", False)

677 + if is_sparse and indexer is not None:

678 + indexer.fill_random_indices(total_q, max_kv_len)

679 +

680 + # Determine which forward method to use

681 + if is_sparse:

682 + # Sparse backends use forward_mqa

683 + forward_fn = lambda: impl.forward_mqa(decode_inputs, kv_cache, metadata, layer)

684 + elif metadata.decode is not None:

685 forward_fn = lambda: impl._forward_decode(

686 decode_inputs, kv_cache, metadata, layer

687 )

@@ -693,11 +732,13 @@ def _run_single_benchmark(

732 def _run_mla_benchmark_batched(

733 backend: str,

734 configs_with_params: list[tuple], # [(config, threshold, num_splits), ...]

735 + index_topk: int = 2048,

736 ) -> list[BenchmarkResult]:

737 """

738 Unified batched MLA benchmark runner for all backends.

739

700 - Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla

740 + Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,

741 + flashinfer_mla_sparse, flashmla_sparse

742

743 This function reuses backend initialization across multiple benchmarks

744 to avoid setup/teardown overhead.

@@ -707,6 +748,7 @@ def _run_mla_benchmark_batched(

748 configs_with_params: List of (config, threshold, num_splits) tuples

749 - threshold: reorder_batch_threshold (FlashAttn/FlashMLA only)

750 - num_splits: num_kv_splits (CUTLASS only)

751 + index_topk: Topk value for sparse MLA backends (default 2048)

752

753 Returns:

754 List of BenchmarkResult objects

@@ -730,19 +772,27 @@ def _run_mla_benchmark_batched(

772 if mla_dims is None:

773 mla_dims = setup_mla_dims("deepseek-v3")

774

775 + # Determine if this is a sparse backend

776 + is_sparse = backend_cfg.get("is_sparse", False)

777 +

778 # Create and set vLLM config for MLA (reused across all benchmarks)

779 vllm_config = create_minimal_vllm_config(

780 model_name="deepseek-v3", # Used only for model path

781 block_size=block_size,

782 mla_dims=mla_dims, # Use custom dims from config or default

783 + index_topk=index_topk if is_sparse else None,

784 )

785

786 results = []

787

788 with set_current_vllm_config(vllm_config):

743 - # Create backend impl, layer, and builder (reused across benchmarks)

744 - impl, layer, builder_instance = _create_backend_impl(

745 - backend_cfg, mla_dims, vllm_config, device

789 + # Create backend impl, layer, builder, and indexer (reused across benchmarks)

790 + impl, layer, builder_instance, indexer = _create_backend_impl(

791 + backend_cfg,

792 + mla_dims,

793 + vllm_config,

794 + device,

795 + index_topk=index_topk if is_sparse else None,

796 )

797

798 # Run each benchmark with the shared impl

@@ -768,6 +818,7 @@ def _run_mla_benchmark_batched(

818 backend_cfg,

819 mla_dims,

820 device,

821 + indexer=indexer,

822 )

823 results.append(result)

824

@@ -793,20 +844,24 @@ def run_mla_benchmark(

844 config,

845 reorder_batch_threshold: int | None = None,

846 num_kv_splits: int | None = None,

847 + index_topk: int = 2048,

848 ) -> BenchmarkResult | list[BenchmarkResult]:

849 """

850 Unified MLA benchmark runner for all backends.

851

800 - Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla

852 + Works for: flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,

853 + flashinfer_mla_sparse, flashmla_sparse

854

855 Always uses batched execution internally for optimal performance.

856

857 Args:

805 - backend: Backend name (flashattn_mla, flashmla, flashinfer_mla, cutlass_mla)

858 + backend: Backend name (flashattn_mla, flashmla, flashinfer_mla, cutlass_mla,

859 + flashinfer_mla_sparse, flashmla_sparse)

860 config: BenchmarkConfig or list of (BenchmarkConfig, param) tuples

861 reorder_batch_threshold: Threshold override for FlashAttn/FlashMLA

862 (single config mode only)

863 num_kv_splits: Number of KV splits for CUTLASS (single config mode only)

864 + index_topk: Topk value for sparse MLA backends (default 2048)

865

866 Returns:

867 BenchmarkResult (single mode) or list of BenchmarkResult (batched mode)

@@ -816,9 +871,9 @@ def run_mla_benchmark(

871 # Already in batched format

872 if len(config) > 0 and isinstance(config[0], tuple):

873 # Format: [(cfg, param), ...] where param is threshold or num_splits

819 - if backend in ("flashattn_mla", "flashmla"):

874 + if backend in ("flashattn_mla", "flashmla", "flashmla_sparse"):

875 configs_with_params = [(cfg, param, None) for cfg, param in config]

821 - else: # cutlass_mla or flashinfer_mla

876 + else: # cutlass_mla, flashinfer_mla, or sparse backends

877 configs_with_params = [(cfg, None, param) for cfg, param in config]

878 else:

879 # Format: [cfg, ...] - just configs

@@ -830,7 +885,7 @@ def run_mla_benchmark(

885 return_single = True

886

887 # Use unified batched execution

833 - results = _run_mla_benchmark_batched(backend, configs_with_params)

888 + results = _run_mla_benchmark_batched(backend, configs_with_params, index_topk)

889

890 # Return single result or list based on input

891 return results[0] if return_single else results

benchmarks/attention_benchmarks/runner.py

+24 -27

@@ -40,29 +40,29 @@ from vllm.v1.kv_cache_interface import FullAttentionSpec

40 # ============================================================================

43 -_BACKEND_CONFIG = {

44 - "flash": {

45 - "module": "vllm.v1.attention.backends.flash_attn",

46 - "backend_class": "FlashAttentionBackend",

47 - },

48 - "triton": {

49 - "module": "vllm.v1.attention.backends.triton_attn",

50 - "backend_class": "TritonAttentionBackend",

51 - },

52 - "flashinfer": {

53 - "module": "vllm.v1.attention.backends.flashinfer",

54 - "backend_class": "FlashInferBackend",

55 - },

56 -}

43 +def _get_backend_config(backend: str) -> dict:

44 + """

45 + Get backend configuration from AttentionBackendEnum.

47 + Args:

48 + backend: Backend name matching AttentionBackendEnum exactly

49 + (e.g., "FLASH_ATTN", "TRITON_ATTN", "FLASHINFER")

59 -def _get_backend_config(backend: str) -> dict:

60 - if backend not in _BACKEND_CONFIG:

51 + Returns:

52 + Dict with backend_class

53 + """

54 + from vllm.v1.attention.backends.registry import AttentionBackendEnum

55 +

56 + try:

57 + backend_enum = AttentionBackendEnum[backend]

58 + backend_class = backend_enum.get_class()

59 + except (KeyError, ValueError) as e:

60 + valid_backends = [b.name for b in AttentionBackendEnum if b.name != "CUSTOM"]

61 raise ValueError(

62 - f"Unknown backend: {backend}. "

63 - f"Available: {', '.join(_BACKEND_CONFIG.keys())}"

64 - )

65 - return _BACKEND_CONFIG[backend]

62 + f"Unknown backend: {backend}. Valid backends: {valid_backends}"

63 + ) from e

64 +

65 + return {"backend_class": backend_class}

68 @contextmanager

@@ -205,10 +205,7 @@ def _create_backend_impl(

205 dtype: torch.dtype,

206 ):

207 """Create backend implementation instance."""

208 - import importlib

209 -

210 - backend_module = importlib.import_module(backend_cfg["module"])

211 - backend_class = getattr(backend_module, backend_cfg["backend_class"])

208 + backend_class = backend_cfg["backend_class"]

209

210 scale = get_attention_scale(config.head_dim)

211

@@ -247,7 +244,7 @@ def _create_metadata_builder(

244

245 # Flashinfer needs get_per_layer_parameters mocked since we don't have

246 # real model layers registered

250 - if backend_name == "flashinfer":

247 + if backend_name == "FLASHINFER":

248 import unittest.mock

249

250 from vllm.v1.attention.backends.utils import PerLayerParameters

@@ -438,7 +435,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:

435 """

436 Run standard attention benchmark with real kernels.

437

441 - Supports: flash, triton, flashinfer

438 + Supports: FLASH_ATTN, TRITON_ATTN, FLASHINFER

439

440 Args:

441 config: Benchmark configuration

@@ -453,7 +450,7 @@ def run_attention_benchmark(config: BenchmarkConfig) -> BenchmarkResult:

450

451 requests = parse_batch_spec(config.batch_spec)

452

456 - if config.backend == "flashinfer":

453 + if config.backend == "FLASHINFER":

454 requests = reorder_for_flashinfer(requests)

455

456 q_lens = [r.q_len for r in requests]

docs/design/attention_backends.md

@@ -128,6 +128,7 @@ Priority is **1 = highest** (tried first).

128 | 4 | `FLASHMLA` |

129 | 5 | `TRITON_MLA` |

130 | 6 | `FLASHMLA_SPARSE` |

131 +| 7 | `FLASHINFER_MLA_SPARSE` |

132

133 **Ampere/Hopper (SM 8.x-9.x):**

134

@@ -204,6 +205,7 @@ configuration.

205 |---------|--------|-----------|-------------|------------|------|--------|-----------|-----|-----------------|--------------|

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

207 | `FLASHINFER_MLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 32, 64 | Any | ❌ | ❌ | ❌ | ❌ | Decoder | 10.x |

208 +| `FLASHINFER_MLA_SPARSE` | fp16, bf16 | `auto`, `bfloat16` | 32, 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 10.x |

209 | `FLASHMLA` | fp16, bf16 | `auto`, `bfloat16`, `fp8`, `fp8_e4m3` | 64 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x-10.x |

210 | `FLASHMLA_SPARSE` | bf16 | `auto`, `bfloat16`, `fp8_ds_mla` | 64 | 576 | ❌ | ✅ | ❌ | ❌ | Decoder | 9.x-10.x |

211 | `FLASH_ATTN_MLA` | fp16, bf16 | `auto`, `bfloat16` | %16 | Any | ❌ | ❌ | ❌ | ✅ | Decoder | 9.x |

tests/v1/attention/test_sparse_mla_backends.py

+180 -70

@@ -1,11 +1,10 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 -"""Unit tests for the FlashMLA sparse backend utilities."""

3 +"""Unit tests for the sparse MLA backends and utilities."""

5 import math

6 from types import MethodType, SimpleNamespace

8 -import numpy as np

8 import pytest

9 import torch

@@ -25,6 +24,9 @@ from vllm.config import set_current_vllm_config

24 from vllm.model_executor.layers.linear import ColumnParallelLinear

25 from vllm.platforms import current_platform

26 from vllm.utils.math_utils import cdiv

27 +from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (

28 + FlashInferMLASparseBackend,

29 +)

30 from vllm.v1.attention.backends.mla.flashmla_sparse import (

31 FlashMLASparseBackend,

32 triton_convert_req_index_to_global_index,

@@ -156,31 +158,47 @@ def _quantize_dequantize_fp8_ds_mla(

158 return dequant_kv_c, dequant_k_pe

159

160

161 +@pytest.mark.parametrize(

162 + "backend_cls",

163 + [FlashMLASparseBackend, FlashInferMLASparseBackend],

164 + ids=["FlashMLA", "FlashInfer"],

165 +)

166 @pytest.mark.parametrize("batch_name", list(SPARSE_BACKEND_BATCH_SPECS.keys()))

160 -@pytest.mark.parametrize("kv_cache_dtype", ["fp8_ds_mla", "auto"])

167 +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_ds_mla"])

168 @pytest.mark.parametrize("tensor_parallel_size", [1, 2, 4])

162 -@pytest.mark.skipif(

163 - torch.cuda.get_device_capability() < (9, 0),

164 - reason="FlashMLASparseBackend requires CUDA 9.0 or higher",

165 -)

169 +@pytest.mark.parametrize("block_size", [32, 64])

170 def test_sparse_backend_decode_correctness(

171 default_vllm_config,

172 dist_init,

173 + backend_cls,

174 batch_name,

175 kv_cache_dtype,

176 tensor_parallel_size,

177 + block_size,

178 workspace_init,

179 ):

174 - if current_platform.is_rocm():

175 - pytest.skip("ROCm does not support fp8_ds_mla data type for kv cache.")

180 + if kv_cache_dtype not in backend_cls.supported_kv_cache_dtypes:

181 + pytest.skip(f"{backend_cls.get_name()} does not support {kv_cache_dtype}")

182

177 - if not torch.cuda.is_available():

178 - pytest.skip("CUDA is required for sparse MLA decode test")

183 + supported_block_sizes = backend_cls.get_supported_kernel_block_sizes()

184 + if block_size not in supported_block_sizes:

185 + pytest.skip(

186 + f"{backend_cls.get_name()} does not support block_size={block_size}"

187 + )

188

180 - device = torch.device("cuda")

181 - dtype = torch.bfloat16

189 + if backend_cls == FlashMLASparseBackend:

190 + ok, reason = flashmla.is_flashmla_sparse_supported()

191 + if not ok:

192 + pytest.skip(reason)

193 + elif backend_cls == FlashInferMLASparseBackend:

194 + if not current_platform.has_device_capability(100):

195 + pytest.skip("FlashInferMLASparseBackend requires SM 10.0 or higher")

196

197 batch_spec = SPARSE_BACKEND_BATCH_SPECS[batch_name]

198 + use_fp8_ds_mla_quantization = kv_cache_dtype == "fp8_ds_mla"

199 +

200 + device = torch.device("cuda")

201 + dtype = torch.bfloat16

202

203 # Model hyper-parameters (kept intentionally small for the unit test)

204 total_num_heads = 128

@@ -192,11 +210,10 @@ def test_sparse_backend_decode_correctness(

210 qk_rope_head_dim = 64

211 v_head_dim = 128

212 head_size = kv_lora_rank + qk_rope_head_dim

195 - topk_tokens = 2048

213 + topk_tokens = 128

214

215 max_seqlen = max(batch_spec.seq_lens)

216 total_cache_tokens = sum(batch_spec.seq_lens)

199 - block_size = 64

217

218 # Note: We use TP=1 to avoid multi-GPU requirements in CI.

219 # The test simulates head partitioning via mocked methods below.

@@ -247,11 +264,55 @@ def test_sparse_backend_decode_correctness(

264 seq_lens = batch_spec.seq_lens

265 query_lens = batch_spec.query_lens

266

267 + # Pre-compute positions and sparse indices for all tokens.

268 + # We need these BEFORE computing the reference to use sparse attention masks.

269 + total_query_tokens = sum(query_lens)

270 + positions = []

271 + for i in range(batch_spec.batch_size):

272 + s_len = seq_lens[i]

273 + q_len = query_lens[i]

274 + ctx_len = s_len - q_len

275 + for q_idx in range(q_len):

276 + positions.append(ctx_len + q_idx)

277 +

278 + # Create sparse indices with UNIQUE per-token offsets to catch bugs where

279 + # the kernel uses wrong indices for some tokens (e.g., due to incorrect

280 + # tensor shapes like [1, num_tokens, ...] instead of [num_tokens, 1, ...]).

281 + # Also include -1 masked indices to verify the kernel handles them correctly.

282 + sparse_indices = torch.empty(

283 + total_query_tokens, topk_tokens, dtype=torch.int32, device=device

284 + )

285 + for tok_idx in range(total_query_tokens):

286 + max_valid_idx = positions[tok_idx]

287 + offset = tok_idx * 7 # Prime number for varied offsets

288 + # Use only half the topk indices as valid, mask the rest with -1

289 + # This tests that the kernel correctly ignores -1 indices

290 + num_valid = min(topk_tokens // 2, max_valid_idx + 1)

291 + if num_valid > 0:

292 + valid_range = torch.arange(num_valid, device=device, dtype=torch.int32)

293 + tok_indices = (valid_range + offset) % (max_valid_idx + 1)

294 + # Pad with -1 for the remaining positions

295 + tok_indices = torch.cat(

296 + [

297 + tok_indices,

298 + torch.full(

299 + (topk_tokens - num_valid,), -1, device=device, dtype=torch.int32

300 + ),

301 + ]

302 + )

303 + else:

304 + tok_indices = torch.full(

305 + (topk_tokens,), -1, device=device, dtype=torch.int32

306 + )

307 + tok_indices[0] = 0 # At least one valid index

308 + sparse_indices[tok_idx] = tok_indices

309 +

310 all_q_vllm, all_kv_c_vllm, all_k_pe_vllm = [], [], []

311 kv_c_contexts, k_pe_contexts = [], []

312 reference_outputs = []

313

314 kv_cache_scale = torch.tensor(1.0, dtype=torch.float32, device=device)

315 + global_token_idx = 0

316

317 for i in range(batch_spec.batch_size):

318 s_len = seq_lens[i]

@@ -268,40 +329,53 @@ def test_sparse_backend_decode_correctness(

329 kv_c_full = torch.rand(s_len, kv_lora_rank, dtype=dtype, device=device)

330 k_pe_full = torch.rand(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device)

331

271 - # SM100 (Blackwell) uses float -> e8m0 -> bf16 scale conversion

272 - # which truncates scales to powers of 2. Simulate this in reference.

273 - is_sm100 = torch.cuda.get_device_capability()[0] >= 10

274 - kv_c_full, k_pe_full = _quantize_dequantize_fp8_ds_mla(

275 - kv_c_full,

276 - k_pe_full.squeeze(1),

277 - block_size=vllm_config.cache_config.block_size,

278 - scale=kv_cache_scale,

279 - simulate_sm100_e8m0_scales=is_sm100,

280 - )

332 + if use_fp8_ds_mla_quantization:

333 + is_sm100 = torch.cuda.get_device_capability()[0] >= 10

334 + kv_c_full, k_pe_squeezed = _quantize_dequantize_fp8_ds_mla(

335 + kv_c_full,

336 + k_pe_full.squeeze(1),

337 + block_size=block_size,

338 + scale=kv_cache_scale,

339 + simulate_sm100_e8m0_scales=is_sm100,

340 + )

341 + k_pe_full = k_pe_squeezed.unsqueeze(1)

342

343 q_nope, q_pe = q_c.split([qk_nope_head_dim, qk_rope_head_dim], dim=-1)

344 ql_nope = torch.einsum("qnh,lnh->qnl", q_nope, W_UK)

345 q_mqa = torch.cat([ql_nope, q_pe], dim=-1)

346

286 - k_mqa = torch.cat([kv_c_full, k_pe_full], dim=-1)

287 - k_mqa = k_mqa.unsqueeze(1).expand(-1, num_heads, -1)

288 - v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_heads, -1)

347 + k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1)

348 + v_mqa = kv_c_full

349

290 - attn_mask = torch.ones(q_len, s_len, dtype=torch.bool, device=device)

291 - causal_mask = torch.tril(torch.ones(q_len, q_len, device=device))

292 - attn_mask[:, ctx_len:] = causal_mask

350 + # Compute sparse SDPA reference per query token using its sparse indices

351 + for q_idx in range(q_len):

352 + tok_sparse_idx = sparse_indices[global_token_idx]

353 + valid_mask = tok_sparse_idx >= 0

354 + valid_indices = tok_sparse_idx[valid_mask].long()

355

294 - q_sdpa_in = q_mqa.unsqueeze(0).transpose(1, 2)

295 - k_sdpa_in = k_mqa.unsqueeze(0).transpose(1, 2)

296 - v_sdpa_in = v_mqa.unsqueeze(0).transpose(1, 2)

356 + q_tok = q_mqa[q_idx : q_idx + 1] # [1, num_heads, head_dim]

357 + k_sparse = k_mqa[valid_indices] # [num_valid, head_dim]

358 + v_sparse = v_mqa[valid_indices] # [num_valid, kv_lora_rank]

359

298 - sdpa_out = torch.nn.functional.scaled_dot_product_attention(

299 - q_sdpa_in, k_sdpa_in, v_sdpa_in, attn_mask=attn_mask, scale=scale

300 - )

301 - sdpa_out = sdpa_out.transpose(1, 2).squeeze(0)

360 + k_sparse = k_sparse.unsqueeze(1).expand(-1, num_heads, -1)

361 + v_sparse = v_sparse.unsqueeze(1).expand(-1, num_heads, -1)

362 +

363 + # SDPA: [1, num_heads, 1, head_dim] x [1, num_heads, num_valid, head_dim]

364 + q_sdpa_in = q_tok.unsqueeze(0).transpose(1, 2)

365 + k_sdpa_in = k_sparse.unsqueeze(0).transpose(1, 2)

366 + v_sdpa_in = v_sparse.unsqueeze(0).transpose(1, 2)

367 +

368 + sdpa_out = torch.nn.functional.scaled_dot_product_attention(

369 + q_sdpa_in, k_sdpa_in, v_sdpa_in, scale=scale

370 + )

371 + sdpa_out = sdpa_out.transpose(1, 2).squeeze(

372 + 0

373 + ) # [1, num_heads, kv_lora_rank]

374

303 - sdpa_out = torch.einsum("qnl,lnv->qnv", sdpa_out, W_UV)

304 - reference_outputs.append(sdpa_out.flatten(start_dim=-2))

375 + sdpa_out = torch.einsum("qnl,lnv->qnv", sdpa_out, W_UV)

376 + reference_outputs.append(sdpa_out.flatten(start_dim=-2))

377 +

378 + global_token_idx += 1

379

380 all_q_vllm.append(q_c)

381 all_kv_c_vllm.append(kv_c_full[ctx_len:])

@@ -334,42 +408,18 @@ def test_sparse_backend_decode_correctness(

408 num_blocks=vllm_config.cache_config.num_gpu_blocks,

409 common_attn_metadata=common_attn_metadata,

410 randomize_blocks=False,

337 - kv_cache_dtype=vllm_config.cache_config.cache_dtype,

411 + kv_cache_dtype=kv_cache_dtype if use_fp8_ds_mla_quantization else "auto",

412 scale=kv_cache_scale,

413 )

414

341 - builder_cls = FlashMLASparseBackend.get_builder_cls()

415 + builder_cls = backend_cls.get_builder_cls()

416 builder = builder_cls(kv_cache_spec, ["placeholder"], vllm_config, device)

417 metadata = builder.build(

418 common_prefix_len=0, common_attn_metadata=common_attn_metadata

419 )

420

347 - starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)

348 - seg_lengths = np.diff(starts)

349 - positions = np.arange(starts[-1], dtype=np.int32) - np.repeat(

350 - starts[:-1], seg_lengths

351 - )

352 - seq_lengths = np.asarray(common_attn_metadata.seq_lens.cpu(), dtype=np.int32)

353 - prefix_lengths = seq_lengths - seg_lengths

354 - positions += np.repeat(prefix_lengths, seg_lengths)

355 -

356 - pos_gpu = torch.as_tensor(positions, device=device, dtype=torch.int32)

357 - topk = metadata.topk_tokens

358 - debug_indices = torch.arange(topk, device=device, dtype=torch.int32).unsqueeze(0)

359 - token_positions = pos_gpu.unsqueeze(1)

360 - causal_mask = debug_indices <= token_positions

361 - debug_indices = torch.where(

362 - causal_mask, debug_indices, torch.full_like(debug_indices, -1)

363 - )

364 -

365 - # FlashMLASparseImpl now reads top-k indices from the indexer-provided

366 - # buffer, so emulate that contract with a simple namespace mock.

367 - debug_indices = debug_indices.expand(metadata.num_actual_tokens, -1).clone()

368 - mock_indexer = SimpleNamespace(topk_indices_buffer=debug_indices)

369 -

370 - ok, reason = flashmla.is_flashmla_sparse_supported()

371 - if not ok:

372 - pytest.skip(reason)

421 + # Use the pre-computed sparse_indices for the mock indexer

422 + mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices)

423

424 kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1)

425 kv_b_proj_weight = kv_b_proj_weight.view(

@@ -383,7 +433,7 @@ def test_sparse_backend_decode_correctness(

433 ).to(device=device, dtype=dtype)

434 mock_kv_b_proj.weight = torch.nn.Parameter(kv_b_proj_weight.T.contiguous())

435

386 - impl_cls = FlashMLASparseBackend.get_impl_cls()

436 + impl_cls = backend_cls.get_impl_cls()

437 with set_current_vllm_config(vllm_config):

438 impl = impl_cls(

439 num_heads=num_heads,

@@ -441,7 +491,7 @@ def test_sparse_backend_decode_correctness(

491

492 # FP8 quantization introduces some error, but should be within reasonable bounds

493 # BF16 (auto) should be very accurate, FP8 allows slightly more tolerance

444 - if kv_cache_dtype == "fp8_ds_mla":

494 + if kv_cache_dtype.startswith("fp8"):

495 torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.05, atol=0.05)

496 else:

497 torch.testing.assert_close(backend_output, sdpa_reference, rtol=0.01, atol=0.01)

@@ -636,3 +686,63 @@ def test_triton_convert_req_index_to_global_index_with_prefill_workspace(block_s

686 def test_split_prefill_chunks(seq_lens, max_buf, expected):

687 out = split_prefill_chunks(seq_lens, max_buf)

688 assert out == expected

689 +

690 +

691 +def test_triton_convert_returns_valid_counts():

692 + """Test that return_valid_counts correctly counts non-negative indices."""

693 + device = torch.device("cuda")

694 + num_tokens = 8

695 + num_requests = 2

696 + max_blocks_per_req = 10

697 + block_size = 64

698 + num_topk_tokens = 128

699 +

700 + req_id = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1], dtype=torch.int32, device=device)

701 + block_table = torch.arange(

702 + num_requests * max_blocks_per_req, dtype=torch.int32, device=device

703 + ).view(num_requests, max_blocks_per_req)

704 +

705 + # Create token indices with varying numbers of valid entries

706 + # Token 0: 64 valid, 64 invalid (-1)

707 + # Token 1: 32 valid, 96 invalid

708 + # Token 2: 128 valid (all)

709 + # Token 3: 1 valid, 127 invalid

710 + # etc.

711 + token_indices = torch.full(

712 + (num_tokens, num_topk_tokens), -1, dtype=torch.int32, device=device

713 + )

714 + expected_valid = []

715 + for i in range(num_tokens):

716 + num_valid = [64, 32, 128, 1, 64, 32, 128, 1][i]

717 + token_indices[i, :num_valid] = torch.arange(

718 + num_valid, dtype=torch.int32, device=device

719 + ) % (block_size * max_blocks_per_req)

720 + expected_valid.append(num_valid)

721 +

722 + expected_valid_tensor = torch.tensor(

723 + expected_valid, dtype=torch.int32, device=device

724 + )

725 +

726 + # Test with return_valid_counts=True

727 + result, valid_counts = triton_convert_req_index_to_global_index(

728 + req_id,

729 + block_table,

730 + token_indices,

731 + BLOCK_SIZE=block_size,

732 + NUM_TOPK_TOKENS=num_topk_tokens,

733 + return_valid_counts=True,

734 + )

735 +

736 + torch.testing.assert_close(valid_counts, expected_valid_tensor, rtol=0, atol=0)

737 +

738 + # Test that return_valid_counts=False returns only the indices

739 + result_only = triton_convert_req_index_to_global_index(

740 + req_id,

741 + block_table,

742 + token_indices,

743 + BLOCK_SIZE=block_size,

744 + NUM_TOPK_TOKENS=num_topk_tokens,

745 + return_valid_counts=False,

746 + )

747 + assert isinstance(result_only, torch.Tensor)

748 + torch.testing.assert_close(result_only, result, rtol=0, atol=0)

tools/pre_commit/generate_attention_backend_docs.py

+42 -2

@@ -901,10 +901,50 @@ def parse_cuda_priority_lists() -> dict[str, list[str]]:

901

902

903 def _get_backends_from_return(stmts: list) -> list[str]:

904 - """Extract backend names from return statements in a list of statements."""

904 + """Extract backend names from return statements in a list of statements.

905 +

906 + Handles starred unpacking (e.g. ``*sparse_backends``) by resolving the

907 + variable from assignments found in the same statement list. When the

908 + variable is conditionally assigned (inside an ``if/else``), the ``else``

909 + branch value is used as the representative default.

910 + """

911 + # Collect variable assignments so we can resolve starred expressions.

912 + # For conditional assignments, last-written (else branch) wins.

913 + var_assigns: dict[str, list[str]] = {}

914 + for stmt in stmts:

915 + if isinstance(stmt, ast.Assign) and isinstance(stmt.value, ast.List):

916 + for target in stmt.targets:

917 + if isinstance(target, ast.Name):

918 + var_assigns[target.id] = [

919 + e.attr for e in stmt.value.elts if isinstance(e, ast.Attribute)

920 + ]

921 + elif isinstance(stmt, ast.If):

922 + for branch in (stmt.body, stmt.orelse):

923 + for branch_stmt in branch:

924 + if isinstance(branch_stmt, ast.Assign) and isinstance(

925 + branch_stmt.value, ast.List

926 + ):

927 + for target in branch_stmt.targets:

928 + if isinstance(target, ast.Name):

929 + var_assigns[target.id] = [

930 + e.attr

931 + for e in branch_stmt.value.elts

932 + if isinstance(e, ast.Attribute)

933 + ]

934 +

935 for stmt in stmts:

936 if isinstance(stmt, ast.Return) and isinstance(stmt.value, ast.List):

907 - return [e.attr for e in stmt.value.elts if isinstance(e, ast.Attribute)]

937 + backends: list[str] = []

938 + for e in stmt.value.elts:

939 + if isinstance(e, ast.Attribute):

940 + backends.append(e.attr)

941 + elif (

942 + isinstance(e, ast.Starred)

943 + and isinstance(e.value, ast.Name)

944 + and e.value.id in var_assigns

945 + ):

946 + backends.extend(var_assigns[e.value.id])

947 + return backends

948 return []

949

950

vllm/model_executor/layers/attention/mla_attention.py

@@ -334,6 +334,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

334 block_size,

335 use_mla=True,

336 use_sparse=use_sparse,

337 + num_heads=self.num_heads,

338 )

339

340 if (

vllm/platforms/cpu.py

@@ -129,6 +129,7 @@ class CpuPlatform(Platform):

129 cls,

130 selected_backend: "AttentionBackendEnum",

131 attn_selector_config: "AttentionSelectorConfig",

132 + num_heads: int | None = None,

133 ) -> str:

134 if selected_backend and selected_backend != AttentionBackendEnum.CPU_ATTN:

135 logger.info("Cannot use %s backend on CPU.", selected_backend)

vllm/platforms/cuda.py

+43 -8

@@ -45,17 +45,29 @@ torch.backends.cuda.enable_cudnn_sdp(False)

45 def _get_backend_priorities(

46 use_mla: bool,

47 device_capability: DeviceCapability,

48 + num_heads: int | None = None,

49 ) -> list[AttentionBackendEnum]:

50 """Get backend priorities with lazy import to avoid circular dependency."""

51 if use_mla:

52 if device_capability.major == 10:

53 + # Prefer FlashInfer at low head counts (FlashMLA uses padding)

54 + if num_heads is not None and num_heads <= 16:

55 + sparse_backends = [

56 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE,

57 + AttentionBackendEnum.FLASHMLA_SPARSE,

58 + ]

59 + else:

60 + sparse_backends = [

61 + AttentionBackendEnum.FLASHMLA_SPARSE,

62 + AttentionBackendEnum.FLASHINFER_MLA_SPARSE,

63 + ]

64 return [

65 AttentionBackendEnum.FLASHINFER_MLA,

66 AttentionBackendEnum.CUTLASS_MLA,

67 AttentionBackendEnum.FLASH_ATTN_MLA,

68 AttentionBackendEnum.FLASHMLA,

69 AttentionBackendEnum.TRITON_MLA,

58 - AttentionBackendEnum.FLASHMLA_SPARSE,

70 + *sparse_backends,

71 ]

72 else:

73 return [

@@ -182,6 +194,8 @@ class CudaPlatformBase(Platform):

194 use_flashmla = False

195 use_cutlass_mla = False

196 use_flashinfer_mla = False

197 + use_flashmla_sparse = False

198 + use_flashinfer_mla_sparse = False

199

200 from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported

201

@@ -217,6 +231,10 @@ class CudaPlatformBase(Platform):

231 use_flashmla = backend == AttentionBackendEnum.FLASHMLA

232 use_cutlass_mla = backend == AttentionBackendEnum.CUTLASS_MLA

233 use_flashinfer_mla = backend == AttentionBackendEnum.FLASHINFER_MLA

234 + use_flashmla_sparse = backend == AttentionBackendEnum.FLASHMLA_SPARSE

235 + use_flashinfer_mla_sparse = (

236 + backend == AttentionBackendEnum.FLASHINFER_MLA_SPARSE

237 + )

238

239 if (

240 use_flashmla

@@ -242,12 +260,24 @@ class CudaPlatformBase(Platform):

260 "Forcing kv cache block size to 64 for FlashInferMLA backend."

261 )

262

245 - # TODO(Chen): remove this hacky code

246 - if use_sparse and cache_config.block_size != 64:

247 - cache_config.block_size = 64

248 - logger.info(

249 - "Forcing kv cache block size to 64 for FlashMLASparse backend."

250 - )

263 + if use_sparse:

264 + if not (use_flashmla_sparse or use_flashinfer_mla_sparse):

265 + use_flashmla_sparse = True

266 +

267 + if use_flashmla_sparse and cache_config.block_size != 64:

268 + cache_config.block_size = 64

269 + logger.info(

270 + "Forcing kv cache block size to 64 for FlashMLASparse backend."

271 + )

272 + elif use_flashinfer_mla_sparse and cache_config.block_size not in (

273 + 32,

274 + 64,

275 + ):

276 + cache_config.block_size = 64

277 + logger.info(

278 + "Forcing kv cache block size to 64 for FlashInferMLASparse "

279 + "backend."

280 + )

281

282 scheduler_config = vllm_config.scheduler_config

283 # Note: model_config may be None during testing

@@ -276,6 +306,7 @@ class CudaPlatformBase(Platform):

306 cls,

307 device_capability: DeviceCapability,

308 attn_selector_config: "AttentionSelectorConfig",

309 + num_heads: int | None = None,

310 ) -> tuple[

311 list[tuple["AttentionBackendEnum", int]],

312 dict["AttentionBackendEnum", list[str]],

@@ -284,7 +315,9 @@ class CudaPlatformBase(Platform):

315 invalid_reasons = {}

316

317 backend_priorities = _get_backend_priorities(

287 - attn_selector_config.use_mla, device_capability

318 + attn_selector_config.use_mla,

319 + device_capability,

320 + num_heads,

321 )

322 for priority, backend in enumerate(backend_priorities):

323 try:

@@ -307,6 +340,7 @@ class CudaPlatformBase(Platform):

340 cls,

341 selected_backend: "AttentionBackendEnum",

342 attn_selector_config: "AttentionSelectorConfig",

343 + num_heads: int | None = None,

344 ) -> str:

345 device_capability = cls.get_device_capability()

346 assert device_capability is not None

@@ -336,6 +370,7 @@ class CudaPlatformBase(Platform):

370 valid_backends_priorities, invalid_reasons = cls.get_valid_backends(

371 device_capability=device_capability,

372 attn_selector_config=attn_selector_config,

373 + num_heads=num_heads,

374 )

375 reasons_str = (

376 "{"

vllm/platforms/interface.py

@@ -233,6 +233,7 @@ class Platform:

233 cls,

234 selected_backend: "AttentionBackendEnum",

235 attn_selector_config: "AttentionSelectorConfig",

236 + num_heads: int | None = None,

237 ) -> str:

238 """Get the attention backend class of a device."""

239 return ""

vllm/platforms/rocm.py

@@ -265,6 +265,7 @@ class RocmPlatform(Platform):

265 cls,

266 selected_backend: "AttentionBackendEnum",

267 attn_selector_config: "AttentionSelectorConfig",

268 + num_heads: int | None = None,

269 ) -> str:

270 from vllm._aiter_ops import rocm_aiter_ops

271

vllm/platforms/xpu.py

@@ -48,6 +48,7 @@ class XPUPlatform(Platform):

48 cls,

49 selected_backend: "AttentionBackendEnum",

50 attn_selector_config: "AttentionSelectorConfig",

51 + num_heads: int | None = None,

52 ) -> str:

53 from vllm.v1.attention.backends.utils import set_kv_cache_layout

vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py new

+353

@@ -0,0 +1,353 @@

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

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

3 +"""FlashInfer MLA Sparse Attention Backend.

4 +

5 +This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k

6 +for models like DeepSeek-V3.2 that use index-based sparse attention.

7 +

8 +For sparse MLA:

9 +- block_tables shape changes from [batch_size, max_num_blocks] (dense)

10 + to [batch_size, q_len_per_request, sparse_mla_top_k] (sparse)

11 +- The sparse indices represent physical cache slot positions to attend to

12 +- sparse_mla_top_k parameter must be set to the topk value

13 +"""

14 +

15 +from dataclasses import dataclass

16 +from typing import TYPE_CHECKING, ClassVar

17 +

18 +import numpy as np

19 +import torch

20 +from flashinfer.decode import trtllm_batch_decode_with_kv_cache_mla

21 +

22 +from vllm.config import VllmConfig

23 +from vllm.config.cache import CacheDType

24 +from vllm.logger import init_logger

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

26 + get_mla_dims,

27 +)

28 +from vllm.platforms.interface import DeviceCapability

29 +from vllm.v1.attention.backend import (

30 + AttentionBackend,

31 + AttentionCGSupport,

32 + AttentionLayer,

33 + AttentionMetadata,

34 + AttentionMetadataBuilder,

35 + AttentionType,

36 + CommonAttentionMetadata,

37 + MultipleOf,

38 + SparseMLAAttentionImpl,

39 +)

40 +from vllm.v1.attention.backends.mla.sparse_utils import (

41 + triton_convert_req_index_to_global_index,

42 +)

43 +from vllm.v1.attention.backends.utils import KVCacheLayoutType

44 +from vllm.v1.kv_cache_interface import AttentionSpec

45 +

46 +if TYPE_CHECKING:

47 + from vllm.model_executor.models.deepseek_v2 import Indexer

48 +

49 +logger = init_logger(__name__)

50 +

51 +FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE = 128 * 1024 * 1024

52 +

53 +

54 +class FlashInferMLASparseBackend(AttentionBackend):

55 + """FlashInfer MLA backend with sparse attention support.

56 +

57 + This backend uses the FlashInfer TRT-LLM MLA kernel with sparse_mla_top_k

58 + for models like DeepSeek-V3.2 that use index-based sparse attention.

59 + """

60 +

61 + accept_output_buffer: bool = True

62 + supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16]

63 + supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [

64 + "auto",

65 + "bfloat16",

66 + ]

67 +

68 + @staticmethod

69 + def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:

70 + return [32, 64]

71 +

72 + @staticmethod

73 + def get_name() -> str:

74 + return "FLASHINFER_MLA_SPARSE"

75 +

76 + @staticmethod

77 + def get_impl_cls() -> type["FlashInferMLASparseImpl"]:

78 + return FlashInferMLASparseImpl

79 +

80 + @staticmethod

81 + def get_builder_cls() -> type["FlashInferMLASparseMetadataBuilder"]:

82 + return FlashInferMLASparseMetadataBuilder

83 +

84 + @classmethod

85 + def get_supported_head_sizes(cls) -> list[int]:

86 + return [576]

87 +

88 + @classmethod

89 + def is_mla(cls) -> bool:

90 + return True

91 +

92 + @classmethod

93 + def is_sparse(cls) -> bool:

94 + return True

95 +

96 + @classmethod

97 + def supports_compute_capability(cls, capability: DeviceCapability) -> bool:

98 + # FlashInfer sparse MLA targets Blackwell (SM 10.x)

99 + return capability.major == 10

100 +

101 + @classmethod

102 + def supports_combination(

103 + cls,

104 + head_size: int,

105 + dtype: torch.dtype,

106 + kv_cache_dtype: CacheDType | None,

107 + block_size: int,

108 + use_mla: bool,

109 + has_sink: bool,

110 + use_sparse: bool,

111 + device_capability: DeviceCapability,

112 + ) -> str | None:

113 + # FlashInfer MLA sparse kernel requires qk_nope_head_dim == 128

114 + from vllm.config import get_current_vllm_config

115 +

116 + vllm_config = get_current_vllm_config()

117 + if vllm_config.model_config is not None:

118 + hf_text_config = vllm_config.model_config.hf_text_config

119 + qk_nope_head_dim = getattr(hf_text_config, "qk_nope_head_dim", 1)

120 + if qk_nope_head_dim != 128:

121 + return (

122 + f"FlashInfer MLA Sparse kernel requires qk_nope_head_dim == 128, "

123 + f"but got {qk_nope_head_dim}"

124 + )

125 + # Check for index_topk which indicates sparse model

126 + if not hasattr(hf_text_config, "index_topk"):

127 + return "FlashInfer MLA Sparse requires model with index_topk config"

128 + return None

129 +

130 + @staticmethod

131 + def get_kv_cache_shape(

132 + num_blocks: int,

133 + block_size: int,

134 + num_kv_heads: int, # assumed to be 1 for MLA

135 + head_size: int,

136 + cache_dtype_str: str = "auto",

137 + ) -> tuple[int, ...]:

138 + return (num_blocks, block_size, head_size)

139 +

140 + @classmethod

141 + def get_required_kv_cache_layout(cls) -> "KVCacheLayoutType | None":

142 + return "HND"

143 +

144 +

145 +@dataclass

146 +class FlashInferMLASparseMetadata(AttentionMetadata):

147 + """Attention metadata for FlashInfer MLA Sparse backend."""

148 +

149 + num_reqs: int

150 + max_query_len: int

151 + max_seq_len: int

152 + num_actual_tokens: int

153 +

154 + # Query start locations

155 + query_start_loc: torch.Tensor

156 + slot_mapping: torch.Tensor

157 + block_table: torch.Tensor

158 + req_id_per_token: torch.Tensor

159 +

160 + # Sequence lengths for all requests (context + query)

161 + seq_lens: torch.Tensor

162 +

163 + # Sparse-specific

164 + block_size: int = 64

165 + topk_tokens: int = 2048

166 +

167 +

168 +class FlashInferMLASparseMetadataBuilder(

169 + AttentionMetadataBuilder[FlashInferMLASparseMetadata]

170 +):

171 + """Builder for FlashInfer MLA Sparse attention metadata."""

172 +

173 + _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

174 +

175 + def __init__(

176 + self,

177 + kv_cache_spec: AttentionSpec,

178 + layer_names: list[str],

179 + vllm_config: VllmConfig,

180 + device: torch.device,

181 + ) -> None:

182 + self.vllm_config = vllm_config

183 + self.layer_names = layer_names

184 + self.kv_cache_spec = kv_cache_spec

185 + self.model_config = vllm_config.model_config

186 + self.device = device

187 +

188 + self.mla_dims = get_mla_dims(self.model_config)

189 + self.topk_tokens = vllm_config.model_config.hf_config.index_topk

190 +

191 + self.req_id_per_token_buffer = torch.empty(

192 + (vllm_config.scheduler_config.max_num_batched_tokens,),

193 + dtype=torch.int32,

194 + device=device,

195 + )

196 +

197 + def build(

198 + self,

199 + common_prefix_len: int,

200 + common_attn_metadata: CommonAttentionMetadata,

201 + fast_build: bool = False,

202 + ) -> FlashInferMLASparseMetadata:

203 + cm = common_attn_metadata

204 + num_tokens = cm.num_actual_tokens

205 +

206 + # Build req_id_per_token mapping

207 + starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)

208 + seg_lengths = np.diff(starts)

209 + req_id_per_token = np.repeat(

210 + np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths

211 + )

212 +

213 + # Zero-fill for cudagraphs

214 + self.req_id_per_token_buffer.fill_(0)

215 + self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

216 + torch.from_numpy(req_id_per_token), non_blocking=True

217 + )

218 + req_id_per_token_tensor = self.req_id_per_token_buffer[:num_tokens]

219 +

220 + return FlashInferMLASparseMetadata(

221 + num_reqs=cm.num_reqs,

222 + max_query_len=cm.max_query_len,

223 + max_seq_len=cm.max_seq_len,

224 + num_actual_tokens=cm.num_actual_tokens,

225 + query_start_loc=cm.query_start_loc,

226 + slot_mapping=cm.slot_mapping,

227 + block_table=cm.block_table_tensor,

228 + req_id_per_token=req_id_per_token_tensor,

229 + seq_lens=cm.seq_lens,

230 + block_size=self.kv_cache_spec.block_size,

231 + topk_tokens=self.topk_tokens,

232 + )

233 +

234 +

235 +# Global workspace buffer (lazily initialized)

236 +_fi_sparse_workspace: torch.Tensor | None = None

237 +

238 +

239 +def _get_workspace_buffer(device: torch.device) -> torch.Tensor:

240 + global _fi_sparse_workspace

241 + if _fi_sparse_workspace is None:

242 + _fi_sparse_workspace = torch.zeros(

243 + FLASHINFER_MLA_SPARSE_WORKSPACE_BUFFER_SIZE,

244 + dtype=torch.uint8,

245 + device=device,

246 + )

247 + return _fi_sparse_workspace

248 +

249 +

250 +class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):

251 + """FlashInfer MLA Sparse implementation.

252 +

253 + Uses the TRT-LLM MLA kernel with sparse_mla_top_k parameter for

254 + sparse attention computation.

255 + """

256 +

257 + def __init__(

258 + self,

259 + num_heads: int,

260 + head_size: int,

261 + scale: float,

262 + num_kv_heads: int,

263 + alibi_slopes: list[float] | None,

264 + sliding_window: int | None,

265 + kv_cache_dtype: str,

266 + logits_soft_cap: float | None,

267 + attn_type: str,

268 + kv_sharing_target_layer_name: str | None,

269 + # MLA Specific Arguments

270 + topk_indice_buffer: torch.Tensor | None = None,

271 + indexer: "Indexer | None" = None,

272 + **mla_args,

273 + ) -> None:

274 + unsupported_features = [alibi_slopes, sliding_window, logits_soft_cap]

275 + if any(unsupported_features):

276 + raise NotImplementedError(

277 + "FlashInferMLASparseImpl does not support one of the following: "

278 + "alibi_slopes, sliding_window, logits_soft_cap"

279 + )

280 +

281 + if attn_type != AttentionType.DECODER:

282 + raise NotImplementedError(

283 + "Encoder self-attention and "

284 + "encoder/decoder cross-attention "

285 + "are not implemented for "

286 + "FlashInferMLASparseImpl"

287 + )

288 +

289 + self.num_heads = num_heads

290 + self.head_size = head_size

291 + self.scale = float(scale)

292 + self.num_kv_heads = num_kv_heads

293 + self.kv_cache_dtype = kv_cache_dtype

294 +

295 + # MLA-specific dimensions

296 + self.kv_lora_rank: int = mla_args["kv_lora_rank"]

297 + self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]

298 + self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]

299 +

300 + assert indexer is not None, "Indexer required for sparse MLA"

301 + self.topk_indices_buffer: torch.Tensor | None = indexer.topk_indices_buffer

302 +

303 + self._workspace_buffer: torch.Tensor | None = None

304 + self.bmm1_scale: float | None = None

305 + self.bmm2_scale: float | None = None

306 +

307 + def forward_mqa(

308 + self,

309 + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

310 + kv_c_and_k_pe_cache: torch.Tensor,

311 + attn_metadata: FlashInferMLASparseMetadata,

312 + layer: AttentionLayer,

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

314 + if isinstance(q, tuple):

315 + q = torch.cat(q, dim=-1)

316 +

317 + num_actual_toks = q.shape[0]

318 +

319 + assert self.topk_indices_buffer is not None

320 + topk_indices = self.topk_indices_buffer[:num_actual_toks]

321 +

322 + topk_indices_physical, seq_lens = triton_convert_req_index_to_global_index(

323 + attn_metadata.req_id_per_token[:num_actual_toks],

324 + attn_metadata.block_table,

325 + topk_indices,

326 + BLOCK_SIZE=attn_metadata.block_size,

327 + NUM_TOPK_TOKENS=topk_indices.shape[1],

328 + return_valid_counts=True,

329 + )

330 +

331 + if self._workspace_buffer is None:

332 + self._workspace_buffer = _get_workspace_buffer(q.device)

333 +

334 + if self.bmm1_scale is None:

335 + self.bmm1_scale = layer._q_scale_float * layer._k_scale_float * self.scale

336 + if self.bmm2_scale is None:

337 + self.bmm2_scale = layer._v_scale_float

338 +

339 + o = trtllm_batch_decode_with_kv_cache_mla(

340 + query=q.unsqueeze(1),

341 + kv_cache=kv_c_and_k_pe_cache.unsqueeze(1),

342 + workspace_buffer=self._workspace_buffer,

343 + qk_nope_head_dim=self.qk_nope_head_dim,

344 + kv_lora_rank=self.kv_lora_rank,

345 + qk_rope_head_dim=self.qk_rope_head_dim,

346 + block_tables=topk_indices_physical.unsqueeze(1),

347 + seq_lens=seq_lens,

348 + max_seq_len=attn_metadata.topk_tokens,

349 + bmm1_scale=self.bmm1_scale,

350 + bmm2_scale=self.bmm2_scale,

351 + sparse_mla_top_k=attn_metadata.topk_tokens,

352 + )

353 + return o.view(-1, o.shape[-2], o.shape[-1]), None

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

+3 -161

@@ -15,7 +15,6 @@ from vllm.model_executor.layers.attention.mla_attention import (

15 )

16 from vllm.platforms import current_platform

17 from vllm.platforms.interface import DeviceCapability

18 -from vllm.triton_utils import tl, triton

18 from vllm.v1.attention.backend import (

19 AttentionBackend,

20 AttentionCGSupport,

@@ -26,6 +25,9 @@ from vllm.v1.attention.backend import (

25 MultipleOf,

26 SparseMLAAttentionImpl,

27 )

28 +from vllm.v1.attention.backends.mla.sparse_utils import (

29 + triton_convert_req_index_to_global_index,

30 +)

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

32 reshape_attn_output_for_spec_decode,

33 reshape_query_for_spec_decode,

@@ -203,166 +205,6 @@ class FlashMLASparseMetadata(AttentionMetadata):

205 fp8_use_mixed_batch: bool = False

206

207

206 -# Kernel with prefill workspace support

207 -@triton.jit

208 -def _convert_req_index_to_global_index_kernel(

209 - req_id_ptr, # int32 [num_tokens]

210 - block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]

211 - token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

212 - out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

213 - prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill

214 - workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr

215 - # shapes (compile-time where possible)

216 - max_num_blocks_per_req: tl.constexpr,

217 - BLOCK_SIZE: tl.constexpr,

218 - BLOCK_N: tl.constexpr, # tile width along columns

219 - HAS_PREFILL: tl.constexpr,

220 - # strides (in elements)

221 - bt_stride0,

222 - bt_stride1,

223 - ti_stride0,

224 - ti_stride1,

225 - out_stride0,

226 - out_stride1,

227 -):

228 - # program_id(0) -> token_id (row)

229 - # program_id(1) -> tile index along columns

230 - token_id = tl.program_id(0)

231 - tile_id = tl.program_id(1)

232 -

233 - # Each program covers BLOCK_N consecutive columns

234 - indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)

235 -

236 - # Load request id for this token (no mask: grid is exact)

237 - req = tl.load(req_id_ptr + token_id)

238 -

239 - # Load token indices for this tile

240 - ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1

241 - tok = tl.load(ti_ptr) # int32

242 -

243 - # Only token == -1 should propagate as -1

244 - is_invalid_tok = tok < 0

245 - is_prefill = False

246 - if HAS_PREFILL:

247 - prefill_req_id = tl.load(prefill_request_id_ptr + token_id)

248 - is_prefill = prefill_req_id >= 0

249 - # Compute block id and in-block offset

250 - block_id = tok // BLOCK_SIZE

251 - inblock_off = tok % BLOCK_SIZE

252 -

253 - # Guard block_table access

254 - valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)

255 - bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1

256 - is_invalid_tok |= ~valid_block

257 - base = tl.load(bt_ptr, mask=valid_block & ~is_prefill, other=0)

258 - out_val = base * BLOCK_SIZE + inblock_off

259 -

260 - # Override with prefill output if prefill is enabled

261 - if HAS_PREFILL:

262 - workspace_start = tl.load(

263 - workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0

264 - )

265 - prefill_out = workspace_start + tok

266 - out_val = tl.where(is_prefill, prefill_out, out_val)

267 - out_val = tl.where(is_invalid_tok, -1, out_val)

268 -

269 - # Store results

270 - out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1

271 - tl.store(out_ptr_ij, out_val)

272 -

273 -

274 -def triton_convert_req_index_to_global_index(

275 - req_id: torch.Tensor, # int32 [num_tokens]

276 - block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req]

277 - token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS]

278 - BLOCK_SIZE: int = 64,

279 - NUM_TOPK_TOKENS: int = 2048,

280 - BLOCK_N: int = 128, # tile width along columns

281 - HAS_PREFILL_WORKSPACE: bool = False,

282 - prefill_workspace_request_ids: torch.Tensor | None = None,

283 - prefill_workspace_starts: torch.Tensor | None = None,

284 -):

285 - """

286 - out[token_id, indice_id] =

287 - block_table[req_id[token_id],

288 - token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE

289 - + token_indices[token_id, indice_id] % BLOCK_SIZE

290 -

291 - Only when token_indices[token_id, indice_id] == -1 do we output -1.

292 - For safety, we also output -1 if the derived block_id would be

293 - out-of-bounds.

294 -

295 - When HAS_PREFILL_WORKSPACE is True, prefill tokens are mapped to workspace offsets

296 - instead of global cache slots. prefill_workspace_request_ids and

297 - prefill_workspace_starts must be provided.

298 -

299 - prefill_workspace_request_ids: int32 [num_tokens], -1 for decode else

300 - prefill request index (maps to prefill_workspace_starts)

301 - prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace

302 - starts for each prefill request

303 - """

304 - assert req_id.dtype == torch.int32

305 - assert block_table.dtype == torch.int32

306 - assert token_indices.dtype == torch.int32

307 - assert token_indices.shape[1] == NUM_TOPK_TOKENS

308 - assert NUM_TOPK_TOKENS % BLOCK_N == 0, (

309 - f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})"

310 - )

311 -

312 - if HAS_PREFILL_WORKSPACE:

313 - assert prefill_workspace_request_ids is not None

314 - assert prefill_workspace_starts is not None

315 - assert prefill_workspace_request_ids.dtype == torch.int32

316 - assert prefill_workspace_starts.dtype == torch.int32

317 -

318 - num_tokens = req_id.shape[0]

319 - max_num_blocks_per_req = block_table.shape[1]

320 - tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N

321 -

322 - # Ensure contiguous tensors on the same device

323 - req_id_c = req_id.contiguous()

324 - block_table_c = block_table.contiguous()

325 - token_indices_c = token_indices.contiguous()

326 - out = torch.empty_like(token_indices_c)

327 -

328 - # Strides in elements

329 - bt_stride0, bt_stride1 = block_table_c.stride()

330 - ti_stride0, ti_stride1 = token_indices_c.stride()

331 - out_stride0, out_stride1 = out.stride()

332 -

333 - # Prepare prefill pointers

334 - if HAS_PREFILL_WORKSPACE:

335 - assert prefill_workspace_request_ids is not None # for mypy

336 - assert prefill_workspace_starts is not None # for mypy

337 - assert prefill_workspace_request_ids.is_contiguous()

338 - assert prefill_workspace_starts.is_contiguous()

339 -

340 - # Exact 2D grid: tokens × column tiles

341 - grid = (num_tokens, tiles_per_row)

342 -

343 - _convert_req_index_to_global_index_kernel[grid](

344 - req_id_c,

345 - block_table_c,

346 - token_indices_c,

347 - out,

348 - prefill_workspace_request_ids,

349 - prefill_workspace_starts,

350 - # shapes / constexprs

351 - max_num_blocks_per_req,

352 - BLOCK_SIZE,

353 - BLOCK_N,

354 - HAS_PREFILL_WORKSPACE,

355 - # strides

356 - bt_stride0,

357 - bt_stride1,

358 - ti_stride0,

359 - ti_stride1,

360 - out_stride0,

361 - out_stride1,

362 - )

363 - return out

364 -

365 -

208 def get_prefill_workspace_size(max_model_len: int):

209 # NOTE(Lucas): 5 is a magic number for controlling the prefill buffer size.

210 # May be tuned later.

vllm/v1/attention/backends/mla/sparse_utils.py new

+191

@@ -0,0 +1,191 @@

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

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

3 +"""Utility functions for sparse MLA backends."""

4 +

5 +import torch

6 +

7 +from vllm.triton_utils import tl, triton

8 +

9 +

10 +# Kernel with prefill workspace support and valid count tracking

11 +@triton.jit

12 +def _convert_req_index_to_global_index_kernel(

13 + req_id_ptr, # int32 [num_tokens]

14 + block_table_ptr, # int32 [num_requests, max_num_blocks_per_req]

15 + token_indices_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

16 + out_ptr, # int32 [num_tokens, NUM_TOPK_TOKENS]

17 + valid_count_ptr, # int32 [num_tokens] - output valid count per row

18 + prefill_request_id_ptr, # int32 [num_tokens], -1 for decode, >=0 for prefill

19 + workspace_starts_ptr, # int32 [num_prefill_reqs+1] or nullptr

20 + # shapes (compile-time where possible)

21 + max_num_blocks_per_req: tl.constexpr,

22 + BLOCK_SIZE: tl.constexpr,

23 + BLOCK_N: tl.constexpr, # tile width along columns

24 + HAS_PREFILL: tl.constexpr,

25 + COUNT_VALID: tl.constexpr, # whether to count valid indices

26 + # strides (in elements)

27 + bt_stride0,

28 + bt_stride1,

29 + ti_stride0,

30 + ti_stride1,

31 + out_stride0,

32 + out_stride1,

33 +):

34 + # program_id(0) -> token_id (row)

35 + # program_id(1) -> tile index along columns

36 + token_id = tl.program_id(0)

37 + tile_id = tl.program_id(1)

38 +

39 + # Each program covers BLOCK_N consecutive columns

40 + indice_id = tile_id * BLOCK_N + tl.arange(0, BLOCK_N)

41 +

42 + # Load request id for this token (no mask: grid is exact)

43 + req = tl.load(req_id_ptr + token_id)

44 +

45 + # Load token indices for this tile

46 + ti_ptr = token_indices_ptr + token_id * ti_stride0 + indice_id * ti_stride1

47 + tok = tl.load(ti_ptr) # int32

48 +

49 + # Only token == -1 should propagate as -1

50 + is_invalid_tok = tok < 0

51 + is_prefill = False

52 + if HAS_PREFILL:

53 + prefill_req_id = tl.load(prefill_request_id_ptr + token_id)

54 + is_prefill = prefill_req_id >= 0

55 + # Compute block id and in-block offset

56 + block_id = tok // BLOCK_SIZE

57 + inblock_off = tok % BLOCK_SIZE

58 +

59 + # Guard block_table access

60 + valid_block = (block_id < max_num_blocks_per_req) & (block_id >= 0)

61 + bt_ptr = block_table_ptr + req * bt_stride0 + block_id * bt_stride1

62 + is_invalid_tok |= ~valid_block

63 + base = tl.load(bt_ptr, mask=valid_block & ~is_prefill, other=0)

64 + out_val = base * BLOCK_SIZE + inblock_off

65 +

66 + # Override with prefill output if prefill is enabled

67 + if HAS_PREFILL:

68 + workspace_start = tl.load(

69 + workspace_starts_ptr + prefill_req_id, mask=is_prefill, other=0

70 + )

71 + prefill_out = workspace_start + tok

72 + out_val = tl.where(is_prefill, prefill_out, out_val)

73 + out_val = tl.where(is_invalid_tok, -1, out_val)

74 +

75 + # Store results

76 + out_ptr_ij = out_ptr + token_id * out_stride0 + indice_id * out_stride1

77 + tl.store(out_ptr_ij, out_val)

78 +

79 + # Count valid indices in this tile and atomically add to row total

80 + if COUNT_VALID:

81 + tile_valid_count = tl.sum((~is_invalid_tok).to(tl.int32))

82 + tl.atomic_add(valid_count_ptr + token_id, tile_valid_count)

83 +

84 +

85 +def triton_convert_req_index_to_global_index(

86 + req_id: torch.Tensor, # int32 [num_tokens]

87 + block_table: torch.Tensor, # int32 [num_requests, max_num_blocks_per_req]

88 + token_indices: torch.Tensor, # int32 [num_tokens, NUM_TOPK_TOKENS]

89 + BLOCK_SIZE: int = 64,

90 + NUM_TOPK_TOKENS: int = 2048,

91 + BLOCK_N: int = 128, # tile width along columns

92 + HAS_PREFILL_WORKSPACE: bool = False,

93 + prefill_workspace_request_ids: torch.Tensor | None = None,

94 + prefill_workspace_starts: torch.Tensor | None = None,

95 + return_valid_counts: bool = False,

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

97 + """

98 + out[token_id, indice_id] =

99 + block_table[req_id[token_id],

100 + token_indices[token_id, indice_id] // BLOCK_SIZE] * BLOCK_SIZE

101 + + token_indices[token_id, indice_id] % BLOCK_SIZE

102 +

103 + Only when token_indices[token_id, indice_id] == -1 do we output -1.

104 + For safety, we also output -1 if the derived block_id would be

105 + out-of-bounds.

106 +

107 + When HAS_PREFILL_WORKSPACE is True, prefill tokens are mapped to workspace offsets

108 + instead of global cache slots. prefill_workspace_request_ids and

109 + prefill_workspace_starts must be provided.

110 +

111 + prefill_workspace_request_ids: int32 [num_tokens], -1 for decode else

112 + prefill request index (maps to prefill_workspace_starts)

113 + prefill_workspace_starts: int32 [num_prefills], 0-indexed workspace

114 + starts for each prefill request

115 +

116 + When return_valid_counts is True, also returns the count of valid (non -1)

117 + indices per row, computed during the same kernel pass (no extra overhead).

118 + """

119 + assert req_id.dtype == torch.int32

120 + assert block_table.dtype == torch.int32

121 + assert token_indices.dtype == torch.int32

122 + assert token_indices.shape[1] == NUM_TOPK_TOKENS

123 + assert NUM_TOPK_TOKENS % BLOCK_N == 0, (

124 + f"NUM_TOPK_TOKENS ({NUM_TOPK_TOKENS}) must be divisible by BLOCK_N ({BLOCK_N})"

125 + )

126 +

127 + if HAS_PREFILL_WORKSPACE:

128 + assert prefill_workspace_request_ids is not None

129 + assert prefill_workspace_starts is not None

130 + assert prefill_workspace_request_ids.dtype == torch.int32

131 + assert prefill_workspace_starts.dtype == torch.int32

132 +

133 + num_tokens = req_id.shape[0]

134 + max_num_blocks_per_req = block_table.shape[1]

135 + tiles_per_row = NUM_TOPK_TOKENS // BLOCK_N

136 +

137 + # Ensure contiguous tensors on the same device

138 + req_id_c = req_id.contiguous()

139 + block_table_c = block_table.contiguous()

140 + token_indices_c = token_indices.contiguous()

141 + out = torch.empty_like(token_indices_c)

142 +

143 + # Allocate valid count buffer if needed (must be zero-initialized for atomics)

144 + valid_counts: torch.Tensor | None = None

145 + if return_valid_counts:

146 + valid_counts = torch.zeros(

147 + num_tokens, dtype=torch.int32, device=token_indices.device

148 + )

149 +

150 + # Strides in elements

151 + bt_stride0, bt_stride1 = block_table_c.stride()

152 + ti_stride0, ti_stride1 = token_indices_c.stride()

153 + out_stride0, out_stride1 = out.stride()

154 +

155 + # Prepare prefill pointers

156 + if HAS_PREFILL_WORKSPACE:

157 + assert prefill_workspace_request_ids is not None # for mypy

158 + assert prefill_workspace_starts is not None # for mypy

159 + assert prefill_workspace_request_ids.is_contiguous()

160 + assert prefill_workspace_starts.is_contiguous()

161 +

162 + # Exact 2D grid: tokens × column tiles

163 + grid = (num_tokens, tiles_per_row)

164 +

165 + _convert_req_index_to_global_index_kernel[grid](

166 + req_id_c,

167 + block_table_c,

168 + token_indices_c,

169 + out,

170 + valid_counts,

171 + prefill_workspace_request_ids,

172 + prefill_workspace_starts,

173 + # shapes / constexprs

174 + max_num_blocks_per_req,

175 + BLOCK_SIZE,

176 + BLOCK_N,

177 + HAS_PREFILL_WORKSPACE,

178 + return_valid_counts,

179 + # strides

180 + bt_stride0,

181 + bt_stride1,

182 + ti_stride0,

183 + ti_stride1,

184 + out_stride0,

185 + out_stride1,

186 + )

187 +

188 + if return_valid_counts:

189 + assert valid_counts is not None

190 + return out, valid_counts

191 + return out

vllm/v1/attention/backends/registry.py

@@ -62,6 +62,10 @@ class AttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):

62 FLASHINFER_MLA = (

63 "vllm.v1.attention.backends.mla.flashinfer_mla.FlashInferMLABackend"

64 )

65 + FLASHINFER_MLA_SPARSE = (

66 + "vllm.v1.attention.backends.mla.flashinfer_mla_sparse."

67 + "FlashInferMLASparseBackend"

68 + )

69 TRITON_MLA = "vllm.v1.attention.backends.mla.triton_mla.TritonMLABackend"

70 CUTLASS_MLA = "vllm.v1.attention.backends.mla.cutlass_mla.CutlassMLABackend"

71 FLASHMLA = "vllm.v1.attention.backends.mla.flashmla.FlashMLABackend"

vllm/v1/attention/selector.py

+5 -2

@@ -53,6 +53,7 @@ def get_attn_backend(

53 use_sparse: bool = False,

54 use_mm_prefix: bool = False,

55 attn_type: str | None = None,

56 + num_heads: int | None = None,

57 ) -> type[AttentionBackend]:

58 """Selects which attention backend to use and lazily imports it."""

@@ -66,7 +67,6 @@ def get_attn_backend(

67 from vllm.config import get_current_vllm_config

69 vllm_config = get_current_vllm_config()

69 - backend_enum = vllm_config.attention_config.backend

71 attn_selector_config = AttentionSelectorConfig(

72 head_size=head_size,

@@ -81,8 +81,9 @@ def get_attn_backend(

81 )

83 return _cached_get_attn_backend(

84 - backend=backend_enum,

84 + backend=vllm_config.attention_config.backend,

85 attn_selector_config=attn_selector_config,

86 + num_heads=num_heads,

87 )

@@ -90,12 +91,14 @@ def get_attn_backend(

91 def _cached_get_attn_backend(

92 backend,

93 attn_selector_config: AttentionSelectorConfig,

94 + num_heads: int | None = None,

95 ) -> type[AttentionBackend]:

96 from vllm.platforms import current_platform

98 attention_cls = current_platform.get_attn_backend_cls(

99 backend,

100 attn_selector_config=attn_selector_config,

101 + num_heads=num_heads,

102 )

103 if not attention_cls:

104 raise ValueError(