@tduncan / Nvidia-Optimization-Stack-VLLM / commits / c4f5cd60da

[1/N] Add dense MHA path for sparse MLA short sequences (#47327)

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com> Co-authored-by: Claude Opus 4.8 <noreply@anthropic.com>

Matthew Bonanni committed Jul 13, 2026 at 20:29 UTC c4f5cd60dae386d106c9b8a12dbab24e2e9dda0b

21 files changed +1991 -647

benchmarks/attention_benchmarks/benchmark.py

+175 -3

@@ -75,7 +75,11 @@ def run_mla_benchmark(config: BenchmarkConfig, **kwargs) -> BenchmarkResult:

75 from mla_runner import run_mla_benchmark as run_mla

77 return run_mla(

78 - config.backend, config, prefill_backend=config.prefill_backend, **kwargs

78 + config.backend,

79 + config,

80 + prefill_backend=config.prefill_backend,

81 + sparse_mla_force_mqa=config.sparse_mla_force_mqa,

82 + **kwargs,

83 )

@@ -592,6 +596,30 @@ def main():

596 default="profile",

597 help="Output file name for ncu profile (default: 'profile').",

598 )

599 + parser.add_argument(

600 + "--torch-profile",

601 + action="store_true",

602 + default=False,

603 + help="Collect a PyTorch profiler Chrome trace for each benchmark run.",

604 + )

605 + parser.add_argument(

606 + "--torch-profile-dir",

607 + default=None,

608 + help="Directory for PyTorch profiler traces.",

609 + )

610 + parser.add_argument(

611 + "--torch-profile-iters",

612 + type=int,

613 + default=3,

614 + help="Number of forward passes to record per PyTorch profiler trace.",

615 + )

616 + parser.add_argument(

617 + "--sparse-mla-mha-variants",

618 + nargs="+",

619 + default=None,

620 + choices=["dense_mha", "mqa"],

621 + help="Sparse MLA variants to run in mha_vs_mqa mode. Defaults to both.",

622 + )

623

624 # Parameter sweep (use YAML config for advanced sweeps)

625 parser.add_argument(

@@ -641,6 +669,7 @@ def main():

669

670 # Prefill backends (e.g., ["fa3", "fa4"])

671 args.prefill_backends = yaml_config.get("prefill_backends", None)

672 + args.prefill_backend = yaml_config.get("prefill_backend", None)

673

674 # FP8 output benchmark knobs; CLI wins.

675 if args.fp8_output_scale is None:

@@ -683,6 +712,9 @@ def main():

712 args.num_q_heads = model.get("num_q_heads", args.num_q_heads)

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

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

715 + args.max_model_len = model.get(

716 + "max_model_len", getattr(args, "max_model_len", None)

717 + )

718 # MLA-specific dimensions

719 args.kv_lora_rank = model.get("kv_lora_rank", args.kv_lora_rank)

720 args.qk_nope_head_dim = model.get("qk_nope_head_dim", args.qk_nope_head_dim)

@@ -701,6 +733,21 @@ def main():

733 args.cuda_graphs = yaml_config["cuda_graphs"]

734 if "ncu_profile" in yaml_config:

735 args.ncu_profile = yaml_config["ncu_profile"]

736 + if "torch_profile" in yaml_config:

737 + args.torch_profile = yaml_config["torch_profile"]

738 + if "torch_profile_dir" in yaml_config:

739 + args.torch_profile_dir = yaml_config["torch_profile_dir"]

740 + if "torch_profile_iters" in yaml_config:

741 + args.torch_profile_iters = yaml_config["torch_profile_iters"]

742 + args.sparse_mla_topk_pattern = yaml_config.get(

743 + "sparse_mla_topk_pattern", "random"

744 + )

745 + args.sparse_mla_dense_mha_max_seq_len = yaml_config.get(

746 + "sparse_mla_dense_mha_max_seq_len", None

747 + )

748 + args.sparse_mla_mha_variants = yaml_config.get(

749 + "sparse_mla_mha_variants", args.sparse_mla_mha_variants

750 + )

751

752 # Parameter sweep configuration

753 if "parameter_sweep" in yaml_config:

@@ -842,8 +889,6 @@ def main():

889 num_kv_heads=args.num_kv_heads,

890 block_size=args.block_size,

891 device=args.device,

845 - repeats=args.repeats,

846 - warmup_iters=args.warmup_iters,

892 profile_memory=args.profile_memory,

893 kv_cache_dtype=args.kv_cache_dtype,

894 use_cuda_graphs=args.cuda_graphs,

@@ -1063,6 +1108,133 @@ def main():

1108 f"\n [yellow]Prefill always faster for batch_size={bs}[/]"

1109 )

1110

1111 + # Handle MHA vs MQA comparison mode for sparse MLA

1112 + elif hasattr(args, "mode") and args.mode == "mha_vs_mqa":

1113 + console.print("[yellow]Mode: MHA vs MQA comparison for sparse MLA[/]")

1114 +

1115 + sparse_mla_topk_pattern = getattr(args, "sparse_mla_topk_pattern", "random")

1116 + dense_mha_max_seq_len = getattr(args, "sparse_mla_dense_mha_max_seq_len", None)

1117 + prefill_backend = getattr(args, "prefill_backend", None)

1118 + if prefill_backend:

1119 + console.print(f"Prefill backend: {prefill_backend}")

1120 + available_variants = [

1121 + ("dense_mha", False, "dense"),

1122 + ("mqa", True, "auto"),

1123 + ]

1124 + requested_variants = getattr(args, "sparse_mla_mha_variants", None)

1125 + if requested_variants is not None:

1126 + valid_variants = {label for label, _, _ in available_variants}

1127 + invalid_variants = sorted(set(requested_variants) - valid_variants)

1128 + if invalid_variants:

1129 + raise ValueError(

1130 + "Invalid sparse_mla_mha_variants entries: "

1131 + f"{invalid_variants}. Valid variants are: "

1132 + f"{sorted(valid_variants)}"

1133 + )

1134 + requested_variant_set = set(requested_variants)

1135 + variants = [

1136 + variant

1137 + for variant in available_variants

1138 + if variant[0] in requested_variant_set

1139 + ]

1140 + else:

1141 + variants = available_variants

1142 + formatter = ResultsFormatter(console)

1143 + total = 0

1144 + for spec in args.batch_specs:

1145 + q_len = max(request.q_len for request in parse_batch_spec(spec))

1146 + for variant_label, _, _ in variants:

1147 + if (

1148 + variant_label == "dense_mha"

1149 + and dense_mha_max_seq_len is not None

1150 + and q_len > dense_mha_max_seq_len

1151 + ):

1152 + continue

1153 + total += len(backends)

1154 +

1155 + with tqdm(total=total, desc="Benchmarking") as pbar:

1156 + for spec in args.batch_specs:

1157 + q_len = max(request.q_len for request in parse_batch_spec(spec))

1158 + for backend in backends:

1159 + for variant_label, force_mqa, mha_mode in variants:

1160 + if (

1161 + variant_label == "dense_mha"

1162 + and dense_mha_max_seq_len is not None

1163 + and q_len > dense_mha_max_seq_len

1164 + ):

1165 + continue

1166 + config = BenchmarkConfig(

1167 + backend=f"{backend}_{variant_label}",

1168 + batch_spec=spec,

1169 + num_layers=args.num_layers,

1170 + head_dim=args.head_dim,

1171 + num_q_heads=args.num_q_heads,

1172 + num_kv_heads=args.num_kv_heads,

1173 + block_size=args.block_size,

1174 + device=args.device,

1175 + max_model_len=getattr(args, "max_model_len", None),

1176 + kv_cache_dtype=args.kv_cache_dtype,

1177 + profile_memory=args.profile_memory,

1178 + use_cuda_graphs=args.cuda_graphs,

1179 + ncu_profile=args.ncu_profile,

1180 + torch_profile=args.torch_profile,

1181 + torch_profile_dir=args.torch_profile_dir,

1182 + torch_profile_iters=args.torch_profile_iters,

1183 + warmup_ms=args.warmup_ms,

1184 + kv_lora_rank=getattr(args, "kv_lora_rank", None),

1185 + qk_nope_head_dim=getattr(args, "qk_nope_head_dim", None),

1186 + qk_rope_head_dim=getattr(args, "qk_rope_head_dim", None),

1187 + v_head_dim=getattr(args, "v_head_dim", None),

1188 + sparse_mla_force_mqa=force_mqa,

1189 + sparse_mla_mha_mode=mha_mode,

1190 + sparse_mla_dense_mha_max_seq_len=dense_mha_max_seq_len,

1191 + sparse_mla_topk_pattern=sparse_mla_topk_pattern,

1192 + prefill_backend=prefill_backend,

1193 + )

1194 +

1195 + # run_mla_benchmark needs the real backend name

1196 + from mla_runner import run_mla_benchmark as run_mla

1197 +

1198 + run_label = f"{backend}_{variant_label} {spec}"

1199 + pbar.set_postfix_str(run_label)

1200 +

1201 + try:

1202 + result = run_mla(

1203 + backend,

1204 + config,

1205 + prefill_backend=prefill_backend,

1206 + sparse_mla_force_mqa=force_mqa,

1207 + )

1208 + except Exception as e:

1209 + result = BenchmarkResult(

1210 + config=config,

1211 + mean_time=float("inf"),

1212 + median_time=float("inf"),

1213 + std_time=0,

1214 + min_time=float("inf"),

1215 + max_time=float("inf"),

1216 + error=str(e),

1217 + )

1218 +

1219 + all_results.append(result)

1220 + if args.output_csv:

1221 + formatter.save_csv(all_results, args.output_csv)

1222 + if args.output_json:

1223 + formatter.save_json(all_results, args.output_json)

1224 +

1225 + if not result.success:

1226 + console.print(

1227 + f"[red]Error {backend}_{variant_label} "

1228 + f"{spec}: {result.error}[/]"

1229 + )

1230 +

1231 + pbar.update(1)

1232 +

1233 + # Display results with variant labels as separate "backends"

1234 + console.print("\n[bold green]MHA vs MQA Results:[/]")

1235 + variant_backends = [f"{b}_{v}" for b in backends for v, _, _ in variants]

1236 + formatter.print_table(all_results, variant_backends)

1237 +

1238 # Handle model parameter sweep mode

1239 elif hasattr(args, "model_parameter_sweep") and args.model_parameter_sweep:

1240 # Model parameter sweep

benchmarks/attention_benchmarks/common.py

+68 -37

@@ -4,8 +4,10 @@

4 """Common utilities for attention benchmarking."""

6 import csv

7 +import gc

8 import json

9 import math

10 +from collections.abc import Sequence

11 from dataclasses import asdict, dataclass

12 from pathlib import Path

13 from typing import Any

@@ -44,10 +46,13 @@ def run_do_bench(

46 kwargs: dict[str, Any] = {"return_mode": "all"}

47 if use_cuda_graphs:

48 result = triton.testing.do_bench_cudagraph(benchmark_fn, **kwargs)

49 + gc.collect()

50 + torch.accelerator.empty_cache()

51 else:

52 if warmup_ms is not None:

53 kwargs["warmup"] = warmup_ms

54 result = triton.testing.do_bench(benchmark_fn, **kwargs)

55 + torch.accelerator.synchronize()

56 return result

@@ -91,42 +96,6 @@ except ImportError:

96 AttentionLayerBase = object # Fallback

94 -class MockKVBProj:

95 - """Mock KV projection layer for MLA prefill mode.

96 -

97 - Mimics ColumnParallelLinear behavior for kv_b_proj in MLA backends.

98 - Projects kv_c_normed to [qk_nope_head_dim + v_head_dim] per head.

99 - """

100 -

101 - def __init__(self, num_heads: int, qk_nope_head_dim: int, v_head_dim: int):

102 - self.num_heads = num_heads

103 - self.qk_nope_head_dim = qk_nope_head_dim

104 - self.v_head_dim = v_head_dim

105 - self.out_dim = qk_nope_head_dim + v_head_dim

106 - self.weight = torch.empty(0, dtype=torch.bfloat16)

107 -

108 - def __call__(self, x: torch.Tensor) -> tuple[torch.Tensor]:

109 - """

110 - Project kv_c_normed to output space.

111 -

112 - Args:

113 - x: Input tensor [num_tokens, kv_lora_rank]

114 -

115 - Returns:

116 - Tuple containing output tensor

117 - [num_tokens, num_heads, qk_nope_head_dim + v_head_dim]

118 - """

119 - num_tokens = x.shape[0]

120 - result = torch.randn(

121 - num_tokens,

122 - self.num_heads,

123 - self.out_dim,

124 - device=x.device,

125 - dtype=x.dtype,

126 - )

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

128 -

129 -

99 class MockIndexer:

100 """Mock Indexer for sparse MLA backends.

101

@@ -158,6 +127,60 @@ class MockIndexer:

127 )

128 self.topk_indices_buffer[:num_tokens] = indices

129

130 + def fill_indices(

131 + self,

132 + num_tokens: int,

133 + max_kv_len: int,

134 + pattern: str = "random",

135 + requests: Sequence[Any] | None = None,

136 + ):

137 + if pattern == "random":

138 + self.fill_random_indices(num_tokens, max_kv_len)

139 + return

140 + if pattern == "prefix":

141 + indices = torch.arange(

142 + self.topk_tokens,

143 + dtype=torch.int32,

144 + device=self.topk_indices_buffer.device,

145 + )

146 + indices = (indices % max_kv_len).expand(num_tokens, -1)

147 + self.topk_indices_buffer[:num_tokens] = indices

148 + return

149 + if pattern == "sliding_window":

150 + if requests is None:

151 + start = max(max_kv_len - self.topk_tokens, 0)

152 + indices = torch.arange(

153 + start,

154 + start + self.topk_tokens,

155 + dtype=torch.int32,

156 + device=self.topk_indices_buffer.device,

157 + )

158 + indices = indices.clamp(max=max_kv_len - 1).expand(num_tokens, -1)

159 + self.topk_indices_buffer[:num_tokens] = indices

160 + return

161 +

162 + rows = []

163 + offsets = torch.arange(

164 + self.topk_tokens,

165 + dtype=torch.int32,

166 + device=self.topk_indices_buffer.device,

167 + ) - (self.topk_tokens - 1)

168 + for request in requests:

169 + q_len = request.q_len

170 + kv_len = request.kv_len

171 + context_len = kv_len - q_len

172 + positions = torch.arange(

173 + context_len,

174 + kv_len,

175 + dtype=torch.int32,

176 + device=self.topk_indices_buffer.device,

177 + )

178 + row_indices = positions[:, None] + offsets[None, :]

179 + rows.append(row_indices.clamp(min=0, max=kv_len - 1))

180 + self.topk_indices_buffer[:num_tokens] = torch.cat(rows, dim=0)

181 + return

182 + raise ValueError(f"Unknown sparse MLA topk pattern: {pattern}")

183 +

184

185 class MockLayer(AttentionLayerBase):

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

@@ -252,10 +275,14 @@ class BenchmarkConfig:

275 num_kv_heads: int

276 block_size: int

277 device: str

278 + max_model_len: int | None = None

279 dtype: torch.dtype = torch.float16

280 profile_memory: bool = False

257 - use_cuda_graphs: bool = False

281 + use_cuda_graphs: bool = True

282 ncu_profile: bool = False

283 + torch_profile: bool = False

284 + torch_profile_dir: str | None = None

285 + torch_profile_iters: int = 3

286 warmup_ms: int | None = None

287

288 # "auto" or "fp8"

@@ -271,6 +298,10 @@ class BenchmarkConfig:

298 # Backend-specific tuning

299 num_kv_splits: int | None = None # CUTLASS MLA

300 reorder_batch_threshold: int | None = None # FlashAttn MLA, FlashMLA

301 + sparse_mla_force_mqa: bool = False # Force MQA path for sparse MLA

302 + sparse_mla_mha_mode: str = "auto" # "auto" or "dense"

303 + sparse_mla_dense_mha_max_seq_len: int | None = None

304 + sparse_mla_topk_pattern: str = "random" # "random", "prefix", "sliding_window"

305 num_splits: int | None = None # FlashAttention split-K (0=auto, 1=disabled)

306

307

benchmarks/attention_benchmarks/configs/mla_sparse_mha_vs_mqa.yaml new

+474

@@ -0,0 +1,474 @@

1 +# Sparse MLA benchmark: forward_mha vs forward_mqa

2 +#

3 +# Usage:

4 +# python benchmark.py --config configs/mla_sparse_mha_vs_mqa.yaml

5 +#

6 +# Heatmap grid:

7 +# - batch_size: 1, 2, 4, 8, 16, 32

8 +# - seq_len: 32, 64, 128, 256, 512, 1024, 2048

9 +# - q_len: powers of two through seq_len

10 +#

11 +# Specs with q_len < seq_len include context; the q_len == seq_len diagonal

12 +# covers pure prefill.

13 +# The model shape below is the DP case. For the TP8 run, manually change

14 +# model.num_q_heads from 128 to 16 before rerunning this benchmark.

15 +

16 +mode: mha_vs_mqa

17 +

18 +model:

19 + name: "deepseek-v3"

20 + num_layers: 60

21 + num_q_heads: 128

22 + num_kv_heads: 1

23 + head_dim: 576

24 + kv_lora_rank: 512

25 + qk_nope_head_dim: 128

26 + qk_rope_head_dim: 64

27 + v_head_dim: 128

28 + block_size: 128

29 + max_model_len: 2048

30 +

31 +batch_specs:

32 + # Batch size 1

33 + # seq_len = 32

34 + - "1q1s32"

35 + - "1q2s32"

36 + - "1q4s32"

37 + - "1q8s32"

38 + - "1q16s32"

39 + - "1q32"

40 + # seq_len = 64

41 + - "1q1s64"

42 + - "1q2s64"

43 + - "1q4s64"

44 + - "1q8s64"

45 + - "1q16s64"

46 + - "1q32s64"

47 + - "1q64"

48 + # seq_len = 128

49 + - "1q1s128"

50 + - "1q2s128"

51 + - "1q4s128"

52 + - "1q8s128"

53 + - "1q16s128"

54 + - "1q32s128"

55 + - "1q64s128"

56 + - "1q128"

57 + # seq_len = 256

58 + - "1q1s256"

59 + - "1q2s256"

60 + - "1q4s256"

61 + - "1q8s256"

62 + - "1q16s256"

63 + - "1q32s256"

64 + - "1q64s256"

65 + - "1q128s256"

66 + - "1q256"

67 + # seq_len = 512

68 + - "1q1s512"

69 + - "1q2s512"

70 + - "1q4s512"

71 + - "1q8s512"

72 + - "1q16s512"

73 + - "1q32s512"

74 + - "1q64s512"

75 + - "1q128s512"

76 + - "1q256s512"

77 + - "1q512"

78 + # seq_len = 1024

79 + - "1q1s1024"

80 + - "1q2s1024"

81 + - "1q4s1024"

82 + - "1q8s1024"

83 + - "1q16s1024"

84 + - "1q32s1024"

85 + - "1q64s1024"

86 + - "1q128s1024"

87 + - "1q256s1024"

88 + - "1q512s1024"

89 + - "1q1024"

90 + # seq_len = 2048

91 + - "1q1s2048"

92 + - "1q2s2048"

93 + - "1q4s2048"

94 + - "1q8s2048"

95 + - "1q16s2048"

96 + - "1q32s2048"

97 + - "1q64s2048"

98 + - "1q128s2048"

99 + - "1q256s2048"

100 + - "1q512s2048"

101 + - "1q1024s2048"

102 + - "1q2048"

103 +

104 + # Batch size 2

105 + # seq_len = 32

106 + - "2q1s32"

107 + - "2q2s32"

108 + - "2q4s32"

109 + - "2q8s32"

110 + - "2q16s32"

111 + - "2q32"

112 + # seq_len = 64

113 + - "2q1s64"

114 + - "2q2s64"

115 + - "2q4s64"

116 + - "2q8s64"

117 + - "2q16s64"

118 + - "2q32s64"

119 + - "2q64"

120 + # seq_len = 128

121 + - "2q1s128"

122 + - "2q2s128"

123 + - "2q4s128"

124 + - "2q8s128"

125 + - "2q16s128"

126 + - "2q32s128"

127 + - "2q64s128"

128 + - "2q128"

129 + # seq_len = 256

130 + - "2q1s256"

131 + - "2q2s256"

132 + - "2q4s256"

133 + - "2q8s256"

134 + - "2q16s256"

135 + - "2q32s256"

136 + - "2q64s256"

137 + - "2q128s256"

138 + - "2q256"

139 + # seq_len = 512

140 + - "2q1s512"

141 + - "2q2s512"

142 + - "2q4s512"

143 + - "2q8s512"

144 + - "2q16s512"

145 + - "2q32s512"

146 + - "2q64s512"

147 + - "2q128s512"

148 + - "2q256s512"

149 + - "2q512"

150 + # seq_len = 1024

151 + - "2q1s1024"

152 + - "2q2s1024"

153 + - "2q4s1024"

154 + - "2q8s1024"

155 + - "2q16s1024"

156 + - "2q32s1024"

157 + - "2q64s1024"

158 + - "2q128s1024"

159 + - "2q256s1024"

160 + - "2q512s1024"

161 + - "2q1024"

162 + # seq_len = 2048

163 + - "2q1s2048"

164 + - "2q2s2048"

165 + - "2q4s2048"

166 + - "2q8s2048"

167 + - "2q16s2048"

168 + - "2q32s2048"

169 + - "2q64s2048"

170 + - "2q128s2048"

171 + - "2q256s2048"

172 + - "2q512s2048"

173 + - "2q1024s2048"

174 + - "2q2048"

175 +

176 + # Batch size 4

177 + # seq_len = 32

178 + - "4q1s32"

179 + - "4q2s32"

180 + - "4q4s32"

181 + - "4q8s32"

182 + - "4q16s32"

183 + - "4q32"

184 + # seq_len = 64

185 + - "4q1s64"

186 + - "4q2s64"

187 + - "4q4s64"

188 + - "4q8s64"

189 + - "4q16s64"

190 + - "4q32s64"

191 + - "4q64"

192 + # seq_len = 128

193 + - "4q1s128"

194 + - "4q2s128"

195 + - "4q4s128"

196 + - "4q8s128"

197 + - "4q16s128"

198 + - "4q32s128"

199 + - "4q64s128"

200 + - "4q128"

201 + # seq_len = 256

202 + - "4q1s256"

203 + - "4q2s256"

204 + - "4q4s256"

205 + - "4q8s256"

206 + - "4q16s256"

207 + - "4q32s256"

208 + - "4q64s256"

209 + - "4q128s256"

210 + - "4q256"

211 + # seq_len = 512

212 + - "4q1s512"

213 + - "4q2s512"

214 + - "4q4s512"

215 + - "4q8s512"

216 + - "4q16s512"

217 + - "4q32s512"

218 + - "4q64s512"

219 + - "4q128s512"

220 + - "4q256s512"

221 + - "4q512"

222 + # seq_len = 1024

223 + - "4q1s1024"

224 + - "4q2s1024"

225 + - "4q4s1024"

226 + - "4q8s1024"

227 + - "4q16s1024"

228 + - "4q32s1024"

229 + - "4q64s1024"

230 + - "4q128s1024"

231 + - "4q256s1024"

232 + - "4q512s1024"

233 + - "4q1024"

234 + # seq_len = 2048

235 + - "4q1s2048"

236 + - "4q2s2048"

237 + - "4q4s2048"

238 + - "4q8s2048"

239 + - "4q16s2048"

240 + - "4q32s2048"

241 + - "4q64s2048"

242 + - "4q128s2048"

243 + - "4q256s2048"

244 + - "4q512s2048"

245 + - "4q1024s2048"

246 + - "4q2048"

247 +

248 + # Batch size 8

249 + # seq_len = 32

250 + - "8q1s32"

251 + - "8q2s32"

252 + - "8q4s32"

253 + - "8q8s32"

254 + - "8q16s32"

255 + - "8q32"

256 + # seq_len = 64

257 + - "8q1s64"

258 + - "8q2s64"

259 + - "8q4s64"

260 + - "8q8s64"

261 + - "8q16s64"

262 + - "8q32s64"

263 + - "8q64"

264 + # seq_len = 128

265 + - "8q1s128"

266 + - "8q2s128"

267 + - "8q4s128"

268 + - "8q8s128"

269 + - "8q16s128"

270 + - "8q32s128"

271 + - "8q64s128"

272 + - "8q128"

273 + # seq_len = 256

274 + - "8q1s256"

275 + - "8q2s256"

276 + - "8q4s256"

277 + - "8q8s256"

278 + - "8q16s256"

279 + - "8q32s256"

280 + - "8q64s256"

281 + - "8q128s256"

282 + - "8q256"

283 + # seq_len = 512

284 + - "8q1s512"

285 + - "8q2s512"

286 + - "8q4s512"

287 + - "8q8s512"

288 + - "8q16s512"

289 + - "8q32s512"

290 + - "8q64s512"

291 + - "8q128s512"

292 + - "8q256s512"

293 + - "8q512"

294 + # seq_len = 1024

295 + - "8q1s1024"

296 + - "8q2s1024"

297 + - "8q4s1024"

298 + - "8q8s1024"

299 + - "8q16s1024"

300 + - "8q32s1024"

301 + - "8q64s1024"

302 + - "8q128s1024"

303 + - "8q256s1024"

304 + - "8q512s1024"

305 + - "8q1024"

306 + # seq_len = 2048

307 + - "8q1s2048"

308 + - "8q2s2048"

309 + - "8q4s2048"

310 + - "8q8s2048"

311 + - "8q16s2048"

312 + - "8q32s2048"

313 + - "8q64s2048"

314 + - "8q128s2048"

315 + - "8q256s2048"

316 + - "8q512s2048"

317 + - "8q1024s2048"

318 + - "8q2048"

319 +

320 + # Batch size 16

321 + # seq_len = 32

322 + - "16q1s32"

323 + - "16q2s32"

324 + - "16q4s32"

325 + - "16q8s32"

326 + - "16q16s32"

327 + - "16q32"

328 + # seq_len = 64

329 + - "16q1s64"

330 + - "16q2s64"

331 + - "16q4s64"

332 + - "16q8s64"

333 + - "16q16s64"

334 + - "16q32s64"

335 + - "16q64"

336 + # seq_len = 128

337 + - "16q1s128"

338 + - "16q2s128"

339 + - "16q4s128"

340 + - "16q8s128"

341 + - "16q16s128"

342 + - "16q32s128"

343 + - "16q64s128"

344 + - "16q128"

345 + # seq_len = 256

346 + - "16q1s256"

347 + - "16q2s256"

348 + - "16q4s256"

349 + - "16q8s256"

350 + - "16q16s256"

351 + - "16q32s256"

352 + - "16q64s256"

353 + - "16q128s256"

354 + - "16q256"

355 + # seq_len = 512

356 + - "16q1s512"

357 + - "16q2s512"

358 + - "16q4s512"

359 + - "16q8s512"

360 + - "16q16s512"

361 + - "16q32s512"

362 + - "16q64s512"

363 + - "16q128s512"

364 + - "16q256s512"

365 + - "16q512"

366 + # seq_len = 1024

367 + - "16q1s1024"

368 + - "16q2s1024"

369 + - "16q4s1024"

370 + - "16q8s1024"

371 + - "16q16s1024"

372 + - "16q32s1024"

373 + - "16q64s1024"

374 + - "16q128s1024"

375 + - "16q256s1024"

376 + - "16q512s1024"

377 + - "16q1024"

378 + # seq_len = 2048

379 + - "16q1s2048"

380 + - "16q2s2048"

381 + - "16q4s2048"

382 + - "16q8s2048"

383 + - "16q16s2048"

384 + - "16q32s2048"

385 + - "16q64s2048"

386 + - "16q128s2048"

387 + - "16q256s2048"

388 + - "16q512s2048"

389 + - "16q1024s2048"

390 + - "16q2048"

391 +

392 + # Batch size 32

393 + # seq_len = 32

394 + - "32q1s32"

395 + - "32q2s32"

396 + - "32q4s32"

397 + - "32q8s32"

398 + - "32q16s32"

399 + - "32q32"

400 + # seq_len = 64

401 + - "32q1s64"

402 + - "32q2s64"

403 + - "32q4s64"

404 + - "32q8s64"

405 + - "32q16s64"

406 + - "32q32s64"

407 + - "32q64"

408 + # seq_len = 128

409 + - "32q1s128"

410 + - "32q2s128"

411 + - "32q4s128"

412 + - "32q8s128"

413 + - "32q16s128"

414 + - "32q32s128"

415 + - "32q64s128"

416 + - "32q128"

417 + # seq_len = 256

418 + - "32q1s256"

419 + - "32q2s256"

420 + - "32q4s256"

421 + - "32q8s256"

422 + - "32q16s256"

423 + - "32q32s256"

424 + - "32q64s256"

425 + - "32q128s256"

426 + - "32q256"

427 + # seq_len = 512

428 + - "32q1s512"

429 + - "32q2s512"

430 + - "32q4s512"

431 + - "32q8s512"

432 + - "32q16s512"

433 + - "32q32s512"

434 + - "32q64s512"

435 + - "32q128s512"

436 + - "32q256s512"

437 + - "32q512"

438 + # seq_len = 1024

439 + - "32q1s1024"

440 + - "32q2s1024"

441 + - "32q4s1024"

442 + - "32q8s1024"

443 + - "32q16s1024"

444 + - "32q32s1024"

445 + - "32q64s1024"

446 + - "32q128s1024"

447 + - "32q256s1024"

448 + - "32q512s1024"

449 + - "32q1024"

450 + # seq_len = 2048

451 + - "32q1s2048"

452 + - "32q2s2048"

453 + - "32q4s2048"

454 + - "32q8s2048"

455 + - "32q16s2048"

456 + - "32q32s2048"

457 + - "32q64s2048"

458 + - "32q128s2048"

459 + - "32q256s2048"

460 + - "32q512s2048"

461 + - "32q1024s2048"

462 + - "32q2048"

463 +

464 +backends:

465 + - FLASHMLA_SPARSE

466 +

467 +device: "cuda:0"

468 +profile_memory: false

469 +sparse_mla_dense_mha_max_seq_len: 2048

470 +sparse_mla_topk_pattern: "random"

471 +

472 +output:

473 + csv: "benchmark_output/mla_sparse_mha_vs_mqa.csv"

474 + json: "benchmark_output/mla_sparse_mha_vs_mqa.json"

benchmarks/attention_benchmarks/mla_runner.py

+171 -17

@@ -9,6 +9,8 @@ needing full VllmConfig integration.

9 """

11 import statistics

12 +import tempfile

13 +from pathlib import Path

15 import numpy as np

16 import torch

@@ -17,7 +19,6 @@ from common import (

19 BenchmarkResult,

20 MockHfConfig,

21 MockIndexer,

20 - MockKVBProj,

22 MockLayer,

23 run_do_bench,

24 run_ncu_profile,

@@ -33,8 +34,59 @@ from vllm.config import (

34 VllmConfig,

35 set_current_vllm_config,

36 )

37 +from vllm.model_executor.layers.linear import ColumnParallelLinear

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

40 +

41 +def _safe_profile_name(value: str) -> str:

42 + return "".join(c if c.isalnum() or c in "._-" else "_" for c in value)

43 +

44 +

45 +def _create_kv_b_proj(

46 + mla_dims: dict,

47 + device: torch.device,

48 +):

49 + kv_b_proj = ColumnParallelLinear(

50 + mla_dims["kv_lora_rank"],

51 + mla_dims["num_q_heads"]

52 + * (mla_dims["qk_nope_head_dim"] + mla_dims["v_head_dim"]),

53 + bias=False,

54 + params_dtype=torch.bfloat16,

55 + quant_config=None,

56 + prefix="benchmark.kv_b_proj",

57 + ).to(device)

58 + with torch.no_grad():

59 + kv_b_proj.weight.copy_(torch.randn_like(kv_b_proj.weight))

60 + return kv_b_proj

61 +

62 +

63 +def _ensure_single_rank_model_parallel() -> None:

64 + import torch.distributed as dist

65 +

66 + from vllm.distributed import (

67 + ensure_model_parallel_initialized,

68 + init_distributed_environment,

69 + model_parallel_is_initialized,

70 + )

71 +

72 + if not dist.is_available():

73 + return

74 + if not dist.is_initialized():

75 + with tempfile.NamedTemporaryFile(

76 + prefix="vllm_bench_dist_", delete=False

77 + ) as init_file:

78 + distributed_init_method = f"file://{init_file.name}"

79 + init_distributed_environment(

80 + world_size=1,

81 + rank=0,

82 + distributed_init_method=distributed_init_method,

83 + local_rank=0,

84 + backend="nccl",

85 + )

86 + if not model_parallel_is_initialized():

87 + ensure_model_parallel_initialized(1, 1)

88 +

89 +

90 # ============================================================================

91 # VllmConfig Creation

92 # ============================================================================

@@ -66,10 +118,12 @@ def create_minimal_vllm_config(

118 block_size: int = 128,

119 max_num_seqs: int = 256,

120 max_num_batched_tokens: int = 8192,

121 + max_model_len: int = 32768,

122 mla_dims: dict | None = None,

123 index_topk: int | None = None,

124 prefill_backend: str | None = None,

125 kv_cache_dtype: str = "auto",

126 + sparse_mla_force_mqa: bool = False,

127 ) -> VllmConfig:

128 """

129 Create minimal VllmConfig for MLA benchmarks.

@@ -86,6 +140,8 @@ def create_minimal_vllm_config(

140 prefill_backend: Prefill backend name (e.g., "fa3", "fa4", "flashinfer",

141 "trtllm"). Configures the attention config to force

142 the specified prefill backend.

143 + sparse_mla_force_mqa: If True, forces all sparse MLA tokens through

144 + forward_mqa (even prefill tokens).

145

146 Returns:

147 VllmConfig for benchmarking

@@ -131,7 +187,7 @@ def create_minimal_vllm_config(

187 trust_remote_code=True,

188 dtype="bfloat16",

189 seed=0,

134 - max_model_len=32768,

190 + max_model_len=max_model_len,

191 quantization=None,

192 enforce_eager=False,

193 max_logprobs=20,

@@ -163,7 +219,7 @@ def create_minimal_vllm_config(

219 scheduler_config = SchedulerConfig(

220 max_num_seqs=max_num_seqs,

221 max_num_batched_tokens=max(max_num_batched_tokens, max_num_seqs),

166 - max_model_len=32768,

222 + max_model_len=max_model_len,

223 is_encoder_decoder=False,

224 enable_chunked_prefill=True,

225 )

@@ -192,6 +248,9 @@ def create_minimal_vllm_config(

248 "flash_attn_version"

249 ]

250

251 + if sparse_mla_force_mqa:

252 + vllm_config.attention_config.sparse_mla_force_mqa = True

253 +

254 return vllm_config

255

256

@@ -548,12 +607,7 @@ def _create_backend_impl(

607 # Calculate scale

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

609

551 - # Create mock kv_b_proj layer for prefill mode

552 - mock_kv_b_proj = MockKVBProj(

553 - num_heads=mla_dims["num_q_heads"],

554 - qk_nope_head_dim=mla_dims["qk_nope_head_dim"],

555 - v_head_dim=mla_dims["v_head_dim"],

556 - )

610 + kv_b_proj = _create_kv_b_proj(mla_dims, device)

611

612 # Create indexer for sparse backends

613 indexer = None

@@ -584,7 +638,7 @@ def _create_backend_impl(

638 "qk_rope_head_dim": mla_dims["qk_rope_head_dim"],

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

640 "v_head_dim": mla_dims["v_head_dim"],

587 - "kv_b_proj": mock_kv_b_proj,

641 + "kv_b_proj": kv_b_proj,

642 }

643

644 # Add indexer for sparse backends

@@ -785,14 +839,35 @@ def _run_single_benchmark(

839 # Fill indexer with random indices for sparse backends

840 is_sparse = backend_cfg.get("is_sparse", False)

841 if is_sparse and indexer is not None:

788 - indexer.fill_random_indices(total_q, max_kv_len)

842 + indexer.fill_indices(

843 + total_q,

844 + max_kv_len,

845 + getattr(config, "sparse_mla_topk_pattern", "random"),

846 + )

847

848 # Determine which forward methods to use based on metadata.

791 - # Sparse MLA backends always use forward_mqa

792 - has_decode = is_sparse or getattr(metadata, "decode", None) is not None

793 - has_prefill = not is_sparse and getattr(metadata, "prefill", None) is not None

849 + # Non-sparse backends use .decode/.prefill sub-objects.

850 + # Sparse backends use num_decode_tokens/num_prefills directly.

851 + #

852 + # sparse_mla_force_mqa overrides: even for prefill metadata, use MQA.

853 + force_mqa = getattr(config, "sparse_mla_force_mqa", False)

854 + force_dense_mha = getattr(config, "sparse_mla_mha_mode", "auto") == "dense"

855 + if force_mqa:

856 + has_decode = True

857 + has_prefill = False

858 + elif is_sparse:

859 + has_decode = metadata.num_decode_tokens > 0

860 + has_prefill = metadata.num_prefills > 0

861 + else:

862 + has_decode = metadata.decode is not None

863 + has_prefill = metadata.prefill is not None

864 if not has_decode and not has_prefill:

865 raise RuntimeError("Metadata has neither decode nor prefill metadata")

866 + if is_sparse and force_dense_mha and not has_prefill:

867 + raise RuntimeError(

868 + "Sparse MLA dense_mha benchmark did not produce prefill metadata. "

869 + "Check reorder_batch_threshold/path forcing."

870 + )

871

872 num_decode = (

873 metadata.num_decode_tokens

@@ -871,7 +946,6 @@ def _run_single_benchmark(

946 metadata,

947 prefill_inputs["k_scale"],

948 prefill_fp8_output if fused_output else prefill_inputs["output"],

874 - prefill_output_scale if fused_output else None,

949 )

950 if fused_output:

951 out = prefill_fp8_output

@@ -898,6 +972,48 @@ def _run_single_benchmark(

972 throughput_tokens_per_sec=0.0,

973 )

974

975 + if config.torch_profile:

976 + profile_dir = Path(

977 + config.torch_profile_dir or "benchmark_outputs/torch_profiles"

978 + )

979 + profile_dir.mkdir(parents=True, exist_ok=True)

980 + trace_name = _safe_profile_name(f"{config.backend}_{config.batch_spec}")

981 + trace_path = profile_dir / f"{trace_name}.json"

982 + iters = max(config.torch_profile_iters, 1)

983 +

984 + forward_fn()

985 + torch.accelerator.synchronize()

986 + with torch.profiler.profile(

987 + activities=[

988 + torch.profiler.ProfilerActivity.CPU,

989 + torch.profiler.ProfilerActivity.CUDA,

990 + ],

991 + record_shapes=True,

992 + profile_memory=True,

993 + with_stack=False,

994 + ) as prof:

995 + for _ in range(iters):

996 + forward_fn()

997 + torch.accelerator.synchronize()

998 + prof.step()

999 + prof.export_chrome_trace(str(trace_path))

1000 + print(f"Saved PyTorch profiler trace to {trace_path}")

1001 + print(

1002 + prof.key_averages().table(

1003 + sort_by="cuda_time_total",

1004 + row_limit=25,

1005 + )

1006 + )

1007 + return BenchmarkResult(

1008 + config=config,

1009 + mean_time=0.0,

1010 + median_time=0.0,

1011 + std_time=0.0,

1012 + min_time=0.0,

1013 + max_time=0.0,

1014 + throughput_tokens_per_sec=0.0,

1015 + )

1016 +

1017 all_ms = run_do_bench(benchmark_fn, config.use_cuda_graphs, config.warmup_ms)

1018

1019 # Convert ms to seconds per layer

@@ -920,6 +1036,7 @@ def _run_mla_benchmark_batched(

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

1037 index_topk: int = 2048,

1038 prefill_backend: str | None = None,

1039 + sparse_mla_force_mqa: bool = False,

1040 output_scale: float | None = None,

1041 fuse_quant_op: bool = False,

1042 ) -> list[BenchmarkResult]:

@@ -940,6 +1057,8 @@ def _run_mla_benchmark_batched(

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

1058 prefill_backend: Prefill backend name (e.g., "fa3", "fa4").

1059 When set, forces the specified FlashAttention version for prefill.

1060 + sparse_mla_force_mqa: If True, forces all sparse MLA tokens through

1061 + forward_mqa (even prefill tokens).

1062

1063 Returns:

1064 List of BenchmarkResult objects

@@ -980,21 +1099,41 @@ def _run_mla_benchmark_batched(

1099 sum(r.q_len for r in parse_batch_spec(cfg.batch_spec))

1100 for cfg, *_ in configs_with_params

1101 )

1102 + max_model_len = max(

1103 + max_total_q,

1104 + max(

1105 + getattr(cfg, "max_model_len", None) or 32768

1106 + for cfg, *_ in configs_with_params

1107 + ),

1108 + )

1109

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

1111 vllm_config = create_minimal_vllm_config(

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

1113 block_size=block_size,

1114 max_num_batched_tokens=max_total_q,

1115 + max_model_len=max_model_len,

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

1117 index_topk=index_topk if is_sparse else None,

1118 prefill_backend=prefill_backend,

1119 kv_cache_dtype=kv_cache_dtype,

1120 + sparse_mla_force_mqa=sparse_mla_force_mqa,

1121 )

1122

1123 results = []

1124

1125 + # Initialize workspace manager (needed by metadata builders)

1126 + from vllm.v1.worker.workspace import (

1127 + init_workspace_manager,

1128 + is_workspace_manager_initialized,

1129 + )

1130 +

1131 + if not is_workspace_manager_initialized():

1132 + init_workspace_manager(device)

1133 +

1134 with set_current_vllm_config(vllm_config):

1135 + _ensure_single_rank_model_parallel()

1136 +

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

1138 impl, layer, builder_instance, indexer = _create_backend_impl(

1139 backend_cfg,

@@ -1040,9 +1179,20 @@ def _run_mla_benchmark_batched(

1179 for config, threshold, num_splits in configs_with_params:

1180 # Set threshold for this benchmark (FlashAttn/FlashMLA only)

1181 original_threshold = None

1043 - if threshold is not None and builder_instance:

1182 + effective_threshold = threshold

1183 + force_dense_mha = (

1184 + is_sparse

1185 + and getattr(config, "sparse_mla_mha_mode", "auto") == "dense"

1186 + and not getattr(config, "sparse_mla_force_mqa", False)

1187 + )

1188 + if force_dense_mha:

1189 + # Sparse MLA normally treats q_len <= 1 as decode. Use an

1190 + # impossible threshold so dense_mha benchmarks actually run

1191 + # the prefill/MHA path, including q_len=1 short extends.

1192 + effective_threshold = -1

1193 + if effective_threshold is not None and builder_instance:

1194 original_threshold = builder_instance.reorder_batch_threshold

1045 - builder_instance.reorder_batch_threshold = threshold

1195 + builder_instance.reorder_batch_threshold = effective_threshold

1196

1197 # Set num_splits for CUTLASS

1198 original_num_splits = None

@@ -1090,6 +1240,7 @@ def run_mla_benchmark(

1240 num_kv_splits: int | None = None,

1241 index_topk: int = 2048,

1242 prefill_backend: str | None = None,

1243 + sparse_mla_force_mqa: bool = False,

1244 output_scale: float | None = None,

1245 fuse_quant_op: bool = False,

1246 ) -> BenchmarkResult | list[BenchmarkResult]:

@@ -1111,6 +1262,8 @@ def run_mla_benchmark(

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

1263 prefill_backend: Prefill backend name (e.g., "fa3", "fa4").

1264 When set, forces the specified FlashAttention version for prefill.

1265 + sparse_mla_force_mqa: If True, forces all sparse MLA tokens through

1266 + forward_mqa (even prefill tokens).

1267 output_scale: Static per-tensor FP8 scale for prefill output (None = bf16).

1268 fuse_quant_op: With output_scale set, fuse the FP8 write into the prefill

1269 kernel vs a standalone post-quant kernel. See _run_single_benchmark.

@@ -1142,6 +1295,7 @@ def run_mla_benchmark(

1295 configs_with_params,

1296 index_topk,

1297 prefill_backend=prefill_backend,

1298 + sparse_mla_force_mqa=sparse_mla_force_mqa,

1299 output_scale=output_scale,

1300 fuse_quant_op=fuse_quant_op,

1301 )

docs/design/attention_backends.md

+1 -1

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

205

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

tests/v1/attention/test_dspark_noncausal_sparse_mla.py

+17

@@ -304,6 +304,23 @@ def _run_sparse_backend_vs_sdpa(

304 scale=kv_cache_scale,

305 )

306

307 + # The sparse builder clones the layer's dense-MHA prefill backend from

308 + # static_forward_context; register a mock layer carrying one.

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

310 +

311 + prefill_backend = get_mla_prefill_backend(vllm_config)(

312 + num_heads=num_heads,

313 + scale=scale,

314 + kv_lora_rank=kv_lora_rank,

315 + qk_nope_head_dim=qk_nope_head_dim,

316 + qk_rope_head_dim=qk_rope_head_dim,

317 + v_head_dim=v_head_dim,

318 + vllm_config=vllm_config,

319 + )

320 + vllm_config.compilation_config.static_forward_context["placeholder"] = (

321 + SimpleNamespace(prefill_backend=prefill_backend)

322 + )

323 +

324 builder = backend_cls.get_builder_cls()(

325 kv_cache_spec, ["placeholder"], vllm_config, device

326 )

tests/v1/attention/test_mla_backends.py

+9 -4

@@ -1031,6 +1031,11 @@ def run_attention_backend(

1031 @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8", "fp8_e4m3"])

1032 @pytest.mark.parametrize(("q_scale", "k_scale"), [(1.0, 1.0), (2.0, 3.0)])

1033 @pytest.mark.parametrize("prefill_backend", PREFILL_BACKENDS_TO_TEST)

1034 +@pytest.mark.parametrize(

1035 + ("qk_nope_head_dim", "v_head_dim"),

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

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

1038 +)

1039 def test_backend_correctness(

1040 default_vllm_config,

1041 dist_init,

@@ -1042,6 +1047,8 @@ def test_backend_correctness(

1047 q_scale: float,

1048 k_scale: float,

1049 prefill_backend: MLAPrefillBackendEnum,

1050 + qk_nope_head_dim: int,

1051 + v_head_dim: int,

1052 ):

1053 """

1054 Test that all backends produce similar outputs to a reference implementation

@@ -1090,9 +1097,9 @@ def test_backend_correctness(

1097 MLAPrefillSelectorConfig(

1098 dtype=torch.bfloat16,

1099 mla_dimensions=MLADimensions(

1093 - qk_nope_head_dim=128,

1100 + qk_nope_head_dim=qk_nope_head_dim,

1101 qk_rope_head_dim=64,

1095 - v_head_dim=128,

1102 + v_head_dim=v_head_dim,

1103 ),

1104 ),

1105 )

@@ -1168,8 +1175,6 @@ def test_backend_correctness(

1175 dtype = _convert_dtype_to_torch(vllm_config.model_config.dtype)

1176 kv_lora_rank = 512

1177 qk_rope_head_dim = 64

1171 - qk_nope_head_dim = 128

1172 - v_head_dim = 128

1178 total_head_size = kv_lora_rank + qk_rope_head_dim

1179 assert kv_lora_rank + qk_rope_head_dim == head_size, (

1180 f"MLA dimensions don't match: {total_head_size} != {head_size}"

tests/v1/attention/test_sparse_mla_backends.py

+269

@@ -442,6 +442,23 @@ def test_sparse_backend_decode_correctness(

442 scale=kv_cache_scale,

443 )

444

445 + # The sparse builder clones the layer's dense-MHA prefill backend from

446 + # static_forward_context; register a mock layer carrying one.

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

448 +

449 + prefill_backend = get_mla_prefill_backend(vllm_config)(

450 + num_heads=num_heads,

451 + scale=scale,

452 + kv_lora_rank=kv_lora_rank,

453 + qk_nope_head_dim=qk_nope_head_dim,

454 + qk_rope_head_dim=qk_rope_head_dim,

455 + v_head_dim=v_head_dim,

456 + vllm_config=vllm_config,

457 + )

458 + vllm_config.compilation_config.static_forward_context["placeholder"] = (

459 + SimpleNamespace(prefill_backend=prefill_backend)

460 + )

461 +

462 builder_cls = backend_cls.get_builder_cls()

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

464 metadata = builder.build(

@@ -722,6 +739,258 @@ def test_split_prefill_chunks(seq_lens, max_buf, expected):

739 assert out == expected

740

741

742 +PREFILL_BATCH_SPECS = {

743 + "short_dense_mha": BatchSpec(seq_lens=[64, 128], query_lens=[64, 128]),

744 + "short_context_dense_mha": BatchSpec(seq_lens=[128, 160], query_lens=[64, 32]),

745 +}

746 +

747 +

748 +@pytest.mark.skipif(

749 + torch.cuda.get_device_capability()[0] < 10,

750 + reason="Sparse MLA forward_mha requires FA4 (SM100+)",

751 +)

752 +@pytest.mark.parametrize("batch_name", list(PREFILL_BATCH_SPECS.keys()))

753 +@pytest.mark.parametrize("kv_cache_dtype", ["auto"])

754 +def test_sparse_backend_prefill_correctness(

755 + default_vllm_config,

756 + dist_init,

757 + batch_name,

758 + kv_cache_dtype,

759 + workspace_init,

760 +):

761 + """Test single-pass FA4 dense forward_mha for sparse MLA prefill."""

762 + backend_cls = FlashMLASparseBackend

763 + batch_spec = PREFILL_BATCH_SPECS[batch_name]

764 +

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

766 + dtype = torch.bfloat16

767 + block_size = 64

768 +

769 + num_heads = 128

770 + kv_lora_rank = 512

771 + qk_nope_head_dim = 128

772 + qk_rope_head_dim = 64

773 + v_head_dim = 128

774 + head_size = kv_lora_rank + qk_rope_head_dim

775 + topk_tokens = 512

776 +

777 + max_seqlen = max(batch_spec.seq_lens)

778 + total_cache_tokens = sum(batch_spec.seq_lens)

779 +

780 + vllm_config = create_vllm_config(

781 + model_name="deepseek-ai/DeepSeek-V2-Lite-Chat",

782 + tensor_parallel_size=1,

783 + max_model_len=max_seqlen,

784 + num_gpu_blocks=max(2048, cdiv(total_cache_tokens, block_size) + 1),

785 + block_size=block_size,

786 + hf_config_override={

787 + "index_topk": topk_tokens,

788 + "attn_module_list_cfg": [{"topk_tokens": topk_tokens}],

789 + },

790 + )

791 + model_config = vllm_config.model_config

792 + model_config.hf_text_config = SimpleNamespace(

793 + q_lora_rank=None,

794 + kv_lora_rank=kv_lora_rank,

795 + qk_nope_head_dim=qk_nope_head_dim,

796 + qk_rope_head_dim=qk_rope_head_dim,

797 + v_head_dim=v_head_dim,

798 + model_type="deepseek_v2",

799 + )

800 + model_config.dtype = dtype

801 + model_config.get_num_attention_heads = MethodType(

802 + lambda self, parallel_config: num_heads, model_config

803 + )

804 + model_config.get_num_kv_heads = MethodType(

805 + lambda self, parallel_config: 1, model_config

806 + )

807 + model_config.get_head_size = MethodType(lambda self: head_size, model_config)

808 + model_config.get_sliding_window = MethodType(lambda self: None, model_config)

809 +

810 + kv_cache_spec = create_standard_kv_cache_spec(vllm_config)

811 + scale = 1.0 / math.sqrt(qk_nope_head_dim + qk_rope_head_dim)

812 +

813 + torch.manual_seed(42)

814 +

815 + W_UK = torch.rand(

816 + kv_lora_rank, num_heads, qk_nope_head_dim, dtype=dtype, device=device

817 + )

818 + W_UV = torch.rand(kv_lora_rank, num_heads, v_head_dim, dtype=dtype, device=device)

819 +

820 + seq_lens = batch_spec.seq_lens

821 + query_lens = batch_spec.query_lens

822 +

823 + # Compute dense reference outputs.

824 + total_query_tokens = sum(query_lens)

825 + sparse_indices = torch.zeros(

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

827 + )

828 +

829 + all_q, all_kv_c_new, all_k_pe_new = [], [], []

830 + kv_c_contexts, k_pe_contexts = [], []

831 + reference_outputs = []

832 +

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

834 + s_len = seq_lens[i]

835 + q_len = query_lens[i]

836 + ctx_len = s_len - q_len

837 +

838 + q_mha = torch.rand(

839 + q_len,

840 + num_heads,

841 + qk_nope_head_dim + qk_rope_head_dim,

842 + dtype=dtype,

843 + device=device,

844 + )

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

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

847 +

848 + # Decompress all KV for reference

849 + kv_b_weight = torch.cat([W_UK, W_UV], dim=-1).view(

850 + kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)

851 + )

852 + kv_decompressed = (kv_c_full @ kv_b_weight).view(

853 + s_len, num_heads, qk_nope_head_dim + v_head_dim

854 + )

855 + k_nope_all, v_all = kv_decompressed.split(

856 + [qk_nope_head_dim, v_head_dim], dim=-1

857 + )

858 + k_pe_expanded = k_pe_full.expand(-1, num_heads, -1)

859 + k_all = torch.cat([k_nope_all, k_pe_expanded], dim=-1)

860 +

861 + for j in range(q_len):

862 + attend_end = ctx_len + j + 1

863 + q_tok = q_mha[j : j + 1] # (1, H, D_qk)

864 + k_attend = k_all[:attend_end] # (N, H, D_qk)

865 + v_attend = v_all[:attend_end] # (N, H, D_v)

866 +

867 + q_sdpa = q_tok.unsqueeze(0).transpose(1, 2).float()

868 + k_sdpa = k_attend.unsqueeze(0).transpose(1, 2).float()

869 + v_sdpa = v_attend.unsqueeze(0).transpose(1, 2).float()

870 +

871 + out = torch.nn.functional.scaled_dot_product_attention(

872 + q_sdpa, k_sdpa, v_sdpa, scale=scale

873 + )

874 + out = out.transpose(1, 2).squeeze(0) # (1, H, D_v)

875 + reference_outputs.append(out.to(dtype).flatten(start_dim=-2))

876 +

877 + all_q.append(q_mha)

878 + all_kv_c_new.append(kv_c_full[ctx_len:])

879 + all_k_pe_new.append(k_pe_full[ctx_len:])

880 + kv_c_contexts.append(kv_c_full)

881 + k_pe_contexts.append(k_pe_full)

882 +

883 + query_cat = torch.cat(all_q, dim=0)

884 + kv_c_cat = torch.cat(all_kv_c_new, dim=0)

885 + k_pe_cat = torch.cat(all_k_pe_new, dim=0)

886 + ref_output = torch.cat(reference_outputs, dim=0)

887 +

888 + vllm_config.cache_config.cache_dtype = kv_cache_dtype

889 + vllm_config.model_config.hf_config.index_topk = topk_tokens

890 +

891 + common_attn_metadata = create_common_attn_metadata(

892 + batch_spec,

893 + vllm_config.cache_config.block_size,

894 + device,

895 + arange_block_indices=True,

896 + )

897 +

898 + kv_cache = create_and_prepopulate_kv_cache(

899 + kv_c_contexts=kv_c_contexts,

900 + k_pe_contexts=k_pe_contexts,

901 + block_size=block_size,

902 + head_size=head_size,

903 + dtype=dtype,

904 + device=device,

905 + num_blocks=vllm_config.cache_config.num_gpu_blocks,

906 + common_attn_metadata=common_attn_metadata,

907 + randomize_blocks=False,

908 + kv_cache_dtype=kv_cache_dtype,

909 + )

910 +

911 + # The sparse builder clones the layer's dense-MHA prefill backend from

912 + # static_forward_context; register a mock layer carrying one.

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

914 +

915 + prefill_backend = get_mla_prefill_backend(vllm_config)(

916 + num_heads=num_heads,

917 + scale=scale,

918 + kv_lora_rank=kv_lora_rank,

919 + qk_nope_head_dim=qk_nope_head_dim,

920 + qk_rope_head_dim=qk_rope_head_dim,

921 + v_head_dim=v_head_dim,

922 + vllm_config=vllm_config,

923 + )

924 + vllm_config.compilation_config.static_forward_context["placeholder"] = (

925 + SimpleNamespace(prefill_backend=prefill_backend)

926 + )

927 +

928 + builder_cls = backend_cls.get_builder_cls()

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

930 + # Drive the queries through the dense-MHA prefill path directly (the routing

931 + # threshold would otherwise classify these short queries as MQA decodes).

932 + builder.reorder_batch_threshold = 1

933 + metadata = builder.build(

934 + common_prefix_len=0, common_attn_metadata=common_attn_metadata

935 + )

936 +

937 + mock_indexer = SimpleNamespace(topk_indices_buffer=sparse_indices)

938 +

939 + kv_b_proj_weight = torch.cat([W_UK, W_UV], dim=-1).view(

940 + kv_lora_rank, num_heads * (qk_nope_head_dim + v_head_dim)

941 + )

942 +

943 + mock_kv_b_proj = ColumnParallelLinear(

944 + input_size=kv_lora_rank,

945 + output_size=num_heads * (qk_nope_head_dim + v_head_dim),

946 + bias=False,

947 + ).to(device=device, dtype=dtype)

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

949 +

950 + impl_cls = backend_cls.get_impl_cls()

951 + with set_current_vllm_config(vllm_config):

952 + impl = impl_cls(

953 + num_heads=num_heads,

954 + head_size=head_size,

955 + scale=scale,

956 + num_kv_heads=1,

957 + alibi_slopes=None,

958 + sliding_window=None,

959 + kv_cache_dtype=kv_cache_dtype,

960 + logits_soft_cap=None,

961 + attn_type="decoder",

962 + kv_sharing_target_layer_name=None,

963 + q_lora_rank=None,

964 + kv_lora_rank=kv_lora_rank,

965 + qk_nope_head_dim=qk_nope_head_dim,

966 + qk_rope_head_dim=qk_rope_head_dim,

967 + qk_head_dim=qk_nope_head_dim + qk_rope_head_dim,

968 + v_head_dim=v_head_dim,

969 + kv_b_proj=mock_kv_b_proj,

970 + indexer=mock_indexer,

971 + )

972 + impl.process_weights_after_loading(dtype)

973 +

974 + out_buffer = torch.empty(

975 + total_query_tokens, num_heads * v_head_dim, dtype=dtype, device=device

976 + )

977 +

978 + with torch.inference_mode():

979 + impl.forward_mha(

980 + q=query_cat,

981 + kv_c_normed=kv_c_cat,

982 + k_pe=k_pe_cat,

983 + kv_c_and_k_pe_cache=kv_cache,

984 + attn_metadata=metadata,

985 + k_scale=torch.tensor(1.0, device=device),

986 + output=out_buffer,

987 + )

988 +

989 + assert out_buffer.shape == ref_output.shape

990 + assert torch.isfinite(out_buffer).all(), "Non-finite values in output"

991 + torch.testing.assert_close(out_buffer, ref_output, rtol=0.01, atol=0.01)

992 +

993 +

994 @pytest.mark.parametrize(

995 "seq_lens,query_lens,workspace_size,max_logits_bytes,expected",

996 [

vllm/config/attention.py

@@ -59,6 +59,11 @@ class AttentionConfig:

59 use_non_causal: bool = False

60 """Whether to use non-causal (bidirectional) attention."""

62 + sparse_mla_force_mqa: bool = False

63 + """Force sparse MLA to use forward_mqa for all requests, including prefill.

64 + When False (default), pure prefill batches use forward_mha when implemented.

65 + Set to True to always use the MQA path."""

66 +

67 flex_attn_block_m: int | None = None

68 """Triton kernel BLOCK_M tile size for flex attention.

69 Must be a power of 2 >= 16. If None and VLLM_BATCH_INVARIANT=1,

vllm/model_executor/layers/attention/mla_attention.py

+315 -255

@@ -259,7 +259,6 @@ from vllm.v1.attention.backend import (

259 AttentionType,

260 CommonAttentionMetadata,

261 MLAAttentionImpl,

262 - SparseMLAAttentionImpl,

262 )

263 from vllm.v1.attention.backends.mla.prefill import (

264 MLAPrefillBackend,

@@ -490,16 +489,30 @@ class MLAAttention(nn.Module, AttentionLayerBase):

489 raise ValueError(f"Duplicate layer name: {prefix}")

490 compilation_config.static_forward_context[prefix] = self

491

493 - prefill_backend_cls = get_mla_prefill_backend(vllm_config)

494 - self.prefill_backend = prefill_backend_cls(

495 - num_heads=self.num_heads,

496 - scale=self.scale,

497 - kv_lora_rank=self.kv_lora_rank,

498 - qk_nope_head_dim=self.qk_nope_head_dim,

499 - qk_rope_head_dim=self.qk_rope_head_dim,

500 - v_head_dim=self.v_head_dim,

501 - vllm_config=vllm_config,

502 - )

492 + self.prefill_backend: MLAPrefillBackend | None

493 + try:

494 + prefill_backend_cls = get_mla_prefill_backend(vllm_config)

495 + except ValueError:

496 + if (

497 + not self.impl.is_sparse

498 + or vllm_config.attention_config.mla_prefill_backend is not None

499 + ):

500 + raise

501 + logger.warning_once(

502 + "No MLA prefill backend supports this model; sparse MLA will use the "

503 + "top-k MQA path only (no dense-MHA prefill)."

504 + )

505 + self.prefill_backend = None

506 + else:

507 + self.prefill_backend = prefill_backend_cls(

508 + num_heads=self.num_heads,

509 + scale=self.scale,

510 + kv_lora_rank=self.kv_lora_rank,

511 + qk_nope_head_dim=self.qk_nope_head_dim,

512 + qk_rope_head_dim=self.qk_rope_head_dim,

513 + v_head_dim=self.v_head_dim,

514 + vllm_config=vllm_config,

515 + )

516

517 self.kv_cache = torch.tensor([])

518

@@ -693,23 +706,28 @@ class MLAAttention(nn.Module, AttentionLayerBase):

706 if fp8_attention and self.kv_cache_dtype != "fp8_ds_mla":

707 kv_cache = kv_cache.view(current_platform.fp8_dtype())

708

696 - # Sparse MLA impls only support forward_mqa (decode-style attention)

697 - is_sparse_impl = isinstance(self.impl, SparseMLAAttentionImpl)

698 -

699 - if is_sparse_impl:

700 - num_mqa_tokens = q.size(0)

701 - num_mha_tokens = 0

702 - else:

703 - assert (

704 - attn_metadata.num_decodes is not None

705 - and attn_metadata.num_prefills is not None

706 - and attn_metadata.num_decode_tokens is not None

709 + assert (

710 + attn_metadata.num_decodes is not None

711 + and attn_metadata.num_prefills is not None

712 + and attn_metadata.num_decode_tokens is not None

713 + )

714 + num_mqa_tokens = attn_metadata.num_decode_tokens

715 + num_mha_tokens = q.size(0) - num_mqa_tokens

716 +

717 + if self.impl.is_sparse and num_mha_tokens > 0:

718 + prefill_max_seq_len = attn_metadata.prefill_max_seq_len # type: ignore[attr-defined]

719 + use_mha = (

720 + self.prefill_backend is not None

721 + and prefill_max_seq_len <= attn_metadata.topk_tokens # type: ignore[attr-defined]

722 + and not self._vllm_config.attention_config.sparse_mla_force_mqa

723 )

708 - num_mqa_tokens = attn_metadata.num_decode_tokens

709 - num_mha_tokens = q.size(0) - num_mqa_tokens

724 + if not use_mha:

725 + num_mqa_tokens = q.size(0)

726 + num_mha_tokens = 0

727

728 mha_use_quant_output = (

729 quant_key is not None

730 + and self.prefill_backend is not None

731 and self.prefill_backend.supports_quant_output(quant_key)

732 and attn_metadata is not None

733 and attn_metadata.prefill is not None

@@ -807,7 +825,7 @@ class MLAAttention(nn.Module, AttentionLayerBase):

825 mqa_q = get_dcp_group().all_gather(mqa_q, dim=1)

826

827 # call decode attn

810 - if not is_sparse_impl:

828 + if not self.impl.is_sparse:

829 assert attn_metadata.decode is not None

830 attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self) # type: ignore[attr-defined]

831

@@ -1398,6 +1416,182 @@ def backend_supports_prefill_query_quantization() -> bool:

1416 )

1417

1418

1419 +def build_mla_chunked_context_metadata(

1420 + *,

1421 + context_lens_cpu: torch.Tensor,

1422 + prefill_query_start_loc_cpu: torch.Tensor,

1423 + num_prefills: int,

1424 + chunked_prefill_workspace: torch.Tensor,

1425 + chunked_prefill_workspace_size: int,

1426 + block_size: int,

1427 + align_chunk_to_block: bool,

1428 + device: torch.device,

1429 + dcp_world_size: int,

1430 + dcp_local_block_size: int,

1431 + dcp_virtual_block_size: int,

1432 +) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None":

1433 + """Build chunked-context metadata for an MLA prefill.

1434 +

1435 + Shared by dense and sparse builders. Splits each prefill's context

1436 + into workspace-sized chunks and, under DCP, plans the per-rank interleaved

1437 + local chunks the all-gather reduction consumes.

1438 +

1439 + Args:

1440 + context_lens_cpu: Per-prefill context length (seq_len - query_len).

1441 + prefill_query_start_loc_cpu: Prefill query cumulative offsets (0-based).

1442 + num_prefills: Number of prefill requests.

1443 + chunked_prefill_workspace: Scratch buffer the context gather writes to.

1444 + chunked_prefill_workspace_size: Row capacity of the workspace.

1445 + block_size: KV cache page size for chunk-start alignment.

1446 + align_chunk_to_block: Round the chunk size down to ``block_size``.

1447 + device: Target device for the returned tensors.

1448 + dcp_world_size: Decode-context-parallel world size (1 if disabled).

1449 + dcp_local_block_size: Per-rank interleave block size for DCP.

1450 + dcp_virtual_block_size: ``dcp_local_block_size * dcp_world_size``.

1451 +

1452 + Returns:

1453 + The chunked-context metadata, or None when no prefill has any context.

1454 + """

1455 + # NOTE: it is recommended you read the `Chunked Prefill` section in the

1456 + # comment at the top of the file before trying to understand this code.

1457 + max_context_len = context_lens_cpu.max().item()

1458 + if max_context_len <= 0:

1459 + return None

1460 + num_prefills_with_context = int((context_lens_cpu > 0).sum().item())

1461 +

1462 + # Currently we allocate an equal amount of workspace for each prefill with

1463 + # context; we could probably use a more advanced algorithm here and allocate

1464 + # more workspace to prefills with longer context lengths.

1465 + max_context_chunk = chunked_prefill_workspace_size // num_prefills_with_context

1466 + if align_chunk_to_block:

1467 + # The `gather_and_maybe_dequant_cache` kernel cannot handle chunk

1468 + # starts that are not aligned to block_size, so round down.

1469 + max_context_chunk = round_down(max_context_chunk, block_size)

1470 + assert max_context_chunk > 0

1471 +

1472 + num_chunks = cdiv(max_context_len, max_context_chunk)

1473 + # e.g. max_context_chunk=256, num_chunks=3, num_prefills=4 ->

1474 + # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]

1475 + # Note(simon): this is done on CPU because of downstream's use of `to_list`.

1476 + chunk_starts = torch.empty(

1477 + num_chunks, num_prefills, dtype=torch.int32, pin_memory=True

1478 + ).copy_(

1479 + torch.arange(num_chunks, dtype=torch.int32)

1480 + .multiply_(max_context_chunk)

1481 + .unsqueeze(1)

1482 + )

1483 + chunk_ends = torch.min(

1484 + context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk

1485 + )

1486 + chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)

1487 +

1488 + cu_seq_lens_cpu = torch.zeros(

1489 + num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True

1490 + )

1491 + torch.cumsum(chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32)

1492 + chunk_total_token = cu_seq_lens_cpu[:, -1]

1493 +

1494 + max_tokens_over_chunk = chunk_total_token.max().item()

1495 + token_to_seq_cpu = torch.zeros(

1496 + (num_chunks, max_tokens_over_chunk), dtype=torch.int32, pin_memory=True

1497 + )

1498 + req_indices = torch.arange(num_prefills, dtype=torch.int32)

1499 + for i in range(num_chunks):

1500 + token_to_seq = torch.repeat_interleave(req_indices, chunk_seq_lens[i])

1501 + token_to_seq_cpu[i, : token_to_seq.shape[0]] = token_to_seq

1502 +

1503 + prefill_tokens_with_context = prefill_query_start_loc_cpu[

1504 + num_prefills_with_context

1505 + ].item()

1506 +

1507 + metadata_cls = MLACommonPrefillMetadata.ChunkedContextMetadata

1508 + if dcp_world_size > 1:

1509 + local_context_lens_allranks = get_dcp_local_seq_lens(

1510 + context_lens_cpu, dcp_world_size, None, dcp_local_block_size

1511 + )

1512 + # Note(qcs): Per-rank local context lengths, padded to

1513 + # `dcp_local_block_size`.

1514 + padded_local_context_lens_cpu: torch.Tensor = (

1515 + cdiv(context_lens_cpu, dcp_virtual_block_size) * dcp_local_block_size

1516 + )

1517 + # Note(hc): The above max_context_chunk already enforces block_size

1518 + # alignment; DCP only requires block_size be divisible by dcp_world_size,

1519 + # because DCP uses cp_gather_cache, which does not require chunk starts

1520 + # aligned to block_size.

1521 + assert max_context_chunk % dcp_world_size == 0

1522 + padded_local_max_context_chunk = (

1523 + cdiv(max_context_chunk, dcp_virtual_block_size) * dcp_local_block_size

1524 + )

1525 + local_chunk_starts = torch.empty(

1526 + num_chunks, num_prefills, dtype=torch.int32, pin_memory=True

1527 + ).copy_(

1528 + torch.arange(num_chunks, dtype=torch.int32)

1529 + .multiply_(padded_local_max_context_chunk)

1530 + .unsqueeze(1)

1531 + )

1532 + local_chunk_ends = torch.min(

1533 + padded_local_context_lens_cpu.unsqueeze(0),

1534 + local_chunk_starts + padded_local_max_context_chunk,

1535 + )

1536 + padded_local_chunk_seq_lens = (local_chunk_ends - local_chunk_starts).clamp(

1537 + min=0

1538 + )

1539 +

1540 + padded_local_cu_seq_lens_cpu = torch.zeros(

1541 + num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True

1542 + )

1543 + torch.cumsum(

1544 + padded_local_chunk_seq_lens,

1545 + dim=1,

1546 + out=padded_local_cu_seq_lens_cpu[:, 1:],

1547 + dtype=torch.int32,

1548 + )

1549 + max_padded_local_tokens = padded_local_cu_seq_lens_cpu[:, -1].max().item()

1550 + padded_local_token_to_seq_cpu = torch.zeros(

1551 + (num_chunks, max_padded_local_tokens), dtype=torch.int32

1552 + )

1553 + for i in range(num_chunks):

1554 + tts = torch.repeat_interleave(req_indices, padded_local_chunk_seq_lens[i])

1555 + padded_local_token_to_seq_cpu[i, : tts.shape[0]] = tts

1556 +

1557 + chunked_context_metadata = metadata_cls(

1558 + cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),

1559 + starts=local_chunk_starts.to(device, non_blocking=True),

1560 + seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),

1561 + max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

1562 + seq_lens=chunk_seq_lens,

1563 + token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),

1564 + chunk_total_token=chunk_total_token.tolist(),

1565 + workspace=chunked_prefill_workspace,

1566 + prefill_tokens_with_context=prefill_tokens_with_context,

1567 + padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),

1568 + local_context_lens_allranks=local_context_lens_allranks.tolist(),

1569 + padded_local_cu_seq_lens=padded_local_cu_seq_lens_cpu.to(

1570 + device, non_blocking=True

1571 + ),

1572 + padded_local_token_to_seq=padded_local_token_to_seq_cpu.to(

1573 + device, non_blocking=True

1574 + ),

1575 + cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),

1576 + chunk_size=padded_local_max_context_chunk,

1577 + )

1578 + else:

1579 + chunked_context_metadata = metadata_cls(

1580 + cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),

1581 + starts=chunk_starts.to(device, non_blocking=True),

1582 + seq_tot=chunk_seq_lens.sum(dim=1).tolist(),

1583 + max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

1584 + seq_lens=chunk_seq_lens,

1585 + token_to_seq=token_to_seq_cpu.to(device, non_blocking=True),

1586 + chunk_total_token=chunk_total_token,

1587 + workspace=chunked_prefill_workspace,

1588 + prefill_tokens_with_context=prefill_tokens_with_context,

1589 + )

1590 +

1591 + assert max(chunked_context_metadata.max_seq_lens) <= chunked_prefill_workspace_size

1592 + return chunked_context_metadata

1593 +

1594 +

1595 class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1596 """

1597 NOTE: Please read the comment at the top of the file before trying to

@@ -1662,8 +1856,6 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1856 context_lens_cpu = (

1857 seq_lens_cpu[reqs_start:num_reqs] - prefill_query_lens_cpu

1858 )

1665 - max_context_len_cpu = context_lens_cpu.max().item()

1666 - num_prefills_with_context_cpu = (context_lens_cpu > 0).sum().item()

1859 prefill_query_start_loc = (

1860 query_start_loc[reqs_start:] - query_start_loc[reqs_start]

1861 )

@@ -1671,188 +1863,19 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]):

1863 query_start_loc_cpu[reqs_start:] - query_start_loc_cpu[reqs_start]

1864 )

1865

1674 - chunked_context_metadata = None

1675 - if max_context_len_cpu > 0:

1676 - # NOTE: it is recommend you read the `Chunked Prefill` section

1677 - # in the comment at the top of the file before trying to

1678 - # understand the following code

1679 -

1680 - # currently we allocate an equal amount of workspace for each

1681 - # prefill in the batch, we could probably use a more advanced

1682 - # algorithm here and allocate more workspace to prefills with

1683 - # longer context lengths

1684 - max_context_chunk = (

1685 - self.chunked_prefill_workspace_size // num_prefills_with_context_cpu

1686 - )

1687 -

1688 - # align max_context_chunk to page_size by rounding down,

1689 - # currently the `gather_and_maybe_dequant_cache` kernel

1690 - # cannot handle `context_chunk_starts` that are not aligned

1691 - # to page_size

1692 - max_context_chunk = round_down(max_context_chunk, self.page_size)

1693 -

1694 - assert max_context_chunk > 0

1695 - num_chunks = cdiv(max_context_len_cpu, max_context_chunk)

1696 -

1697 - # if `max_context_chunk = 256`, `num_chunks = 3`, and

1698 - # `num_prefills_with_context = 4`, create a tensor that looks

1699 - # like

1700 - # [[0, 0, 0, 0], [256, 256, 256, 256], [512, 512, 512, 512]]

1701 - # Note(simon): this is done in CPU because of downstream's

1702 - # of `to_list`.

1703 - chunk_starts = torch.empty(

1704 - num_chunks, num_prefills, dtype=torch.int32, pin_memory=True

1705 - ).copy_(

1706 - torch.arange(num_chunks, dtype=torch.int32)

1707 - .multiply_(max_context_chunk)

1708 - .unsqueeze(1)

1709 - )

1710 - chunk_ends = torch.min(

1711 - context_lens_cpu.unsqueeze(0), chunk_starts + max_context_chunk

1712 - )

1713 - chunk_seq_lens = (chunk_ends - chunk_starts).clamp(min=0)

1714 -

1715 - cu_seq_lens_cpu = torch.zeros(

1716 - num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True

1717 - )

1718 - torch.cumsum(

1719 - chunk_seq_lens, dim=1, out=cu_seq_lens_cpu[:, 1:], dtype=torch.int32

1720 - )

1721 - chunk_total_token = cu_seq_lens_cpu[:, -1]

1722 -

1723 - max_token_num_over_chunk = chunk_total_token.max().item()

1724 - token_to_seq_tensor_cpu = torch.zeros(

1725 - [num_chunks, max_token_num_over_chunk],

1726 - dtype=torch.int32,

1727 - pin_memory=True,

1728 - )

1729 - range_idx = torch.arange(num_prefills, dtype=torch.int32)

1730 - for i in range(num_chunks):

1731 - chunk_token_to_seq_tensor = torch.repeat_interleave(

1732 - range_idx, chunk_seq_lens[i]

1733 - )

1734 - chunk_len = chunk_token_to_seq_tensor.shape[0]

1735 - token_to_seq_tensor_cpu[i, :chunk_len] = chunk_token_to_seq_tensor

1736 -

1737 - if self.dcp_world_size > 1:

1738 - local_context_lens_allranks = get_dcp_local_seq_lens(

1739 - context_lens_cpu,

1740 - self.dcp_world_size,

1741 - None,

1742 - self.dcp_local_block_size,

1743 - )

1744 - # Note(qcs): The max local context lengths

1745 - # padded to `dcp_local_block_size`.

1746 - padded_local_context_lens_cpu: torch.Tensor = (

1747 - cdiv(

1748 - context_lens_cpu,

1749 - self.dcp_virtual_block_size,

1750 - )

1751 - * self.dcp_local_block_size

1752 - )

1753 - # Note(hc): The above max_context_chunk already enforces

1754 - # block_size alignment, DCP just need the block_size can

1755 - # be divisible by dcp_world_size, because DCP use

1756 - # cp_gather_cache which not require `cp_chunk_starts`

1757 - # aligned to page_size.

1758 - assert max_context_chunk % self.dcp_world_size == 0

1759 - padded_local_max_context_chunk_across_ranks = (

1760 - cdiv(

1761 - max_context_chunk,

1762 - self.dcp_virtual_block_size,

1763 - )

1764 - * self.dcp_local_block_size

1765 - )

1766 - local_chunk_starts = torch.empty(

1767 - num_chunks, num_prefills, dtype=torch.int32, pin_memory=True

1768 - ).copy_(

1769 - torch.arange(num_chunks, dtype=torch.int32)

1770 - .multiply_(padded_local_max_context_chunk_across_ranks)

1771 - .unsqueeze(1)

1772 - )

1773 - local_chunk_ends = torch.min(

1774 - padded_local_context_lens_cpu.unsqueeze(0),

1775 - local_chunk_starts

1776 - + padded_local_max_context_chunk_across_ranks,

1777 - )

1778 - padded_local_chunk_seq_lens = (

1779 - local_chunk_ends - local_chunk_starts

1780 - ).clamp(min=0)

1781 -

1782 - padded_local_cu_chunk_seq_lens_cpu = torch.zeros(

1783 - num_chunks, num_prefills + 1, dtype=torch.int32, pin_memory=True

1784 - )

1785 - torch.cumsum(

1786 - padded_local_chunk_seq_lens,

1787 - dim=1,

1788 - out=padded_local_cu_chunk_seq_lens_cpu[:, 1:],

1789 - dtype=torch.int32,

1790 - )

1791 - max_padded_local_tokens_over_chunk = (

1792 - padded_local_cu_chunk_seq_lens_cpu[:, -1].max().item()

1793 - )

1794 - padded_local_token_to_seq_tensor_cpu = torch.zeros(

1795 - [num_chunks, max_padded_local_tokens_over_chunk],

1796 - dtype=torch.int32,

1797 - )

1798 - for i in range(num_chunks):

1799 - chunk_token_to_seq_tensor = torch.repeat_interleave(

1800 - range_idx, padded_local_chunk_seq_lens[i]

1801 - )

1802 - chunk_len = chunk_token_to_seq_tensor.shape[0]

1803 - padded_local_token_to_seq_tensor_cpu[i, :chunk_len] = (

1804 - chunk_token_to_seq_tensor

1805 - )

1806 -

1807 - prefill_tokens_with_context = None

1808 - if num_prefills_with_context_cpu > 0:

1809 - prefill_tokens_with_context = prefill_query_start_loc_cpu[

1810 - num_prefills_with_context_cpu

1811 - ].item()

1812 - _ChunkedMetadata = MLACommonPrefillMetadata.ChunkedContextMetadata

1813 - if self.dcp_world_size > 1:

1814 - chunked_context_metadata = _ChunkedMetadata(

1815 - cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),

1816 - starts=local_chunk_starts.to(device, non_blocking=True),

1817 - seq_tot=padded_local_chunk_seq_lens.sum(dim=1).tolist(),

1818 - max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

1819 - seq_lens=chunk_seq_lens,

1820 - token_to_seq=token_to_seq_tensor_cpu.to(

1821 - device, non_blocking=True

1822 - ),

1823 - chunk_total_token=chunk_total_token.tolist(),

1824 - workspace=self.chunked_prefill_workspace,

1825 - padded_local_chunk_seq_lens=padded_local_chunk_seq_lens.tolist(),

1826 - local_context_lens_allranks=local_context_lens_allranks.tolist(),

1827 - padded_local_cu_seq_lens=padded_local_cu_chunk_seq_lens_cpu.to(

1828 - device, non_blocking=True

1829 - ),

1830 - padded_local_token_to_seq=padded_local_token_to_seq_tensor_cpu.to(

1831 - device, non_blocking=True

1832 - ),

1833 - cu_seq_lens_lst=cu_seq_lens_cpu.tolist(),

1834 - chunk_size=padded_local_max_context_chunk_across_ranks,

1835 - prefill_tokens_with_context=prefill_tokens_with_context,

1836 - )

1837 - else:

1838 - chunked_context_metadata = _ChunkedMetadata(

1839 - cu_seq_lens=cu_seq_lens_cpu.to(device, non_blocking=True),

1840 - starts=chunk_starts.to(device, non_blocking=True),

1841 - seq_tot=chunk_seq_lens.sum(dim=1).tolist(),

1842 - max_seq_lens=chunk_seq_lens.max(dim=1).values.tolist(),

1843 - seq_lens=chunk_seq_lens,

1844 - token_to_seq=token_to_seq_tensor_cpu.to(

1845 - device, non_blocking=True

1846 - ),

1847 - chunk_total_token=chunk_total_token,

1848 - workspace=self.chunked_prefill_workspace,

1849 - prefill_tokens_with_context=prefill_tokens_with_context,

1850 - )

1851 -

1852 - assert (

1853 - max(chunked_context_metadata.max_seq_lens)

1854 - <= self.chunked_prefill_workspace_size

1855 - )

1866 + chunked_context_metadata = build_mla_chunked_context_metadata(

1867 + context_lens_cpu=context_lens_cpu,

1868 + prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,

1869 + num_prefills=num_prefills,

1870 + chunked_prefill_workspace=self.chunked_prefill_workspace,

1871 + chunked_prefill_workspace_size=self.chunked_prefill_workspace_size,

1872 + block_size=self.page_size,

1873 + align_chunk_to_block=True,

1874 + device=device,

1875 + dcp_world_size=self.dcp_world_size,

1876 + dcp_local_block_size=self.dcp_local_block_size,

1877 + dcp_virtual_block_size=self.dcp_virtual_block_size,

1878 + )

1879

1880 prefill_metadata = MLACommonPrefillMetadata(

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

@@ -1985,19 +2008,14 @@ def reorg_kvcache(

2008 return reorganized_kv_c_normed, reorganized_k_pe

2009

2010

1988 -class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2011 +class MLACommonBaseImpl(MLAAttentionImpl[A], Generic[A]):

2012 """

1990 - NOTE: Please read the comment at the top of the file before trying to

1991 - understand this class

2013 + Shared MLA base providing dense-MHA prefill (via the selected

2014 + MLAPrefillBackend) for both dense and sparse impls; subclasses add decode

2015 + (``forward_mqa``).

2016 """

2017

1994 - def fused_output_quant_supported(self, quant_key):

1995 - return quant_key in (

1996 - kFp8StaticTensorSym,

1997 - kNvfp4Dynamic,

1998 - kFp8Dynamic128Sym,

1999 - kFp8Dynamic64Sym,

2000 - )

2018 + _use_flashinfer_concat_mla_k: bool

2019

2020 def __init__(

2021 self,

@@ -2005,60 +2023,25 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2023 head_size: int,

2024 scale: float,

2025 num_kv_heads: int,

2008 - alibi_slopes: list[float] | None,

2009 - sliding_window: int | None,

2026 kv_cache_dtype: str,

2011 - logits_soft_cap: float | None,

2012 - attn_type: str,

2013 - kv_sharing_target_layer_name: str | None,

2014 - # MLA Specific Arguments

2015 - q_lora_rank: int | None,

2027 kv_lora_rank: int,

2028 qk_nope_head_dim: int,

2029 qk_rope_head_dim: int,

2030 qk_head_dim: int,

2031 v_head_dim: int,

2032 kv_b_proj: ColumnParallelLinear,

2022 - # DSV3.2 MLA Specific Arguments

2023 - indexer: object | None = None,

2024 - q_pad_num_heads: int | None = None,

2033 ) -> None:

2026 - if kv_sharing_target_layer_name is not None:

2027 - raise NotImplementedError("KV sharing is not supported for MLA")

2028 -

2034 self.num_heads = num_heads

2035 self.head_size = head_size

2036 self.scale = float(scale)

2037 self.num_kv_heads = num_kv_heads

2038 self.kv_cache_dtype = kv_cache_dtype

2034 -

2035 - self.q_lora_rank = q_lora_rank

2039 self.kv_lora_rank = kv_lora_rank

2040 self.qk_nope_head_dim = qk_nope_head_dim

2041 self.qk_rope_head_dim = qk_rope_head_dim

2042 self.qk_head_dim = qk_head_dim

2043 self.v_head_dim = v_head_dim

2044 self.kv_b_proj = kv_b_proj

2042 - self.indexer = indexer

2043 - self.q_pad_num_heads = q_pad_num_heads

2044 - self.supports_quant_query_input = True

2045 - self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()

2046 -

2047 - # Use flashinfer's optimized concat_mla_k kernel when available.

2048 - # The kernel is optimized for DeepSeek V3 dimensions:

2049 - # num_heads=128, nope_dim=128, rope_dim=64

2050 - self._use_flashinfer_concat_mla_k = (

2051 - has_flashinfer()

2052 - and (self.num_heads == 128)

2053 - and (self.qk_nope_head_dim == 128)

2054 - and (self.qk_rope_head_dim == 64)

2055 - )

2056 -

2057 - self.dcp_world_size: int = -1

2058 -

2059 - self.cp_kv_cache_interleave_size: int = (

2060 - get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size

2061 - )

2045

2046 def _concat_k_nope_k_pe(

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

@@ -2341,7 +2324,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2324

2325 return output, output_lse

2326

2344 - def forward_mha(

2327 + def forward_mha( # type: ignore[override]

2328 self,

2329 q: torch.Tensor,

2330 kv_c_normed: torch.Tensor,

@@ -2428,6 +2411,83 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):

2411 output_prefill = output_prefill.flatten(start_dim=-2)

2412 output.copy_(output_prefill)

2413

2414 +

2415 +class MLACommonImpl(MLACommonBaseImpl[M], Generic[M]):

2416 + """

2417 + NOTE: Please read the comment at the top of the file before trying to

2418 + understand this class

2419 + """

2420 +

2421 + def fused_output_quant_supported(self, quant_key):

2422 + return quant_key in (

2423 + kFp8StaticTensorSym,

2424 + kNvfp4Dynamic,

2425 + kFp8Dynamic128Sym,

2426 + kFp8Dynamic64Sym,

2427 + )

2428 +

2429 + def __init__(

2430 + self,

2431 + num_heads: int,

2432 + head_size: int,

2433 + scale: float,

2434 + num_kv_heads: int,

2435 + alibi_slopes: list[float] | None,

2436 + sliding_window: int | None,

2437 + kv_cache_dtype: str,

2438 + logits_soft_cap: float | None,

2439 + attn_type: str,

2440 + kv_sharing_target_layer_name: str | None,

2441 + # MLA Specific Arguments

2442 + q_lora_rank: int | None,

2443 + kv_lora_rank: int,

2444 + qk_nope_head_dim: int,

2445 + qk_rope_head_dim: int,

2446 + qk_head_dim: int,

2447 + v_head_dim: int,

2448 + kv_b_proj: ColumnParallelLinear,

2449 + # DSV3.2 MLA Specific Arguments

2450 + indexer: object | None = None,

2451 + q_pad_num_heads: int | None = None,

2452 + ) -> None:

2453 + if kv_sharing_target_layer_name is not None:

2454 + raise NotImplementedError("KV sharing is not supported for MLA")

2455 +

2456 + super().__init__(

2457 + num_heads,

2458 + head_size,

2459 + scale,

2460 + num_kv_heads,

2461 + kv_cache_dtype,

2462 + kv_lora_rank,

2463 + qk_nope_head_dim,

2464 + qk_rope_head_dim,

2465 + qk_head_dim,

2466 + v_head_dim,

2467 + kv_b_proj,

2468 + )

2469 + self.q_lora_rank = q_lora_rank

2470 + self.indexer = indexer

2471 + self.q_pad_num_heads = q_pad_num_heads

2472 + self.supports_quant_query_input = True

2473 + self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled()

2474 +

2475 + # Use flashinfer's optimized concat_mla_k kernel when available.

2476 + # The kernel is optimized for DeepSeek V3 dimensions:

2477 + # num_heads=128, nope_dim=128, rope_dim=64

2478 + self._use_flashinfer_concat_mla_k = (

2479 + has_flashinfer()

2480 + and (self.num_heads == 128)

2481 + and (self.qk_nope_head_dim == 128)

2482 + and (self.qk_rope_head_dim == 64)

2483 + )

2484 +

2485 + self.dcp_world_size: int = -1

2486 +

2487 + self.cp_kv_cache_interleave_size: int = (

2488 + get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size

2489 + )

2490 +

2491 @abstractmethod

2492 def forward_mqa(

2493 self,

vllm/model_executor/layers/attention/sparse_mla_attention.py new

+316

@@ -0,0 +1,316 @@

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

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

3 +"""Shared forward_mha implementation and metadata builder for sparse MLA backends."""

4 +

5 +from shutil import which

6 +from typing import TYPE_CHECKING, Generic, TypeVar

7 +

8 +import numpy as np

9 +import torch

10 +

11 +from vllm.distributed import get_dcp_group

12 +from vllm.logger import init_logger

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

14 + MLACommonBaseImpl,

15 + MLACommonPrefillMetadata,

16 + build_mla_chunked_context_metadata,

17 + get_mla_dims,

18 +)

19 +from vllm.platforms import current_platform

20 +from vllm.utils.flashinfer import has_flashinfer

21 +from vllm.utils.torch_utils import np_to_pinned_tensor

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

23 + AttentionMetadata,

24 + AttentionMetadataBuilder,

25 +)

26 +from vllm.v1.attention.backends.utils import split_decodes_and_prefills

27 +

28 +if TYPE_CHECKING:

29 + from vllm.config import VllmConfig

30 + from vllm.model_executor.layers.linear import ColumnParallelLinear

31 + from vllm.v1.attention.backend import CommonAttentionMetadata

32 + from vllm.v1.kv_cache_interface import AttentionSpec

33 +

34 +logger = init_logger(__name__)

35 +

36 +T = TypeVar("T", bound=AttentionMetadata)

37 +

38 +

39 +class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):

40 + metadata_cls: type[T]

41 +

42 + def __init__(

43 + self,

44 + kv_cache_spec: "AttentionSpec",

45 + layer_names: list[str],

46 + vllm_config: "VllmConfig",

47 + device: torch.device,

48 + ) -> None:

49 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)

50 + self.vllm_config = vllm_config

51 + self.device = device

52 + self.model_config = vllm_config.model_config

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

54 + self.topk_tokens: int = vllm_config.model_config.hf_config.index_topk

55 + self.req_id_per_token_buffer = torch.empty(

56 + (vllm_config.scheduler_config.max_num_batched_tokens,),

57 + dtype=torch.int32,

58 + device=device,

59 + )

60 + parallel_config = vllm_config.parallel_config

61 + try:

62 + self.dcp_world_size = get_dcp_group().world_size

63 + self.dcp_rank = get_dcp_group().rank_in_group

64 + except AssertionError:

65 + # DCP might not be initialized in testing

66 + self.dcp_world_size = 1

67 + self.dcp_rank = 0

68 + self.cp_kv_cache_interleave_size = parallel_config.cp_kv_cache_interleave_size

69 + self.dcp_local_block_size = self.cp_kv_cache_interleave_size

70 + self.dcp_virtual_block_size = self.dcp_local_block_size * self.dcp_world_size

71 +

72 + self.chunked_prefill_workspace_size = (

73 + self.determine_chunked_prefill_workspace_size(vllm_config)

74 + )

75 + workspace_head_size = (

76 + self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim

77 + )

78 + workspace_rows = self.chunked_prefill_workspace_size

79 + if self.dcp_world_size > 1:

80 + # DCP gathers each rank's local KV shard into the workspace, so it

81 + # needs an extra 1/DCP rows beyond the TP allocation.

82 + assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0

83 + workspace_rows += self.chunked_prefill_workspace_size // self.dcp_world_size

84 + self.chunked_prefill_workspace = torch.empty(

85 + (workspace_rows, workspace_head_size),

86 + dtype=self.model_config.dtype,

87 + device=device,

88 + )

89 + layer_prefill_backend = vllm_config.compilation_config.static_forward_context[

90 + layer_names[0]

91 + ].prefill_backend

92 + self._prefill_backend = (

93 + layer_prefill_backend.clone() if layer_prefill_backend is not None else None

94 + )

95 +

96 + @staticmethod

97 + def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int:

98 + scheduler_config = vllm_config.scheduler_config

99 + cache_config = vllm_config.cache_config

100 + model_config = vllm_config.model_config

101 + topk_tokens = model_config.hf_config.index_topk

102 +

103 + workspace_size = min(

104 + max(

105 + 8 * model_config.max_model_len,

106 + 4 * scheduler_config.max_num_seqs * cache_config.block_size,

107 + ),

108 + 64 * 1024,

109 + scheduler_config.max_num_seqs * topk_tokens,

110 + )

111 + return max(

112 + workspace_size,

113 + scheduler_config.max_num_seqs * cache_config.block_size,

114 + )

115 +

116 + def _build_req_id_per_token(

117 + self,

118 + common_attn_metadata: "CommonAttentionMetadata",

119 + ) -> torch.Tensor:

120 + num_tokens = common_attn_metadata.num_actual_tokens

121 + starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)

122 + seg_lengths = np.diff(starts)

123 + req_id_per_token = np.repeat(

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

125 + )

126 + self.req_id_per_token_buffer.fill_(0)

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

128 + np_to_pinned_tensor(req_id_per_token), non_blocking=True

129 + )

130 + return self.req_id_per_token_buffer[:num_tokens]

131 +

132 + def _build_chunked_context_fields(

133 + self,

134 + common_attn_metadata: "CommonAttentionMetadata",

135 + num_decodes: int,

136 + num_prefills: int,

137 + prefill_query_lens_cpu: torch.Tensor | None,

138 + ) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None":

139 + if num_prefills == 0 or prefill_query_lens_cpu is None:

140 + return None

141 +

142 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

143 + assert seq_lens_cpu is not None

144 + context_lens_cpu = (

145 + seq_lens_cpu[num_decodes : num_decodes + num_prefills]

146 + - prefill_query_lens_cpu

147 + )

148 + qsl_cpu = common_attn_metadata.query_start_loc_cpu

149 + prefill_query_start_loc_cpu = qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]

150 +

151 + return build_mla_chunked_context_metadata(

152 + context_lens_cpu=context_lens_cpu,

153 + prefill_query_start_loc_cpu=prefill_query_start_loc_cpu,

154 + num_prefills=num_prefills,

155 + chunked_prefill_workspace=self.chunked_prefill_workspace,

156 + chunked_prefill_workspace_size=self.chunked_prefill_workspace_size,

157 + block_size=self.kv_cache_spec.block_size,

158 + align_chunk_to_block=current_platform.is_cuda(),

159 + device=self.device,

160 + dcp_world_size=self.dcp_world_size,

161 + dcp_local_block_size=self.dcp_local_block_size,

162 + dcp_virtual_block_size=self.dcp_virtual_block_size,

163 + )

164 +

165 + def build(

166 + self,

167 + common_prefix_len: int,

168 + common_attn_metadata: "CommonAttentionMetadata",

169 + fast_build: bool = False,

170 + ) -> T:

171 + req_id_per_token = self._build_req_id_per_token(common_attn_metadata)

172 +

173 + num_decodes, num_prefills, num_decode_tokens, _ = split_decodes_and_prefills(

174 + common_attn_metadata,

175 + decode_threshold=self.reorder_batch_threshold or 1,

176 + )

177 + (

178 + prefill_query_start_loc,

179 + prefill_max_query_len,

180 + prefill_query_lens_cpu,

181 + ) = self._build_prefill_fields(common_attn_metadata, num_decodes, num_prefills)

182 +

183 + prefill_max_seq_len = 0

184 + prefill: MLACommonPrefillMetadata | None = None

185 + if num_prefills > 0 and self._prefill_backend is not None:

186 + seq_lens_cpu = common_attn_metadata.seq_lens_cpu_upper_bound

187 + assert seq_lens_cpu is not None

188 + prefill_max_seq_len = int(

189 + seq_lens_cpu[num_decodes : num_decodes + num_prefills].max().item()

190 + )

191 + prefill = MLACommonPrefillMetadata(

192 + block_table=common_attn_metadata.block_table_tensor[num_decodes:, ...],

193 + query_start_loc=prefill_query_start_loc,

194 + max_query_len=prefill_max_query_len,

195 + chunked_context=self._build_chunked_context_fields(

196 + common_attn_metadata,

197 + num_decodes,

198 + num_prefills,

199 + prefill_query_lens_cpu,

200 + ),

201 + q_data_type=self.model_config.dtype,

202 + output_dtype=self.model_config.dtype,

203 + prefill_backend=self._prefill_backend,

204 + )

205 + self._prefill_backend.prepare_metadata(prefill)

206 +

207 + return self.metadata_cls( # type: ignore[call-arg]

208 + num_reqs=common_attn_metadata.num_reqs,

209 + max_query_len=common_attn_metadata.max_query_len,

210 + max_seq_len=common_attn_metadata.max_seq_len,

211 + num_actual_tokens=common_attn_metadata.num_actual_tokens,

212 + query_start_loc=common_attn_metadata.query_start_loc,

213 + slot_mapping=common_attn_metadata.slot_mapping,

214 + block_table=common_attn_metadata.block_table_tensor,

215 + req_id_per_token=req_id_per_token,

216 + seq_lens=common_attn_metadata.seq_lens,

217 + block_size=self.kv_cache_spec.block_size,

218 + topk_tokens=self.topk_tokens,

219 + num_decodes=num_decodes,

220 + num_prefills=num_prefills,

221 + num_decode_tokens=num_decode_tokens,

222 + prefill_max_seq_len=prefill_max_seq_len,

223 + prefill=prefill,

224 + cp_kv_cache_interleave_size=self.cp_kv_cache_interleave_size,

225 + )

226 +

227 + @staticmethod

228 + def _build_prefill_fields(

229 + common_attn_metadata: "CommonAttentionMetadata",

230 + num_decodes: int,

231 + num_prefills: int,

232 + ) -> tuple[

233 + torch.Tensor | None, # prefill_query_start_loc

234 + int, # prefill_max_query_len

235 + torch.Tensor | None, # prefill_query_lens_cpu

236 + ]:

237 + if num_prefills == 0:

238 + return None, 0, None

239 +

240 + offset = common_attn_metadata.query_start_loc[num_decodes]

241 + prefill_query_start_loc = (

242 + common_attn_metadata.query_start_loc[num_decodes:] - offset

243 + )

244 +

245 + qsl_cpu = common_attn_metadata.query_start_loc_cpu

246 + prefill_qsl_cpu = qsl_cpu[num_decodes:] - qsl_cpu[num_decodes]

247 + prefill_query_lens = prefill_qsl_cpu[1:] - prefill_qsl_cpu[:-1]

248 + prefill_max_query_len = int(prefill_query_lens.max().item())

249 +

250 + return (

251 + prefill_query_start_loc,

252 + prefill_max_query_len,

253 + prefill_query_lens,

254 + )

255 +

256 +

257 +class SparseMLACommonImpl(MLACommonBaseImpl[T], Generic[T]):

258 + """Sparse MLA base: shared dense-MHA prefill (from MLACommonBaseImpl) plus the

259 + sparse top-k MQA decode path. Subclasses implement forward_mqa."""

260 +

261 + is_sparse = True

262 +

263 + def __init__(

264 + self,

265 + num_heads: int,

266 + head_size: int,

267 + scale: float,

268 + num_kv_heads: int,

269 + alibi_slopes: list[float] | None,

270 + sliding_window: int | None,

271 + kv_cache_dtype: str,

272 + logits_soft_cap: float | None,

273 + attn_type: str,

274 + kv_sharing_target_layer_name: str | None,

275 + # MLA-specific

276 + q_lora_rank: int | None,

277 + kv_lora_rank: int,

278 + qk_nope_head_dim: int,

279 + qk_rope_head_dim: int,

280 + qk_head_dim: int,

281 + v_head_dim: int,

282 + kv_b_proj: "ColumnParallelLinear",

283 + indexer: object | None = None,

284 + topk_indices_buffer: torch.Tensor | None = None,

285 + q_pad_num_heads: int | None = None,

286 + ) -> None:

287 + super().__init__(

288 + num_heads,

289 + head_size,

290 + scale,

291 + num_kv_heads,

292 + kv_cache_dtype,

293 + kv_lora_rank,

294 + qk_nope_head_dim,

295 + qk_rope_head_dim,

296 + qk_head_dim,

297 + v_head_dim,

298 + kv_b_proj,

299 + )

300 +

301 + # The indexer carries the shared buffer for normal layers and tests;

302 + # the explicitly-passed buffer covers backbone skip layers, whose

303 + # indexer is not constructed (see deepseek_v2.py).

304 + self.topk_indices_buffer: torch.Tensor | None = (

305 + indexer.topk_indices_buffer # type: ignore[attr-defined]

306 + if indexer is not None

307 + else topk_indices_buffer

308 + )

309 +

310 + self._use_flashinfer_concat_mla_k = (

311 + has_flashinfer()

312 + and which("ninja") is not None

313 + and (self.num_heads == 128)

314 + and (self.qk_nope_head_dim == 128)

315 + and (self.qk_rope_head_dim == 64)

316 + )

vllm/v1/attention/backend.py

+4 -81

@@ -781,6 +781,10 @@ class AttentionImplBase(ABC, Generic[T]):

781 method - subclasses define their own forward interfaces.

782 """

783

784 + # Whether this impl uses a sparse (top-k) attention path. Used by MLA to

785 + # route between the dense-MHA prefill and sparse-MQA paths.

786 + is_sparse: ClassVar[bool] = False

787 +

788 # Required attributes that all impls should have

789 num_heads: int

790 head_size: int

@@ -976,7 +980,6 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):

980 ) -> None:

981 raise NotImplementedError

982

979 - @abstractmethod

983 def forward_mha(

984 self,

985 q: torch.Tensor,

@@ -1038,86 +1041,6 @@ class MLAAttentionImpl(AttentionImplBase[T], Generic[T]):

1041 )

1042

1043

1041 -class SparseMLAAttentionImpl(AttentionImplBase[T], Generic[T]):

1042 - """Sparse MLA attention implementation with only forward_mqa method.

1043 -

1044 - Sparse MLA implementations only support decode (MQA-style) attention.

1045 - They do not support prefill (MHA-style) attention.

1046 - """

1047 -

1048 - def fused_output_quant_supported(self, quant_key: "QuantKey"):

1049 - """

1050 - Does this attention implementation support fused output quantization.

1051 - Since MLA quantization is done manually in forward_impl (common code),

1052 - all MLA backends support it by default.

1053 - """

1054 - return quant_key in (

1055 - kFp8StaticTensorSym,

1056 - kNvfp4Dynamic,

1057 - kFp8Dynamic128Sym,

1058 - kFp8Dynamic64Sym,

1059 - )

1060 -

1061 - @abstractmethod

1062 - def __init__(

1063 - self,

1064 - num_heads: int,

1065 - head_size: int,

1066 - scale: float,

1067 - num_kv_heads: int,

1068 - alibi_slopes: list[float] | None,

1069 - sliding_window: int | None,

1070 - kv_cache_dtype: str,

1071 - logits_soft_cap: float | None,

1072 - attn_type: str,

1073 - kv_sharing_target_layer_name: str | None,

1074 - # MLA Specific Arguments

1075 - q_lora_rank: int | None,

1076 - kv_lora_rank: int,

1077 - qk_nope_head_dim: int,

1078 - qk_rope_head_dim: int,

1079 - qk_head_dim: int,

1080 - v_head_dim: int,

1081 - kv_b_proj: "ColumnParallelLinear",

1082 - indexer: object | None = None,

1083 - q_pad_num_heads: int | None = None,

1084 - ) -> None:

1085 - raise NotImplementedError

1086 -

1087 - @abstractmethod

1088 - def forward_mqa(

1089 - self,

1090 - q: torch.Tensor | tuple[torch.Tensor, torch.Tensor],

1091 - kv_c_and_k_pe_cache: torch.Tensor,

1092 - attn_metadata: T,

1093 - layer: AttentionLayer,

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

1095 - """MQA-style decode forward pass."""

1096 - raise NotImplementedError

1097 -

1098 - def do_kv_cache_update(

1099 - self,

1100 - kv_c_normed: torch.Tensor,

1101 - k_pe: torch.Tensor,

1102 - kv_cache: torch.Tensor,

1103 - slot_mapping: torch.Tensor,

1104 - kv_cache_dtype: str,

1105 - k_scale: torch.Tensor,

1106 - ) -> None:

1107 - if kv_cache.numel() == 0:

1108 - return

1109 - from vllm import _custom_ops as ops

1110 -

1111 - ops.concat_and_cache_mla(

1112 - kv_c_normed,

1113 - k_pe.squeeze(1),

1114 - kv_cache,

1115 - slot_mapping.flatten(),

1116 - kv_cache_dtype=kv_cache_dtype,

1117 - scale=k_scale,

1118 - )

1119 -

1120 -

1044 def subclass_attention_backend(

1045 name_prefix: str,

1046 attention_backend_cls: type[AttentionBackend],

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

+36 -63

@@ -4,22 +4,23 @@

4 from dataclasses import dataclass

5 from typing import Any, ClassVar

7 -import numpy as np

7 import torch

9 from vllm.config import VllmConfig

10 from vllm.config.cache import CacheDType

11 +from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata

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

13 + SparseMLACommonImpl,

14 + SparseMLACommonMetadataBuilder,

15 +)

16 from vllm.platforms.interface import DeviceCapability

13 -from vllm.utils.torch_utils import np_to_pinned_tensor

17 from vllm.v1.attention.backend import (

18 AttentionBackend,

19 AttentionCGSupport,

20 AttentionLayer,

21 AttentionMetadata,

19 - AttentionMetadataBuilder,

20 - CommonAttentionMetadata,

22 + MLAAttentionImpl,

23 MultipleOf,

22 - SparseMLAAttentionImpl,

24 )

25 from vllm.v1.attention.backends.fa_utils import flash_attn_supports_mla

26 from vllm.v1.attention.backends.mla.sparse_utils import (

@@ -50,7 +51,7 @@ class FlashAttnMLASparseBackend(AttentionBackend):

51 return FlashAttnMLASparseMetadataBuilder

53 @staticmethod

53 - def get_impl_cls() -> type[SparseMLAAttentionImpl[Any]]:

54 + def get_impl_cls() -> type[MLAAttentionImpl[Any]]:

55 return FlashAttnMLASparseImpl

57 @classmethod

@@ -125,13 +126,21 @@ class FlashAttnMLASparseMetadata(AttentionMetadata):

126

127 block_table: torch.Tensor

128 req_id_per_token: torch.Tensor

129 + seq_lens: torch.Tensor

130 block_size: int = 64

131 topk_tokens: int = 2048

132 + num_decodes: int = 0

133 + num_prefills: int = 0

134 + num_decode_tokens: int = 0

135 + prefill_max_seq_len: int = 0

136 + prefill: MLACommonPrefillMetadata | None = None

137 + cp_kv_cache_interleave_size: int = 1

138

139

140 class FlashAttnMLASparseMetadataBuilder(

133 - AttentionMetadataBuilder[FlashAttnMLASparseMetadata]

141 + SparseMLACommonMetadataBuilder[FlashAttnMLASparseMetadata]

142 ):

143 + metadata_cls = FlashAttnMLASparseMetadata

144 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

145

146 def __init__(

@@ -141,55 +150,16 @@ class FlashAttnMLASparseMetadataBuilder(

150 vllm_config: VllmConfig,

151 device: torch.device,

152 ) -> None:

144 - self.vllm_config = vllm_config

145 - self.layer_names = layer_names

146 - self.kv_cache_spec = kv_cache_spec

147 - self.model_config = vllm_config.model_config

148 - self.device = device

149 -

150 - self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)

151 -

152 - self.topk_tokens = vllm_config.model_config.hf_config.index_topk

153 - self.req_id_per_token_buffer = torch.empty(

154 - (vllm_config.scheduler_config.max_num_batched_tokens,),

155 - dtype=torch.int32,

156 - device=device,

157 - )

158 -

159 - def build(

160 - self,

161 - common_prefix_len: int,

162 - common_attn_metadata: CommonAttentionMetadata,

163 - fast_build: bool = False,

164 - ) -> FlashAttnMLASparseMetadata:

165 - cm = common_attn_metadata

166 - num_tokens = cm.num_actual_tokens

167 - starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)

168 - seg_lengths = np.diff(starts)

169 - req_id_per_token = np.repeat(

170 - np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths

171 - )

172 -

173 - self.req_id_per_token_buffer.fill_(0)

174 - self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

175 - np_to_pinned_tensor(req_id_per_token), non_blocking=True

176 - )

153 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)

154

178 - return FlashAttnMLASparseMetadata(

179 - num_reqs=cm.num_reqs,

180 - max_query_len=cm.max_query_len,

181 - max_seq_len=cm.max_seq_len,

182 - num_actual_tokens=cm.num_actual_tokens,

183 - query_start_loc=cm.query_start_loc,

184 - slot_mapping=cm.slot_mapping,

185 - block_table=cm.block_table_tensor,

186 - req_id_per_token=self.req_id_per_token_buffer[:num_tokens],

187 - block_size=self.kv_cache_spec.block_size,

188 - topk_tokens=self.topk_tokens,

155 + num_q_heads = self.model_config.get_num_attention_heads(

156 + vllm_config.parallel_config

157 )

158 + threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256)

159 + self._init_reorder_batch_threshold(threshold, supports_spec_as_decode=True)

160

161

192 -class FlashAttnMLASparseImpl(SparseMLAAttentionImpl[FlashAttnMLASparseMetadata]):

162 +class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):

163 def __init__(

164 self,

165 num_heads: int,

@@ -217,22 +187,25 @@ class FlashAttnMLASparseImpl(SparseMLAAttentionImpl[FlashAttnMLASparseMetadata])

187 "FlashAttnMLASparseImpl currently supports only FP16/BF16 KV cache."

188 )

189

220 - self.num_heads = num_heads

221 - self.head_size = head_size

222 - self.scale = float(scale)

223 - self.num_kv_heads = num_kv_heads

224 - self.kv_cache_dtype = kv_cache_dtype

225 - self.kv_lora_rank: int = mla_args["kv_lora_rank"]

226 - self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]

227 - self.topk_indices_buffer: torch.Tensor | None = (

228 - indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer

190 + super().__init__(

191 + num_heads,

192 + head_size,

193 + scale,

194 + num_kv_heads,

195 + alibi_slopes,

196 + sliding_window,

197 + kv_cache_dtype,

198 + logits_soft_cap,

199 + attn_type,

200 + kv_sharing_target_layer_name,

201 + indexer=indexer,

202 + topk_indices_buffer=topk_indices_buffer,

203 + **mla_args,

204 )

205 assert self.topk_indices_buffer is not None, (

206 "Indexer or topk_indices_buffer required for sparse MLA"

207 )

208 self.supports_quant_query_input = False

234 - self.dcp_world_size = -1

235 - self.q_pad_num_heads = None

209

210 def forward_mqa(

211 self,

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

+37 -94

@@ -5,37 +5,33 @@

5 from dataclasses import dataclass

6 from typing import TYPE_CHECKING, ClassVar

8 -import numpy as np

8 import torch

10 from vllm import envs

11 from vllm.config import VllmConfig

12 from vllm.config.cache import CacheDType

13 from vllm.logger import init_logger

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

16 - get_mla_dims,

14 +from vllm.model_executor.layers.attention.mla_attention import MLACommonPrefillMetadata

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

16 + SparseMLACommonImpl,

17 + SparseMLACommonMetadataBuilder,

18 )

19 from vllm.platforms.interface import DeviceCapability

19 -from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor

20 +from vllm.utils.torch_utils import is_quantized_kv_cache

21 from vllm.v1.attention.backend import (

22 AttentionBackend,

23 AttentionCGSupport,

24 AttentionLayer,

25 AttentionMetadata,

25 - AttentionMetadataBuilder,

26 AttentionType,

27 - CommonAttentionMetadata,

27 + MLAAttentionImpl,

28 MultipleOf,

29 - SparseMLAAttentionImpl,

29 )

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

31 triton_convert_req_index_to_global_index,

32 triton_filter_and_convert_dcp_index,

33 )

35 -from vllm.v1.attention.backends.utils import (

36 - KVCacheLayoutType,

37 - split_decodes_and_prefills,

38 -)

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

35 from vllm.v1.kv_cache_interface import AttentionSpec

37 if TYPE_CHECKING:

@@ -85,7 +81,7 @@ class FlashInferMLASparseTRTLLMBackend(_FlashInferMLASparseBackendBase):

81 return [32, 64]

83 @staticmethod

88 - def get_impl_cls() -> type[SparseMLAAttentionImpl]:

84 + def get_impl_cls() -> type[MLAAttentionImpl]:

85 return FlashInferMLASparseImpl

87 @classmethod

@@ -162,7 +158,7 @@ class FlashInferMLASparseSM120Backend(_FlashInferMLASparseBackendBase):

158 return [64, 256]

159

160 @staticmethod

165 - def get_impl_cls() -> type[SparseMLAAttentionImpl]:

161 + def get_impl_cls() -> type[MLAAttentionImpl]:

162 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse_sm120 import (

163 FlashInferMLASparseSM120Impl,

164 )

@@ -256,7 +252,10 @@ class FlashInferMLASparseMetadata(AttentionMetadata):

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

253 seq_lens: torch.Tensor

254 num_decodes: int

255 + num_prefills: int

256 num_decode_tokens: int

257 + prefill_max_seq_len: int = 0

258 + prefill: MLACommonPrefillMetadata | None = None

259

260 # Sparse-specific

261 block_size: int = 64

@@ -265,10 +264,11 @@ class FlashInferMLASparseMetadata(AttentionMetadata):

264

265

266 class FlashInferMLASparseMetadataBuilder(

268 - AttentionMetadataBuilder[FlashInferMLASparseMetadata]

267 + SparseMLACommonMetadataBuilder[FlashInferMLASparseMetadata]

268 ):

269 """Builder for FlashInfer MLA Sparse attention metadata."""

270

271 + metadata_cls = FlashInferMLASparseMetadata

272 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

273

274 def __init__(

@@ -278,75 +278,20 @@ class FlashInferMLASparseMetadataBuilder(

278 vllm_config: VllmConfig,

279 device: torch.device,

280 ) -> None:

281 - self.vllm_config = vllm_config

282 - self.layer_names = layer_names

283 - self.kv_cache_spec = kv_cache_spec

284 - self.model_config = vllm_config.model_config

285 - self.device = device

286 -

287 - self.mla_dims = get_mla_dims(self.model_config)

288 - self.topk_tokens = vllm_config.model_config.hf_config.index_topk

281 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)

282

283 + num_q_heads = vllm_config.model_config.get_num_attention_heads(

284 + vllm_config.parallel_config

285 + )

286 + threshold = {8: 128, 16: 128, 32: 128, 64: 256, 128: 1024}.get(

287 + num_q_heads, 1024

288 + )

289 self._init_reorder_batch_threshold(

291 - 1,

290 + threshold,

291 supports_spec_as_decode=True,

292 supports_dcp_with_varlen=True,

293 )

294

296 - self.req_id_per_token_buffer = torch.empty(

297 - (vllm_config.scheduler_config.max_num_batched_tokens,),

298 - dtype=torch.int32,

299 - device=device,

300 - )

301 -

302 - def build(

303 - self,

304 - common_prefix_len: int,

305 - common_attn_metadata: CommonAttentionMetadata,

306 - fast_build: bool = False,

307 - ) -> FlashInferMLASparseMetadata:

308 - cm = common_attn_metadata

309 - num_tokens = cm.num_actual_tokens

310 - assert self.reorder_batch_threshold is not None

311 - num_decodes, _, num_decode_tokens, _ = split_decodes_and_prefills(

312 - cm,

313 - decode_threshold=self.reorder_batch_threshold,

314 - treat_short_extends_as_decodes=True,

315 - )

316 -

317 - # Build req_id_per_token mapping

318 - starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)

319 - seg_lengths = np.diff(starts)

320 - req_id_per_token = np.repeat(

321 - np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths

322 - )

323 -

324 - # Zero-fill for cudagraphs

325 - self.req_id_per_token_buffer.fill_(0)

326 - self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

327 - np_to_pinned_tensor(req_id_per_token), non_blocking=True

328 - )

329 - req_id_per_token_tensor = self.req_id_per_token_buffer[:num_tokens]

330 -

331 - return FlashInferMLASparseMetadata(

332 - num_reqs=cm.num_reqs,

333 - max_query_len=cm.max_query_len,

334 - max_seq_len=cm.max_seq_len,

335 - num_actual_tokens=cm.num_actual_tokens,

336 - query_start_loc=cm.query_start_loc,

337 - slot_mapping=cm.slot_mapping,

338 - block_table=cm.block_table_tensor,

339 - req_id_per_token=req_id_per_token_tensor,

340 - seq_lens=cm.seq_lens,

341 - num_decodes=num_decodes,

342 - num_decode_tokens=num_decode_tokens,

343 - block_size=self.kv_cache_spec.block_size,

344 - topk_tokens=self.topk_tokens,

345 - cp_kv_cache_interleave_size=(

346 - self.vllm_config.parallel_config.cp_kv_cache_interleave_size

347 - ),

348 - )

349 -

295

296 # Global workspace buffer (lazily initialized)

297 _fi_sparse_workspace: torch.Tensor | None = None

@@ -365,7 +310,7 @@ def _get_workspace_buffer(device: torch.device) -> torch.Tensor:

310 return _fi_sparse_workspace

311

312

368 -class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):

313 +class FlashInferMLASparseImpl(SparseMLACommonImpl[FlashInferMLASparseMetadata]):

314 """FlashInfer MLA Sparse implementation.

315

316 Uses the TRT-LLM MLA kernel with sparse_mla_top_k parameter for

@@ -407,22 +352,20 @@ class FlashInferMLASparseImpl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata

352 "FlashInferMLASparseImpl"

353 )

354

410 - self.num_heads = num_heads

411 - self.head_size = head_size

412 - self.scale = float(scale)

413 - self.num_kv_heads = num_kv_heads

414 - self.kv_cache_dtype = kv_cache_dtype

415 -

416 - # MLA-specific dimensions

417 - self.kv_lora_rank: int = mla_args["kv_lora_rank"]

418 - self.qk_nope_head_dim: int = mla_args["qk_nope_head_dim"]

419 - self.qk_rope_head_dim: int = mla_args["qk_rope_head_dim"]

420 -

421 - # The indexer carries the shared buffer for normal layers and tests;

422 - # the explicitly-passed buffer covers backbone skip layers, whose

423 - # indexer is not constructed (see deepseek_v2.py).

424 - self.topk_indices_buffer: torch.Tensor | None = (

425 - indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer

355 + super().__init__(

356 + num_heads,

357 + head_size,

358 + scale,

359 + num_kv_heads,

360 + alibi_slopes,

361 + sliding_window,

362 + kv_cache_dtype,

363 + logits_soft_cap,

364 + attn_type,

365 + kv_sharing_target_layer_name,

366 + indexer=indexer,

367 + topk_indices_buffer=topk_indices_buffer,

368 + **mla_args,

369 )

370

371 self._workspace_buffer: torch.Tensor | None = None

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

+4 -2

@@ -9,7 +9,7 @@ import torch

9 from vllm.v1.attention.backend import (

10 AttentionLayer,

11 AttentionType,

12 - SparseMLAAttentionImpl,

12 + MLAAttentionImpl,

13 )

14 from vllm.v1.attention.backends.mla.flashinfer_mla_sparse import (

15 FlashInferMLASparseMetadata,

@@ -29,9 +29,11 @@ def _kv_scale_format_for_model(model_type: str | None) -> str:

29 return "pow2_fp32"

32 -class FlashInferMLASparseSM120Impl(SparseMLAAttentionImpl[FlashInferMLASparseMetadata]):

32 +class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata]):

33 """SM120 FlashInfer sparse-MLA implementation."""

35 + is_sparse = True

36 +

37 def __init__(

38 self,

39 num_heads: int,

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

+55 -73

@@ -1,31 +1,30 @@

1 # SPDX-License-Identifier: Apache-2.0

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

3 from dataclasses import dataclass

4 -from typing import TYPE_CHECKING, Any, ClassVar

4 +from typing import TYPE_CHECKING, ClassVar

6 -import numpy as np

6 import torch

8 from vllm import _custom_ops as ops

9 from vllm.config import VllmConfig, get_current_vllm_config

10 from vllm.config.cache import CacheDType

11 from vllm.logger import init_logger

13 -from vllm.model_executor.layers.attention.mla_attention import (

14 - get_mla_dims,

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

13 +from vllm.model_executor.layers.attention.sparse_mla_attention import (

14 + SparseMLACommonImpl,

15 + SparseMLACommonMetadataBuilder,

16 )

17 from vllm.platforms import current_platform

18 from vllm.platforms.interface import DeviceCapability

19 from vllm.utils.platform_utils import num_compute_units

19 -from vllm.utils.torch_utils import is_quantized_kv_cache, np_to_pinned_tensor

20 +from vllm.utils.torch_utils import is_quantized_kv_cache

21 from vllm.v1.attention.backend import (

22 AttentionBackend,

23 AttentionCGSupport,

24 AttentionLayer,

25 AttentionMetadata,

25 - AttentionMetadataBuilder,

26 CommonAttentionMetadata,

27 MultipleOf,

28 - SparseMLAAttentionImpl,

28 )

29 from vllm.v1.attention.backends.mla.sparse_utils import (

30 triton_convert_req_index_to_global_index,

@@ -109,7 +108,7 @@ class FlashMLASparseBackend(AttentionBackend):

108 return FlashMLASparseMetadataBuilder

109

110 @staticmethod

112 - def get_impl_cls() -> type[SparseMLAAttentionImpl[Any]]:

111 + def get_impl_cls() -> type["FlashMLASparseImpl"]:

112 return FlashMLASparseImpl

113

114 @classmethod

@@ -159,6 +158,14 @@ class FlashMLASparseMetadata(AttentionMetadata):

158 block_size: int = 64

159 topk_tokens: int = 2048

160

161 + num_decodes: int = 0

162 + num_prefills: int = 0

163 + num_decode_tokens: int = 0

164 + seq_lens: torch.Tensor | None = None

165 + prefill_max_seq_len: int = 0

166 + prefill: MLACommonPrefillMetadata | None = None

167 + cp_kv_cache_interleave_size: int = 1

168 +

169 @dataclass

170 class FP8KernelMetadata:

171 scheduler_metadata: FlashMLASchedMeta

@@ -228,8 +235,11 @@ def get_prefill_workspace_size(max_model_len: int):

235 return max_model_len * 5

236

237

231 -class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetadata]):

238 +class FlashMLASparseMetadataBuilder(

239 + SparseMLACommonMetadataBuilder[FlashMLASparseMetadata]

240 +):

241 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

242 + metadata_cls = FlashMLASparseMetadata

243

244 def __init__(

245 self,

@@ -238,29 +248,30 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

248 vllm_config: VllmConfig,

249 device: torch.device,

250 ) -> None:

241 - self.vllm_config = vllm_config

242 - self.layer_names = layer_names

251 + super().__init__(kv_cache_spec, layer_names, vllm_config, device)

252 cache_config = vllm_config.cache_config

244 - self.kv_cache_spec = kv_cache_spec

245 - self.model_config = vllm_config.model_config

253 parallel_config = vllm_config.parallel_config

247 - self.device = device

254

249 - # Classify single-token queries (plus num_speculative_tokens via

250 - # supports_spec_as_decode=True) as decodes; longer queries go to

251 - # prefill.

252 - self._init_reorder_batch_threshold(1, supports_spec_as_decode=True)

255 + num_q_heads = self.model_config.get_num_attention_heads(parallel_config)

256 + if current_platform.is_device_capability_family(100):

257 + threshold = {8: 128, 16: 128, 32: 128, 64: 256, 128: 1024}.get(

258 + num_q_heads, 1024

259 + )

260 + else:

261 + threshold = {16: 128, 32: 128, 64: 256, 128: 256}.get(num_q_heads, 256)

262 + self._init_reorder_batch_threshold(

263 + threshold,

264 + supports_spec_as_decode=True,

265 + )

266

267 sm_count = num_compute_units(device.index)

268

269 self.num_heads = self.model_config.get_num_attention_heads(parallel_config)

257 - self.mla_dims = get_mla_dims(self.model_config)

270 # FP8 decode kernel only supports h_q = 64 or 128, so we need to pad

271 self.fp8_decode_padded_heads = (

272 FlashMLASparseImpl._compute_fp8_decode_padded_heads(self.num_heads)

273 )

274

263 - self.topk_tokens = vllm_config.model_config.hf_config.index_topk

275 self.use_fp8_kv_cache = cache_config.cache_dtype == "fp8_ds_mla"

276 max_num_seqs = vllm_config.scheduler_config.max_num_seqs

277 # Shape: [max_num_seqs], all elements = topk_tokens (constant for full-CG)

@@ -306,11 +317,6 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

317 dtype=torch.int32,

318 device=device,

319 )

309 - self.req_id_per_token_buffer = torch.empty(

310 - (vllm_config.scheduler_config.max_num_batched_tokens,),

311 - dtype=torch.int32,

312 - device=device,

313 - )

320

321 def _build_fp8_mixed_decode_prefill(

322 self,

@@ -493,51 +499,24 @@ class FlashMLASparseMetadataBuilder(AttentionMetadataBuilder[FlashMLASparseMetad

499 common_attn_metadata: CommonAttentionMetadata,

500 fast_build: bool = False,

501 ) -> FlashMLASparseMetadata:

496 - cm = common_attn_metadata

497 - num_tokens = cm.num_actual_tokens

498 - starts = np.asarray(cm.query_start_loc_cpu, dtype=np.int32)

499 - seg_lengths = np.diff(starts)

500 - req_id_per_token = np.repeat(

501 - np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths

502 - )

503 - # Zero-fill for cudagraphs

504 - self.req_id_per_token_buffer.fill_(0)

505 - self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(

506 - np_to_pinned_tensor(req_id_per_token), non_blocking=True

507 - )

508 - req_id_per_token = self.req_id_per_token_buffer[:num_tokens]

502 + metadata = super().build(common_prefix_len, common_attn_metadata, fast_build)

503

510 - fp8_extra_metadata: (

511 - FlashMLASparseMetadata.FP8SeparatePrefillDecode

512 - | FlashMLASparseMetadata.FP8KernelMetadata

513 - | None

514 - ) = None

504 fp8_use_mixed_batch = self.num_heads < MIN_HEADS_FOR_BF16_PREFILL

505 + metadata.fp8_use_mixed_batch = fp8_use_mixed_batch

506 if self.use_fp8_kv_cache:

507 if fp8_use_mixed_batch:

518 - fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(cm)

508 + metadata.fp8_extra_metadata = self._build_fp8_mixed_decode_prefill(

509 + common_attn_metadata

510 + )

511 else:

520 - fp8_extra_metadata = self._build_fp8_separate_prefill_decode(cm)

521 -

522 - metadata = FlashMLASparseMetadata(

523 - num_reqs=cm.num_reqs,

524 - max_query_len=cm.max_query_len,

525 - max_seq_len=cm.max_seq_len,

526 - num_actual_tokens=cm.num_actual_tokens,

527 - query_start_loc=cm.query_start_loc,

528 - slot_mapping=cm.slot_mapping,

529 - block_table=cm.block_table_tensor,

530 - req_id_per_token=req_id_per_token,

531 - block_size=self.kv_cache_spec.block_size,

532 - topk_tokens=self.topk_tokens,

533 - fp8_extra_metadata=fp8_extra_metadata,

534 - fp8_use_mixed_batch=fp8_use_mixed_batch,

535 - )

512 + metadata.fp8_extra_metadata = self._build_fp8_separate_prefill_decode(

513 + common_attn_metadata

514 + )

515

516 return metadata

517

518

540 -class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):

519 +class FlashMLASparseImpl(SparseMLACommonImpl[FlashMLASparseMetadata]):

520 @staticmethod

521 def _compute_fp8_decode_padded_heads(num_heads: int) -> int:

522 # FP8 decode kernel only supports h_q = 64 or 128

@@ -561,19 +540,22 @@ class FlashMLASparseImpl(SparseMLAAttentionImpl[FlashMLASparseMetadata]):

540 indexer: "Indexer | None" = None,

541 **mla_args,

542 ) -> None:

564 - self.num_heads = num_heads

565 - self.head_size = head_size

566 - self.scale = float(scale)

567 - self.num_kv_heads = num_kv_heads

568 - self.kv_cache_dtype = kv_cache_dtype

569 - self.kv_lora_rank: int = mla_args["kv_lora_rank"]

570 - self.softmax_scale = scale

571 - # The indexer carries the shared buffer for normal layers and tests;

572 - # the explicitly-passed buffer covers backbone skip layers, whose

573 - # indexer is not constructed (see deepseek_v2.py).

574 - self.topk_indices_buffer: torch.Tensor | None = (

575 - indexer.topk_indices_buffer if indexer is not None else topk_indices_buffer

543 + super().__init__(

544 + num_heads,

545 + head_size,

546 + scale,

547 + num_kv_heads,

548 + alibi_slopes,

549 + sliding_window,

550 + kv_cache_dtype,

551 + logits_soft_cap,

552 + attn_type,

553 + kv_sharing_target_layer_name,

554 + indexer=indexer,

555 + topk_indices_buffer=topk_indices_buffer,

556 + **mla_args,

557 )

558 + self.softmax_scale = scale

559 # Prefill BF16 kernel requires 64 on Hopper, 128 on Blackwell

560 self.prefill_padding = (

561 128 if current_platform.is_device_capability_family(100) else 64

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

+6 -3

@@ -238,7 +238,9 @@ def get_max_prefill_buffer_size(vllm_config: VllmConfig):

238

239

240 class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

241 - reorder_batch_threshold: int = 1

241 + # The indexer opts out of the shared reorder-threshold vote (see __init__),

242 + # so this is None; its own split uses self.decode_threshold.

243 + reorder_batch_threshold: int | None = None

244

245 @classmethod

246 def get_cudagraph_support(

@@ -285,7 +287,8 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

287 )

288

289 next_n = self.num_speculative_tokens + 1

288 - self.reorder_batch_threshold += self.num_speculative_tokens

290 + self.decode_threshold = next_n

291 + self.reorder_batch_threshold = None

292 # NOTE: SM100 datacenter GPUs support any next_n natively via the

293 # multi-atom paged MQA logits kernels (FP8 and FP4 indexer

294 # caches). Outside the SM100 family the FP8

@@ -565,7 +568,7 @@ class DeepseekV32IndexerMetadataBuilder(AttentionMetadataBuilder):

568 num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens = (

569 split_decodes_and_prefills(

570 common_attn_metadata,

568 - decode_threshold=self.reorder_batch_threshold,

571 + decode_threshold=self.decode_threshold,

572 require_uniform=not self.use_flattening,

573 )

574 )

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

+16 -3

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

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

5 import functools

6 -from typing import TYPE_CHECKING

6 +from typing import TYPE_CHECKING, ClassVar

8 import torch

@@ -25,7 +25,10 @@ from vllm.v1.attention.backends.fa_utils import (

25 get_flash_attn_version,

26 is_flash_attn_varlen_func_available,

27 )

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

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

29 + MLADimensions,

30 + MLAPrefillBackend,

31 +)

33 if TYPE_CHECKING:

34 from vllm.config import VllmConfig

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

43 class FlashAttnPrefillBackend(MLAPrefillBackend):

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

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

47 + MLADimensions(

48 + qk_nope_head_dim=128,

49 + qk_rope_head_dim=64,

50 + v_head_dim=128,

51 + ),

52 + ]

53 +

54 @staticmethod

55 def get_name() -> str:

56 return "FLASH_ATTN"

@@ -76,7 +87,9 @@ class FlashAttnPrefillBackend(MLAPrefillBackend):

87 )

88 qk_head_dim = qk_nope_head_dim + qk_rope_head_dim

89 self.flash_attn_varlen_func = flash_attn_varlen_func

79 - self.vllm_flash_attn_version = get_flash_attn_version(head_size=qk_head_dim)

90 + self.vllm_flash_attn_version = get_flash_attn_version(

91 + head_size=qk_head_dim, head_size_v=v_head_dim

92 + )

93 if self.vllm_flash_attn_version is not None:

94 self.flash_attn_varlen_func = functools.partial(

95 flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version

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

+4 -2

@@ -24,8 +24,8 @@ from vllm.v1.attention.backend import (

24 AttentionMetadata,

25 AttentionMetadataBuilder,

26 CommonAttentionMetadata,

27 + MLAAttentionImpl,

28 MultipleOf,

28 - SparseMLAAttentionImpl,

29 )

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

31 AiterMLAHelper,

@@ -626,7 +626,9 @@ def reference_mla_sparse_prefill(

626 return (result, lse)

627

628

629 -class ROCMAiterMLASparseImpl(SparseMLAAttentionImpl[ROCMAiterMLASparseMetadata]):

629 +class ROCMAiterMLASparseImpl(MLAAttentionImpl[ROCMAiterMLASparseMetadata]):

630 + is_sparse = True

631 +

632 def __init__(

633 self,

634 num_heads: int,

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

+5 -7

@@ -292,8 +292,7 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

292 - Chunked prefill (aligns with the indexer's chunking)

293 """

294

295 - # Base threshold: query_len <= 1 is decode

296 - reorder_batch_threshold: int = 1

295 + reorder_batch_threshold: int | None = None

296 _cudagraph_support: ClassVar[AttentionCGSupport] = AttentionCGSupport.UNIFORM_BATCH

297

298 def __init__(self, *args, **kwargs):

@@ -315,14 +314,13 @@ class DeepseekSparseSWAMetadataBuilder(AttentionMetadataBuilder):

314 )

315 # Decode can have query_len up to

316 # 1 + (2 if parallel drafting else 1) * num_speculative_tokens.

318 - # This MUST match the flashmla_sparse / indexer threshold so that

319 - # all backends agree on the decode/prefill split.

317 + # sparse_swa has no MQA-vs-dense-MHA routing, so multi-token queries take

318 + # the prefill path and the decode/prefill split stays at that width.

319 spec_mult = (

320 2 if (spec_config is not None and spec_config.parallel_drafting) else 1

321 )

323 - self.decode_threshold = (

324 - self.reorder_batch_threshold + spec_mult * self.num_speculative_tokens

325 - )

322 + self.decode_threshold = 1 + spec_mult * self.num_speculative_tokens

323 + self.reorder_batch_threshold = None

324

325 hf_config = self.vllm_config.model_config.hf_config

326 assert hasattr(hf_config, "sliding_window")

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

+4 -2

@@ -21,7 +21,7 @@ from vllm.v1.attention.backend import (

21 AttentionMetadata,

22 AttentionMetadataBuilder,

23 CommonAttentionMetadata,

24 - SparseMLAAttentionImpl,

24 + MLAAttentionImpl,

25 )

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

27 triton_convert_req_index_to_global_index,

@@ -170,7 +170,9 @@ class XPUMLASparseMetadataBuilder(AttentionMetadataBuilder[XPUMLASparseMetadata]

170 return metadata

171

172

173 -class XPUMLASparseImpl(SparseMLAAttentionImpl[XPUMLASparseMetadata]):

173 +class XPUMLASparseImpl(MLAAttentionImpl[XPUMLASparseMetadata]):

174 + is_sparse = True

175 +

176 def __init__(

177 self,

178 num_heads: int,